TensorRT部署手把手教您，跟我一起来学习。
这东西就是 NVidia 在自家显卡上做了一个深度学习 inference 加速的框架，只要你把训练好的模型参数和结构告诉他，他就能自动帮你优化（硬件相关），以达到最快速度。
这涉及两个问题：
应该以什么模型格式把模型喂给 TensorRT？
如何使用 TensorRT 优化后的模型？
对于第一个问题：现在的深度学习框架非常多，不止常用的 pytorch/tensorflow，而即使是同一种框架还可以使用不同的编程语言实现。让 TensorRT 对每一个框架都直接支持，显然是不可能的。
TensorRT 只需要知道网络的结构和参数即可，它支持三种转换入口：
TF-TRT，要求是 TensorFlow 模型
ONNX 模型格式
使用 TensorRT API 手动把模型搭起来，然后把参数加载进去
第一种不够灵活，第三种比较麻烦，所以最省事方便的就是第二种方法。本文介绍第二种。
ONNX 就是一个通用的神经网络格式，一个 .onnx 文件内包含了网络的结构和参数。甭管是用什么深度学习框架写的网络，只要把模型导出成 ONNX 格式，就跟原本的代码没有关系了。
转成 ONNX 格式还没有被优化，需要再使用 TensorRT 读取它并优化成 TensorRT Engine。优化参数也在这一步指定。
对于第二个问题：得到的 TensorRT Engine 是硬件相关的，之后跑模型只需要运行这个 Engine 即可。调用 TensorRT Engine 需要使用 TensorRT Runtime API。
所以整个逻辑就是：
把你的模型导出成 ONNX 格式。
把 ONNX 格式模型输入给 TensorRT，并指定优化参数。
使用 TensorRT 优化得到 TensorRT Engine。
使用 TensorRT Engine 进行 inference。
屏蔽 nouveau 驱动
nouveau 是系统自带的一个显示驱动程序，需要先将其禁用，然后再进行下一步操作，否则在安装显卡驱动时，会提示：You appear to be running an X server …，然后安装失败。分别打开如下两个文件（如果没有就创建一个），并在其中输入如下两句，然后保存。
# vim /etc/modprobe.d/nvidia-installer-disable-nouveau.conf
# vim /lib/modprobe.d/nvidia-installer-disable-nouveau.conf
...
blacklist nouveau
options nouveau modeset=0
重做 initramfs镜像
重做镜像之后启动才会屏蔽驱动，否则无效，重做时应先rm已有驱动，否则会提示无法覆盖。
这一步需要确保 boot 文件目录的空间足够，否则会失败。建议大于 400 MB
# cp /boot/initramfs-$(uname -r).img /boot/initramfs-$(uname -r).img.bak
# dracut /boot/initramfs-$(uname -r).img $(uname -r) --force
# rm /boot/initramfs-$(uname -r).img.bak ; 这一步可不执行
如果想用YOLOv5对图像做目标检测，在将图像输入给模型之前还需要做一定的预处理操作，预处理操作应该与模型训练时所做的操作一致。YOLOv5的输入是RGB格式的3通道图像，图像的每个像素需要除以255来做归一化，并且数据要按照CHW的顺序进行排布。所以YOLOv5的预处理大致可以分为两个步骤：
将原始输入图像缩放到模型需要的尺寸，比如640x640。这一步需要注意的是，原始图像是按照等比例进行缩放的，如果缩放后的图像某个维度上比目标值小，那么就需要进行填充。举个例子：假设输入图像尺寸为768x576，模型输入尺寸为640x640，按照等比例缩放的原则缩放后的图像尺寸为640x480，那么在y方向上还需要填充640-480=160（分别在图像的顶部和底部各填充80）。来看一下实现代码：
cv::Mat input_image = cv::imread("dog.jpg");
cv::Mat resize_image;
const int model_width = 640;
const int model_height = 640;
const float ratio = std::min(model_width / (input_image.cols * 1.0f),
model_height / (input_image.rows * 1.0f));
// 等比例缩放
const int border_width = input_image.cols * ratio;
const int border_height = input_image.rows * ratio;
// 计算偏移值
const int x_offset = (model_width - border_width) / 2;
const int y_offset = (model_height - border_height) / 2;
cv::resize(input_image, resize_image, cv::Size(border_width, border_height));
cv::copyMakeBorder(resize_image, resize_image, y_offset, y_offset, x_offset,
x_offset, cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114));
// 转换为RGB格式
cv::cvtColor(resize_image, resize_image, cv::COLOR_BGR2RGB);
格式转化
要将tensorflow的pb文件转化为uff格式的文件，首先找到convert_to_uff文件，看自己用的是哪个版本的python，如果是python3，则在/usr/lib/python3.5/dist-packages/uff/bin文件夹下，如果是python2，则在/usr/lib/python2.7/dist-packages/uff/bin文件夹下
我们在终端中进入end_to_end_tensorflow_mnist，运行以下指令
首先使用 Pytorch 实现一个和上文一致的模型，即只对输入做一次池化并输出；然后将 Pytorch 模型转换成 ONNX 模型；最后将 ONNX 模型转换成 TensorRT 模型。这里主要使用了 TensorRT 的 OnnxParser 功能，它可以将 ONNX 模型解析到 TensorRT 的网络中。最后我们同样可以得到一个 TensorRT 模型，其功能与上述方式实现的模型功能一致。
import torch
import onnx
import tensorrt as trt
onnx_model = 'model.onnx'
class NaiveModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.pool = torch.nn.MaxPool2d(2, 2)
def forward(self, x):
return self.pool(x)
device = torch.device('cuda:0')
# generate ONNX model
torch.onnx.export(NaiveModel(), torch.randn(1, 3, 224, 224), onnx_model, input_names=['input'], output_names=['output'], opset_version=11)
onnx_model = onnx.load(onnx_model)
# create builder and network
logger = trt.Logger(trt.Logger.ERROR)
builder = trt.Builder(logger)
EXPLICIT_BATCH = 1 << (int)(
trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
network = builder.create_network(EXPLICIT_BATCH)
# parse onnx
parser = trt.OnnxParser(network, logger)
if not parser.parse(onnx_model.SerializeToString()):
error_msgs = ''
for error in range(parser.num_errors):
error_msgs += f'{parser.get_error(error)}\n'
raise RuntimeError(f'Failed to parse onnx, {error_msgs}')
config = builder.create_builder_config()
config.max_workspace_size = 1<<20
profile = builder.create_optimization_profile()
profile.set_shape('input', [1,3 ,224 ,224], [1,3,224, 224], [1,3 ,224 ,224])
config.add_optimization_profile(profile)
# create engine
with torch.cuda.device(device):
engine = builder.build_engine(network, config)
with open('model.engine', mode='wb') as f:
f.write(bytearray(engine.serialize()))
print("generating file done!")
有疑问加站长微信联系（非本文作者）
