腾讯云 T4 GPU 上 PyTorch 模型 TensorRT 推理加速实战：从 ONNX 导出到动态 Shape 与 FP16 调优

# 建议使用 NVIDIA 官方 PyTorch 镜像（已适配腾讯云 T4 驱动）
docker pull nvcr.io/nvidia/pytorch:23.01-py3
docker run --gpus all -it --rm \
  -v $(pwd):/workspace \
  nvcr.io/nvidia/pytorch:23.01-py3 bash

# 内部已预装：torch 1.13.1, tensorrt 8.5.3, onnx 1.13.0
# 额外安装 onnx-simplifier 与 polygraphy（调试神器）
pip install onnx-simplifier polygraphy

import torch
import torchvision.models as models
from timm import create_model

# 固定随机种子，确保可复现
torch.manual_seed(42)

# ---------- CNN: ResNet50 ----------
resnet50 = models.resnet50(pretrained=True).cuda().eval()

# ---------- Transformer: ViT-B/16 ----------
vit_base = create_model('vit_base_patch16_224', pretrained=True).cuda().eval()

# 构建虚拟输入（batch=1，后续动态调整）
dummy_input = torch.randn(1, 3, 224, 224, device='cuda')

def export_to_onnx(model, dummy_input, onnx_path, dynamic_batch=True, dynamic_size=False):
    # 动态轴配置
    dynamic_axes = {'input': {0: 'batch'}, 'output': {0: 'batch'}}
    if dynamic_size:
        # 仅对支持可变分辨率的模型启用（如 ResNet，ViT 因位置编码固定 224 不建议）
        dynamic_axes['input'][2] = 'height'
        dynamic_axes['input'][3] = 'width'
    
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        input_names=['input'],
        output_names=['output'],
        opset_version=14,          # 推荐 14，支持 BERT 等复杂算子
        dynamic_axes=dynamic_axes,
        do_constant_folding=True,  # 折叠常量，减小模型大小
        verbose=False,
        export_params=True,
    )
    print(f"ONNX exported to {onnx_path}")

# 导出 ResNet50（支持动态宽高）
export_to_onnx(resnet50, dummy_input, 'resnet50.onnx', dynamic_batch=True, dynamic_size=True)

# 导出 ViT（仅动态 batch）
export_to_onnx(vit_base, dummy_input, 'vit_base.onnx', dynamic_batch=True, dynamic_size=False)

import onnx
from onnxsim import simplify

def simplify_onnx(input_path, output_path):
    model = onnx.load(input_path)
    model_simp, check = simplify(model, 
                                  dynamic_input_shape=True,  # 保留动态轴
                                  input_shapes=None)         # 不固定具体尺寸
    assert check, "Simplification failed"
    onnx.save(model_simp, output_path)
    print(f"Simplified ONNX saved to {output_path}")

simplify_onnx('resnet50.onnx', 'resnet50_sim.onnx')
simplify_onnx('vit_base.onnx', 'vit_base_sim.onnx')

import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine_static(onnx_path, engine_path, fp16=False, batch_size=1):
    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 error in range(parser.num_errors):
                print(parser.get_error(error))
            raise RuntimeError("ONNX parse failed")
    
    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 30  # 1GB
    if fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    
    # 设置固定优化形状（batch=1, 3, 224, 224）
    profile = builder.create_optimization_profile()
    profile.set_shape('input', (batch_size, 3, 224, 224), 
                               (batch_size, 3, 224, 224), 
                               (batch_size, 3, 224, 224))
    config.add_optimization_profile(profile)
    
    engine = builder.build_serialized_network(network, config)
    with open(engine_path, 'wb') as f:
        f.write(engine)
    print(f"Static engine saved to {engine_path}")

build_engine_static('resnet50_sim.onnx', 'resnet50_fp16.engine', fp16=True)

