onnx设置动态batch/修改onnx的batch

• 一. 训练结束后（例如pytorch）导出onnx模型时， 设置动态batch

• def export_onnx(model, input_hwc, output_file, input_names=['input'],
              output_names=['output'], show=False, opset_version=9, dynamic=True):
      if dynamic:
          dynamic_axes = {
              'input': {
                  0: 'batch',
              },
              'output': {
                  0: 'batch'
              }
          }
      else:
          dynamic_axes = {}
  
      h, w, c = input_hwc
      input_shape = [1, c, h, w]
      one_img = torch.randn(input_shape)
  
      # register_extra_symbolics(opset_version)
      with torch.no_grad():
          torch.onnx.export(
              model.cpu().eval(),
              one_img,
              output_file,
              input_names=input_names,
              output_names=output_names,
              export_params=True,
              keep_initializers_as_inputs=True,
              verbose=show,
              dynamic_axes=dynamic_axes,
              opset_version=opset_version)
      print('>>> finish export onnx:', output_file)
  

  二.

• 通过onnx库修改onnx模型的batch

• # 安装onnx：pip install onnx
  import onnx
  def change_input_dim(model):
      # Use some symbolic name not used for any other dimension
      sym_batch_dim = "N"
      # or an actal value
      actual_batch_dim = "4" 
  
      # The following code changes the first dimension of every input to be batch-dim
      # Modify as appropriate ... note that this requires all inputs to
      # have the same batch_dim 
      inputs = model.graph.input
      for input in inputs:
          # Checks omitted.This assumes that all inputs are tensors and have a shape with first dim.
          # Add checks as needed.
          dim1 = input.type.tensor_type.shape.dim[0]
          # update dim to be a symbolic value
          dim1.dim_param = sym_batch_dim
          # or update it to be an actual value:
          # dim1.dim_value = actual_batch_dim
  
  
  def apply(transform, infile, outfile):
      model = onnx.load(infile)
      transform(model)
      onnx.save(model, outfile)
  
  apply(change_input_dim, onnx_pth, save_pth)
  
  

三.动态推理。

# 安装onnxruntime： pip install onnxruntime
import onnxruntime as ort
import numpy as np
import torch 
x1 = torch.rand(1,3,112,112)
ort_sess1 = ort.InferenceSession(path_to_onnx_model)
outputs1 = ort_sess1.run(None, {'input.1': x2.numpy()})