Spaces:
Sleeping
Sleeping
File size: 742 Bytes
383bfb8 6f74e93 383bfb8 6f74e93 383bfb8 6f74e93 383bfb8 6f74e93 383bfb8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
from .resnet import *
from .mobilenet_v2 import *
visual_backbone_list = ['ResNet10', 'mobilenet_v2']
def get_visual_backbone(args):
if args.visual_backbone in visual_backbone_list:
model = eval(args.visual_backbone)()
if hasattr(args, 'pretrain_vis_path') and args.pretrain_vis_path != "":
model.load_model(pretrain=args.pretrain_vis_path)
if hasattr(args, 'logger'):
args.logger.info("Visual backbone has been loaded...")
else:
if hasattr(args, 'logger'):
args.logger.info("Visual backbone choose to train from scratch")
return model
else:
raise NotImplementedError("Unsupported Backbone: {}".format(args.visual_backbone)) |