def model_fn(model_dir):
print('Loading model.')
# First, load the parameters used to create the model.
model_info = {}
model_info_path = os.path.join(model_dir, 'model_info.pth')
with open(model_info_path, 'rb') as f:
model_info = torch.load(f)
print('model_info: {}'.format(model_info))
# Determine the device and construct the model.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = LSTMPredictor(
model_info['input_dim'],
model_info['hidden_dim'],
model_info['output_dim'],
model_info['n_layers']
)
# Load the stored model parameters.
model_path = os.path.join(model_dir, 'model.pth')
with open(model_path, 'rb') as f:
model.load_state_dict(torch.load(f))
model.to(device).eval()
print('Done loading model.')
return model
def load_model(pickled_model, args_dict, test_data):
lstm_settings_dict = { # this works because these variables are set in locals from json_dict
'no_subnets': args_dict['no_subnets'],
'hidden_dims': {
'master': args_dict['hidden_nodes_master'],
'acous': args_dict['hidden_nodes_acous'],
'visual': args_dict['hidden_nodes_visual'],
},
'uses_master_time_rate': {},
'time_step_size': {},
'is_irregular': {},
'layers': num_layers,
'dropout': args_dict['dropout_dict'],
'freeze_glove': args_dict['freeze_glove_embeddings']
}
lstm_settings_dict = test_data.get_lstm_settings_dict(
lstm_settings_dict) # add some extra items to the lstm settings related to the dataset
model = LSTMPredictor(lstm_settings_dict=lstm_settings_dict, feature_size_dict=test_data.get_feature_size_dict(),
batch_size=train_batch_size, seq_length=args_dict['sequence_length'],
prediction_length=prediction_length, embedding_info=test_data.get_embedding_info())
with open(pickled_model, "rb") as model_file:
if torch.cuda.is_available():
model.load_state_dict(torch.load(model_file))
else:
model.load_state_dict(torch.load(model_file, map_location=torch.device('cpu')))
return model
#import glob
#PATH="./model_save/checkpoint_{0}.pkl".format(prediction_length)
PATH = "./model_save/checkpoint.pkl"
fileexists = os.path.isfile(PATH)
optimizer_list = []
if fileexists:
model = LSTMPredictor(lstm_settings_dict=lstm_settings_dict,
feature_size_dict=feature_size_dict,
batch_size=train_batch_size,
seq_length=sequence_length,
prediction_length=prediction_length,
embedding_info=embedding_info)
checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer_list.append(
optim.Adam(model.out.parameters(),
lr=learning_rate,
weight_decay=l2_dict['out']))
for lstm_key in model.lstm_dict.keys():
optim_instance = optim.Adam(model.lstm_dict[lstm_key].parameters())
# optim_instance.add_param_group({"params": torch.tensor([0])})
optimizer_list.append(optim_instance)
for i, optimizer in enumerate(optimizer_list):
optimizer.load_state_dict(checkpoint['optimizer_state_dict' + str(i)])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
model.eval()