Преобразование grid_sample PyTorch в CoreML (через coremltools)

torch.nn.functional.grid_sample (источник здесь, щелкните документы для документации ) в настоящее время не поддерживается CoreML (и их библиотекой утилит преобразования: coremltools).

Я ищу способ экспортировать слой, показанный ниже, из torchscript PyTorch (документы здесь) в CoreML (либо с использованием пользовательского op, созданного с помощью Swift, либо с помощью эффективной переписывания grid_sample PyTorch).

Подробные сведения и советы по началу работы см. в разделе «Советы».

Минимальный проверяемый пример

import coremltools as ct
import torch


class GridSample(torch.nn.Module):
    def forward(self, inputs, grid):
        # Rest could be the default behaviour, e.g. bilinear
        return torch.nn.functional.grid_sample(inputs, grid, align_corners=True)


# Image could also have more in_channels, different dimension etc.,
# for example (2, 32, 64, 64)
image = torch.randn(2, 3, 32, 32)  # (batch, in_channels, width, height)
grid = torch.randint(low=-1, high=2, size=(2, 64, 64, 2)).float()

layer = GridSample()
# You could use `torch.jit.script` if preferable
scripted = torch.jit.trace(layer, (image, grid))

# Sanity check
print(scripted(image, grid).shape)


# Error during conversion
coreml_layer = ct.converters.convert(
    scripted,
    source="pytorch",
    inputs=[
        ct.TensorType(name="image", shape=image.shape),
        ct.TensorType(name="grid", shape=grid.shape),
    ],
)

что вызывает следующую ошибку:

Traceback (most recent call last):
  File "/home/REDACTED/Downloads/sample.py", line 23, in <module>
    coreml_layer = ct.converters.convert(
  File "/home/REDACTED/.conda/envs/REDACTED/lib/python3.9/site-packages/coremltools/converters/_converters_entry.py", line 175, in convert
    mlmodel = mil_convert(
  File "/home/REDACTED/.conda/envs/REDACTED/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 128, in mil_convert
    proto = mil_convert_to_proto(, convert_from, convert_to,
  File "/home/REDACTED/.conda/envs/REDACTED/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 171, in mil_convert_to_proto
    prog = frontend_converter(, **kwargs)
  File "/home/REDACTED/.conda/envs/REDACTED/lib/python3.9/site-packages/coremltools/converters/mil/converter.py", line 85, in __call__
    return load(*args, **kwargs)
  File "/home/REDACTED/.conda/envs/REDACTED/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 81, in load
    raise e
  File "/home/REDACTED/.conda/envs/REDACTED/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 73, in load
    prog = converter.convert()
  File "/home/REDACTED/.conda/envs/REDACTED/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 227, in convert
    convert_nodes(self.context, self.graph)
  File "/home/REDACTED/.conda/envs/REDACTED/lib/python3.9/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 54, in convert_nodes
    raise RuntimeError(
RuntimeError: PyTorch convert function for op 'grid_sampler' not implemented.

Зависимости

Python (conda):

• coremltools==4.1
• torch==1.8.0

Вы также можете использовать _10 _ / _ 11_ билдов (по крайней мере, на день написания: 2021-03-20)

подсказки

Они были разделены на два возможных решения, которые я сейчас вижу:

Только PyTorch

Перепишите torch.nn.functional.grid_sample с нуля.

• Для этого потребуется придерживаться только операций PyTorch с тензорами, поскольку циклы (например, тройная вложенность) могут повесить преобразователь и будут слишком неэффективными.
• Вы не можете использовать __getitem__ для list или связанных типов - похоже, работает с torch.Tensor, но были проблемы с этим, поэтому вам следует помнить об этом, если вы получите RuntimeError: PyTorch convert function for op '__getitem__' not implemented

Плюсы:

• Нет необходимости в двух языках и использовании единой технологии

Минусы:

• Ограничено циклами и потребует использования векторизованных операций (большую часть времени).

Swift и CoreML

Зарегистрируйте пользовательский уровень, который отвечает за работу grid_sample. Подойдет реализация только CPU (хотя было бы здорово использовать Apple Metal для ускорения графического процессора).

Поскольку мне не нравится Swift, я собрал несколько ресурсов, которые могут вам помочь:

• https://coremltools.readme.io/docs/custom-operators - отправная точка , Только Python, довольно просто, только уровень регистрации для преобразования
• https://developer.apple.com/documentation/coreml/mlcustomlayer - API слой, который нужно было бы закодировать в Swift
• https://developer.apple.com/documentation/coreml/core_ml_api/creating_a_custom_layer - подробнее о вышеупомянутом (но не много)
• https://machinethink.net/blog/coreml-custom-layers/ - сообщение в блоге с примером и уровнем отправки на устройства (GPU, CPU). Требуется Swift (версия CPU), Metal (реализация GPU). Возможная реализация Metal может быть основана на PyTorch CUDA impl, CPU и Swift также могут быть связаны. 3 года, так что имейте в виду, что слой активации swish, кажется, является хорошей отправной точкой (и другие сообщения того же автора проливают свет на сам CoreML).
• https://github.com/hollance/CoreML-Custom-Layers - репо для над

Плюсы:

• Возможность использовать циклы и более тонкий контроль над алгоритмом
• Может быть проще, поскольку мы не ограничены операциями, которые CoreML в настоящее время может читать.

Минусы:

• Два языка
• Скудная документация