(self)
| 994 | |
| 995 | |
| 996 | def train(self): |
| 997 | if self.ngpus > 1: |
| 998 | dummy_batch_data = next(iter(self.train_loader)) |
| 999 | dummy_batch_data.to(self.device, non_blocking=True) |
| 1000 | with torch.no_grad(): |
| 1001 | if self.flag_use_ema_model: |
| 1002 | _ = self.model(dummy_batch_data, num_loop_infer=1) |
| 1003 | _ = self.model(dummy_batch_data, num_loop_infer=1, flag_use_ema_infer=True) |
| 1004 | else: |
| 1005 | _ = self.model(dummy_batch_data, num_loop_infer=1) |
| 1006 | |
| 1007 | if (not self.cfg['distributed']) or (self.cfg['distributed'] and dist.get_rank() == 0): |
| 1008 | logger.info(f'==================== Total number of parameters: {count_parameters(self.model):.3f}M') |
| 1009 | |
| 1010 | local_rank = int(os.environ["LOCAL_RANK"]) |
| 1011 | self.model = torch.nn.parallel.DistributedDataParallel( |
| 1012 | self.model, |
| 1013 | device_ids=[local_rank], |
| 1014 | output_device=local_rank, |
| 1015 | find_unused_parameters=True, |
| 1016 | # find_unused_parameters=False |
| 1017 | ) |
| 1018 | else: |
| 1019 | dummy_batch_data = next(iter(self.train_loader)) |
| 1020 | dummy_batch_data.to(self.device, non_blocking=True) |
| 1021 | with torch.no_grad(): |
| 1022 | # _ = self.model(dummy_batch_data, num_loop_infer=1) |
| 1023 | if self.flag_use_ema_model: |
| 1024 | _ = self.model(dummy_batch_data, num_loop_infer=1) |
| 1025 | _ = self.model(dummy_batch_data, num_loop_infer=1, flag_use_ema_infer=True) |
| 1026 | else: |
| 1027 | _ = self.model(dummy_batch_data, num_loop_infer=1) |
| 1028 | logger.info(f'==================== Total number of parameters: {count_parameters(self.model):.3f}M') |
| 1029 | |
| 1030 | |
| 1031 | if not self.resume_training: |
| 1032 | self.perform_best = np.Infinity |
| 1033 | self.perform_best_ep = -1 |
| 1034 | self.start_epoch = 0 |
| 1035 | self.perform_best_metrics = {} |
| 1036 | else: |
| 1037 | self.perform_best, self.perform_best_ep, self.start_epoch, self.perform_best_metrics = self._init_training_wt_checkpoint(self.checkpt_path) |
| 1038 | |
| 1039 | local_best = self.perform_best |
| 1040 | local_best_ep = self.perform_best_ep |
| 1041 | local_best_metrics = self.perform_best_metrics |
| 1042 | if self.flag_use_ema_model: |
| 1043 | local_best_ema = self.perform_best |
| 1044 | local_best_ep_ema = self.perform_best_ep |
| 1045 | local_best_metrics_ema =self.perform_best_metrics |
| 1046 | for epoch in range(self.start_epoch, self.cfg['train']['epochs']): |
| 1047 | with Timer(rest_epochs=self.cfg['train']['epochs'] - (epoch + 1)) as timer: |
| 1048 | train_loss, train_matrix = self.exec_epoch(epoch, flag='train') |
| 1049 | valid_loss, val_matrix = self.exec_epoch(epoch, flag='valid') |
| 1050 | if self.flag_use_ema_model: |
| 1051 | valid_loss_ema, valid_matrix_ema = self.exec_epoch(epoch, flag='valid', |
| 1052 | flag_infer_ema=True) |
| 1053 | if self.scheduler: |
no test coverage detected