(self, epoch, flag, flag_infer_ema=False)
| 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 |
| 896 | print_freq_lst = [i * print_freq for i in range(1, print_period)] + [total_batch - 1] |
| 897 | |
| 898 | # start loops |
| 899 | for batch_id, batch in enumerate(self.train_loader): |
| 900 | # data |
| 901 | batch.to(self.device, non_blocking=True) |
| 902 | |
| 903 | # forward |
| 904 | self.optim.zero_grad() |
| 905 | if flag_return_losses: |
| 906 | pred, loss, record_losses = self.model(batch, flag_return_losses=True) |
| 907 | else: |
| 908 | pred, loss = self.model(batch) |
| 909 | |
| 910 | # records |
| 911 | cur_matrix = self.matrix(pred) |
| 912 | if (not self.cfg['distributed']) or (self.cfg['distributed'] and dist.get_rank() == 0): |
| 913 | # logger.info(f"Iter:{batch_id}/{total_batch} - {str(cur_matrix)}") |
| 914 | # print(cur_matrix) |
| 915 | pass |
| 916 | if batch_id == 0: |
| 917 | for key in cur_matrix: |
| 918 | train_matrix[key] = AVGMeter() |
| 919 | |
| 920 | for key in cur_matrix: |
| 921 | train_matrix[key].update(cur_matrix[key]) |
| 922 | |
| 923 | # backwards |
| 924 | loss.backward() |
| 925 | clip_grad_norm_(self.model.parameters(), 1.0) |
| 926 | self.optim.step() |
| 927 | train_loss.update(loss.item()) |
| 928 | |
| 929 | # update ema |
| 930 | if self.flag_use_ema_model: |
| 931 | if self.cfg['distributed']: |
| 932 | self.model.module.update_ema_model(epoch, batch_id + epoch * total_batch, total_batch) |
| 933 | else: |
| 934 | self.model.update_ema_model(epoch, batch_id + epoch * total_batch, total_batch) |
| 935 | |
| 936 | # print stats |
| 937 | if (batch_id in print_freq_lst) or ((batch_id + 1) == total_batch): |
| 938 | if self.cfg['distributed']: |
no test coverage detected