推理引擎

# TensorRT基本流程
import tensorrt as trt

# 1. 创建builder和network
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 2. 解析ONNX模型
parser = trt.OnnxParser(network, logger)
parser.parse_from_file("model.onnx")

# 3. 配置优化
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30  # 1GB

# 4. 构建engine
engine = builder.build_serialized_network(network, config)

# INT8量化配置
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = MyCalibrator(calibration_data)

# 使用FP16
config.set_flag(trt.BuilderFlag.FP16)

// Plugin示例：自定义激活函数
class MyActivationPlugin : public nvinfer1::IPluginV2 {
public:
    // 实现enqueue执行推理
    int enqueue(int batchSize, const void* const* inputs,
                void* const* outputs, void* workspace,
                cudaStream_t stream) override;
};

# Triton模型仓库结构
model_repository/
├── model_a/
│   ├── config.pbtxt
│   └── 1/
│       └── model.onnx
├── model_b/
│   ├── config.pbtxt
│   └── 1/
│       └── model.plan  # TensorRT engine