PyTorch model export to ONNX file example (LeNet-5)

      从PyTorch模型导出到ONNXfile is called byPyTorch的torch.onnx.export接口实现.

      torch.onnx.export：如果pytorchThe model is neithertorch.jit.ScriptModule也不是orch.jit.ScriptFunction,它(torch.nn.Module)会run一次pytorch模型,以便将其转换为TorchScript graph被导出(相当于torch.jit.trace,track its implementation,Then export the tracked model to onnx文件).生成的onnxThe file contains a binaryprotocol buffer,It contains the network structure and parameters of the model you exported.

      参数说明：

      (1).model：要导出的pytorch模型,可以为torch.nn.Module, torch.jit.ScriptModule或torch.jit.ScriptFunction.

      (2).args：模型的输入,可以为tuple或torch.Tensor.

      (3).f：A file-like object or a string containing the filename.A binary protocol buffer will be written to this file.

      (4).export_params=True：默认值为True.如果为True,then all parameters will be exported.如果要导出未经训练的模型,This parameter needs to be set to False.如果为True,导出的模型将首先将其所有参数作为参数,Its order is given bymodel.stat_dict().values()指定.

      (5).verbose=False：默认值为False.如果为True,then prints a description of the model being exported to standard output.此外,最终的ONNX graphWill contain fields from the exported model"doc_string",其中提到了"model"source code location.

      (6).training=TrainingMode.EVAL：默认值为TrainingMode.EVAL.TrainingMode.EVAL：以推理模式导出模型.TrainingMode.PRESERVE：如果model.training为False,The model is exported in inference mode;如果model.training为True,Export the model in training mode.TrainingMode.TRAINING：Export the model in training mode,Disable optimizations that might interfere with training.

      (7).input_names=None：类型为str的列表,默认为空列表.按顺序分配给graphThe name of the input node.如果不设置的话,Some simple names are automatically assigned,如input.1.ONNXEach input and output of the modeltensor都有一个名字.

      (8).output_names=None：类型为str的列表,默认为空列表.按顺序分配给graphThe name of the output node.如果不设置的话,Some simple names or numbers are automatically assigned,如logits、25.ONNXEach input and output of the modeltensor都有一个名字.

      (9).operator_export_type=None：enum类型,默认为None.None通常表示"`OperatorExportTypes.ONNX",但是,如果PyTorch是用"DPYTORCH_ONNX_CAFFE2_BUNDLE"构建的,则None表示"OperatorExportTypes.ONNX_ATEN_FALLBACK".OperatorExportTypes.ONNX：Export all operations as generalONNX操作(in the default action domain(opset domain)).OperatorExportTypes.ONNX_FALLTHROUGH：Attempt to convert all operations to the standard in the default operation domainONNX操作.OperatorExportTypes.ONNX_ATEN：所有的ATen操作(ops)都导出为ATen操作.OperatorExportTypes.ONNX_ATEN_FALLBACK:尝试将每个ATenActions are exported as regularONNX操作.

      (10).opset_version=None：int类型,在PyTorch 1.11.0版本中,默认值为9,This value range must be in[7, 15]范围内.每个PyTorchVersions correspond to different value ranges.ONNX算子集版本.参考：https://github.com/onnx/onnx/blob/main/docs/Operators.md

      (11).do_constant_folding=True：默认为True.应用constant-folding优化.constant-foldingSome operations with all constant inputs will be replaced with precomputed constant nodes.

      (12).dynamic_axes=None：字典类型,默认为空字典.默认情况下,The exported model will have all inputs and outputstensors的shape设置为与argsExact match given in .指定输入输出tensorwhich dimensions are dynamic,ONNXBy default all involved in the operationtensor都是静态的(tensor的shape不发生改变).

      (13).keep_initializers_as_inputs=None：bool类型,默认为None.如果为True,则导出的graphAll initializers in (Usually corresponds to a parameter)will also be added as an input to graph.如果为False,then the initializer is not added as an inputgraph,并且仅将非参数输入添加为输入.

      (14).custom_opsets=None：字典类型,默认为空字典.schema字典：Key(str)：opset域名;Value(int)：opset版本.

      (15).export_modules_as_functions=False：bool类型或set of type of nn.Module,默认为False.将所有nn.Module forwardCall export asONNX中的本地函数(local function).or instructions to be inONNXA specific module type that is exported as a local function in .

      以下是将LeNet-5.pth导出到LeNet-5.onnx的示例：

      1.加载LeNet-5.pth：https://blog.csdn.net/fengbingchun/article/details/125462001

      See the link above for the production of this model,Because this sample code will use another directorypython脚本中的函数,需导入,代码段如下：

