(self)
| 404 | self.logger.info(log_msg) |
| 405 | |
| 406 | def save(self): |
| 407 | if ((self.tmp.current_step + 1) % self.config.get('ckpt_interval', 1000) == 0 or |
| 408 | self.tmp.current_step + 1 == self.config.max_iter |
| 409 | ) and self.ginfo.task_rank == 0: |
| 410 | save_state({ |
| 411 | 'step': self.tmp.current_step+1, |
| 412 | 'state_dict': self.model.state_dict(), |
| 413 | 'optimizer': self.optimizer.state_dict(), |
| 414 | }, '{}/ckpt_task{}'.format(self.ckpt_path, self.ginfo.task_id), 'newest') |
| 415 | if self.config.get('save_interval', -1) > 0 and (self.tmp.current_step+1) % self.config.save_interval == 0 and self.ginfo.task_rank == 0: |
| 416 | save_state({ |
| 417 | 'step': self.tmp.current_step+1, |
| 418 | 'state_dict': self.model.state_dict(), |
| 419 | 'optimizer': self.optimizer.state_dict(), |
| 420 | }, '{}/ckpt_task{}'.format(self.ckpt_path, self.ginfo.task_id), self.tmp.current_step+1) |
| 421 | |
| 422 | if (self.tmp.current_step+1) in self.config.get('save_iters', []) and self.ginfo.task_rank == 0: |
| 423 | save_state({ |
| 424 | 'step': self.tmp.current_step+1, |
| 425 | 'state_dict': self.model.state_dict(), |
| 426 | 'optimizer': self.optimizer.state_dict(), |
| 427 | }, '{}/ckpt_task{}'.format(self.ckpt_path, self.ginfo.task_id), self.tmp.current_step+1) |
| 428 | |
| 429 | if (self.tmp.current_step + 1) % self.config.get('ckpt_interval', 1000) == 0 and self.config.get('cuda_empty_cache', False): |
| 430 | torch.cuda.empty_cache() |
| 431 | |
| 432 | def prepare_data(self): |
| 433 | self.tmp.input_var = dict() |
no test coverage detected