Convert pythorch pth model to onnx with fixed width and variable height

Alan_S: <p>I also posted the question <a href="https://stackoverflow.com/questions/79291415/convert-pythorch-pth-model-to-onnx-with-fixed-width-and-variable-height" rel="noopener nofollow ugc">here</a>.</p>
<p>My code uses PyTorch to perform segmentation annotations on PNG images. The input images have a width of 512 pixels or a multiple of this, but the height can range from 400 to 900 pixels. The code, along with the PyTorch model (*.pth file), works as expected.</p>
<p>I am currently attempting to convert my *.pth model to *.onnx. The code itself hasn’t changed much (only modifications related to ONNX, naturally), but the issue I am encountering is with the model conversion.</p>
<p>Here is my code for the model conversion:</p>
<pre data-code-wrap="python"><code class="lang-python">import onnx
import torch
import torch.nn as nn
import torch.nn.functional as F

# pip install torch onnx


class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        super().__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)
        self.down1 = Down(64, 128)
        self.down2 = Down(128, 256)
        self.down3 = Down(256, 512)
        factor = 2 if bilinear else 1
        self.down4 = Down(512, 1024 // factor)
        self.up1 = Up(1024, 512, bilinear)
        self.up2 = Up(512, 256, bilinear)
        self.up3 = Up(256, 128, bilinear)
        self.up4 = Up(128, 64 * factor, bilinear)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits


class DoubleConv(nn.Module):
    """(convolution =&gt; [BN] =&gt; ReLU) * 2"""

    def __init__(self, in_channels, out_channels, mid_channels=None):
        super().__init__()
        if not mid_channels:
            mid_channels = out_channels
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, x):
        return self.double_conv(x)


class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(nn.MaxPool2d(2), DoubleConv(in_channels, out_channels))

    def forward(self, x):
        return self.maxpool_conv(x)


class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()
        self.up: nn.Upsample | nn.ConvTranspose2d

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels // 2, in_channels // 2)
        else:
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
            self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is CHW
        diffY = torch.tensor([x2.size()[2] - x1.size()[2]])
        diffX = torch.tensor([x2.size()[3] - x1.size()[3]])

        x1 = F.pad(
            x1,
            [
                torch.div(diffX, 2, rounding_mode="floor"),
                diffX - torch.div(diffX, 2, rounding_mode="floor"),
                torch.div(diffY, 2, rounding_mode="floor"),
                diffY - torch.div(diffY, 2, rounding_mode="floor"),
            ],
        )
        # if you have padding issues, see
        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)


def convert_pytorch_to_onnx(pytorch_model_path, onnx_model_path):
    # Load the PyTorch model
    model = UNet(n_channels=1, n_classes=1)
    model.load_state_dict(torch.load(pytorch_model_path, map_location="cpu"))
    model.eval()

    # Create dummy input with dynamic size
    dummy_input = torch.randn(1, 1, 512, 512)  # Height [400 to 900], Width fixed at 512

    # Export the model
    torch.onnx.export(
        model,  # model being run
        dummy_input,  # model input (or a tuple for multiple inputs)
        onnx_model_path,  # where to save the model
        export_params=True,  # store the trained parameter weights inside the model file
        opset_version=20,  # the ONNX version to export the model to
        do_constant_folding=True,  # whether to execute constant folding for optimization
        input_names=["input"],  # the model's input names
        output_names=["output"],  # the model's output names
        dynamic_axes={
            "input": {0: "batch_size", 2: "height", 3: "width"},  # variable length axes
            "output": {0: "batch_size", 2: "height", 3: "width"},
        },
    )

    # Verify the model
    onnx_model = onnx.load(onnx_model_path)
    onnx.checker.check_model(onnx_model)

    print(f"Model {pytorch_model_path} converted to {onnx_model_path}")


# List of models to convert
models = [
    "models/case1.pth",
    "models/case2.pth",
]

# Convert each model
for model_path in models:
    onnx_path = model_path.replace(".pth", ".onnx")
    convert_pytorch_to_onnx(model_path, onnx_path)
</code></pre>
<p>Using the ONNX models created with:</p>
<pre data-code-wrap="python"><code class="lang-python">dummy_input = torch.randn(1, 1, 512, 512)  # Height [400 to 900], Width fixed at 512
</code></pre>
<p>Inference on images with heights between 400 and 600 pixels seems to work similarly to my PyTorch code. However, images with a height of 800 pixels produce incorrect results when they work at all.</p>
<p>Conversely, if I convert the model using:</p>
<pre data-code-wrap="python"><code class="lang-python">dummy_input = torch.randn(1, 1, 885, 512)
</code></pre>
<p>An image with a height of 885 pixels works perfectly.</p>
<p>I’m not an expert in PyTorch or ONNX. For now, the only “workable” solution I’ve found is to use 512x512 in the <code>dummy_input</code> and add a few lines to my Python code to crop the top and bottom of input images if they are taller than 512 pixels. However, the results are not identical to those of the original PyTorch code.</p>
<p>I’m unsure exactly what PyTorch does with variable-height inputs — I didn’t build the model — but from inspecting the UNet class, it’s clear that it is not cropping the input. It’s more likely “reducing” the height until it fits the model’s structure.</p>
<p>If that’s the case, how can I replicate this behaviour with an ONNX model?</p>
<p>I’m also wondering if torch-onnx: github . com / justinchuby / torch-onnx or torch-dynamo: pytorch . org / docs / stable/ onnx.html can help me here.</p>

Alan_S: <p>With some hints from a colleague I at least got want I wanted.<br>
In the end I did two changes on the code above:</p>
<p>I simplified the <code>Up.forward()</code> so it would be properly translated to ONNX, I believe:</p>
<pre data-code-wrap="python"><code class="lang-python">    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is Channel Height Width
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]

        padding_left = diffX // 2
        padding_right = diffX - padding_left
        padding_top = diffY // 2
        padding_bottom = diffY - padding_top

        x1 = F.pad(x1, [padding_left, padding_right, padding_top, padding_bottom])

        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)
</code></pre>
<p>And I did:</p>
<pre data-code-wrap="python"><code class="lang-python">    dummy_input = torch.randn(1, 1, 768, 768)  # H[400 to 900] x W (512, 768, 1024, 1536)
</code></pre>
<p>As I realised the model was trained with that input shape.<br>
<code>torch.dynamo</code> did not work for me.</p>