pytorch模型导出到onnx时预测固定标签

Question

我在 pytorch 中训练了 resnet-18 模型。而且在pytorch中效果很好
但是，当我将其转换为 onnx 并在 cv2 中进行预测时，模型仅预测 1~2 个标签（它应该预测 0~17 个标签）。
这是我的模型导出代码

    model.eval()
    x = torch.randn(1, 3, 512, 384, requires_grad=True)

    # export model
    torch.onnx.export(model, x, "model.onnx", export_params=True, opset_version=10, do_constant_folding=True, input_names = ['input'], output_names = ['output'])

这是我在 cv2 中进行推理的代码

self.transform = albumentations.Compose([
        albumentations.Resize(512, 384, cv2.INTER_LINEAR),
        albumentations.GaussianBlur(3, sigma_limit=(0.1, 2)),
        albumentations.Normalize(mean=(0.5), std=(0.2)),
        albumentations.ToFloat(max_value=255)
        ])
...
#image crop code: works fine in pytorch
image = frame[ymin:ymax, xmin:xmax]  #type(frame)==numpy.array, RGB form
augmented = self.transform(image=image)
image = augmented["image"]
...
#inference code: does not work well
net=cv2.dnn.readNet("Model.onnx")
blob = cv2.dnn.blobFromImage(image, swapRB=False, crop=False)
net.setInput(blob)
label = np.array(net.forward())
text  =  'Label: '+str(np.argmax(label[0]))

所有转换设置在 pytorch 中运行良好。
这段代码可能有什么问题？

Answer 1

您的代码的问题可能与以不同方式预处理图像有关：self.transform 重新缩放图像，但是当您阅读 blob 时，您并没有这样做。要验证这一点，您可以读取相同的图像并检查 image 和 blob 是否相等（例如，使用 torch.allclose），当实际（随机）增强被禁用时。

pytorch模型导出到onnx时预测固定标签

pytorch model predicts fixed label when it exports to onnx

cv2

pytorch

onnx