345345t43

3t

ta t 

65u

ysrxysxhy

000000

’‘’

'''

import torch

from model import Model

old_net_path = 'lenet.pth'

new_net_path = 'lenet.onnx'

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu' )

# 导入模型

net = Model().to(device)

net.load_state_dict(torch.load(old_net_path, map_location=device))

net.eval()

input = torch.randn(1, 1, 30, 30).to(device)   # BCHW  其中Batch必须为1，因为测试时一般为1，尺寸HW必须和训练时的尺寸一致

torch.onnx.export(net, input, new_net_path, verbose=False)

'''

wrw4r4t