NNAPI Serializer & PixelShuffle

PyTorch Version: 1.8.1+cu102

Using this tutorial for NNAPI ((Prototype) Convert MobileNetV2 to NNAPI — PyTorch Tutorials 1.8.1+cu102 documentation), I’m trying to convert my QAT trained model with qnnpack backend, using torch.backends._nnapi.prepare.convert_model_to_nnapi(traced, input_tensor)

However in my model I make use of torch.nn.functional.pixel_shuffle, and when I run the conversion function I receive the following error:

File ".../pytorch/lib/python3.8/site-packages/torch/backends/_nnapi/serializer.py", line 641, in add_node                                                                                                                             
    raise Exception("Unsupported node kind (%r) in node %r" % (node.kind(), node))                                                                                                                                                                                                                             
Exception: Unsupported node kind ('aten::pixel_shuffle') in node %99 : Tensor = aten::pixel_shuffle(%x_hat, %98)

It seems like PixelShuffle is not supported; is that correct?

However I can see here DEPTH_TO_SPACE and SPACE_TO_DEPTH are NNAPI operation codes: pytorch/serializer.py at cf7d56d8f223d4b10a624a5f7af886a2b67f0955 · pytorch/pytorch · GitHub

Is there any way to use PixelShuffle like behaviour with the NNAPI?

Apologies PixelShuffle is not currently supported by the NNAPI converter.

You can see list of supported ops at this time here: pytorch/serializer.py at master · pytorch/pytorch · GitHub

We’re adding more ops soon!