From af69351b76e3b8193d71b1b710e04f2e8f3f19d7 Mon Sep 17 00:00:00 2001 From: ExtReMLapin <3909752+ExtReMLapin@users.noreply.github.com> Date: Sun, 25 Oct 2026 14:52:16 +0100 Subject: [PATCH] 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 | 2 +- python/rapidocr/ch_ppocr_rec/utils.py | 7 +++++-- python/rapidocr/inference_engine/base.py | 7 ++++++- .../rapidocr/inference_engine/pytorch/main.py | 17 +++++++++++++++-- 4 files changed, 27 insertions(+), 6 deletions(-) diff --git a/python/rapidocr/ch_ppocr_rec/main.py b/python/rapidocr/ch_ppocr_rec/main.py index 11ee322..f6b7564 100644 --- a/python/rapidocr/ch_ppocr_rec/main.py +++ b/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, diff --git a/python/rapidocr/ch_ppocr_rec/utils.py b/python/rapidocr/ch_ppocr_rec/utils.py index 76c5dd0..8c1243b 100644 --- a/python/rapidocr/ch_ppocr_rec/utils.py +++ b/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) diff --git a/python/rapidocr/inference_engine/base.py b/python/rapidocr/inference_engine/base.py index 83b1fc2..a4c96a8 100644 --- a/python/rapidocr/inference_engine/base.py +++ b/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: diff --git a/python/rapidocr/inference_engine/pytorch/main.py b/python/rapidocr/inference_engine/pytorch/main.py index 8357d9f..f7e2a69 100644 --- a/python/rapidocr/inference_engine/pytorch/main.py +++ b/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()