jetson TX2 NX 使用 jetson inference 推理自己的Pytorch模型

本文全程参考NVIDIA 官方教程。

简介

针对某个具体任务，合理构建数据集后，总能快速获得一些具有实际效果的深度学习模型，但该模型仅能在深度学习框架（如pytorch）下运行，无法实际运用到生产环境中。此时，模型部署工作势在必行，ubuntu pytorch+NVIDIA Jetson TX2 是一个不错的实现路径。
本文将使用将在个人PC上使用pytorch训练一个分类模型，并将模型部署到Jetson TX2 NX，实现实时推理。

jetson-inference 安装

jetson-inference 是NVIDIA提供的开源项目，用于jetson系列的设备进行快速部署任务，支持设备包括Jetson Nano/TX1/TX2/Xavier NX/AGX Xavier/AGX Ori。

本人手上的jetson TX2在刷机时已经安装好jetpack4.5.0和cuda库，所以直接进行工程构建即可。进入项目内的Building the Project from Source，按照步骤执行。 具体步骤包括：

$ sudo apt-get update
$ sudo apt-get install git cmake libpython3-dev python3-numpy
$ git clone --recursive https://github.com/dusty-nv/jetson-inference #克隆完整工程，包括分支
$ cd jetson-inference
$ mkdir build
$ cd build
$ cmake ../
$ make -j$(nproc) # 根据核心数量进行多线程编译
$ sudo make install
$ sudo ldconfig

其中，在git clone 工程时可能会因为网络原因无法克隆分支，此时可以选择手动下载分支，添加到具体目录下。所有下载分支的具体目录和下载链接在jetson-inference目录下的.gitmodules文件里。
jetson-inference/.gitmodules：

[submodule "utils"]
	path = utils
	url = https://github.com/dusty-nv/jetson-utils
	branch = master
[submodule "tools/camera-capture"]
	path = tools/camera-capture
	url = https://github.com/dusty-nv/camera-capture
	branch = master
[submodule "python/training/classification"]
	path = python/training/classification
	url = https://github.com/dusty-nv/pytorch-classification
	branch = master
[submodule "python/training/detection"]
	path = python/training/detection
	url = https://github.com/dusty-nv/pytorch-detection
	branch = master
[submodule "python/training/segmentation"]
	path = python/training/segmentation
	url = https://github.com/dusty-nv/pytorch-segmentation
	branch = master
[submodule "plugins/pose"]
	path = plugins/pose
	url = https://github.com/dusty-nv/trt_pose
[submodule "docker/containers"]
	path = docker/containers
	url = https://github.com/dusty-nv/jetson-containers
	branch = master

cmake 过程中会联网下载一些文件，其中下载模型可能因为网络原因下载失败，本文是部署自己训练的模型，不需要执行开源项目自带的模型，所以注释掉该部分。
注释jetson-inference/CMakePreBuild.sh的52行./download-models.sh，如果不需要在tx2上安装pytorch，则53行的./install-pytorch.sh也可以注释掉（部署是使用NVIDIA的jetson.inference，无需使用pytorch框架，这也正是部署的意义所在）

cmake成功：
make成功：

sudo make install 安装到/usr/lib/python3.6/dist-packages/目录下

至此，部署环境已经配置完毕，后续将部署自己的模型。

pytorch模型转换

该步骤在自己的电脑上完成即可。
分类网络选择一个简单基础的Resnet34,训练和测试Demo在此不多赘述。网络经过训练后可以获得一个Resnet34模型resnet34.pth，pth格式是pytorch保存模型格式，无法在其他深度学习框架下直接使用，所以需要将pth格式转换为ONNX格式。
ONNX是微软和Facebook提出的一种表示深度学习模型的开放格式，定义了一套独立于环境和平台的标准格式。
参考/jetson-inference/python/training/classification/onnx_export.py提供的导出ONNX的例程。
my_onnx_export.py：

#
# converts a saved PyTorch model to ONNX format
# 
import os
import argparse

import torch
import torchvision.models as models

from reshape import reshape_model
from src.resnet import resnet50,resnet34

model_names = sorted(name for name in models.__dict__
    if name.islower() and not name.startswith("__")
    and callable(models.__dict__[name]))

# parse command line
parser = argparse.ArgumentParser()
parser.add_argument('--input', type=str, default='best_val_acc_d0.pth', help="path to input PyTorch model (default: model_best.pth.tar)")
parser.add_argument('--output', type=str, default='', help="desired path of converted ONNX model (default: <ARCH>.onnx)")
parser.add_argument('--model-dir', type=str, default='weights', help="directory to look for the input PyTorch model in, and export the converted ONNX model to (if --output doesn't specify a directory)")
parser.add_argument('--no-softmax', type=bool, default=False, help="disable adding nn.Softmax layer to model (default is to add Softmax)")

opt = parser.parse_args() 
print(opt)

# format input model path
if opt.model_dir:
	opt.model_dir = os.path.expanduser(opt.model_dir)
	opt.input = os.path.join(opt.model_dir, opt.input)

# set the device
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('running on device ' + str(device))

# load the model checkpoint
print('loading checkpoint:  ' + opt.input)
checkpoint = torch.load(opt.input)
# load the model weights
model = resnet34(num_classes=5).to(device)
model.load_state_dict(checkpoint)

# add softmax layer
if not opt.no_softmax:
	print('adding nn.Softmax layer to model...')
	model = torch.nn.Sequential(model, torch.nn.Softmax(1))

model.to(device)
model.eval()

print(model)

# create example image data
input = torch.ones((1, 3, 256, 320)).cuda()
print('input size:  {:d}x{:d}'.format(256, 320))

# format output model path
if not opt.output:
	opt.output = 'resnet34.onnx'

if opt.model_dir and opt.output.find('/') == -1 and opt.output.find('\\') == -1:
	opt.output = os.path.join(opt.model_dir, opt.output)

# export the model
input_names = [ "input_0" ]
output_names = [ "output_0" ]

print('exporting model to ONNX...')
torch.onnx.export(model, input, opt.output, verbose=True, input_names=input_names, output_names=output_names)
print('model exported to:  {:s}'.format(opt.output))

weights目录下获得resnet34.onnx，模型转换成功。

部署模型

该步骤在TX2进行。
将resnet34.onnx放到TX2内，写一个简单的程序调用模型预测一张图片。准备文件包括
推理脚本：recognition.py
onnx模型：resnet34.onnx
测试图片：1.bmp
标签txt文件：labels.txt

recognition.py：

#!/usr/bin/python3

import jetson.inference
import jetson.utils
import argparse

# parse the command line
parser = argparse.ArgumentParser()
parser.add_argument("--filename", default='./1.bmp', type=str, help="filename of the image to process")
args = parser.parse_args()
# load an image (into shared CPU/GPU memory)
img = jetson.utils.loadImage(args.filename)
# load the recognition network
net=jetson.inference.imageNet('resnet-34', ['--model=./resnet34.onnx','--input_blob=input_0','--output_blob=output_0', '--labels=./labels.txt'])
# classify the image
class_idx, confidence = net.Classify(img)
# find the object description
class_desc = net.GetClassDesc(class_idx)
# print out the result
print("image is recognized as '{:s}' (class #{:d}) with {:f}% confidence".format(class_desc, class_idx, confidence * 100))

labels.txt：

bg
corn
rice
soybean
wheat

运行推理脚本：

python3 recognition.py

获得预测结果，类别为corn，置信度100%。

至此，完成简单分类网络部署。