Flexible shapes not working with PyTorch unified converter

[3DTOPO]

🐞Describe the bug

I made changes to my model so I could use the recommended unified convertor. Conversion is successful without issue and shows that flexible shapes are supported (in both Python and Xcode).

Running prediction with a shape in the supported ranges (and any shape other than the fixed shape) will fail with an error. The fixed shape input works as expected. I've tried both GPU and CPU only.

Trace

[espresso] [Espresso::handle_ex_plan] exception=Espresso exception: "Not implemented": axis -4 not implemented status=-9
[coreml] Failure dynamically resizing for sequence length.
[coreml] Failure in resetSizes.
prediction error: Error Domain=com.apple.CoreML Code=0 "Failure dynamically resizing for sequence length." UserInfo={NSLocalizedDescription=Failure dynamically resizing for sequence length.}

To Reproduce

The source code and model is in the attached archive.

import torch
import torch.nn as nn
import coremltools as ct
import coremltools.proto.FeatureTypes_pb2 as ft
from coremltools.models.neural_network import flexible_shape_utils
from model import TransformerNet

channels = 3
width = 1024
height = 1024

torch_model = TransformerNet()
#torch_model.load_state_dict(torch.load('TrainedModel.pth', map_location=torch.device('cpu')))
torch_model.eval()

example_input = torch.rand(1, channels, width, height)
traced_model = torch.jit.trace(torch_model, example_input)

mlmodel = ct.convert(
    traced_model,
    inputs=[ct.ImageType(name="input_1", shape=example_input.shape)], 
    minimum_ios_deployment_target='13'
)

#note if "input" is used for the name it creates a name collision

spec = mlmodel.get_spec()

# needed because documentation states:
# outputs must not be specified for PyTorch
output = spec.description.output[0]
output.type.imageType.colorSpace = ft.ImageFeatureType.RGB
output.type.imageType.height = height
output.type.imageType.width = width

ct.utils.rename_feature(spec, '782', 'output')

img_size_ranges = flexible_shape_utils.NeuralNetworkImageSizeRange(height_range=(256, 3072), width_range=(256, 3072))
flexible_shape_utils.update_image_size_range(spec, feature_name='input_1', size_range=img_size_ranges)
flexible_shape_utils.update_image_size_range(spec, feature_name='output', size_range=img_size_ranges)

ct.utils.save_spec(spec, "TransformerNet.mlmodel")

model.py:

import torch
import torch.nn as nn

class TransformerNet(torch.nn.Module):

    def __init__(self):
        super(TransformerNet, self).__init__()
        # Initial convolution layers
        self.conv1 = ConvLayer(3, 8, kernel_size=9, stride=1)
        self.in1 = torch.nn.InstanceNorm2d(8, affine=True)
        self.conv2 = ConvLayer(8, 16, kernel_size=3, stride=2)
        self.in2 = torch.nn.InstanceNorm2d(16, affine=True)
        self.conv3 = ConvLayer(16, 32, kernel_size=3, stride=2)
        self.in3 = torch.nn.InstanceNorm2d(32, affine=True)
        # Residual layers
        self.res1 = ResidualBlock(32)
        self.res2 = ResidualBlock(32)
        self.res3 = ResidualBlock(32)
        self.res4 = ResidualBlock(32)
        self.res5 = ResidualBlock(32)
        # Upsampling Layers
        self.deconv1 = UpsampleConvLayer(32, 16, kernel_size=3, stride=1, upsample=2)
        self.in4 = torch.nn.InstanceNorm2d(16, affine=True)
        self.deconv2 = UpsampleConvLayer(16, 8, kernel_size=3, stride=1, upsample=2)
        self.in5 = torch.nn.InstanceNorm2d(8, affine=True)
        self.deconv3 = ConvLayer(8, 3, kernel_size=9, stride=1)
        # Non-linearities
        self.relu = torch.nn.ReLU()

    def forward(self, X):
        y = self.relu(self.in1(self.conv1(X)))
        y = self.relu(self.in2(self.conv2(y)))
        y = self.relu(self.in3(self.conv3(y)))
        y = self.res1(y)
        y = self.res2(y)
        y = self.res3(y)
        y = self.res4(y)
        y = self.res5(y)
        y = self.relu(self.in4(self.deconv1(y)))
        y = self.relu(self.in5(self.deconv2(y)))
        y = self.deconv3(y)
        return y

class ConvLayer(torch.nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, stride):
        super(ConvLayer, self).__init__()
        reflection_padding = kernel_size // 2
        #self.reflection_pad = torch.nn.ReflectionPad2d(reflection_padding)
        self.reflection_pad = ReflectPad2d_rev(reflection_padding)
        self.conv2d = torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride)

    def forward(self, x):
        out = self.reflection_pad(x)
        out = self.conv2d(out)
        return out

class ResidualBlock(torch.nn.Module):

    def __init__(self, channels):
        super(ResidualBlock, self).__init__()
        self.conv1 = ConvLayer(channels, channels, kernel_size=3, stride=1)
        self.in1 = torch.nn.InstanceNorm2d(channels, affine=True)
        self.conv2 = ConvLayer(channels, channels, kernel_size=3, stride=1)
        self.in2 = torch.nn.InstanceNorm2d(channels, affine=True)
        self.relu = torch.nn.ReLU()

    def forward(self, x):
        residual = x
        out = self.relu(self.in1(self.conv1(x)))
        out = self.in2(self.conv2(out))
        out = out + residual
        return out

class UpsampleConvLayer(torch.nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, stride, upsample=None):
        super(UpsampleConvLayer, self).__init__()
        self.upsample = upsample
        reflection_padding = kernel_size // 2
        #self.reflection_pad = torch.nn.ReflectionPad2d(reflection_padding)
        self.reflection_pad = ReflectPad2d_rev(reflection_padding)
        self.conv2d = torch.nn.Conv2d(in_channels, out_channels, kernel_size, stride)

    def forward(self, x):
        x_in = x
        if self.upsample:
            x_in = torch.nn.functional.interpolate(x_in, mode='nearest', scale_factor=self.upsample)
        out = self.reflection_pad(x_in)
        out = self.conv2d(out)
        return out

class ReflectPad2d_rev(nn.Module):

    def __init__(self, size):
        super().__init__()
        self.size = size

    def forward(self, x):
        a = self.size
        L_list, R_list = [], []
        U_list, D_list = [], []
        for i in range(a):#i:0, 1
            l = x[:, :, :, (a-i):(a-i+1)]
            L_list.append(l)
            r = x[:, :, :, (i-a-1):(i-a)]
            R_list.append(r)
        L_list.append(x)
        x = torch.cat(L_list+R_list[::-1], dim=3)
        for i in range(a):
            u = x[:, :, (a-i):(a-i+1), :]
            U_list.append(u)
            d = x[:, :, (i-a-1):(i-a), :]
            D_list.append(d)
        U_list.append(x)
        x = torch.cat(U_list+D_list[::-1], dim=2)
        return x

System environment (please complete the following information):

• coremltools 4.0
• OS MacOS
• macOS 10.15.7 (19H2)
• Version 12.1 (12A7403)
• virtualenv
• python version 3.7
  • pytorch 1.70

Additional context

This issue severely restricts deploying MLModels across my workflow.

[Mstronach]

Hi Jeshua, thank you for submitting this issue. Is this different than issue #992 ? Thanks!

[3DTOPO]

Hi Mstronach, you're quite welcome. Yes, it appears to be a duplicate. I hadn't noticed the double-post before. It must have gotten posted twice by accident.

[TobyRoseman]

Yes, it appears to be a duplicate.

Ok, closing as a duplicate.