fix cls train bug

Author: WenmuZhou

torchocr/data/simple_dataset.py

@@ -247,6 +247,6 @@ def __getitem__(self, properties):
             outs = None
         if outs is None:
             # during evaluation, we should fix the idx to get same results for many times of evaluation.
-            rnd_idx = (idx + 1) % self.__len__()
+            rnd_idx = np.random.randint(self.__len__()) if self.mode == "train" else (idx + 1) % self.__len__()
             return self.__getitem__([img_width, img_height, rnd_idx, wh_ratio])
         return outs

torchocr/engine/trainer.py

@@ -227,9 +227,6 @@ def train(self):
                     train_reader_cost = 0.0
                     train_batch_cost = 0.0
                 reader_start = time.time()
-            if self.local_rank == 0:
-                save_ckpt(self.model, self.cfg, self.optimizer, self.lr_scheduler, epoch, global_step, best_metric,
-                          is_best=False)
             # eval
             if self.local_rank == 0 and epoch > start_eval_epoch and (epoch - start_eval_epoch) % eval_epoch_step == 0:
                 cur_metric = self.eval()
@@ -245,14 +242,17 @@ def train(self):
                 if cur_metric[self.eval_class.main_indicator] >= best_metric[self.eval_class.main_indicator]:
                     best_metric.update(cur_metric)
                     best_metric['best_epoch'] = epoch
-                    save_ckpt(self.model, self.cfg, self.optimizer, self.lr_scheduler, epoch, global_step, best_metric,
-                              is_best=True)
-
                     if self.writer is not None:
                         self.writer.add_scalar(f'EVAL/best_{self.eval_class.main_indicator}',
                                                best_metric[self.eval_class.main_indicator], global_step)
+                    save_ckpt(self.model, self.cfg, self.optimizer, self.lr_scheduler, epoch, global_step, best_metric,
+                              is_best=True)
                 best_str = f"best metric, {', '.join(['{}: {}'.format(k, v) for k, v in best_metric.items()])}"
                 self.logger.info(best_str)
+
+            if self.local_rank == 0:
+                save_ckpt(self.model, self.cfg, self.optimizer, self.lr_scheduler, epoch, global_step, best_metric,
+                          is_best=False)
         best_str = f"best metric, {', '.join(['{}: {}'.format(k, v) for k, v in best_metric.items()])}"
         self.logger.info(best_str)
         if self.writer is not None:
@@ -274,7 +274,6 @@ def eval(self):
             for idx, batch in enumerate(self.valid_dataloader):
                 batch = [t.to(self.device) for t in batch]
                 start = time.time()
-                images = batch[0].to(self.device)
                 if self.scaler:
                     with torch.cuda.amp.autocast():
                         preds = self.model(batch[0], data=batch[1:])

torchocr/postprocess/cls_postprocess.py

@@ -22,5 +22,5 @@ def __call__(self, preds, batch=None, *args, **kwargs):
         decode_out = [(label_list[idx], preds[i, idx]) for i, idx in enumerate(pred_idxs)]
         if batch is None:
             return decode_out
-        label = [(label_list[idx], 1.0) for idx in batch[1].numpy()]
+        label = [(label_list[idx], 1.0) for idx in batch[1].cpu().numpy()]
         return decode_out, label