coremltools: как правильно использовать NeuralNetworkMultiArrayShapeRange?

У меня есть сеть PyTorch, и я хочу развернуть ее на устройствах iOS. Короче говоря, мне не удалось добавить гибкости к форме входного тензора в CoreML.

Сеть - это свертка, которая принимает изображение RGB (сохраненное как тензор) в качестве входных данных и возвращает изображение RGB того же размера. Используя PyTorch, я могу вводить изображения любого размера, который захочу, например тензор размера (1, 3, 300, 300) для изображения 300x300.

Чтобы преобразовать модель PyTorch в модель CoreML, я сначала конвертирую ее в модель ONNX, используя torch.onnx.export. Эта функция требует передать фиктивный ввод, чтобы он мог выполнить график. Итак, я использовал:

input = torch.rand(1, 3, 300, 300)

Я предполагаю, что модель ONNX принимает только изображения / тензоры размера (1, 3, 300, 300). Теперь я могу использовать функцию onnx_coreml.convert для преобразования модели ONNX в модель CoreML. Распечатав описание спецификации модели CoreML с помощью Python, я получаю что-то вроде:

input {
  name: "my_image"
  type {
    multiArrayType {
      shape: 1
      shape: 3
      shape: 300
      shape: 300
      dataType: FLOAT32
    }
  }
}
output {
  name: "my_output"
  type {
    multiArrayType {
      shape: 1
      shape: 3
      shape: 300
      shape: 300
      dataType: FLOAT32
    }
  }
}
metadata {
  userDefined {
    key: "coremltoolsVersion"
    value: "3.1"
  }
}

Входные данные модели должны иметь размер multiArrayType (1, 3, 300, 300). Копируя эту модель в XCode, я могу видеть во время проверки модели, что my_name указано в разделе «Входные данные», и ожидается, что это будет MultiArray (Float32 1 x 3 x 300 x 300). Пока все связно.

Моя проблема - добавить гибкости форме ввода. Я безуспешно пытался использовать coremltools. Это моя проблема. Вот мой код:

import coremltools
from coremltools.models.neural_network import flexible_shape_utils

spec = coremltools.utils.load_spec('my_model.mlmodel')

shape_range = flexible_shape_utils.NeuralNetworkMultiArrayShapeRange()
shape_range.add_channel_range((3,3))
shape_range.add_height_range((64, 5000))
shape_range.add_width_range((64, 5000))
flexible_shape_utils.update_multiarray_shape_range(spec, feature_name='my_image', shape_range=shape_range)
coremltools.models.utils.save_spec(spec, 'my_flexible_model.mlmodel')

Используя Python, я получаю следующее описание спецификации:

input {
  name: "my_image"
  type {
    multiArrayType {
      shape: 1
      shape: 3
      shape: 300
      shape: 300
      dataType: FLOAT32
      shapeRange {
        sizeRanges {
          lowerBound: 3
          upperBound: 3
        }
        sizeRanges {
          lowerBound: 64
          upperBound: 5000
        }
        sizeRanges {
          lowerBound: 64
          upperBound: 5000
        }
      }
    }
  }
}

Указано только 3 диапазона, что имеет смысл, поскольку я определил диапазон только для канала, высоты и ширины, но не для размера пакета. В XCode я получаю следующую ошибку при проверке гибкой модели CoreML:

There was a problem decoding this CoreML document
validator error: Description of multiarray feature 'my_image' has a default 4-d shape but a 3-d shape range

Я почти уверен, что он работал над другим проектом, когда я был на macOS X Mojave, но на данный момент я ни в чем не уверен.

Я использую:

• macOS X Catalina
• conda 4.7.12
• Python 3.7.5
• pytorch 1.3.1
• оннкс 1.6.0
• onnx-coreml 1.1
• coremltools 3.1

Спасибо за помощь