@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]