Оптимизация инференса нейросетей: конвертация в TensorRT на NVIDIA GPU

Конвертация моделей в TensorRT

При инференсе BERT-base на T4 с batch=8 задержка в PyTorch FP32 составляет ~12ms — для real-time систем это многовато. TensorRT снижает latency до 2.9ms в FP16 и 1.8ms в INT8. Разбираемся, как этого достичь на практике.

Мы специализируемся на конвертации моделей любой сложности: от BERT до LLaMA. Выполнили 20+ проектов. Гарантируем совместимость и ускорение. Наши инженеры сертифицированы NVIDIA. Закажите конвертацию — мы подготовим оптимизированный engine за 3–10 дней. Свяжитесь с нами для оценки вашего проекта.

Как TensorRT ускоряет инференс?

TensorRT оптимизирует граф вычислений: фьюзинг слоёв, удаление dead-операций, выбор оптимальных kernel-реализаций под конкретную GPU. Дополнительно применяет precision scaling (FP16, INT8) с минимальной потерей точности. Результат — до 8x ускорение по сравнению с FP32 PyTorch. NVIDIA TensorRT Documentation подтверждает эти показатели.

Конвертация из ONNX в TensorRT

import tensorrt as trt TRT_LOGGER = trt.Logger(trt.Logger.WARNING) def build_engine(onnx_path: str, engine_path: str, fp16: bool = True, int8: bool = False): builder = trt.Builder(TRT_LOGGER) network = builder.create_network( 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) ) parser = trt.OnnxParser(network, TRT_LOGGER) with open(onnx_path, "rb") as f: if not parser.parse(f.read()): for i in range(parser.num_errors): print(f"ONNX parse error: {parser.get_error(i)}") raise RuntimeError("Failed to parse ONNX") config = builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 * 1024 ** 3) # 4 GB if fp16 and builder.platform_has_fast_fp16: config.set_flag(trt.BuilderFlag.FP16) if int8: config.set_flag(trt.BuilderFlag.INT8) # Нужен calibrator для INT8 config.int8_calibrator = MyCalibrator(calibration_data) # Dynamic shapes — важно для переменной длины входа profile = builder.create_optimization_profile() profile.set_shape( "input_ids", min=(1, 1), opt=(8, 128), max=(32, 512) ) config.add_optimization_profile(profile) serialized_engine = builder.build_serialized_network(network, config) with open(engine_path, "wb") as f: f.write(serialized_engine) print(f"Engine saved to {engine_path}") 

Инференс с TensorRT

import tensorrt as trt import numpy as np import pycuda.driver as cuda import pycuda.autoinit class TRTInferenceSession: def __init__(self, engine_path: str): runtime = trt.Runtime(TRT_LOGGER) with open(engine_path, "rb") as f: self.engine = runtime.deserialize_cuda_engine(f.read()) self.context = self.engine.create_execution_context() self.inputs = [] self.outputs = [] self.bindings = [] for binding in self.engine: shape = self.engine.get_tensor_shape(binding) size = trt.volume(shape) * np.dtype(np.float32).itemsize device_mem = cuda.mem_alloc(size) self.bindings.append(int(device_mem)) if self.engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT: self.inputs.append({"name": binding, "mem": device_mem, "shape": shape}) else: self.outputs.append({"name": binding, "mem": device_mem, "shape": shape}) self.stream = cuda.Stream() def infer(self, inputs: dict[str, np.ndarray]) -> dict[str, np.ndarray]: for inp in self.inputs: data = inputs[inp["name"]].astype(np.float32).ravel() cuda.memcpy_htod_async(inp["mem"], data, self.stream) self.context.execute_async_v2(self.bindings, self.stream.handle) results = {} for out in self.outputs: output = np.empty(out["shape"], dtype=np.float32) cuda.memcpy_dtoh_async(output, out["mem"], self.stream) results[out["name"]] = output self.stream.synchronize() return results 

INT8 Calibration

INT8 требует калибровочных данных для определения диапазонов значений:

