def save_model(self, epoch, directory, is_best=False):
names = self.get_names_if_none()
for name in names:
save_checkpoint(
{
'state_dict': self._models[name].state_dict(),
'epoch': epoch + 1,
'optimizer': self._optims[name].state_dict(),
'scheduler': self._scheds[name].state_dict()
},
osp.join(directory, name),
is_best=is_best)
def save_model(self, epoch, directory, is_best=False):
names = self.get_model_names()
for name in names:
model_dict = self._models[name].state_dict()
optim_dict = None
if self._optims[name] is not None:
optim_dict = self._optims[name].state_dict()
sched_dict = None
if self._scheds[name] is not None:
sched_dict = self._scheds[name].state_dict()
save_checkpoint(
{
'state_dict': model_dict,
'epoch': epoch + 1,
'optimizer': optim_dict,
'scheduler': sched_dict
},
osp.join(directory, name),
is_best=is_best)