第四章 Jetson Nano Unet TensorRT模型推理

Jetson Nano系列学习

第一章 Jetson Nano烧录镜像及jtop安装
第二章 Jetson Nano安装Archiconda、PyTorch、torchvision
第三章 Jetson Nano 虚拟环境中使用PyQt5及TensorRT

---

---

前言

  由于在Jetson nano上运行pytorch生成的.pth模型文件，推理速度很慢，故需要将模型转化为tensorrt的引擎文件，使用Jetson nano自带的tensorrt进行加速推理。由于tensorrt中未实现Unet网络中上采样采取的bilinear双线性插值操作，故上采样使用ConvTranspose2d实现，将上采样中bilinear值设为False，Unet模型参考Pytorch-Unet
  首先使用Unet模型训练好的.pth模型文件利用onnx库转化为.onnx模型文件，然后利用onnx-tensorrt将onnx模型文件转换为tensorrt的engine文件即.trt模型文件，最终使用tensorrt进行推理。参考

---

一、使用onnx将.pth文件转化为.onnx文件

from src import UNet
import torch
import onnx

# 模型路径
model_path = "best_unet.pth"

if __name__ == "__main__":
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))
    # input shape尽量选择能被2整除的输入大小
    dummy_input = torch.randn(1, 3, 584, 565, device=device)
    # 1. Unet模型创建
    model = UNet(in_channels=3, num_classes=1, bilinear=False).to(device)
    # model = model.cuda()
    print("create U-Net model finised ...")
    # 2. 加载权重
    state_dict = torch.load(model_path, map_location=device)['model']
    model.load_state_dict(state_dict)
    print("load weight to model finised ...")
    # 3. 将torch格式转化为onnx
    input_names = ["input"]
    output_names = ["output"]
    torch.onnx.export(model,
                      dummy_input,
                      "best_unet.onnx",
                      verbose=True,
                      input_names=input_names