Py手电模型预测出口到onnx时的固定标签。

Question

我在pytorch训练了resnet-18模型。它能很好地应用于火把。

但是，当我将其转换为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]))

所有转换设置都能很好地工作在移相器中。

这段代码有什么问题？

Answer 1

代码的问题可能与图像的预处理不同有关：self.transform重新调整图像的标度，但是当您读取blob时，并不是这样做的。要验证这一点，您可以读取相同的图像，并在禁用实际(随机)增强时检查image和blob是否相等(例如使用torch.allclose)。