Pytorch利用TensorRT在NVIDIA Jetson Nano上加速

之前做了个项目，就是在Jetson Nano上做实时图像分割。输入是一个HDMI的采集卡，输出是显示器。用pytorch做了个分割模型在上面直接跑。Pytorch用的1.3，具体环境配置略过。

经过代码赤果果的搬运之后，发现inference需要70ms，加上数据采集/数据前处理和后处理和图像显示花120ms。而且原始图像和处理过的图像之间误差了1.5s以上。

首先，我们把inference的cost降下来，首先我们改成TensorRT实现。用pip3 list查找了一下，官方给的系统默认装了TensorRT 5。看了TensorRT的sample，感觉要死了，还要重新把模型撸一遍，我不答应。GitHub看看有没有人写了torch2trt。果真还有，https://github.com/NVIDIA-AI-IOT/torch2trt，里面的readme写的好人性化，我好稀罕。根据它的Setup提示安装好了。接下来根据他的说明改几行代码就OK了

from torch2trt import *
import torch
from torchvision.models import *

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

#随便加个模型，加载好参数
model_org = resnet18(pretrained=True)
model_org.to(device)
model_org.eval() #这步需要

test1 = torch.rand(1,3,224,224).to(device)

#一定要设max_workspace_size，否则会报空间不足的错误
#fp16_mode精度高一些，int8_mode精度不太好，但是速度快
model = torch2trt(model_org, [test1], fp16_mode=True, max_workspace_size=100)
model.eval()

#第一次推断耗时较长
test2 = model(test1.cuda())

...

上面调的代码跟它的说明有些不全的地方，我在注释里面都写了。fp16的推断耗时降低到30ms，int8的推断降低到25ms。但是我不在乎这5ms。但是采集的延时依旧没有降低，还是卡在1s多。究竟是怎么回事呢？

就是怎样的我也不知道了，反正根据网友的解决办法采集2次就能大大降低延时，即：

import cv2

cap = cv2.VideoCapture(0)

while True:
    ret, frame = cap.read()
    ret, frame = cap.read()

我猜估计是总线被占了