Pytroch转Onnx实现（Unet的模型转Onnx）

最新推荐文章于 2025-09-10 01:58:46 发布

原创

最新推荐文章于 2025-09-10 01:58:46 发布 · 2.3k 阅读

标签

#pytorch #python #Unet #Onnx

收录于

主要是使用torch.onnx.export()这个方法来实现。
Unet的实现参考：
链接: https://blog.csdn.net/weixin_44791964/article/details/108866828.
这位博主写的很详细，b站还有实现视频，手把手教学！！！

Onnx转换实现代码

import onnx
import torch.onnx
from unet import Unet
unet=Unet()
print(unet)
model = torch.load('D:/PycharmProjects/test/unet-pytorch-main/model_data/Epoch19-Total_Loss0.5564-Val_Loss0.5045.pth',map_location='cpu')
unet.net.load_state_dict(model)
unet.net.eval()
print(unet.net)
print(torch.__version__)