def build_engine_dynamic(onnx_path, engine_path, fp16=False, 
                         min_batch=1, opt_batch=8, max_batch=32,
                         min_size=224, opt_size=224, max_size=448):  # 仅 ResNet 支持动态尺寸
    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:
        parser.parse(f.read())
    
    config = builder.create_builder_config()
    config.max_workspace_size = 2 << 30  # 2GB
    if fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    
    profile = builder.create_optimization_profile()
    # 动态 batch 和宽高（需确保 ONNX 中对应轴已设为 dynamic）
    profile.set_shape('input', 
                      (min_batch, 3, min_size, min_size),
                      (opt_batch, 3, opt_size, opt_size),
                      (max_batch, 3, max_size, max_size))
    config.add_optimization_profile(profile)
    
    engine = builder.build_serialized_network(network, config)
    with open(engine_path, 'wb') as f:
        f.write(engine)
    print(f"Dynamic engine saved to {engine_path}")

build_engine_dynamic('resnet50_sim.onnx', 'resnet50_dynamic.engine', 
                     fp16=True, min_batch=1, opt_batch=8, max_batch=32)

import os
import numpy as np
from torchvision import transforms, datasets
from torch.utils.data import DataLoader, Subset

class Calibrator(trt.IInt8Calibrator):
    def __init__(self, dataloader, cache_file='calibration.cache'):
        super().__init__()
        self.dataloader = iter(dataloader)
        self.cache_file = cache_file
        self.batch_size = dataloader.batch_size
        self.device = 'cuda'
    
    def get_batch_size(self):
        return self.batch_size
    
    def get_batch(self, names):
        try:
            batch = next(self.dataloader)
            # 假设返回 (images, labels)，仅取 images
            images = batch[0].numpy().astype(np.float32)
            # TensorRT 需要 NHWC？默认 NCHW，若模型为 NCHW 则直接返回
            return [trt.Volume(images)]
        except StopIteration:
            return None
    
    def read_calibration_cache(self):
        if os.path.exists(self.cache_file):
            with open(self.cache_file, 'rb') as f:
                return f.read()
    
    def write_calibration_cache(self, cache):
        with open(self.cache_file, 'wb') as f:
            f.write(cache)

# 准备校准数据集（从 ImageNet 验证集取 100 张）
val_dataset = datasets.ImageNet(root='/data/imagenet', split='val', 
                                transform=transforms.Compose([
                                    transforms.Resize(256),
                                    transforms.CenterCrop(224),
                                    transforms.ToTensor(),
                                    transforms.Normalize(mean=[0.485,0.456,0.406],
                                                         std=[0.229,0.224,0.225])
                                ]))
subset = Subset(val_dataset, indices=range(100))
calib_loader = DataLoader(subset, batch_size=10, shuffle=False, num_workers=4)

calibrator = Calibrator(calib_loader, cache_file='resnet50_int8.cache')

# 构建 INT8 引擎时增加以下配置（在 builder 中）
# config.set_flag(trt.BuilderFlag.INT8)
# config.int8_calibrator = calibrator

import time
import numpy as np

class TRTInference:
    def __init__(self, engine_path, input_names=None, output_names=None):
        self.logger = trt.Logger(trt.Logger.WARNING)
        with open(engine_path, 'rb') as f:
            self.runtime = trt.Runtime(self.logger)
            self.engine = self.runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
        self.input_names = input_names or ['input']
        self.output_names = output_names or ['output']
        # 预分配 CUDA 内存（固定 pool 可提升多 Stream 性能）
        self.bindings = []
        for binding in self.engine:
            shape = self.engine.get_binding_shape(binding)
            dtype = trt.nptype(self.engine.get_binding_dtype(binding))
            size = trt.volume(shape)
            # 若动态轴，shape 可能含 -1，需根据实际输入设置
            self.bindings.append((binding, shape, dtype, size))
    
    def infer(self, input_data, batch_size=1):
        # input_data: numpy array (N, C, H, W)
        # 设置动态 shape（若 engine 支持）
        for i, (name, shape, dtype, size) in enumerate(self.bindings):
            if name == self.input_names[0]:
                # 根据实际输入调整 shape
                current_shape = (batch_size, *input_data.shape[1:])
                self.context.set_binding_shape(i, current_shape)
                # 分配 device memory
                # 简化版：实际可复用预分配 buffer
        # 执行推理
        outputs = []
        for out_name in self.output_names:
            output_shape = tuple(self.context.get_binding_shape(i))
            outputs.append(np.empty(output_shape, dtype=np.float32))
        # 实际需用 cuda.memcpy，此处仅示意
        # 完整实现见下节 benchmark
        return outputs

