def restore(self, model: SimNet, loss_fn: nn.Module, optimizer: Optimizer, current_loss: str):
checkpoint = torch.load(self.path, map_location=common.DEVICE)
model.load_common_state_dict(checkpoint['common_state_dict'])
if current_loss == checkpoint['trained_loss']:
loss_fn.load_state_dict(checkpoint['loss_state_dict'])
if model.loss_module is not None:
model.loss_module.load_state_dict(checkpoint['loss_module_state_dict'])
optimizer.load_state_dict(checkpoint['optim_state_dict'])
print(f"Recovered Model. Epoch {checkpoint['epoch']}. Test Metric {checkpoint['accuracy']}")
return checkpoint
def restore(self, model: MetricNet, loss_fn: nn.Module, optimizer: Optimizer, current_loss: str):
checkpoint = torch.load(self.path, map_location=constants.DEVICE)
model.load_encoder_state_dict(checkpoint['common_state_dict'])
if current_loss == checkpoint['trained_loss']:
loss_fn.load_state_dict(checkpoint['loss_state_dict'])
if model.classifier is not None:
model.classifier.load_state_dict(checkpoint['loss_module_state_dict'])
if self.restore_optimizer:
optimizer.load_state_dict(checkpoint['optim_state_dict'])
print(f"Recovered Model. Epoch {checkpoint['epoch']}. Dev Metric {checkpoint['accuracy']}")
return checkpoint