diff --git a/mineru/model/utils/pytorchocr/base_ocr_v20.py b/mineru/model/utils/pytorchocr/base_ocr_v20.py index 610c4f36..d4177bd0 100644 --- a/mineru/model/utils/pytorchocr/base_ocr_v20.py +++ b/mineru/model/utils/pytorchocr/base_ocr_v20.py @@ -6,16 +6,50 @@ import torch from .modeling.architectures.base_model import BaseModel + +# OCR 推理精度开关：auto 表示 CPU 使用 fp32，非 CPU 自动使用 fp16。 +OCR_INFERENCE_PRECISION = "auto" + + class BaseOCRV20: def __init__(self, config, **kwargs): self.config = config self.build_net(**kwargs) + self.ocr_inference_dtype = torch.float32 self.net.eval() def build_net(self, **kwargs): self.net = BaseModel(self.config, **kwargs) + def _resolve_inference_dtype(self, device): + """根据常量和设备类型解析 OCR 网络推理使用的浮点精度。""" + precision = OCR_INFERENCE_PRECISION.lower() + device_name = str(device).lower() + is_cpu = device_name.startswith("cpu") + + if precision not in {"auto", "fp32", "fp16"}: + raise ValueError( + "OCR_INFERENCE_PRECISION must be one of: auto, fp32, fp16" + ) + if precision == "fp32" or is_cpu: + return torch.float32 + return torch.float16 + + def _apply_inference_precision(self, device): + """将 OCR 网络移动到目标设备，并在非 CPU 半精度场景下切到 fp16。""" + self.net.to(device) + self.ocr_inference_dtype = self._resolve_inference_dtype(device) + if self.ocr_inference_dtype == torch.float16: + self.net.to(dtype=torch.float16) + + def _to_inference_dtype(self, tensor): + """将浮点输入 tensor 转为 OCR 推理精度，整型/布尔辅助输入保持原 dtype。""" + if torch.is_tensor(tensor) and torch.is_floating_point(tensor): + inference_dtype = getattr(self, "ocr_inference_dtype", torch.float32) + return tensor.to(dtype=inference_dtype) + return tensor + @staticmethod def _is_safetensors_path(weights_path): """判断权重文件是否为 safetensors 格式。""" diff --git a/mineru/model/utils/tools/infer/predict_cls.py b/mineru/model/utils/tools/infer/predict_cls.py index 5350c73e..5e08151d 100644 --- a/mineru/model/utils/tools/infer/predict_cls.py +++ b/mineru/model/utils/tools/infer/predict_cls.py @@ -34,7 +34,7 @@ class TextClassifier(BaseOCRV20): self.load_pytorch_weights(self.weights_path) self.net.eval() - self.net.to(self.device) + self._apply_inference_precision(self.device) def resize_norm_img(self, img): imgC, imgH, imgW = self.cls_image_shape @@ -93,8 +93,9 @@ class TextClassifier(BaseOCRV20): with torch.no_grad(): inp = torch.from_numpy(norm_img_batch) inp = inp.to(self.device) + inp = self._to_inference_dtype(inp) prob_out = self.net(inp) - prob_out = prob_out.cpu().numpy() + prob_out = prob_out.float().cpu().numpy() cls_result = self.postprocess_op(prob_out) elapse += time.time() - starttime diff --git a/mineru/model/utils/tools/infer/predict_det.py b/mineru/model/utils/tools/infer/predict_det.py index 29981ef5..f846c740 100644 --- a/mineru/model/utils/tools/infer/predict_det.py +++ b/mineru/model/utils/tools/infer/predict_det.py @@ -119,7 +119,7 @@ class TextDetector(BaseOCRV20): super(TextDetector, self).__init__(network_config, **kwargs) self.load_pytorch_weights(self.weights_path) self.net.eval() - self.net.to(self.device) + self._apply_inference_precision(self.device) for module in self.net.modules(): if hasattr(module, 'rep'): module.rep() @@ -184,23 +184,24 @@ class TextDetector(BaseOCRV20): with torch.no_grad(): inp = torch.from_numpy(batch_tensor) inp = inp.to(self.device) + inp = self._to_inference_dtype(inp) outputs = self.net(inp) # 处理输出 preds = {} if self.det_algorithm == "EAST": - preds['f_geo'] = outputs['f_geo'].cpu().numpy() - preds['f_score'] = outputs['f_score'].cpu().numpy() + preds['f_geo'] = outputs['f_geo'].float().cpu().numpy() + preds['f_score'] = outputs['f_score'].float().cpu().numpy() elif self.det_algorithm == 'SAST': - preds['f_border'] = outputs['f_border'].cpu().numpy() - preds['f_score'] = outputs['f_score'].cpu().numpy() - preds['f_tco'] = outputs['f_tco'].cpu().numpy() - preds['f_tvo'] = outputs['f_tvo'].cpu().numpy() + preds['f_border'] = outputs['f_border'].float().cpu().numpy() + preds['f_score'] = outputs['f_score'].float().cpu().numpy() + preds['f_tco'] = outputs['f_tco'].float().cpu().numpy() + preds['f_tvo'] = outputs['f_tvo'].float().cpu().numpy() elif self.det_algorithm in ['DB', 'PSE', 'DB++']: - preds['maps'] = outputs['maps'].cpu().numpy() + preds['maps'] = outputs['maps'].float().cpu().numpy() elif self.det_algorithm == 'FCE': for i, (k, output) in enumerate(outputs.items()): - preds['level_{}'.format(i)] = output.cpu().numpy() + preds['level_{}'.format(i)] = output.float().cpu().numpy() else: raise NotImplementedError @@ -323,22 +324,23 @@ class TextDetector(BaseOCRV20): with torch.no_grad(): inp = torch.from_numpy(img) inp = inp.to(self.device) + inp = self._to_inference_dtype(inp) outputs = self.net(inp) preds = {} if self.det_algorithm == "EAST": - preds['f_geo'] = outputs['f_geo'].cpu().numpy() - preds['f_score'] = outputs['f_score'].cpu().numpy() + preds['f_geo'] = outputs['f_geo'].float().cpu().numpy() + preds['f_score'] = outputs['f_score'].float().cpu().numpy() elif self.det_algorithm == 'SAST': - preds['f_border'] = outputs['f_border'].cpu().numpy() - preds['f_score'] = outputs['f_score'].cpu().numpy() - preds['f_tco'] = outputs['f_tco'].cpu().numpy() - preds['f_tvo'] = outputs['f_tvo'].cpu().numpy() + preds['f_border'] = outputs['f_border'].float().cpu().numpy() + preds['f_score'] = outputs['f_score'].float().cpu().numpy() + preds['f_tco'] = outputs['f_tco'].float().cpu().numpy() + preds['f_tvo'] = outputs['f_tvo'].float().cpu().numpy() elif self.det_algorithm in ['DB', 'PSE', 'DB++']: - preds['maps'] = outputs['maps'].cpu().numpy() + preds['maps'] = outputs['maps'].float().cpu().numpy() elif self.det_algorithm == 'FCE': for i, (k, output) in enumerate(outputs.items()): - preds['level_{}'.format(i)] = output + preds['level_{}'.format(i)] = output.float().cpu().numpy() else: raise NotImplementedError diff --git a/mineru/model/utils/tools/infer/predict_rec.py b/mineru/model/utils/tools/infer/predict_rec.py index 8ace3943..944671bc 100644 --- a/mineru/model/utils/tools/infer/predict_rec.py +++ b/mineru/model/utils/tools/infer/predict_rec.py @@ -94,13 +94,13 @@ class TextRecognizer(BaseOCRV20): self.load_state_dict(weights) self.net.eval() - self.net.to(self.device) for module in self.net.modules(): if isinstance(module, ConvBNAct): if module.use_act: torch.quantization.fuse_modules(module, ['conv', 'bn', 'act'], inplace=True) else: torch.quantization.fuse_modules(module, ['conv', 'bn'], inplace=True) + self._apply_inference_precision(self.device) def resize_norm_img(self, img, max_wh_ratio): imgC, imgH, imgW = self.rec_image_shape @@ -388,6 +388,11 @@ class TextRecognizer(BaseOCRV20): gsrm_word_pos_inp = gsrm_word_pos_inp.to(self.device) gsrm_slf_attn_bias1_inp = gsrm_slf_attn_bias1_inp.to(self.device) gsrm_slf_attn_bias2_inp = gsrm_slf_attn_bias2_inp.to(self.device) + inp = self._to_inference_dtype(inp) + encoder_word_pos_inp = self._to_inference_dtype(encoder_word_pos_inp) + gsrm_word_pos_inp = self._to_inference_dtype(gsrm_word_pos_inp) + gsrm_slf_attn_bias1_inp = self._to_inference_dtype(gsrm_slf_attn_bias1_inp) + gsrm_slf_attn_bias2_inp = self._to_inference_dtype(gsrm_slf_attn_bias2_inp) backbone_out = self.net.backbone(inp) # backbone_feat prob_out = self.net.head(backbone_out, [encoder_word_pos_inp, gsrm_word_pos_inp, gsrm_slf_attn_bias1_inp, gsrm_slf_attn_bias2_inp]) @@ -405,6 +410,7 @@ class TextRecognizer(BaseOCRV20): with torch.no_grad(): inp = torch.from_numpy(norm_img_batch) inp = inp.to(self.device) + inp = self._to_inference_dtype(inp) preds = self.net(inp) elif self.rec_algorithm == "CAN": @@ -415,6 +421,7 @@ class TextRecognizer(BaseOCRV20): inp = [torch.from_numpy(e_i) for e_i in inputs] inp = [e_i.to(self.device) for e_i in inp] + inp = [self._to_inference_dtype(e_i) for e_i in inp] with torch.no_grad(): outputs = self.net(inp) outputs = [v.cpu().numpy() for k, v in enumerate(outputs)] @@ -427,6 +434,7 @@ class TextRecognizer(BaseOCRV20): with torch.no_grad(): inp = torch.from_numpy(norm_img_batch) inp = inp.to(self.device) + inp = self._to_inference_dtype(inp) preds = self.net(inp) with torch.no_grad():