diff --git a/python/rapidocr/inference_engine/tensorrt/main.py b/python/rapidocr/inference_engine/tensorrt/main.py index 5968cdf..3aebb80 100644 --- a/python/rapidocr/inference_engine/tensorrt/main.py +++ b/python/rapidocr/inference_engine/tensorrt/main.py @@ -1,6 +1,7 @@ # -*- encoding: utf-8 -*- # @Author: SWHL # @Contact: liekkaskono@163.com +import os import traceback from pathlib import Path from typing import Any, Dict, List, Optional @@ -32,7 +33,7 @@ class TRTInferSession(InferSession): engine_path = self._get_engine_path(cfg) self.engine = self._load_or_build_engine(cfg, engine_path) - self.context = self.engine.create_execution_context() + self.context = self._create_execution_context(engine_path) # Allocate memory buffers (pre-allocated with max shape) self.inputs, self.outputs, self.bindings, self.stream = allocate_buffers( @@ -222,8 +223,11 @@ class TRTInferSession(InferSession): model_name = self._get_model_name(cfg) gpu_arch = self._get_gpu_arch() precision = "fp16" if self.engine_cfg.get("use_fp16", True) else "fp32" + tf32_override = os.environ.get("NVIDIA_TF32_OVERRIDE", "unset") - return cache_dir / f"{model_name}_{gpu_arch}_{precision}.engine" + return cache_dir / ( + f"{model_name}_{gpu_arch}_{precision}_tf32{tf32_override}.engine" + ) def _get_model_name(self, cfg: Dict[str, Any]) -> str: """Extract model name from config for engine filename.""" @@ -332,7 +336,23 @@ class TRTInferSession(InferSession): runtime = trt.Runtime(self.trt_logger) with open(engine_path, "rb") as f: engine_data = f.read() - return runtime.deserialize_cuda_engine(engine_data) + + engine = runtime.deserialize_cuda_engine(engine_data) + if engine is None: + raise TensorRTError(f"Failed to deserialize TensorRT engine: {engine_path}") + return engine + + def _create_execution_context(self, engine_path: Path) -> trt.IExecutionContext: + context = self.engine.create_execution_context() + if context is None: + tf32_override = os.environ.get("NVIDIA_TF32_OVERRIDE", "unset") + raise TensorRTError( + "Failed to create TensorRT execution context for " + f"{engine_path}. NVIDIA_TF32_OVERRIDE={tf32_override}. " + "Remove the cached engine or rebuild it under the current " + "TensorRT/CUDA environment." + ) + return context def have_key(self, key: str = "character") -> bool: return False diff --git a/python/tests/test_tensorrt.py b/python/tests/test_tensorrt.py index fe548e1..18f919a 100644 --- a/python/tests/test_tensorrt.py +++ b/python/tests/test_tensorrt.py @@ -111,9 +111,7 @@ def import_tensorrt_main(monkeypatch): return tensorrt_main -def test_tensorrt_cache_dir_initializes_model_root_dir_for_build( - monkeypatch, tmp_path -): +def test_tensorrt_cache_dir_initializes_model_root_dir_for_build(monkeypatch, tmp_path): tensorrt_main = import_tensorrt_main(monkeypatch) downloaded_paths = [] @@ -304,28 +302,6 @@ def test_tensorrt_falls_back_to_rebuild_when_cached_engine_load_fails( assert build_count == 1 -def test_tensorrt_square_padding_and_crop(monkeypatch): - tensorrt_main = import_tensorrt_main(monkeypatch) - - session = object.__new__(tensorrt_main.TRTInferSession) - session._closed = True - session._max_square_size = 64 - - import numpy as np - - img = np.ones((1, 3, 32, 48), dtype=np.float32) - padded, original_hw = session._pad_to_square(img) - - assert padded.shape == (1, 3, 64, 64) - assert original_hw == (32, 48) - np.testing.assert_array_equal(padded[:, :, :32, :48], img) - assert padded[:, :, 32:, :].sum() == 0 - assert padded[:, :, :, 48:].sum() == 0 - - cropped = session._crop_output(padded, original_hw) - assert cropped.shape == (1, 3, 32, 48) - - def test_tensorrt_get_dict_key_url_uses_fallback_engines(monkeypatch): tensorrt_main = import_tensorrt_main(monkeypatch)