Pytorch pth 格式转ONNX 格式

目录

背景介绍

 安装依赖库

转换成onnx 格式

ONNX 转换成TensorRT 格式

---

 下面是Pytorch配套的视频教程：

Pytorch 快速实战教程：0_Pytorch实战前言_哔哩哔哩_bilibili

Pytorch 分割实战教程：介绍一个图像分割的网络搭建利器，Segmentation model PyTorch_哔哩哔哩_bilibili

---

 

背景介绍

PyTorch 训练的模型，需要在Jetson nano 上部署，jetson 原生提供了TensorRT 的支持，所以一个比较好的方式是把它转换成ONNX 格式，然后在通过ONNX 转换成TensorRT 格式

 安装依赖库

所以就需要安装ONNX， 安装ONNX 具体的版本需要依赖于环境中的protobuf 和python版本，我的python版本是3.6.9

pip install onnx==1.11.0
pip install onnx-simplifier

 安装ONNX 没问题，但是安装onnx-simplifier的时候，Jetson xaviar 上死活安装不成功

各种报错

 但是，我换到windows 上和ubuntu 服务器上，就很正常的安装成功了

 谁知道Jetson 上为什么安装失败的，可以私信我或评论区告诉我，感谢~

转换成onnx 格式

with torch.no_grad(): 
	    torch.onnx.export( 
		model, 
		example, 
		"init.onnx", 
		opset_version=11, 
		input_names=['input'], 
		output_names=['output'])

这样就能正常导出了

在Jetson XAVIAR 上运行时可能会报如下错误

Illegal instruction (core dumped)

执行下如下命令即可

 export OPENBLAS_CORETYPE=ARMV8

ONNX 转换成TensorRT 格式

在Jetson 上通过trtexec 工具做转换

trtexec --onnx=init.onnx --saveEngine=init.trt --explicitBatch

发现转换时报如下错误：Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.

 所以需要通过onnx-simplifier 转换一下，转换命令如下

python -m onnxsim init.onnx init_sim.onnx

 这样，就转换成功了，成功之后，后面就可以在TensorRT 上推理了

---