def load_model(args, verbose=False): if args.command != 'train' and not os.path.isfile(args.model): raise RuntimeError('Model file {} does not exist!'.format(args.model)) model = None state = {} _, ext = os.path.splitext(args.model) if args.command == 'train' and (not os.path.exists(args.model) or args.override): if verbose: print('Initializing model...') model = Model(args.backbone, args.classes) model.initialize(args.fine_tune) if verbose: print(model) elif ext == '.pth' or ext == '.torch': if verbose: print('Loading model from {}...'.format(os.path.basename(args.model))) model, state = Model.load(args.model) if verbose: print(model) elif args.command == 'infer' and ext in ['.engine', '.plan']: model = None else: raise RuntimeError('Invalid model format "{}"!'.format(args.ext)) state['path'] = args.model return model, state
def load_model(args, verbose=False): config = {} if args.config: with open(args.config, 'r') as config_file: config = json.load(config_file) if args.command != 'train' and not os.path.isfile(args.model): raise RuntimeError('Model file {} does not exist!'.format(args.model)) model = None state = {} model_name, ext = os.path.splitext(args.model) if args.command == 'train' and (not os.path.exists(args.model) or args.override): if verbose: print('Initializing model...') model = Model(backbones=args.backbone, classes=args.classes, rotated_bbox=args.rotated_bbox, config=config) model.initialize(args.fine_tune) if verbose: print(model) elif ext == '.pth' or ext == '.torch': if verbose: print('Loading model from {}...'.format( os.path.basename(args.model))) exporting = False if args.command == 'eval': exporting = True model, state = Model.load(filename=args.model, rotated_bbox=args.rotated_bbox, config=config, exporting=exporting) if verbose: print(model) elif args.command == 'infer' and ext in ['.engine', '.plan']: model = None else: raise RuntimeError('Invalid model format "{}"!'.format(args.ext)) state['path'] = model_name return model, state