pytorch pth 模型转 onnx模型，并验证结果正确性

刺骨的言语ヽ痛彻心扉 2022-12-23 00:37 461阅读 0赞

> pytorch 模型部署很重要的一步是转存pth模型为ONNX，本文记录方法。

### 转存 onnx ###

*  建立自己的pytorch模型，并加载权重

model = create_model(num_classes=2)
    model.load_state_dict(load(model_path, map_location='cpu')["model"])

*  转存onnx文件

dummy_input = torch.randn(1, 3, 256, 256, device='cpu')
    torch.onnx._export(model, dummy_input, "faster_rcnn.onnx", verbose=True, opset_version=11)

> 将模型保存在了当前目录的 `faster_rcnn.onnx`文件内

#### 验证 onnx 有效性 ####

*  安装 `onnxruntime` 库

pip install onnxruntime

*  加载onnx模型并测试

import onnxruntime
    from onnxruntime.datasets import get_example
    
    def to_numpy(tensor):
        return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
    
    # 测试数据
    dummy_input = torch.randn(1, 3, 256, 256, device='cpu')
    
    example_model = get_example(<absolute_root_to_your_onnx_model_file>)
    # netron.start(example_model) 使用 netron python 包可视化网络
    sess = onnxruntime.InferenceSession(example_model)
    
    # onnx 网络输出
    onnx_out = sess.run(None, { <input_layer_name_of_your_network>: to_numpy(dummy_input)})
    print(onnx_out)
    
    model.eval()
    with torch.no_grad():
        # pytorch model 网络输出
        torch_out = model(dummy_input)
    	print(torch_out)

*  输出：

onnx_out
    [array([[  0.       ,  93.246    , 228.95842  , 256.       ],
           [  0.       ,   2.6370468, 209.39705  , 148.17822  ]],
          dtype=float32), array([1, 1], dtype=int64), array([0.1501071 , 0.07568519], dtype=float32)]
    
    torch_out
    [{'boxes': tensor([[  0.0000,  93.2459, 228.9584, 256.0000],
            [  0.0000,   2.6370, 209.3971, 148.1782]]), 'labels': tensor([1, 1]), 'scores': tensor([0.1501, 0.0757])}]

#### 获取自己网络输入层名称 ####

*  有时对网络不熟悉的情况下不清楚模型输入层的名称，可以使用Netron可视化自己的网络，获取输入层名称，喂入onnx的sess中。

## 注意 ！！！ ##

*  pytorch 模型在转 ONNX 模型的过程中，使用的导出器是一个基于轨迹的导出器,这意味着它执行时需要运行一次模型,然后导出实际参与运算的运算符. 这也意味着, 如果你的模型是动态的,例如,改变一些依赖于输入数据的操作,这时的导出结果是不准确的.同样,一 个轨迹可能只对一个具体的输入尺寸有效 (这是为什么我们在轨迹中需要有明确的输入的原因之一.) 我们建议检查 模型的轨迹,确保被追踪的运算符是合理的. ——— [pytorch 文档][pytorch]
 *  也就是说，如果网络模块中存在 if… else… 类似的分支，在生成ONNX模型时会依据所使用的初始数据来选择其中某一个分支，这样所生成的ONNX模型仅会保留这一个分支的结构，在原始pytorch模型中的其他逻辑能力在该模型中不复存在。

### 参考资料 ###

*  https://pytorch.apachecn.org/docs/0.3/onnx.html
 *  https://www.jianshu.com/p/5a0a09fbdeba

[pytorch]: https://pytorch.apachecn.org/docs/0.3/onnx.html