Pytorch 转TensorRT5
知乎上，说的也不错：
https://zhuanlan.zhihu.com/p/88318324
https://blog.csdn.net/qq_38003892/article/details/89314108
1.torch2trthttps://www.ctolib.com/amp/NVIDIA-AI-IOT-torch2trt.html
https://github.com/DocF/YOLOv3-Torch2TRT
https://github.com/traveller59/torch2trt
像官方代码:
https://github.com/NVIDIA-AI-IOT/torch2trt
import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet
# create some regular pytorch model...
model = alexnet(pretrained=True).eval().cuda()
# create example data
x = torch.ones((1, 3, 224, 224)).cuda()
# convert to TensorRT feeding sample data as input
model_trt = torch2trt(model, [x])
Execute
We can execute the returned TRTModule just like the original PyTorch model
y = model(x)
y_trt = model_trt(x)
# check the output against PyTorch
print(torch.max(torch.abs(y - y_trt)))
Save and load
We can save the model as a state_dict.
torch.save(model_trt.state_dict(), 'alexnet_trt.pth')
We can load the saved model into a TRTModule
from torch2trt import TRTModule
model_trt = TRTModule()
model_trt.load_state_dict(torch.load('alexnet_trt.pth'))
Models
We tested the converter against these models using the test.sh script. You can generate the results by calling
2.Pytorch通过保存为ONNX模型转TensorRT5
这是另一篇博客：有开源代码
https://www.cnblogs.com/darkknightzh/p/11332155.html
https://github.com/darkknightzh/TensorRT_pytorch
1 Pytorch以ONNX方式保存模型
def saveONNX(model, filepath):
'''
保存ONNX模型
:param model: 神经网络模型
:param filepath: 文件保存路径
'''
# 神经网络输入数据类型
dummy_input = torch.randn(self.config.BATCH_SIZE, 1, 28, 28, device='cuda')
torch.onnx.export(model, dummy_input, filepath, verbose=True)
2 利用TensorRT5中ONNX解析器构建Engine
def ONNX_build_engine(onnx_file_path):
'''
通过加载onnx文件，构建engine
:param onnx_file_path: onnx文件路径
:return: engine
'''
# 打印日志
G_LOGGER = trt.Logger(trt.Logger.WARNING)
with trt.Builder(G_LOGGER) as builder, builder.create_network() as network, trt.OnnxParser(network, G_LOGGER) as parser:
builder.max_batch_size = 100
builder.max_workspace_size = 1 << 20
print('Loading ONNX file from path {}...'.format(onnx_file_path))
with open(onnx_file_path, 'rb') as model:
print('Beginning ONNX file parsing')
parser.parse(model.read())
print('Completed parsing of ONNX file')
print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))
engine = builder.build_cuda_engine(network)
print("Completed creating Engine")
# 保存计划文件
# with open(engine_file_path, "wb") as f:
# f.write(engine.serialize())
return engine
3 构建TensorRT运行引擎进行预测
def loadONNX2TensorRT(filepath):
'''
通过onnx文件，构建TensorRT运行引擎
:param filepath: onnx文件路径
'''
# 计算开始时间
Start = time()
engine = self.ONNX_build_engine(filepath)
# 读取测试集
datas = DataLoaders()
test_loader = datas.testDataLoader()
img, target = next(iter(test_loader))
img = img.numpy()
target = target.numpy()
img = img.ravel()
context = engine.create_execution_context()
output = np.empty((100, 10), dtype=np.float32)
# 分配内存
d_input = cuda.mem_alloc(1 * img.size * img.dtype.itemsize)
d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)
bindings = [int(d_input), int(d_output)]
# pycuda操作缓冲区
stream = cuda.Stream()
# 将输入数据放入device
cuda.memcpy_htod_async(d_input, img, stream)
# 执行模型
context.execute_async(100, bindings, stream.handle, None)
# 将预测结果从从缓冲区取出
cuda.memcpy_dtoh_async(output, d_output, stream)
# 线程同步
stream.synchronize()
print("Test Case: " + str(target))
print("Prediction: " + str(np.argmax(output, axis=1)))
print("tensorrt time:", time() - Start)
del context
del engine
作者：ShellCollector
