perf(torch): reduce rec logits on device before copying them to the CPU (#757)

* perf(torch): reduce rec logits on device before copying them to the CPU

Signed-off-by: ExtReMLapin <3909752+ExtReMLapin@users.noreply.github.com>

* refactor: keep __call__ returning raw outputs, add InferSession.infer_argmax for rec

Signed-off-by: ExtReMLapin <3909752+ExtReMLapin@users.noreply.github.com>

---------

Signed-off-by: ExtReMLapin <3909752+ExtReMLapin@users.noreply.github.com>

python/rapidocr/ch_ppocr_rec/main.py

@@ -128,7 +128,7 @@ class TextRecognizer:
                 norm_img_batch.append(norm_img[np.newaxis, :])
             norm_img_batch = np.concatenate(norm_img_batch).astype(np.float32)
 
-            preds = self.session(norm_img_batch)
+            preds = self.session.infer_argmax(norm_img_batch)
             line_results, word_results = self.postprocess_op(
                 preds,
                 return_word_box,

python/rapidocr/ch_ppocr_rec/utils.py

@@ -22,8 +22,11 @@ class CTCLabelDecode:
     def __call__(
         self, preds: np.ndarray, return_word_box: bool = False, **kwargs
     ) -> Tuple[List[Tuple[str, float]], List[Any]]:
-        preds_idx = preds.argmax(axis=2)
-        preds_prob = preds.max(axis=2)
+        if isinstance(preds, tuple):
+            preds_idx, preds_prob = preds
+        else:
+            preds_idx = preds.argmax(axis=2)
+            preds_prob = preds.max(axis=2)
 
         wh_ratio_list = kwargs.get("wh_ratio_list", (1.0,))
         max_wh_ratio = kwargs.get("max_wh_ratio", 1.0)

python/rapidocr/inference_engine/base.py

@@ -5,7 +5,7 @@ import abc
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
-from typing import Dict, List, Union
+from typing import Dict, List, Tuple, Union
 
 import numpy as np
 from omegaconf import OmegaConf
@@ -93,6 +93,11 @@ class InferSession(abc.ABC):
     def __call__(self, input_content: np.ndarray) -> np.ndarray:
         pass
 
+    def infer_argmax(self, input_content: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        """argmax/max over the last axis of the output (rec logits)."""
+        preds = self(input_content)
+        return preds.argmax(axis=2), preds.max(axis=2)
+
     @staticmethod
     def _verify_model(model_path: Union[str, Path, None]):
         if model_path is None:

python/rapidocr/inference_engine/pytorch/main.py

@@ -2,7 +2,7 @@
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
 from pathlib import Path
-from typing import List
+from typing import List, Tuple
 
 import numpy as np
 import torch
@@ -22,10 +22,23 @@ class TorchInferSession(InferSession):
         self.predictor = ModelLoader(cfg, self.device).predictor
 
     def __call__(self, img: np.ndarray):
+        return self._run(img, reduce=False)
+
+    def infer_argmax(self, img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        """For rec: argmax/max of the (B, T, vocab) logits on the device, so only
+        (B, T) is copied to the CPU."""
+        return self._run(img, reduce=True)
+
+    def _run(self, img: np.ndarray, reduce: bool):
         with torch.no_grad():
             inp = torch.from_numpy(img)
             inp = inp.to(self.device)
-            outputs = self.predictor(inp).cpu().numpy()
+            outputs = self.predictor(inp)
+            if reduce:
+                prob, idx = outputs.max(dim=2)
+                outputs = (idx.cpu().numpy(), prob.cpu().numpy())
+            else:
+                outputs = outputs.cpu().numpy()
             if self.device.type == "mps":
                 del inp
                 torch.mps.empty_cache()