【pytorch基础】基于训练的pytorch模型转换为onnx模型并测试
前言
模型部署的过程中,不同的硬件可能支持不同的模型框架,本文介绍pytorch模型文件转换为onnx模型文件的实现过程,主要是基于Pytorch_Unet的实现过程,训练模型转换为onnx模型,并测试onnx的效果;
操作步骤
1. 基于训练完成的pth文件转换为onnx模型;
2. check和验证onnx模型;
3. 基于输入数据测试onnx模型;
实现过程
1. 基于训练完成的pth文件转换为onnx模型;
模型是基于Unet网络构建,基于Carvana数据集进行训练;
![](https://images.cnblogs.com/OutliningIndicators/ContractedBlock.gif)
import io import torch import torch.onnx from unet import UNet import onnx import onnxruntime import numpy as np from PIL import Image import torchvision.transforms as transforms from utils.dataset import BasicDataset
转换过程
![](https://images.cnblogs.com/OutliningIndicators/ContractedBlock.gif)
def test(): model = UNet(n_channels=3, n_classes=1) batch_size = 1 input_shape = (3, 640, 959) # Initialize model with the pretrained weights map_location = lambda storage, loc: storage if torch.cuda.is_available(): map_location = None loaded_model = torch.load(pthfile, map_location=map_location) model.load_state_dict(loaded_model) # set the model to inference mode model.eval() # data type nchw x = torch.rand(batch_size, *input_shape) input_names = ['input'] output_names = ['output'] # # Export the model torch.onnx.export(model, # model being run x, # model input (or a tuple for multiple inputs) onnxpath, # where to save the model (can be a file or file-like object) export_params=True, # store the trained parameter weights inside the model file opset_version=12, # the ONNX version to export the model to do_constant_folding=True, # whether to execute constant folding for optimization input_names = ['input'], # the model's input names output_names = ['output'], # the model's output names dynamic_axes={'input' : {0 : 'batch_size'}, # variable lenght axes 'output' : {0 : 'batch_size'}})
输入数据等
![](https://images.cnblogs.com/OutliningIndicators/ContractedBlock.gif)
pthfile = 'xxx/Pytorch-UNet/checkpoints/CP_epoch5.pth' onnxpath = './unet.onnx' device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
最后会得到onnx模型文件;
注意,模型的输入大小和测试的输入数据一致;
注意,在导出模型之前,请先调用torch_model.eval()
或torch_model.train(False)
,以将模型转换为推理模式,这一点很重要。 这是必需的,因为像dropout
或batchnorm
这样的运算符在推断和训练模式下的行为会有所不同。
注意,除非指定为动态轴,否则输入尺寸将在导出的 ONNX 图中固定为所有输入尺寸。
在此示例中,我们使用输入batch_size=1
导出模型,但随后在torch.onnx.export()
的dynamic_axes
参数中将第一维指定为动态。 因此,导出的模型将接受大小为[batch_size,
3, 640, 959]
的输入,其中batch_size
可以是可变的。
2. check和验证onnx模型;
check模型:
onnx.checker.check_model(onnx_model)
验证模型的结构并确认模型具有有效的架构。
通过检查模型的版本,图的结构以及节点及其输入和输出,可以验证 ONNX 图的有效性。如果有效,则输出为None。
1 2 3 4 | # check model onnx_model = onnx.load(onnxpath) check = onnx.checker.check_model(onnx_model) print ( 'check: ' , check) |
验证模型是否匹配:
验证 ONNX 运行时和 PyTorch 正在为网络计算相同的值。
1 2 3 4 5 6 7 8 9 10 | # check model whether match ort_session = onnxruntime.InferenceSession(onnxpath) # compute ONNX Runtime output prediction ort_inputs = {ort_session.get_inputs()[ 0 ].name:to_numpy(x)} ort_outs = ort_session.run( None , ort_inputs) # compare ONNX Runtime and PyTorch results torch_out = model(x) print ( 'tor_out: ' , torch_out.shape) np.testing.assert_allclose(to_numpy(torch_out), ort_outs[ 0 ], rtol = 1e - 03 , atol = 1e - 05 ) print ( "Exported model has been tested with ONNXRuntime, and the result looks good!" ) |
PyTorch 和 ONNX 运行时的输出在数值上与给定的精度(rtol/ atol)匹配。
注意,测试数据时和模型的输入大小一致的。
问题,为什么模型的输出是 ort_outs[0],比模型预想的输出多出一个维度呢????
验证转换前后模型数据是否一致,注意的是模型是否使用正确;
# Initialize model with the pretrained weights map_location = lambda storage, loc: storage if torch.cuda.is_available(): map_location = None loaded_model = torch.load(pthfile, map_location=map_location) model.load_state_dict(loaded_model) # set the model to inference mode model.eval()
3. 基于输入数据测试onnx模型;
![](https://images.cnblogs.com/OutliningIndicators/ContractedBlock.gif)
import io import torch import torch.onnx from unet import UNet import onnx import onnxruntime import numpy as np from PIL import Image import torchvision.transforms as transforms from utils.dataset import BasicDataset pthfile = 'xxx/Pytorch-UNet/checkpoints_carvana/CP_epoch5.pth' onnxpath = './unet.onnx' imgpath = 'xxx/Pytorch-UNet/output/0cdf5b5d0ce1_01.jpg' device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') def to_numpy(tensor): return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy() def test_onnx(): full_img = Image.open(imgpath) ort_session = onnxruntime.InferenceSession(onnxpath) scale_factor = 0.5 img = torch.from_numpy(BasicDataset.preprocess(full_img, scale_factor)) img = img.to(device=device, dtype=torch.float32) img.unsqueeze_(0) # ONNX RUNTIME ort_inputs = {ort_session.get_inputs()[0].name:to_numpy(img)} ort_outs = ort_session.run(None, ort_inputs) # list. # post process. img_out = ort_outs[0] img_out = torch.from_numpy(img_out) # probs = torch.nn.functional.softmax(img_out, dim=1) probs = torch.sigmoid(img_out) probs = probs.squeeze(0) tf = transforms.Compose( [ transforms.ToPILImage(), transforms.Resize(full_img.size[1]), transforms.ToTensor() ] ) probs = tf(probs.cpu()) full_mask = probs.squeeze().cpu().numpy() mask_thres = 0.5; mask_out = (full_mask > mask_thres) # save image img_out = Image.fromarray((mask_out*255).astype(np.uint8)) img_out.save('./img/onnx_img.jpg') if __name__ == '__main__': test_onnx()
问题
1. ONNX版本问题;参考here;
File "xxx/miniconda3/envs/open_mmlab/lib/python3.8/site-packages/torch/onnx/symbolic_helper.py", line 80, in _parse_arg raise RuntimeError("Failed to export an ONNX attribute '" + v.node().kind() + RuntimeError: Failed to export an ONNX attribute 'onnx::Cast', since it's not constant, please try to make things (e.g., kernel size) static if possible
注意事项
查了pytorch官方文档后发现,这里的upsample只支持nearest一种模式,而我用的是bilinear,在改变了这个之后,结果就对的齐了。
建议:先去官方文档看一下哪些算子支持哪些算子不支持,以及别用Function函数,得用torch.nn里面的层。
我猜想可能是网络中的某些操作过程在pytorch和onnxruntime中实现不一致吧。。。还没解决,怎么溯源呢??
参考
2. Pytorch_ONNX_doc;
3. Carvana_dataset;
4. Unet;
5. github_onnx_q;
完
心正意诚,做自己该做的事情,做自己喜欢做的事情,安静做一枚有思想的技术媛。
版权声明,转载请注明出处:https://www.cnblogs.com/happyamyhope/
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】凌霞软件回馈社区,博客园 & 1Panel & Halo 联合会员上线
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】博客园社区专享云产品让利特惠,阿里云新客6.5折上折
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步