xuebinqin/U-2-Net

onnx result wrong

NingNanXin opened this issue · 5 comments

Hi, Thanks for your work. Recently I retrain U2net for my task and export onnx. The export code as follows:

model = U2NET(3, 1)
state_dict = torch.load("u2net_bce_itr_12000_train_0.127958_tar_0.012812.pth")
model.load_state_dict(state_dict)

model.eval()

input_names = ["input"]
output_names = ["output"]
input_tensor = torch.zeros([1, 3, 320, 320])

# I change the network to get 1 output: d1
torch.onnx.export(model,
                  input_tensor,
                  "human_seg.onnx",
                  opset_version=12,
                  export_params=True,
                  verbose=True,
                  input_names=input_names,
                  output_names=output_names,
                  dynamic_axes={"input": {0: "batch"},
                                "output": {0: "batch"}})

After successfully export onnx, I try to infer an image with onnx. The preprocess code as follows, but the result is wrong.

image = io.imread("xiaozheng.png")
size = (image.shape[1], image.shape[0])
bg = np.ones_like(image) * 255

src = transform.resize(image, (320, 320), mode='constant')

mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
src = src / np.max(src)
src = src - mean
src = src / std

src = np.transpose(src, [2, 0, 1])
src = np.expand_dims(src, 0)
src = src.astype(np.float32)

sess = onnxruntime.InferenceSession("human_seg.onnx")
d1 = sess.run([], {"input": src})[0]

image

Thanks for your help!

save_result func, I comment the "255". Speechless, I am so, emmm.

@NingNanXin I have the same problem. Did you solve this?

@NingNanXin I have the same problem. Did you solve this?

Debug your preprocess and post-process, use torch.onnx.export() can get the right onnx model.

@NingNanXin Thanks for your reply. Can you explain more detail? How did you solve that?

@ducnguyen998

  1. use a random generate image as an input
  2. use torch and onnx to infer your image
  3. use numpy.testing.assert_allclose to compare torch result and onnx result
    For me, 3 is right, then compare your process code with the origin code

@NingNanXin Thank you.
I have my own model that produces accurate results when run in PyTorch format. However, when I convert it to ONNX format, the model gives incorrect results. I attempted to convert the author's pre-trained model to ONNX, and it successfully produced the same results as the PyTorch model. However, when I applied the same process to my own model, it did not work as expected.