CoreML Modelに柔軟な入力形状を設定する

#色々な入力形状に対応できるCoreMLモデルに変換する方法です

#固定ではなくさまざまなサイズの入力を使いたい

デフォルト変換では、CoreMLモデルは入力サイズを固定します。
しかし、実際のユースケースでは、色々なサイズの入力を使いたい場合があります。

#柔軟な入力を使う方法

CoreMLToolsは、柔軟な入力の設定が可能です。

###変換時に設定する方法

import coremltools as ct

input_shape = ct.Shape(shape=(1,3, ct.RangeDim(512, 1024), ct.RangeDim(512, 1024)))
model_input = ct.ImageType(shape=input_shape)

mlmodel = ct.convert(model, inputs=[model_input])

これで、入力画像の縦と横で512~1024の範囲を受け入れられるようになりました。
ct.RangeDim()とすると、無制限の入力形状を受け入れるようになります。

範囲指定ではなく、いくつかのデフォルト値を設定して、メモリを事前に割り当てることで、範囲指定よりも高速になります。

import coremltools as ct

shapes = [(1,3, 512, 512),(1,3, 768, 768),(1,3, 1024, 1024)]
input_shape = ct.EnumeratedShapes(shapes=shapes, 
                                  default=(1,3, 512, 512))
model_input = ct.TensorType(shape=input_shape)

mlmodel = ct.convert(model, inputs=[model_input])

これで、３つの選択肢で入力形状を指定できます。

###変換した後で設定する方法

from coremltools.models.neural_network import flexible_shape_utils
spec = coremltools.utils.load_spec('mymodel.mlmodel')

input_name = spec.description.input[0].name

img_size_ranges = flexible_shape_utils.NeuralNetworkImageSizeRange()
img_size_ranges.add_height_range((256, 4096))
img_size_ranges.add_width_range((256, 4096))
flexible_shape_utils.update_image_size_range(spec, feature_name=input_name,
size_range=img_size_ranges)

from coremltools.models.utils import save_spec
save_spec(spec, 'updatedMyModel.mlmodel')

🐣

---

フリーランスエンジニアです。
お仕事のご相談こちらまで
[e-mail]

Core MLやARKitを使ったアプリを作っています。
機械学習／AR関連の情報を発信しています。

Twitter
Medium