mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Save & load amp state

Browse Source
This commit is contained in:
James Betker 2020-06-18 11:38:48 -06:00
parent 2e3b6bad77
commit efc80f041c
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
codes/models/base_model.py
View File

@ -4,6 +4,7 @@ import torch
import torch.nn as nn import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel from torch.nn.parallel import DistributedDataParallel
import utils.util import utils.util
from apex import amp
class BaseModel(): class BaseModel():
@ -109,6 +110,7 @@ class BaseModel():
state['schedulers'].append(s.state_dict()) state['schedulers'].append(s.state_dict())
for o in self.optimizers: for o in self.optimizers:
state['optimizers'].append(o.state_dict()) state['optimizers'].append(o.state_dict())
state['amp'] = amp.state_dict()
save_filename = '{}.state'.format(iter_step) save_filename = '{}.state'.format(iter_step)
save_path = os.path.join(self.opt['path']['training_state'], save_filename) save_path = os.path.join(self.opt['path']['training_state'], save_filename)
torch.save(state, save_path) torch.save(state, save_path)
@ -129,3 +131,4 @@ class BaseModel():
self.optimizers[i].load_state_dict(o) self.optimizers[i].load_state_dict(o)
for i, s in enumerate(resume_schedulers): for i, s in enumerate(resume_schedulers):
self.schedulers[i].load_state_dict(s) self.schedulers[i].load_state_dict(s)
amp.load_state_dict(resume_state['amp'])