PyTorch导出ONNX格式分割模型及在C#中调用预测

PyTorch导出ONNX格式分割模型及在C#中调用预测

• 一、ONNX简介
• 二、PyTorch导出ONNX格式分割模型
• • 2.1导出流程及实现代码
  • • 2.1.1导出流程
    • 2.1.2实现代码
    • 2.1.3代码说明
    • • 2.1.3.1模型权重加载
      • 2.1.3.2torch.onnx.export函数参数
  • [2.2RuntimeError: ONNX export failed: Couldn't export operator aten::adaptive_avg_pool2d报错解决](#2.2RuntimeError: ONNX export failed: Couldn't export operator aten::adaptive_avg_pool2d报错解决)
  • • 2.2.1尝试提高opset_version算子版本（失败）
    • 2.2.2adaptive_avg_pool2d替换为F.interpolate（失败）
    • 2.2.3使用avg_pool2d替换adaptive_avg_pool2d（成功）
  • 2.3导出模型与原始模型对比
  • • 2.3.1对比验证代码
    • 2.3.2对比结果
• 三、C#中调用ONNX模型进行分割预测
• • 3.1OnnxRuntime软件包安装
  • 3.2调用ONNX模型进行分割预测代码
  • 3.3预测结果验证

一、ONNX简介

ONNX（Open Neural Network Exchange）是一种开放的神经网络交换格式，旨在促进不同深度学习框架之间的互操作性。它使用protobuf二进制格式来序列化模型，从而提供更好的传输性能。

ONNX的主要作用是作为不同深度学习框架之间的中间表示。例如，可以将PyTorch或TensorFlow模型转换为ONNX模型，然后在C#中调用ONNX模型进行预测。

使用C#中调用PyTorch图像分割模型，本质上是跨语言/跨运行时集成。通过ONNX模型导出，再在C#中通过ONNX Runtime调用模型预测，可以不依赖于Python实现高性能预测。

二、PyTorch导出ONNX格式分割模型

2.1导出流程及实现代码

2.1.1导出流程

导出流程如下：

1. 首先需要安装onnx与onnxruntime，使用pip安装命令为：
2. 加载模型结构，并且加载训练好的模型权重，如果不加载权重，推理结果会是随机噪声；
3. 将模型切换为推理模式，如果不切换为推理模式，那么如果模型中存在Dropout层会导致推理结果每次都不一样；
4. 构造一个输入张量，设置张量的shape，应为1×N×H×W，N为输入的图片通道数，一般为3，H和W为图片高宽，应跟训练图片保持一致；
5. 调用torch.onnx.export函数完成导出。

2.1.2实现代码

示例代码如下：

def export_model():
    import torch
    from self_model import self_net
    model = self_net()
    model.load_state_dict(torch.load("model.pth", map_location="cpu"), strict=False)
    model.eval()
    
    dummy = torch.randn(1, 3, 256, 256)
    
    torch.onnx.export(
        model,
        dummy,
        "pred.onnx",
        input_names=["input"],
        output_names=["output"],
        export_params=True,
        do_constant_folding=False,
        opset_version=17,
        dynamic_axes={
            "input": {0: "batch", 2: "height", 3: "width"},
            "output": {0: "batch", 2: "height", 3: "width"}
        }
    )

2.1.3代码说明

2.1.3.1模型权重加载

代码实现与上述流程说明一致，需要说明的是加载模型权重与torch.onnx.export函数的几个参数。

加载模型权重有两种方式：

1. model = self_net()，model.load_state_dict(torch.load("best_model.pth", map_location="cpu"))
2. model=torch.load("best_model.pth")

第二种方式似乎实现更为简洁，但并不推荐使用这种方式，原因是该方式非常依赖原始代码环境，要求完全相同的类路径和相同的代码结构，一般发生改动或者重命名就会加载失败。

而第一种方式只加载权重参数，模型结构可以进行小改，因此推荐使用第一种方式也就是上述代码中的加载方式。

