От FastAPI predict к Triton:
гид по оптимизации
ML-инференса
в продакшене

@app.post("/predict")
def predict(request):
    data = preprocess(request)
    prediction = model(data) 
    return {"result": prediction.tolist()}

import torch

# 1. Загружаем обученную модель
model = MyModel().cuda().eval()

# 2. Создаем dummy input (нужен для трассировки графа)
dummy_input = torch.randn(1, 3, 224, 224, device='cuda')

# 3. Экспорт
# dynamic_axes позволяет менять размер батча в рантайме
torch.onnx.export(
    model, 
    dummy_input, 
    "model.onnx", 
    input_names=["input"], 
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
    opset_version=17  # Версия набора операторов, критически важный параметр
)

Затем пишем код инференса (он не требует PyTorch, только легкую библиотеку `onnxruntime`):

import onnxruntime as ort
import numpy as np

# Выбираем провайдер: по классике сначала пытаемся на GPU, если нет - CPU
providers = [
    ('CUDAExecutionProvider', {
        'device_id': 0,
        'arena_extend_strategy': 'kNextPowerOfTwo',
        'cudnn_conv_algo_search': 'EXHAUSTIVE',
    }),
    'CPUExecutionProvider',
]

# Загрузка сессии (тяжелая операция, делается один раз при старте сервиса)
sess = ort.InferenceSession("model.onnx", providers=providers)

# Инференс
input_name = sess.get_inputs()[0].name
data = np.random.randn(8, 3, 224, 224).astype(np.float32) # батч 8
result = sess.run(None, {input_name: data})

# Пример команды для сборки движка с поддержкой FP16

trtexec --onnx=model.onnx \
        --saveEngine=model.engine \
        --fp16 \
        --minShapes=input:1x3x224x224 \
        --optShapes=input:32x3x224x224 \
        --maxShapes=input:64x3x224x224 \
        --workspace=4096 # Разрешаем использовать 4ГБ памяти для поиска алгоритмов

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit  # Автоматическая инициализация CUDA контекста
import numpy as np

# 1. Загрузка движка из файла
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
with open("model.engine", "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())

# 2. Создание контекста исполнения и выделение памяти
context = engine.create_execution_context()

# Модель может иметь несколько входов/выходов, берем первый вход и выход
h_input = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=np.float32)
h_output = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=np.float32)

# Выделяем память на GPU (Device memory)
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)

# Создаем поток для асинхронного выполнения
stream = cuda.Stream()

def predict(batch_input):
    # Копируем входные данные в "закрепленную" (pagelocked) память, затем на GPU
    np.copyto(h_input, batch_input.ravel())
    cuda.memcpy_htod_async(d_input, h_input, stream)
    
    # Исполнение: передаем адреса буферов на GPU
    context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
    
    # Перенос результата обратно в RAM хоста
    cuda.memcpy_dtoh_async(h_output, d_output, stream)
    
    # Ждем завершения операций в этом потоке
    stream.synchronize()
    return h_output.reshape(engine.get_binding_shape(1))

res = predict(my_numpy_array)

Пример конфига для модели классификации:

name: "resnet50_onnx"
platform: "onnxruntime_onnx" // для TensorRT будет "tensorrt_plan"
max_batch_size: 128 // для TensorRT должен быть >= любому значению в dynamic_batching.preferred_batch_size

// Настройка входов
input [
  {
    name: "input_tensor"
    data_type: TYPE_FP32
    dims: [ 3, 224, 224 ]
  }
]

// Настройка выходов
output [
  {
    name: "softmax_output"
    data_type: TYPE_FP32
    dims: [ 1000 ]
  }
]

// Батчинг
dynamic_batching {
  preferred_batch_size: [ 16, 32 ]
  max_queue_delay_microseconds: 2000 // окно ожидания 2 мс
}

// Запустить 2 копии модели на GPU с ID 0
instance_group [
  {
    count: 2
    kind: KIND_GPU
    gpus: [ 0 ]
  }
]

import torch
import torchvision.models as models
import openvino as ov

# 1. Загружаем модель
weights = models.ResNet50_Weights.DEFAULT
model = models.resnet50(weights=weights)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

# 2. Конвертируем
ov_model = ov.convert_model(model, example_input=dummy_input)

# 3. Сохраненяем модель и веса (квантизирует модель в FP16 по умолчанию)
ov.save_model(ov_model, "resnet50.xml", compress_to_fp16=True)

from openvino.runtime import Core

core = Core()
# 1. Загружаем модель и веса
model = core.read_model(model="resnet50.xml", weights="resnet50.bin")
compiled_model = core.compile_model(model, device_name="CPU") 
# device_name может быть "GPU" (интегрированная графика Intel Iris), "NPU" (новые чипы Meteor Lake) или "AUTO"

input_layer = compiled_model.input(0)
output_layer = compiled_model.output(0)

# 2. Инференс
result = compiled_model([input_data])[output_layer]