Save training state during training, which will be used for resuming
(self, epoch, iter_step)
| 98 | print("Wrong!!!!! pretrained path not exists") |
| 99 | |
| 100 | def save_training_state(self, epoch, iter_step): |
| 101 | """Save training state during training, which will be used for resuming""" |
| 102 | state = {'epoch': epoch, 'iter': iter_step, 'schedulers': [], 'optimizers': []} |
| 103 | for s in self.schedulers: |
| 104 | state['schedulers'].append(s.state_dict()) |
| 105 | for o in self.optimizers: |
| 106 | state['optimizers'].append(o.state_dict()) |
| 107 | save_filename = '{}.state'.format(iter_step) |
| 108 | save_path = os.path.join(self.opt['path']['training_state'], save_filename) |
| 109 | torch.save(state, save_path) |
| 110 | |
| 111 | def resume_training(self, resume_state): |
| 112 | """Resume the optimizers and schedulers for training""" |