【pytorch基础】基于训练的pytorch模型转换为onnx模型并测试

前言

  模型部署的过程中，不同的硬件可能支持不同的模型框架，本文介绍pytorch模型文件转换为onnx模型文件的实现过程，主要是基于Pytorch_Unet的实现过程，训练模型转换为onnx模型，并测试onnx的效果；

操作步骤

1. 基于训练完成的pth文件转换为onnx模型；

2. check和验证onnx模型；

3. 基于输入数据测试onnx模型；

实现过程

1. 基于训练完成的pth文件转换为onnx模型；

模型是基于Unet网络构建，基于Carvana数据集进行训练；

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

转换过程

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'}})

输入数据等

pthfile = 'xxx/Pytorch-UNet/checkpoints/CP_epoch5.pth'
onnxpath = './unet.onnx'
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

最后会得到onnx模型文件；

注意，模型的输入大小和测试的输入数据一致；

2. check和验证onnx模型；

check模型：

onnx.checker.check_model(onnx_model)验证模型的结构并确认模型具有有效的架构。

通过检查模型的版本，图的结构以及节点及其输入和输出，可以验证 ONNX 图的有效性。

 # check model
  onnx_model = onnx.load(onnxpath)
  check = onnx.checker.check_model(onnx_model)
  print('check: ', check)

验证模型是否匹配：

验证 ONNX 运行时和 PyTorch 正在为网络计算相同的值。

  # 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)匹配。

注意，测试数据时和模型的输入大小一致的，

3. 基于输入数据测试onnx模型；

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

查询ONNX版本：

import onnx（或onnxruntime）onnx .__ version __（或onnxruntime .__ version__）

主要原因是ONNX的版本或者pytorch版本的问题；

注意事项

1）转换onnx模型的时候，模型的输入大小是你的输入数据的大小；

2）注意export模型的ONNX的版本是否正确；

3）转换完成之后记得check是否正确，以及验证和pytorch模型的结果是否匹配；

4）测试数据的完成需要基于训练过程完成对应的预处理和后处理过程；

5）注意转换过程中的数据类型，特别是tensor/numpy等；

参考

1. Pytorch_Unet_github;

2. Pytorch_ONNX_doc;

3. Carvana_dataset;

4. Unet;

5. github_onnx_q;

完