Change Upsample Layer to support direct export to CoreML

[dlawrences]

🚀 Feature

The default Upsample Layer seems to not be supported for CoreML conversion through ONNX. It needs to be changed to a simpler Upsample Layer that fixes the Slice operation.

Motivation

When converting exported ONNX model that has resulted from models/onnx_export.py, the following issue is encountered:

179/280: Converting Node Type Slice
Traceback (most recent call last):
  File "models/coreml_export.py", line 11, in <module>
    minimum_ios_deployment_target='13')
  File "/home/innaite/Applications/miniconda3/envs/pytorch_15/lib/python3.7/site-packages/onnx_coreml/converter.py", line 627, in convert
    _convert_node_nd(builder, node, graph, err)
  File "/home/innaite/Applications/miniconda3/envs/pytorch_15/lib/python3.7/site-packages/onnx_coreml/_operators_nd.py", line 2389, in _convert_node_nd
    return converter_fn(builder, node, graph, err)
  File "/home/innaite/Applications/miniconda3/envs/pytorch_15/lib/python3.7/site-packages/onnx_coreml/_operators_nd.py", line 1958, in _convert_slice
    err.unsupported_op_configuration(builder, node, graph, "Input shape not available")
  File "/home/innaite/Applications/miniconda3/envs/pytorch_15/lib/python3.7/site-packages/onnx_coreml/_error_utils.py", line 60, in unsupported_op_configuration
    self.rerun_suggestion)
TypeError: Error while converting op of type: Slice. Error message: Input shape not available 
 Please try converting with higher minimum_ios_deployment_target.
You can also provide custom function/layer to convert the model.

What I also tried so far

Changed default Upsample Layer params

I have also changed the Upsample Layer params from

[-2, 1, nn.Upsample, [None, 2, 'nearest']]

to

[-2, 1, nn.Upsample, [None, 2, 'bilinear']]

This does not fix the issue above.

Added custom Upsample Layer and replaced model definition

I have defined a custom Upsample Layer:

class UpsampleLayer(nn.Module):
    def __init__(self, scale_factor):
        super(UpsampleLayer, self).__init__()
        self.scale_factor = scale_factor

    def forward(self, x):
        assert (x.data.dim() == 4)
        B = x.data.size(0)
        C = x.data.size(1)
        H = x.data.size(2)
        W = x.data.size(3)
        ws = self.scale_factor
        hs = self.scale_factor
        #x = x.view(B, C, H, 1, W, 1).expand(B, C, H, hs, W, ws).contiguous().view(B, C, H * hs, W * ws)

        # coreml does not support tensors of rank 6, so expanding by scale factor has to be done in 2 steps
        x = x.view(B, C, H, 1, W).expand(B, C, H, hs, W).contiguous().view(B, C, H * hs, W)
        # tensor is now temporarily of shape B, C, H * hs, W
        x = x.view(B, C, H * hs, W, 1).expand(B, C, H * hs, W, ws).contiguous().view(B, C, H * hs, W * ws)
        # tensor has final shape of B, C, H * hs, W * ws
        return x

and updated model definition to:

[-2, 1, UpsampleLayer, [2]]

This does fix the conversion issue above from ONNX to CoreML, however it introduces a different issue which I don't how to fix. The loss on the validation dataset sky rockets, thus the performance quickly goes to zero.

image shows TensorBoard logs of model that uses torch Upsample Layer (bue) and of model that uses custom Upsample Layer defined above (orange)

Any thoughts?

Thanks

[github-actions]

Hello @dlawrences, thank you for your interest in our work! Please visit our Custom Training Tutorial to get started, and see our Jupyter Notebook , Docker Image, and Google Cloud Quickstart Guide for example environments.

If this is a bug report, please provide screenshots and minimum viable code to reproduce your issue, otherwise we can not help you.

If this is a custom model or data training question, please note that Ultralytics does not provide free personal support. As a leader in vision ML and AI, we do offer professional consulting, from simple expert advice up to delivery of fully customized, end-to-end production solutions for our clients, such as:

• Cloud-based AI systems operating on hundreds of HD video streams in realtime.
• Edge AI integrated into custom iOS and Android apps for realtime 30 FPS video inference.
• Custom data training, hyperparameter evolution, and model exportation to any destination.

For more information please visit https://www.ultralytics.com.

[glenn-jocher]

@dlawrences thanks for the detailed problem breakdown and very clear details.

I'm not sure that the .expand() operator is accomplishing the effect you are looking for. To verify your changes to the model, I would not train, I would simply run detect.py for anecdotal evidence, and run test.py for statistical comparison with the default model.

In any case, we actually don't have any problem exporting upsample operations. iDetection works great with the current upsample operators.

We do have a problem with a new layer in yolov5, the Focus module in the first layer. The stride 2 operations seem to be interpreted fine by onnx, but coreml is applying stride 2 to all dimensions rather than just the intended ones, though the error does not appear at export time, only at runtime on an iOS device.

[dlawrences]

Thanks @glenn-jocher

Would you be able to share the environment details using which you are able to do full export through ONNX to CoreML and if you are doing any pre-conversion clean-up steps?

Really appreciate your feedback.

Regards,
Laurentiu

[dlawrences]

Hi @glenn-jocher

I appreciate you sharing this with the community and your active development.

Would you be able to share the environment details that you used when exporting to ONNX and CoreML?

Thanks,
Laurentiu

[glenn-jocher]

@dlawrences we've updated export.py a bit to better support the 3 export channels (ONNX, TorchScript and CoremL). As we talked about however, these are simple first steps to point users in the right direction, they are not full export pipelines. See https://docs.ultralytics.com/yolov5/tutorials/model_export

[dlawrences]

Cheers

I had a look - wouldn't you need to specify image scaling as well?

inputs=[ct.ImageType(name="input", shape=example_input.shape, scale=1/255.0, bias=[0, 0, 0])]

I know you do specify it during inference:

img /= 255.0  # 0 - 255 to 0.0 - 1.0

[glenn-jocher]

Yes, I believe you are correct! Can you submit a PR for this bug fix please?

[dlawrences]

Sure. Check out #321

[phamdan]

Sure. Check out #321

@dlawrences
have you exported to coreml. can you share it?