其次torch.load("best_model.pth", map_location="cpu")，该函数设置map_location为cpu，仅是把权重加载到CPU内存中，并不会影响后续onnx模型推理时使用CPU或GPU。

2.1.3.2torch.onnx.export函数参数

torch.onnx.export函数的几个重要参数说明如下：

参数	说明
export_params	是否把权重写入ONNX模型，必须设置为True
do_constant_folding	常量折叠优化，提前合并计算减少运行时计算，建议先设置为False，后续进行验证后可设置为True
opset_version	ONNX算子版本，这是最重要的参数之一，不同算子版本支持的torch函数不一样，一般设置的越高支持的torch函数也越多，但可能存在兼容性问题

2.2RuntimeError: ONNX export failed: Couldn't export operator aten::adaptive_avg_pool2d报错解决

首次尝试导出时使用的pytorch版本为1.10.1，使用的opset_version算子版本为12，结果报错：RuntimeError: ONNX export failed: Couldn't export operator aten::adaptive_avg_pool2d，此外还有类似于如下的一些提示信息：

  %output : Float(*, *, *, *, strides=[160000, 40000, 200, 1], requires_grad=1, device=cpu) = onnx::Resize[coordinate_transformation_mode="align_corners", cubic_coeff_a=-0.75, mode="linear", nearest_mode="floor"](%796, %805, %806, %804) # C:\Miniconda3\envs\pytorch1.10.1\lib\site-packages\torch\nn\functional.py:3731:0
  return (%output)

2.2.1尝试提高opset_version算子版本（失败）

根据报错信息，说明报错是因为当前onnx不支持adaptive_avg_pool2d函数。

由于opset_version算子版本越高，支持的torch函数也越多，因此尝试提高opset_version，从12提高至14，结果仍然报出上述错误，而pytorch1.10.1最高支持14。

所以尝试在pytorch2.5.0中继续提高opset_version，提高至20（pytorch2.5.0最高支持20），依然报出如下错误：

torch.onnx.errors.SymbolicValueError: Unsupported: ONNX export of operator adaptive_avg_pool2d, input size not accessible. Please feel free to request support or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues  [Caused by the value '675 defined in (%675 : Float(*, 128, *, *, strides=[131072, 1024, 32, 1], requires_grad=1, device=cpu) = onnx::HardSwish(%input.547), scope: C:\Miniconda3\envs\pytorch2.5.0\Lib\site-packages\torch\nn\functional.py:2424:0
)' (type 'Tensor') in the TorchScript graph. The containing node has kind 'onnx::HardSwish'.]

错误原因依然是onnx不支持adaptive_avg_pool2d函数。

提高opset_version算子版本失败。

2.2.2adaptive_avg_pool2d替换为F.interpolate（失败）

使用pytorch版本为2.5.0，使用的opset_version算子版本为17，尝试将adaptive_avg_pool2d替换为F.interpolate函数。

当将F.interpolate函数中的参数mode设置为area时就是对区域取平均操作，那么几乎等价于nn.AdaptiveAvgPool2d函数，因此进行替换，结果依然报出上述错误信息。

经过查找torch\nn\functional.py源代码，发现F.interpolate函数实现逻辑如下：

if input.dim() == 3 and mode == "area":
    assert output_size is not None
    return adaptive_avg_pool1d(input, output_size)
if input.dim() == 4 and mode == "area":
    assert output_size is not None
    return adaptive_avg_pool2d(input, output_size)
if input.dim() == 5 and mode == "area":
    assert output_size is not None
    return adaptive_avg_pool3d(input, output_size)

上述代码表明，当mode设置为area时，F.interpolate函数会去调用adaptive_avg_pool2d函数实现平均池化功能，因此尝试失败。

2.2.3使用avg_pool2d替换adaptive_avg_pool2d（成功）

经过查找资料得知，之所以onnx不支持adaptive_avg_pool2d函数，是因为不支持算子参数如kernel size或stride是动态的，而adaptive_avg_pool2d函数需要在运行过程中根据输入图的尺寸动态计算kernel_size和stride，因此报错。

