Pytorch 模型导出为 ONNX 或 TensorRT 格式

将 Pytorch 模型导出为 ONNX 或 TensorRT 格式的具体步骤如下：

### 导出为 ONNX 格式
1. 安装 onnx 包：`pip install onnx`
2. 加载 Pytorch 模型并将其转换为 ONNX 模型：

import torch
import torchvision
import onnx

# 加载 Pytorch 模型
model = torchvision.models.resnet18(pretrained=True)
# 转换为 ONNX 模型
dummy_input = torch.randn(1, 3, 224, 224)
input_names = ["input"]
output_names = ["output"]
onnx_path = "resnet18.onnx"
torch.onnx.export(model, dummy_input, onnx_path, verbose=True, input_names=input_names, output_names=output_names)

3. 导入 ONNX 模型：

import onnx

# 加载 ONNX 模型
onnx_path = "resnet18.onnx"
model = onnx.load(onnx_path)

### 导出为 TensorRT 格式
1. 安装 TensorRT 并设置环境变量：

# 安装 TensorRT
!pip install nvidia-pyindex
!pip install nvidia-tensorrt
# 设置 TensorRT 环境变量
import os
os.environ["LD_LIBRARY_PATH"] += ":/usr/local/cuda/lib64:/usr/lib/x86_64-linux-gnu"

2. 加载 Pytorch 模型并将其转换为 TensorRT 模型：

import tensorrt as trt
import pycuda.driver as cuda
import torch
import torchvision

# 加载 Pytorch 模型
model = torchvision.models.resnet18(pretrained=True)
# 转换为 TensorRT 模型
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
trt_runtime = trt.Runtime(TRT_LOGGER)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
    builder.max_workspace_size = 1 << 30
    builder.max_batch_size = 1
    # 加载 ONNX 模型
    onnx_path = "resnet18.onnx"
    with open(onnx_path, "rb") as f:
        parser.parse(f.read())
    # 构建 TensorRT 引擎
    engine = builder.build_cuda_engine(network)
    # 保存 TensorRT 引擎
    with open("resnet18.trt", "wb") as f:
        f.write(engine.serialize())

3. 导入 TensorRT 模型：

import tensorrt as trt

# 加载 TensorRT 模型
trt_path = "resnet18.trt"
with open(trt_path, "rb") as f, trt.Runtime(trt.Logger(trt.Logger.WARNING)) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())