pytorch模型如何转成TensorRT

PyTorch模型可以通过以下步骤转换为TensorRT：

1. 定义PyTorch模型并加载预训练权重。
2. 使用ONNX将PyTorch模型转换为ONNX格式。
3. 使用TensorRT的ONNX解析器将ONNX模型转换为TensorRT引擎。
4. 在TensorRT引擎上执行推理。

下面是一个简单的示例代码：

import torch
import tensorrt as trt
import numpy as np
import onnx
import onnx_tensorrt.backend as backend

# 定义PyTorch模型并加载预训练权重
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = torch.nn.Linear(10, 5)
        self.fc2 = torch.nn.Linear(5, 2)
    
    def forward(self, x):
        x = self.fc1(x)
        x = torch.nn.functional.relu(x)
        x = self.fc2(x)
        return x

model = MyModel()
model.load_state_dict(torch.load('model.pth'))

# 使用ONNX将PyTorch模型转换为ONNX格式
dummy_input = torch.randn(1, 10)
input_names = ['input']
output_names = ['output']
torch.onnx.export(model, dummy_input, 'model.onnx', verbose=False, input_names=input_names, output_names=output_names)

# 使用TensorRT的ONNX解析器将ONNX模型转换为TensorRT引擎
onnx_model = onnx.load('model.onnx')
engine = backend.prepare(onnx_model, device='CUDA:0')

# 在TensorRT引擎上执行推理
input_data = np.random.randn(1, 10).astype(np.float32)
output_data = engine.run(input_data)[0]

在上面的代码中，我们首先定义了一个简单的PyTorch模型，并加载了预训练权重。然后，我们使用`torch.onnx.export()`将PyTorch模型转换为ONNX格式，并定义了输入和输出的名称。接着，我们使用ONNX解析器将ONNX模型转换为TensorRT引擎。最后，我们使用TensorRT引擎执行推理，得到输出数据。