所以尝试提前计算出kernel_size和stride，然后使用avg_pool2d替换adaptive_avg_pool2d。

例如输入图尺寸为32×32，而想得到1×1输出图尺寸，计算stride = H // k，kernel = H - (k - 1) * stride，得出kernel与stride都为32，则将nn.AdaptiveAvgPool2d(1)替换为nn.AvgPool2d(kernel_size=32, stride=32)。

将模型结构中所有adaptive_avg_pool2d都替换为avg_pool2d，然后尝试导出，结果成功导出。

2.3导出模型与原始模型对比

2.3.1对比验证代码

可以构造一个随机tensor作为输入数据，然后分别使用原始pytorch模型与onnx模型进行预测，对输出结果进行对比，验证代码如下：

def infer_onnx(input_data):
    import onnx
    import onnxruntime as ort
    onnx_model = onnx.load("pred.onnx")
    onnx.checker.check_model(onnx_model)

    ort_session = ort.InferenceSession("pred.onnx")

    #input_data = np.random.randn(1,3,200,200).astype(np.float32)
    outputs = ort_session.run(None, {"input": input_data})
    #print(outputs[0].shape)
    np.save('onnx_output_data.npy', outputs[0])

def infer_origin(input_data):
    import torch
    from self_model import self_net
    model = self_net()
    model.load_state_dict(torch.load("model.pth", map_location="cpu", weights_only=True), strict=False)
    model.eval()
    with torch.no_grad():
        outputs = model(torch.from_numpy(input_data))[0]
        pred_np = outputs.numpy()
        #print(pred_np.shape)
        np.save('origin_output_data.npy', pred_np)

def compare_originAndonnx():
    origin_out=np.load('origin_output_data.npy')
    onnx_out=np.load('onnx_output_data.npy')
    print(np.allclose(onnx_out, origin_out, atol=1e-2))
    #print(onnx_out[0,0,100:110,100:110])
    #print(origin_out[0,0,100:110,100:110])

if __name__ == '__main__':
    input_data = np.random.randn(1,3,200,200).astype(np.float32)
    np.save('test_input_data.npy', input_data)
    input_data=np.load('test_input_data.npy')
    infer_onnx(input_data)
    infer_origin(input_data)
    compare_originAndonnx()

2.3.2对比结果

onnx模型与原始torch模型输出数据如下：

[[2.2444808 2.556353  2.8660774 2.8883095 2.8982658 2.9018817 3.0587878
  3.1235728 3.0752256 3.132615 ]
 [2.282033  2.7556944 3.3516119 3.3406177 3.2286854 3.051212  3.2959785
  3.4294393 3.4262075 3.3748384]
 [2.306106  3.095516  4.18652   4.094642  3.9580243 3.794877  3.976505
  4.1066055 4.173426  4.0021214]
 [2.3726444 3.2268372 4.389828  4.256221  4.1885242 4.168866  4.2012367
  4.143094  3.9737947 3.884371 ]
 [2.3362284 2.9494276 3.972338  3.8259952 3.789136  3.8312256 3.9066408
  3.812707  3.5107985 3.4080713]
 [2.2282205 2.388075  3.1222186 2.9816258 2.9476242 2.9932737 3.2456686
  3.2231762 2.8630984 2.6808338]
 [2.254444  2.3769994 3.0225296 2.7452154 2.670063  2.74163   2.9335928
  2.8601842 2.460877  2.3203175]
 [2.2024732 2.27948   2.8722925 2.5166574 2.402021  2.4606936 2.563119
  2.4566424 2.0936174 2.0083768]
 [2.0544739 2.0758    2.6599777 2.2977428 2.145669  2.1437838 2.1209278
  2.003303  1.7692943 1.7560914]
 [1.7041425 1.7208269 2.3349662 2.0115452 1.8203009 1.724907  1.7033668
  1.5785713 1.3269699 1.2522831]]

