Search code examples
pytorchonnxonnx-coreml

Error while converting pytorch model to core-ml


C = torch.cat((A,B),1)

shape of tensors:

A is (1, 128, 128, 256)
B is (1, 1, 128, 256)

Expected C value is (1, 129, 128, 256)

This code is working on pytorch, but while converting to core-ml it gives me below error:

"Error while converting op of type: {}. Error message: {}\n".format(node.op_type, err_message, )
TypeError: Error while converting op of type: Concat. Error message: unable to translate constant array shape to CoreML shape"

Solution

  • It was coremltools version related issue. Tried with latest beta coremltools 3.0b2.

    Following works without any error with latest beta.

    import torch
    
    class cat_model(torch.nn.Module):
        def __init__(self):
            super(cat_model, self).__init__()
    
        def forward(self, a, b):
            c = torch.cat((a, b), 1)
            # print(c.shape)
            return c
    
    a = torch.randn((1, 128, 128, 256))
    b = torch.randn((1, 1, 128, 256))
    
    model = cat_model()
    torch.onnx.export(model, (a, b), 'cat_model.onnx')
    
    import onnx
    model = onnx.load('cat_model.onnx')
    onnx.checker.check_model(model)
    print(onnx.helper.printable_graph(model.graph))
    
    from onnx_coreml import convert
    mlmodel = convert(model)