mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-12 05:50:49 +08:00
131 lines
4.9 KiB
Python
131 lines
4.9 KiB
Python
# -*- encoding: utf-8 -*-
|
|
# @Author: SWHL
|
|
# @Contact: liekkaskono@163.com
|
|
from pathlib import Path
|
|
from typing import Any, Dict
|
|
|
|
import tensorrt as trt
|
|
|
|
from ...utils.log import logger
|
|
from ...utils.typings import OCRVersion
|
|
|
|
|
|
class TRTEngineBuilder:
|
|
"""Build TensorRT engine from ONNX model"""
|
|
|
|
def __init__(
|
|
self,
|
|
onnx_path: Path,
|
|
engine_path: Path,
|
|
cfg: Dict[str, Any],
|
|
task_type: str,
|
|
trt_logger: trt.Logger,
|
|
ocr_version: OCRVersion,
|
|
):
|
|
self.onnx_path = onnx_path
|
|
self.engine_path = engine_path
|
|
self.cfg = cfg
|
|
self.task_type = task_type
|
|
self.trt_logger = trt_logger
|
|
self.ocr_version = ocr_version
|
|
|
|
def build(self) -> trt.ICudaEngine:
|
|
"""Build TensorRT engine from ONNX"""
|
|
builder = trt.Builder(self.trt_logger)
|
|
network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
|
|
network = builder.create_network(network_flags)
|
|
parser = trt.OnnxParser(network, self.trt_logger)
|
|
|
|
# Parse ONNX model
|
|
with open(self.onnx_path, "rb") as f:
|
|
if not parser.parse(f.read()):
|
|
for i in range(parser.num_errors):
|
|
logger.error(f"ONNX parse error: {parser.get_error(i)}")
|
|
raise RuntimeError("Failed to parse ONNX model")
|
|
|
|
# Configure builder
|
|
config = builder.create_builder_config()
|
|
|
|
# Set workspace size
|
|
workspace_size = self.cfg.get("workspace_size", 1 << 30) # 1GB default
|
|
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_size)
|
|
|
|
# Set precision
|
|
if self.cfg.get("use_fp16", True) and builder.platform_has_fast_fp16:
|
|
config.set_flag(trt.BuilderFlag.FP16)
|
|
logger.info("Using FP16 precision")
|
|
else:
|
|
logger.info("Using FP32 precision")
|
|
|
|
if self.cfg.get("use_int8", False) and builder.platform_has_fast_int8:
|
|
config.set_flag(trt.BuilderFlag.INT8)
|
|
logger.info("Using INT8 precision")
|
|
|
|
# Add optimization profile for dynamic shapes
|
|
profile = builder.create_optimization_profile()
|
|
self._set_dynamic_shapes(network, profile)
|
|
config.add_optimization_profile(profile)
|
|
|
|
# Build engine
|
|
logger.info("Building TensorRT engine (this may take a few minutes)...")
|
|
serialized_engine = builder.build_serialized_network(network, config)
|
|
|
|
if serialized_engine is None:
|
|
raise RuntimeError("Failed to build TensorRT engine")
|
|
|
|
# Save engine to cache
|
|
self.engine_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(self.engine_path, "wb") as f:
|
|
f.write(serialized_engine)
|
|
logger.info(f"TensorRT engine saved to {self.engine_path}")
|
|
|
|
# Deserialize and return
|
|
runtime = trt.Runtime(self.trt_logger)
|
|
return runtime.deserialize_cuda_engine(serialized_engine)
|
|
|
|
def _set_dynamic_shapes(
|
|
self, network: trt.INetworkDefinition, profile: trt.IOptimizationProfile
|
|
):
|
|
"""Set dynamic shape optimization profiles"""
|
|
profile_key = f"{self.task_type}_profile"
|
|
profile_cfg = self.cfg.get(profile_key, {})
|
|
|
|
# Default profiles based on task type
|
|
if self.task_type == "det":
|
|
min_shape = profile_cfg.get("min_shape", (1, 3, 32, 32))
|
|
opt_shape = profile_cfg.get("opt_shape", (1, 3, 736, 736))
|
|
max_shape = profile_cfg.get("max_shape", (1, 3, 2048, 2048))
|
|
elif self.task_type == "rec":
|
|
min_shape = profile_cfg.get("min_shape", (1, 3, 48, 32))
|
|
opt_shape = profile_cfg.get("opt_shape", (6, 3, 48, 320))
|
|
max_shape = profile_cfg.get("max_shape", (6, 3, 48, 2048))
|
|
elif self.task_type == "cls":
|
|
if self.ocr_version == OCRVersion.PPOCRV4:
|
|
min_shape = profile_cfg.get("min_shape", (1, 3, 48, 192))
|
|
opt_shape = profile_cfg.get("opt_shape", (6, 3, 48, 192))
|
|
max_shape = profile_cfg.get("max_shape", (6, 3, 48, 192))
|
|
elif self.ocr_version == OCRVersion.PPOCRV5:
|
|
min_shape = profile_cfg.get("min_shape", (1, 3, 80, 160))
|
|
opt_shape = profile_cfg.get("opt_shape", (6, 3, 80, 160))
|
|
max_shape = profile_cfg.get("max_shape", (6, 3, 80, 160))
|
|
else:
|
|
# Generic fallback
|
|
min_shape = (1, 3, 32, 32)
|
|
opt_shape = (1, 3, 224, 224)
|
|
max_shape = (1, 3, 2048, 2048)
|
|
|
|
# Set shapes for input tensor
|
|
input_tensor = network.get_input(0)
|
|
input_name = input_tensor.name
|
|
|
|
profile.set_shape(
|
|
input_name,
|
|
min=tuple(min_shape),
|
|
opt=tuple(opt_shape),
|
|
max=tuple(max_shape),
|
|
)
|
|
logger.info(
|
|
f"Set optimization profile for {input_name}: "
|
|
f"min={min_shape}, opt={opt_shape}, max={max_shape}"
|
|
)
|