pytorch 转onnx

三只熊999

大佬们，渣渣最近要把pytorch 的模型转成onnx 的，转之前模型是可以预测的可视

转之后框都飞了。
下面是我转换的脚本
在预测的时候回警告说CleanUnusedInitializers] Removing initializer 'layer3.1.bn1.num_batches_tracked'. It is not used by any node and should be removed from the model
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
pt_model_path = ’model_best.pth'
onnx_model_path = 'model_latest.onnx
model = resnet18(pretrained=True)
model.eval().cuda()
model.load_state_dict(torch.load(pt_model_path), False
input_tensor = torch.randn(1,3,640,640).cuda()
input_names = ['input']
output_names = ['output']
print(type(input_tensor))
torch.onnx.export(model, (input_tensor,), onnx_model_path, verbose=True, input_names=input_names, output_names=output_names,opset_version=10)