opensource/MinerU
3
0
Fork 0
mirror of https://github.com/opendatalab/MinerU.git synced 2026-09-21 12:42:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: implement dynamic OCR inference precision handling for improved performance

Browse Source
This commit is contained in:
myhloli
2026-06-17 01:04:06 +08:00
parent 10fdbe02ca
commit 735db71ea6
4 changed files with 65 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mineru/model/utils/pytorchocr/base_ocr_v20.py
+34
View File
@@ -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 格式。"""
mineru/model/utils/tools/infer/predict_cls.py
+3 -2
View File
@@ -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
mineru/model/utils/tools/infer/predict_det.py
+19 -17
View File
@@ -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
mineru/model/utils/tools/infer/predict_rec.py
+9 -1
View File
@@ -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():