Error in converting pytorch model to mlmodel with single input image input

1

I am trying to convert a Pytorch model to .mlmodel. so I am first converting .pth file to .onnx file then .onnx file to .mlmodel file.

My Input size is : (1,3,299,299)

Here is my code for conversion:

class DenseNet(nn.Module):
    def __init__(self):
        super(DenseNet, self).__init__()
        
        # get the pretrained DenseNet model
        densenet = models.densenet161(pretrained=True)
        
        # freeze the gradients to avoid training
        for i, param in enumerate(densenet.parameters()):
            param.requires_grad_(False)
        
        # transfer learning procedure
        # take everything before the 12th layer of the densenet
        modules = list(densenet.children())[0][:12]
        self.densenet = nn.Sequential(*modules)
        
        # transfer the classifier
        self.classifier1 = nn.Linear(in_features=2208, out_features=4096)
        self.classifier2 = nn.Linear(in_features=4096, out_features=2887)
        
        # relu activation
        self.prelu = nn.PReLU()

        # dropout
        self.dropout = nn.Dropout(p=0.5)
        
        # max pool
        self.avg_pool = nn.AvgPool2d(kernel_size=7)
        
    def forward(self, x):
        
        # get the features from the original VGG
        features = self.densenet(x)
        # flat the features
        features = self.avg_pool(features).squeeze()
        # out
        features = self.dropout(self.prelu(self.classifier1(features)))
        logits = self.classifier2(features)
        return logits


def load_checkpoint(checkpoint_path):
    checkpoint = torch.load(checkpoint_path,map_location="cpu")
    model = DenseNet()
    model.load_state_dict(checkpoint)
    return model

model = load_checkpoint('model.pth')
dummy_input = torch.randn(1, 3, 299, 299, device='cpu')
torch.onnx.export(model, dummy_input, "model_output.onnx")

This script is running successfully while .onnx to .mlmodel script gives me the error.

Here is the script that I am using :

import sys
from onnx import onnx_pb
from onnx_coreml import convert

model_in = 'model_output.onnx'
model_out = 'model_output.mlmodel'

model_file = open(model_in, 'rb')
model_proto = onnx_pb.ModelProto()
model_proto.ParseFromString(model_file.read())
coreml_model = convert(model_proto)
coreml_model.save(model_out)

This script is giving me the following errors :

570/574: Converting Node Type MatMul
Traceback (most recent call last):
File “densenet.py”, line 84, in
coreml_model = convert(model_proto)
File “/home/ubuntu/anaconda3/envs/py36/lib/python3.6/site-packages/onnx_coreml/converter.py”, line 458, in convert
_convert_node(builder, node, graph, err)
File “/home/ubuntu/anaconda3/envs/py36/lib/python3.6/site-packages/onnx_coreml/_operators.py”, line 1755, in _convert_node
return converter_fn(builder, node, graph, err)
File “/home/ubuntu/anaconda3/envs/py36/lib/python3.6/site-packages/onnx_coreml/_operators.py”, line 1091, in _convert_matmul
return err.unsupported_op_configuration(builder, node, graph, “CoreML incompatible axis placement”)
File “/home/ubuntu/anaconda3/envs/py36/lib/python3.6/site-packages/onnx_coreml/_error_utils.py”, line 56, in unsupported_op_configuration
“Error while converting op of type: {}. Error message: {}\n”.format(node.op_type, err_message, )
TypeError: Error while converting op of type: MatMul. Error message: CoreML incompatible axis placement

But when I change the dummy input size in pth to onnx model to

dummy_input = torch.randn(10, 3, 299, 299, device=‘cpu’)

or any number greater than 1.It works prefectly fine.
But I need the input to be single color image with size (299,299).