| 836 | |
| 837 | # torch.autograd.set_detect_anomaly(True) |
| 838 | class SubclassOven(Oven): |
| 839 | def __init__(self, cfg, log_dir, out_dir): |
| 840 | super(SubclassOven,self).__init__(cfg) |
| 841 | self.cfg = cfg |
| 842 | self.ngpus = cfg.get('ngpus', 1) |
| 843 | if self.ngpus == 0: |
| 844 | self.device = 'cpu' |
| 845 | else: |
| 846 | self.device = 'cuda' |
| 847 | if (not self.cfg['distributed']) or (self.cfg['distributed'] and dist.get_rank() == 0): |
| 848 | self.reporter = Reporter(cfg, log_dir) |
| 849 | self.out_dir = out_dir |
| 850 | self.matrix = self._init_matrix() |
| 851 | self.train_loader, self.valid_loader = self._init_data() |
| 852 | self.criterion = self._init_criterion() |
| 853 | self.model = self._init_model() |
| 854 | self.optim, self.scheduler = self._init_optim() |
| 855 | checkpt_path = self.cfg['model'].get("resume_ckpt_path", "") |
| 856 | # self.resume_training = True if os.path.exists(os.path.join(self.cfg['log_path'], 'ckpt_latest.pt')) else False |
| 857 | self.resume_training = True if os.path.exists(checkpt_path) else False |
| 858 | self.checkpt_path = checkpt_path |
| 859 | # using ema info |
| 860 | self.flag_use_ema_model = self.cfg['model'].get("flag_use_ema", False) |
| 861 | |
| 862 | def _init_matrix(self): |
| 863 | if self.cfg['model']['matrix'] == 'vm_va': |
| 864 | return vm_va_matrix |
| 865 | else: |
| 866 | raise TypeError(f"No such of matrix {self.cfg['model']['matrix']}") |
| 867 | |
| 868 | def _init_model(self): |
| 869 | model = IterGCN(**self.cfg['model']) |
| 870 | model = model.to(self.device) |
| 871 | return model |
| 872 | |
| 873 | def _init_criterion(self): |
| 874 | if self.cfg['loss']['type'] == "deltapq_loss": |
| 875 | return deltapq_loss |
| 876 | elif self.cfg['loss']['type'] == "bi_deltapq_loss": |
| 877 | return bi_deltapq_loss |
| 878 | else: |
| 879 | raise TypeError(f"No such of loss {self.cfg['loss']['type']}") |
| 880 | |
| 881 | def exec_epoch(self, epoch, flag, flag_infer_ema=False): |
| 882 | flag_return_losses = self.cfg.get("flag_return_losses", False) |
| 883 | if flag == 'train': |
| 884 | if (not self.cfg['distributed']) or (self.cfg['distributed'] and dist.get_rank() == 0): |
| 885 | logger.info(f'-------------------- Epoch: {epoch+1} --------------------') |
| 886 | self.model.train() |
| 887 | if self.cfg['distributed']: |
| 888 | self.train_loader.sampler.set_epoch(epoch) |
| 889 | |
| 890 | # record vars |
| 891 | train_loss = AVGMeter() |
| 892 | train_matrix = dict() |
| 893 | total_batch = len(self.train_loader) |
| 894 | print_period = self.cfg['train'].get('logs_freq', 8) |
| 895 | print_freq = total_batch // print_period |