import torch
import torchvision as tv
from u2net import U2NET
def deeplabv32onnx():
model = U2NET()
model.load_state_dict(torch.load(r'D:\Dely\unet.pt'))
model.eval()
dummy_input = torch.randn(1, 3, 512,512)
model(dummy_input)
im = torch.zeros(1, 3,512,512)#.to("cpu")
torch.onnx.export(model, im,
"u2net.onnx",
verbose=False,
opset_version=11,
training=torch.onnx.TrainingMode.EVAL,
do_constant_folding=True,
input_names=['input1'],
output_names=['o0', 'o1','o2','o3','o4','o5','o6'],
dynamic_axes={'input': {0: 'batch', 2: 'height', 3: 'width'}}
)
deeplabv32onnx()
u2net导onnx
最新推荐文章于 2024-06-25 10:17:20 发布