def load_execution_engine(path, verbose=True):
checkpoint = load_cpu(path)
kwargs = checkpoint['execution_engine_kwargs']
state = checkpoint['execution_engine_state']
kwargs['verbose'] = verbose
model = ModuleNet(**kwargs)
cur_state = model.state_dict()
model.load_state_dict(state)
return model, kwargs
def get_execution_engine(args):
vocab = utils.load_vocab(args.vocab_json)
if args.execution_engine_start_from is not None:
ee, kwargs = utils.load_execution_engine(
args.execution_engine_start_from)
# TODO: Adjust vocab?
else:
kwargs = {
'vocab': vocab,
'feature_dim': parse_int_list(args.feature_dim),
'stem_batchnorm': args.module_stem_batchnorm == 1,
'stem_num_layers': args.module_stem_num_layers,
'module_dim': args.module_dim,
'module_residual': args.module_residual == 1,
'module_batchnorm': args.module_batchnorm == 1,
'classifier_proj_dim': args.classifier_proj_dim,
'classifier_downsample': args.classifier_downsample,
'classifier_fc_layers': parse_int_list(args.classifier_fc_dims),
'classifier_batchnorm': args.classifier_batchnorm == 1,
'classifier_dropout': args.classifier_dropout,
}
ee = ModuleNet(**kwargs)
ee.cuda()
ee.train()
return ee, kwargs