[[2.2445092 2.556378  2.8660927 2.8883183 2.8982725 2.9018898 3.0587964
  3.123582  3.0752337 3.1326241]
 [2.2820568 2.7557082 3.351613  3.3406167 3.2286835 3.0512104 3.295968
  3.4294264 3.426199  3.3748393]
 [2.3061204 3.0955148 4.186507  4.094637  3.958023  3.794877  3.9764862
  4.106581  4.1734123 4.002126 ]
 [2.3726614 3.2268379 4.389816  4.2562065 4.188513  4.1688614 4.2012234
  4.1430774 3.9737802 3.8843684]
 [2.3362448 2.9494247 3.9723227 3.825974  3.7891164 3.8312113 3.9066267
  3.8126924 3.5107827 3.4080613]
 [2.2282343 2.388065  3.1221967 2.981601  2.947598  2.993248  3.2456493
  3.2231605 2.8630822 2.6808164]
 [2.2544525 2.3769825 3.0224996 2.7451882 2.670036  2.7416017 2.9335737
  2.860169  2.4608595 2.3203013]
 [2.2024744 2.279457  2.87226   2.5166295 2.401994  2.4606667 2.5630996
  2.4566264 2.093599  2.0083604]
 [2.054464  2.0757723 2.659951  2.2977154 2.1456437 2.1437619 2.1209073
  2.0032835 1.7692753 1.7560725]
 [1.7041312 1.7207947 2.3349338 2.0115201 1.8202796 1.7248867 1.7033452
  1.5785491 1.3269478 1.2522633]]

可以直观地看到输出结果只在小数点后4位开始才产生不同，因此模型导出的正确性可以得到验证。

使用np.allclose() 函数可以给出一个定量的结果。

三、C#中调用ONNX模型进行分割预测

使用ONNX的CPU版进行预测，使用GPU版进行预测的过程与此类似。

3.1OnnxRuntime软件包安装

以vs2019中安装为例，依次打开工具------NuGet包管理器------管理解决方案的NuGet程序包，搜索OnnxRuntime，安装如下图两个软件包：

3.2调用ONNX模型进行分割预测代码

安装之后，通过如下代码添加引用：

using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;

调用预测代码如下：

public void MeasureSingleImage()
{
    string modelPath = @"D:\pred.onnx";
    string imagePath = @"D:\000012.jpg";

    var session = new InferenceSession(modelPath);
    var inputTensor = Preprocess(imagePath);
    var inputName = session.InputMetadata.Keys.First();
    var inputs = new List<NamedOnnxValue>
    {
        NamedOnnxValue.CreateFromTensor(inputName, inputTensor)
    };
    var results = session.Run(inputs);
    var output = results.First().AsTensor<float>();
    var mask = Argmax(output, 3, 256, 256);
    SaveMask(mask, 256, 256, @"D:\result.jpg");
    return result;
}

public DenseTensor<float> Preprocess(string imagePath)
{
    int W = 256, H = 256;

    Bitmap bmp = new Bitmap(imagePath);

    float[] data = new float[1 * 3 * H * W];

    for (int y = 0; y < H; y++)
    {
        for (int x = 0; x < W; x++)
        {
            System.Drawing.Color c = bmp.GetPixel(x, y);
            data[0 * H * W + y * W + x] = c.R / 255f;
            data[1 * H * W + y * W + x] = c.G / 255f;
            data[2 * H * W + y * W + x] = c.B / 255f;
        }
    }

    return new DenseTensor<float>(data, new[] { 1, 3, H, W });
}

public int[] Argmax(Tensor<float> output, int C, int H, int W)
{
    float[] data = output.ToArray();
    int[] mask = new int[H * W];

    for (int y = 0; y < H; y++)
    {
        for (int x = 0; x < W; x++)
        {
            float maxVal = float.MinValue;
            int maxIdx = 0;

            for (int c = 0; c < C; c++)
            {
                int idx = c * H * W + y * W + x;

                if (data[idx] > maxVal)
                {
                    maxVal = data[idx];
                    maxIdx = c;
                }
            }

            mask[y * W + x] = maxIdx;
        }
    }

    return mask;
}

