TensorRT学习笔记

TensorRT是Nvidia开发的用于加速模型推理的SDK

pytorch –> onnx –> tensorrt

一、 pytorch –> onnx


---

步骤

1. 加载pytorch模型

import torchvision.models as models

resnext50_32x4d = models.resnext50_32x4d(pretrained=True)

2. 保存onnx模型

import torch

BATCH_SIZE = 64
dummy_input=torch.randn(BATCH_SIZE, 3, 224, 224)

import torch.onnx
torch.onnx.export(resnext50_32x4d, dummy_input, "resnet50_onnx_model.onnx", verbose=False)

二、onnx –> tensorrt engine

有两种转换方法：

1. 使用 trtexec
2. 使用 TensorRT API

这里我们使用第一种方法：

trtexec --onnx=resnet50_onnx_model.onnx --saveEngine=resnet_engine.trt  --explicitBatch

这会把resnet50_onnx_model.onnx转换为resnet_engine.trt。