(self)
| 85 | self.train() |
| 86 | |
| 87 | def train(self): |
| 88 | |
| 89 | for epoch in range(self.args.start_epoch, self.args.start_epoch + self.args.epochs): |
| 90 | if self.args.distributed: |
| 91 | self.train_sampler.set_epoch(epoch) |
| 92 | self.train_epoch(epoch) |
| 93 | if is_main_process(): |
| 94 | torch.save({ |
| 95 | 'epoch': epoch, |
| 96 | 'model': self.network_without_ddp.state_dict(), |
| 97 | 'optimizer': self.optimizer.state_dict(), |
| 98 | "lr_scheduler": self.lr_scheduler.state_dict()}, |
| 99 | "{}/model_{:0>6}.ckpt".format(self.args.log_dir, epoch)) |
| 100 | |
| 101 | if (epoch % self.args.eval_freq == 0) or (epoch == self.args.epochs - 1): |
| 102 | self.validate(epoch) |
| 103 | torch.cuda.empty_cache() |
| 104 | |
| 105 | def train_epoch(self, epoch): |
| 106 | self.network.train() |
no test coverage detected