import sys
sys.path.append("..") # 为了导入pytorch目录中的内容
from pytorch.lenet5.test_lenet5_mnist import LeNet5, list_files, get_image_label

      加载LeNet-5.pthThe code snippet for the model is as follows：需将model设置为评估模式

def load_pytorch_model(model_name):
    model = LeNet5(n_classes=10).to('cpu') # 实例化一个LeNet5网络对象
    model.load_state_dict(torch.load(model_name)) # 加载pytorch模型
    model.eval() # 将网络设置为评估模式

    return model

      2.导出到onnx文件,and verify thisonnx文件的正确性,代码段如下：

def export_model_from_pytorch_to_onnx(pytorch_model, onnx_model_name):
    batch_size = 1
    # input to the model
    x = torch.randn(batch_size, 1, 32, 32)
    out = pytorch_model(x)
    #print("out:", out)

    # export the model
    torch.onnx.export(pytorch_model,             # model being run
                      x,                         # model input (or a tuple for multiple inputs)
                      onnx_model_name,           # 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=9,          # 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 length axes
                                    'output' : {0 : 'batch_size'}})

def verify_onnx_model(onnx_model_name):
    # model is an in-memory ModelProto
    model = onnx.load(onnx_model_name)
    #print("the model is:\n{}".format(model))

    # check the model
    try:
        onnx.checker.check_model(model)
    except onnx.checker.ValidationError as e:
        print("  the model is invalid: %s" % e)
        exit(1)
    else:
        print("  the model is valid")

      3.准备测试图像,This test image is also from the link above：一共10幅,0到9各一幅,如下图所示,注意：训练图像背景色为黑色,而测试图像背景色为白色：

def image_preprocess(image_names, image_name_suffix):
    input_data = []
    labels = []

    for image_name in image_names:
        label = get_image_label(image_name, image_name_suffix)
        labels.append(label)

        img = cv2.imread(image_name, cv2.IMREAD_GRAYSCALE)
        img = cv2.resize(img, (32, 32))
        # MNIST图像背景为黑色,而测试图像的背景色为白色,识别前需要做转换
        img = cv2.bitwise_not(img)
        norm_img = cv2.normalize(img, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
        norm_img = norm_img.reshape(1, 1, 32, 32).astype('float32')
        #print(f"img type: {type(norm_img)}, shape: {norm_img.shape}")
        input_data.append(norm_img)

    return input_data, labels

       4.通过ONNX Runtime进行推理,验证LeNet-5.onnx,代码段如下：

def softmax(x):
    x = x.reshape(-1)
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum(axis=0)

def postprocess(result):
    return softmax(np.array(result)).tolist()

def inference(model_name, image_names, input_data, labels):
    session = onnxruntime.InferenceSession(model_name, None)
    # get the name of the first input of the model
    input_name = session.get_inputs()[0].name
    count = 0

    for data in input_data:
        raw_result = session.run([], {input_name: data})

        res = postprocess(raw_result)
        idx = np.argmax(res)

        image_name = image_names[count][image_names[count].rfind("/")+1:]
        print(f"  image name: {image_name}, actual value: {labels[count]}, predict value: {idx}, percentage: {round(res[idx]*100, 4)}%")
        count += 1

      执行结果如下：Consistent with the results in the link above

       GitHub：https://github.com/fengbingchun/PyTorch_Test