| 544 | self.logger.info(log_msg) |
| 545 | |
| 546 | def save(self): |
| 547 | config = self.config |
| 548 | tmp = self.tmp |
| 549 | ginfo = self.ginfo |
| 550 | if config.save_interval > 0 and (tmp.current_step+1) % 1000 == 0 and ginfo.task_rank == 0: |
| 551 | save_state({ |
| 552 | 'step': tmp.current_step+1, |
| 553 | 'backbone_args': config.get('backbone', None), |
| 554 | 'neck_args': config.get('neck', None), |
| 555 | 'decoder_args': config.get('decoder', None), |
| 556 | 'state_dict': self.model.state_dict(), |
| 557 | 'optimizer': self.optimizer.state_dict(), |
| 558 | }, '{}/ckpt_task{}'.format(self.ckpt_path, ginfo.task_id), 'newest') |
| 559 | if config.save_interval > 0 and (tmp.current_step+1) % config.save_interval == 0 and ginfo.task_rank == 0: |
| 560 | save_state({ |
| 561 | 'step': tmp.current_step+1, |
| 562 | 'backbone_args': config.get('backbone', None), |
| 563 | 'neck_args': config.get('neck', None), |
| 564 | 'decoder_args': config.get('decoder', None), |
| 565 | 'state_dict': self.model.state_dict(), |
| 566 | 'optimizer': self.optimizer.state_dict(), |
| 567 | }, '{}/ckpt_task{}'.format(self.ckpt_path, ginfo.task_id), tmp.current_step+1) |
| 568 | if config.save_interval > 0 and tmp.current_step+1 == len(self.loader) and ginfo.task_rank == 0: |
| 569 | save_state({ |
| 570 | 'step': tmp.current_step+1, |
| 571 | 'backbone_args': config.get('backbone', None), |
| 572 | 'neck_args': config.get('neck', None), |
| 573 | 'decoder_args': config.get('decoder', None), |
| 574 | 'state_dict': self.model.state_dict(), |
| 575 | 'optimizer': self.optimizer.state_dict(), |
| 576 | }, '{}/ckpt_task{}'.format(self.ckpt_path, ginfo.task_id), 'final') |
| 577 | |
| 578 | def post_run(self): |
| 579 | pass |