public void SaveMask(int[] mask, int W, int H, string path)
{
    Bitmap bmp = new Bitmap(W, H);

    System.Drawing.Color[] colors = new System.Drawing.Color[]
    {
        System.Drawing.Color.Black,   // 类0
        System.Drawing.Color.Red,     // 类1
        System.Drawing.Color.Green,   // 类2
    };

    for (int y = 0; y < H; y++)
    {
        for (int x = 0; x < W; x++)
        {
            int cls = mask[y * W + x];
            bmp.SetPixel(x, y, colors[cls]);
        }
    }

    bmp.Save(path);
}

上述代码实现了完整的图片分割预测过程，并将预测结果图片保存到本地。

3.3预测结果验证

将C#预测输出的原始结果保存到本地，然后与python调用onnx模型预测的输出结果进行对比，对比结果如下：

[[2.2444808 2.556353  2.8660774 2.8883095 2.8982658 2.9018817 3.0587878
  3.1235728 3.0752256 3.132615 ]
 [2.282033  2.7556944 3.3516119 3.3406177 3.2286854 3.051212  3.2959785
  3.4294393 3.4262075 3.3748384]
 [2.306106  3.095516  4.18652   4.094642  3.9580243 3.794877  3.976505
  4.1066055 4.173426  4.0021214]
 [2.3726444 3.2268372 4.389828  4.256221  4.1885242 4.168866  4.2012367
  4.143094  3.9737947 3.884371 ]
 [2.3362284 2.9494276 3.972338  3.8259952 3.789136  3.8312256 3.9066408
  3.812707  3.5107985 3.4080713]
 [2.2282205 2.388075  3.1222186 2.9816258 2.9476242 2.9932737 3.2456686
  3.2231762 2.8630984 2.6808338]
 [2.254444  2.3769994 3.0225296 2.7452154 2.670063  2.74163   2.9335928
  2.8601842 2.460877  2.3203175]
 [2.2024732 2.27948   2.8722925 2.5166574 2.402021  2.4606936 2.563119
  2.4566424 2.0936174 2.0083768]
 [2.0544739 2.0758    2.6599777 2.2977428 2.145669  2.1437838 2.1209278
  2.003303  1.7692943 1.7560914]
 [1.7041425 1.7208269 2.3349662 2.0115452 1.8203009 1.724907  1.7033668
  1.5785713 1.3269699 1.2522831]]

[[2.2411404 2.551806  2.8599524 2.8858123 2.8971422 2.9004815 3.0548823
  3.1190667 3.0724592 3.1328373]
 [2.2809253 2.749422  3.3415031 3.335905  3.2258985 3.0478518 3.291164
  3.4245348 3.4228873 3.3722012]
 [2.3083029 3.0885808 4.1736574 4.087658  3.9529195 3.7888155 3.971172
  4.1028733 4.172366  4.0009   ]
 [2.3687706 3.2184021 4.377794  4.249323  4.1821823 4.1598463 4.192381
  4.1344895 3.965546  3.876351 ]
 [2.3266954 2.9367476 3.9588048 3.8145552 3.778482  3.8204818 3.8995743
  3.8074884 3.5051818 3.402461 ]
 [2.213361  2.3691175 3.1053224 2.9624617 2.9306092 2.9816453 3.24404
  3.2270458 2.8669393 2.6840363]
 [2.2366161 2.3475637 2.9898705 2.7265735 2.6562533 2.7262516 2.9277108
  2.859814  2.461125  2.3203077]
 [2.1836388 2.2493584 2.8394773 2.4996498 2.390599  2.4476776 2.557334
  2.4550443 2.0924668 2.0083396]
 [2.0370405 2.0570173 2.6461732 2.2837353 2.1356308 2.140637  2.1205826
  2.004176  1.7694407 1.7599378]
 [1.6873618 1.697017  2.31027   1.9969264 1.8129518 1.7227268 1.7007617
  1.5763365 1.326082  1.2520404]]

可以直观地看到输出结果只在小数点后3位开始才产生不同，因此C#模型预测的正确性可以得到验证。