Hi, my custom model looks like this:
mask_conv_model = models.resnet50(pretrained=True)
for param in mask_conv_model.parameters():
param.requires_grad = False
num_ftrs = mask_conv_model.fc.in_features
mask_conv_model.fc = nn.Sequential(nn.Linear(num_ftrs, masked_output_ftrs), nn.ReLU(), nn.Dropout())
for name, param in mask_conv_model.named_parameters():
if param.requires_grad == True:
mask_params.append(param)
mask_conv_model_optimizer = optim.Adam(mask_params, lr=masked_learning_rate)
Then later on, I save my model:
torch.save({
'epoch': epoch,
'target_seq_len': hype['target_seq_len'],
'input_seq_len': hype['input_seq_len'],
'masked_output_ftrs': hype['masked_output_ftrs'],
'mask_conv': dict_of_models['MASKED CONV MODEL'].state_dict(),
}, os.path.join(save_dir, "MASK_CONV.tar"))
Now, to load the model I have to write a lot of code:
mask_conv_model = models.resnet50(pretrained=True)
mask_checkpoint = torch.load(mask_dir)
epoch = mask_checkpoint['epoch']
masked_output_ftrs = mask_checkpoint['masked_output_ftrs']
num_ftrs = mask_conv_model.fc.in_features
mask_conv_model.fc = nn.Linear(num_ftrs, masked_output_ftrs)
mask_conv_model.load_state_dict(mask_checkpoint['mask_conv'])
Is there a one-line method of loading models? Something like this:
mask_checkpoint = torch.load(mask_dir)
mask_conv_model = mask_checkpoint.load_state_dict(mask_checkpoint['mask_conv'])