import torch
import onnxruntime as ort
import pycuda.driver as cuda
import pycuda.autoinit  # 需要 pycuda 库

# 准备测试数据（batch=8）
batch = 8
test_input = torch.randn(batch, 3, 224, 224, device='cuda')

# ---------- PyTorch ----------
def bench_torch(model, input_tensor, iterations=100):
    torch.cuda.synchronize()
    start = time.time()
    with torch.no_grad():
        for _ in range(iterations):
            out = model(input_tensor)
    torch.cuda.synchronize()
    return (time.time() - start) / iterations

# ---------- ONNX Runtime ----------
ort_session = ort.InferenceSession('resnet50_sim.onnx', 
                                   providers=['CUDAExecutionProvider'])
def bench_ort(input_numpy, iterations=100):
    start = time.time()
    for _ in range(iterations):
        ort_session.run(['output'], {'input': input_numpy})
    return (time.time() - start) / iterations

# ---------- TensorRT（封装完整） ----------
class TRTInferenceOptimized:
    def __init__(self, engine_path):
        with open(engine_path, 'rb') as f:
            self.runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
            self.engine = self.runtime.deserialize_cuda_engine(f.read())
        self.context = self.engine.create_execution_context()
        self.bindings = []
        self.buffer_shapes = {}
        for i, binding in enumerate(self.engine):
            shape = self.engine.get_binding_shape(binding)
            if shape[0] == -1:  # 动态 batch
                shape = (batch, *shape[1:])  # 暂设 batch=8
            size = trt.volume(shape)
            dtype = trt.nptype(self.engine.get_binding_dtype(binding))
            # 分配 device memory
            device_mem = cuda.mem_alloc(size * dtype().itemsize)
            self.bindings.append(int(device_mem))
            self.buffer_shapes[binding] = shape
        self.output_idx = self.engine.get_binding_index('output')
    
    def infer(self, input_np):
        # 拷贝 input 到 device
        cuda.memcpy_htod(self.bindings[0], input_np.ravel().astype(np.float32))
        self.context.execute_v2(self.bindings)
        # 从 device 读取 output
        output_shape = self.buffer_shapes['output']
        output_np = np.empty(output_shape, dtype=np.float32)
        cuda.memcpy_dtoh(output_np, self.bindings[self.output_idx])
        return output_np

trt_engine = TRTInferenceOptimized('resnet50_dynamic.engine')
def bench_trt(input_np, iterations=100):
    start = time.time()
    for _ in range(iterations):
        _ = trt_engine.infer(input_np)
    return (time.time() - start) / iterations

# 执行测试
input_np = test_input.cpu().numpy()
iterations = 200  # 预热 100 次后取平均

# 预热
for _ in range(50):
    _ = resnet50(test_input)
    ort_session.run(['output'], {'input': input_np})
    _ = trt_engine.infer(input_np)

latency_torch = bench_torch(resnet50, test_input, iterations)
latency_ort = bench_ort(input_np, iterations)
latency_trt = bench_trt(input_np, iterations)

print(f"PyTorch GPU  : {latency_torch*1000:.2f} ms")
print(f"ONNX Runtime : {latency_ort*1000:.2f} ms")
print(f"TensorRT FP16: {latency_trt*1000:.2f} ms")
print(f"Speedup vs PyTorch: {latency_torch/latency_trt:.2f}x")

streams = [cuda.Stream() for _ in range(4)]
contexts = [engine.create_execution_context() for _ in range(4)]
# 每个 Stream 绑定独立 device 内存，执行异步推理