def handle_load(self, load_cmd):
response = Response()
checkpoint = torch.load(load_cmd.filename)
self.model.load_state_dict(checkpoint['model_state_dict'])
self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
response.msg = "Model and optimizer loaded"
response.success = True
return response
def handle_init(self, init_cmd):
response = Response()
response.success = False
self.model = GenericModel(self.model_cfg).to(self.device)
self.optimizer = optim.Adam(self.model.parameters(), lr=self.optim_cfg['learning_rate'])
response.msg = "Model and optimizer reinitialized"
response.success = True
return response
def handle_save(self, save_cmd):
response = Response()
checkpoint = {
'model_state_dict': self.model.state_dict(),
'optimizer_state_dict': self.optimizer.state_dict()
}
torch.save(checkpoint, save_cmd.filename)
response.msg = f"Model and optimizer saved at {save_cmd.filename}"
response.success = True
return response