mirror of
https://github.com/wassname/Castor.git
synced 2026-08-21 11:09:28 +08:00
Michael Tu
d7a631b0a9
* Refactor datasets * Update evaluators * Update trainers * Update main and MP-CNN model * Add serialization util * Fix bugs * Refactoring for NCE to use new parent class
24 lines
630 B
Python
24 lines
630 B
Python
"""
|
|
Utils for serialization
|
|
"""
|
|
import torch
|
|
|
|
|
|
def save_checkpoint(epoch, arch, state_dict, optimizer_state, eval_metric, filename):
|
|
for k, tensor in state_dict.items():
|
|
state_dict[k] = tensor.cpu()
|
|
|
|
state = {
|
|
'epoch': epoch,
|
|
'arch': arch,
|
|
'state_dict': state_dict,
|
|
'optimizer_state': None, # currently do not save optimizer state
|
|
'eval_metric': eval_metric
|
|
}
|
|
torch.save(state, filename)
|
|
|
|
|
|
def load_checkpoint(filename):
|
|
state = torch.load(filename)
|
|
return state['epoch'], state['arch'], state['state_dict'], state['optimizer_state'], state['eval_metric']
|