def load_model(exp, save_dir):
# load trained model
tags_path = exp.get_data_path(exp.name, exp.version)
tags_path = os.path.join(tags_path, 'meta_tags.csv')
checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
weights_dir = os.path.join(save_dir, checkpoints[0])
trained_model = LightningTemplateModel.load_from_metrics(
weights_path=weights_dir, tags_csv=tags_path, on_gpu=True)
assert trained_model is not None, 'loading model failed'
return trained_model
def get_model():
# set up model with these hyperparams
root_dir = os.path.dirname(os.path.realpath(__file__))
hparams = Namespace(
**{
'drop_prob': 0.2,
'batch_size': 32,
'in_features': 28 * 28,
'learning_rate': 0.001 * 8,
'optimizer_name': 'adam',
'data_root': os.path.join(root_dir, 'mnist'),
'out_features': 10,
'hidden_dim': 1000
})
model = LightningTemplateModel(hparams)
return model, hparams
def get_model():
# set up model with these hyperparams
hparams = get_hparams()
model = LightningTemplateModel(hparams)
return model, hparams