Getting unexpected output after conversion of model to ONNX

[prateek9623]

Hi, I am trying to convert "timm-efficientnet-b0" to ONNX with mixed-precision and facing issues while validating the onnx model. In full-precision, outputs are approximately matching, but in mixed-precision, there are a lot of differences in the o/p and in the graph (compared to the full-precision graph).

Testing code:

# %%
# import sys
# !{sys.executable} -m pip install --upgrade segmentation-models-pytorch
# !{sys.executable} -m pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu117
# !{sys.executable} -m pip install --upgrade onnxruntime-gpu

# %%
import torch
import segmentation_models_pytorch as smp

device = torch.device("cuda")
model = smp.Unet(encoder_name="timm-efficientnet-b0", encoder_weights="noisy-student", in_channels=3, classes=2).to(
    device
)
model.eval()

# %%
import numpy as np
input = np.random.rand(1, 3, 256, 256).astype(np.float32)

# %%

pytorch_input = torch.from_numpy(input).cuda()
enable_mixed_precision = True
with torch.no_grad():
    with torch.cuda.amp.autocast(enable_mixed_precision):
        pytorch_output = model(pytorch_input)
pytorch_output.cpu().detach().numpy()[0][0][0][0]

# %%
export_path = "Model.onnx"
with torch.no_grad():
    with torch.cuda.amp.autocast(enabled=enable_mixed_precision):
        torch.onnx.export(
            model,
            pytorch_input,
            export_path,
            verbose=False,
            do_constant_folding=True,
            opset_version=17,
            export_params=True,
            input_names=["input"],
            output_names=["output1"],
            dynamic_axes={
                "input": {0: "batch_size"},
                "output1": {0: "batch_size"}
            },
        )

# %%
import onnxruntime as ort

session_options = ort.SessionOptions()
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
session_options.log_severity_level = 0
providers = [
    (
        "CUDAExecutionProvider",
        {
            "device_id": 0,
            "arena_extend_strategy": "kNextPowerOfTwo",
            "gpu_mem_limit": 2 * 1024 * 1024 * 1024,
            "cudnn_conv_algo_search": "EXHAUSTIVE",
            # 'enable_cuda_graph': True,
            "do_copy_in_default_stream": True,
        },
    )
]
ort_session = ort.InferenceSession(export_path, session_options, providers)


# %%
ort_output = ort_session.run(None, {"input": input})
ort_output[0][0][0][0][0]

# %%
print(np.abs(pytorch_output[0].cpu().numpy() - ort_output[0]).mean())

Full-precision graph: https://github.com/qubvel/segmentation_models.pytorch/assets/17100687/6ae20ac7-56c1-4391-82ce-6390fc81077c

Half-precision graph: https://github.com/qubvel/segmentation_models.pytorch/assets/17100687/5cd3b19f-3610-4c17-8022-362a32e29561

[github-actions[bot]]

This issue is stale because it has been open 60 days with no activity. Remove stale label or comment or this will be closed in 7 days.

[prateek9623]

adding comment so that issue doesn't become stale

[github-actions[bot]]

This issue is stale because it has been open 60 days with no activity. Remove stale label or comment or this will be closed in 7 days.

[github-actions[bot]]

This issue was closed because it has been stalled for 7 days with no activity.