class BertCalibrator(trt.IInt8EntropyCalibrator2): def __init__(self, calibration_texts: list[str], cache_file: str = "calibration.cache"): super().__init__() self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") self.cache_file = cache_file self.data = iter( self.tokenizer(calibration_texts, padding="max_length", truncation=True, max_length=128, return_tensors="np") ) self.device_input = cuda.mem_alloc(128 * 4) # input_ids buffer def get_batch_size(self) -> int: return 16 def get_batch(self, names: list[str]) -> list | None: try: batch = {k: next(self.data) for k in ["input_ids", "attention_mask"]} cuda.memcpy_htod(self.device_input, batch["input_ids"].astype(np.int32).ravel()) return [int(self.device_input)] except StopIteration: return None def read_calibration_cache(self) -> bytes | None: if os.path.exists(self.cache_file): with open(self.cache_file, "rb") as f: return f.read() return None def write_calibration_cache(self, cache: bytes) -> None: with open(self.cache_file, "wb") as f: f.write(cache) 

Использование torch-tensorrt

Более простой путь для PyTorch моделей:

import torch_tensorrt trt_model = torch_tensorrt.compile( model, inputs=[ torch_tensorrt.Input( min_shape=[1, 1], opt_shape=[8, 128], max_shape=[32, 512], dtype=torch.int32 ) ], enabled_precisions={torch.float16}, workspace_size=4 * 1024 ** 3, truncate_long_and_double=True ) torch.jit.save(trt_model, "bert_trt.ts") 

Типичный прирост производительности

На T4 GPU, BERT-base, batch=8, seq=128:

Почему стоит использовать INT8 калибровку?

INT8 даёт максимальный прирост скорости (6-8x), но требует аккуратной калибровки. Без неё точность может упасть на 2-5%. Мы используем Entropy Calibrator 2, который минимизирует потери. Если точность критична, оставляем FP16 — он безопаснее и даёт 4x ускорение. TensorRT INT8 также экономит GPU-часы: задача, выполнявшаяся час на PyTorch FP32, с INT8 решается за 9 минут.

Типичные ошибки при конвертации

Этапы конвертации: пошаговый how-to

1. Экспорт модели в ONNX: Используйте torch.onnx.export с dynamic_axes для трансформеров.
2. Проверка совместимости: Прогнать ONNX через onnx.checker и onnxsim.
3. Сбор калибровочного датасета (для INT8): 500–2000 сэмплов из вашей предметной области.
4. Построение engine: Выбор precision, настройка dynamic shapes, запуск билдера.
5. Валидация: Сравнить выходы TensorRT и оригинальной модели на тестовых данных.
6. Интеграция: Адаптировать пайплайн под TensorRT runtime (Python или C++).
7. Профилирование: Измерить latency p99, проверить стабильность.

Что входит в работу

Результат конвертации — не просто engine, а полный пакет:

• Оптимизированный TensorRT engine в формате .plan.
• Код инференс-сессии на Python или C++ с примерами.
• Калибровочный датасет и скрипты калибровки (при INT8).
• Интеграция в ваш пайплайн: контейнеризация, CI/CD.
• Документация с описанием конфигураций и рекомендаций по масштабированию.
• Поддержка после интеграции — консультации по профилированию и тюнингу.

Как мы работаем

• Анализ: Изучаем модель, выявляем узкие места (например, неэффективные view-операции).
• Экспорт: Конвертируем в ONNX, устраняем ошибки совместимости.
• Калибровка: Подбираем калибровочные данные, настраиваем Entropy Calibrator 2.
• Генерация engine: Собираем TensorRT с dynamic shapes и выбранной точностью.
• Интеграция: Встраиваем engine в ваш инференс-пайплайн (Python/C++).
• Документация: Передаём конфиги, код, инструкции.

Срок выполнения — от 3 до 10 дней в зависимости от сложности модели. Стоимость рассчитывается индивидуально. Закажите конвертацию — мы подготовим оптимизированный engine за 3–10 дней.

Наш опыт

Многолетний опыт в AI/ML, 20+ проектов по оптимизации инференса. Работали с моделями BERT, GPT-2/3, LLaMA, YOLOv8, ResNet. Гарантируем ускорение до заявленных значений. Обращайтесь за консультацией — оценим ваш проект.