DeepLabをCore MLフォーマットに変換する
環境
- OS: Ubuntu 20.04
- pyenv 2.1.0
fronzen graphを作る
Pythonの環境を整えます。coremltoolsは5系だとエラーが出ます。tensorflowは1系を使うのですが、Pythonは3.7を使いました。
$ mkdir deeplab-coreml
$ cd deeplab-coreml
$ pyenv local 3.7.10
$ python -m venv .venv
$ poetry init
$ poetry add tensorflow=1.15.5
$ poetry add tf-slim=1.1.0
$ poetry add coremltools=4.1
$ poetry add jupyter matplotlib
frozen graphの置き場を作っておきます。
$ mkdir -p weights/deeplabv3
deeplabを取得とPYTHONPATH
の設定をします。
$ git clone https://github.com/tensorflow/models.git
$ cd models
$ git checkout d7eabefa595f325611798a6d88d0f415986b8152
$ cd research
$ export PYTHONPATH=`pwd`:`pwd`/slim
deeplabの学習済みモデルを取得し、frozen graphを生成します。設定により多少変わりますが、バックボーンがmobilenetv2の場合とxception_65の場合を紹介します。
$ cd deeplab
$ mkdir checkpoint
$ curl http://download.tensorflow.org/models/deeplabv3_mnv2_pascal_train_aug_2018_01_29.tar.gz --output ./checkpoint/deeplabv3_mnv2_pascal_train_aug.tar.gz
$ tar xzvf checkpoint/deeplabv3_mnv2_pascal_train_aug.tar.gz -C checkpoint/
$ python export_model.py --model_variant="mobilenet_v2" \
--checkpoint_path=./checkpoint/deeplabv3_mnv2_pascal_train_aug/model.ckpt-30000\
--export_path=../../../weights/deeplabv3/frozen_graph.pb
$ cd deeplab
$ mkdir checkpoint
$ curl http://download.tensorflow.org/models/deeplabv3_pascal_train_aug_2018_01_04.tar.gz --output checkpoint/deeplabv3_pascal_train_aug_2018_01_04.tar.gz
$ python export_model.py \
--model_variant="xception_65" \
--atrous_rates=6 \
--atrous_rates=12 \
--atrous_rates=18 \
--output_stride=16 \
--decoder_output_stride=4 \
--num_classes=21 \
--checkpoint_path=./checkpoint/deeplabv3_pascal_train_aug/model.ckpt \
--export_path=../../../weights/deeplabv3/frozen_graph.pb
どのようなオプションを指定すべきかはTensorFlow DeepLab Model Zooの表とexport_model.py
のオプションのコメントを見ると大体分かると思います。
coremltoolsで変換
ここからcoremltools
を使ってCore ML形式にfrozen graphを変換します。以下に登場するPythonのコード片はプロジェクトルートにjupyterノートブックを作ったと思ってください。
import coremltools as ct
import json
json
モジュールは、XCodeから見たとき分かりやすいように設定するために使います。
model = ct.convert('./weights/deeplabv3/frozen_graph.pb', inputs=[ct.ImageType(shape=(1, 513, 513, 3))])
WARNING:root:Output var SemanticPredictions of type i32 in function main is cast to type fp32
という警告が出ますが、後でint32
に修正するので大丈夫です。画像を入力にするモデルの変換は大抵ct.ImageType
を使います。shape
のはじめの1はバッチ数なので、ここはいつでも1だと思って良いと思います。
spec = model.get_spec()
builder = ct.models.neural_network.NeuralNetworkBuilder(spec=spec)
変換には特にbuilder
は要らないですが、入力や出力の状態を確認するのに使います。
builder.inspect_input_features()
[Id: 0] Name: ImageTensor
Type: imageType {
width: 513
height: 513
colorSpace: RGB
}
RGBの画像になってますね。でも名前が気に入らないので、変更します。
ct.utils.rename_feature(spec, "ImageTensor", "image")
builder.inspect_input_features()
[Id: 0] Name: image
Type: imageType {
width: 513
height: 513
colorSpace: RGB
}
inputのnameがimage
に変わりました。次に出力を確認します。
builder.inspect_output_features()
[Id: 1] Name: SemanticProbabilities
Type: multiArrayType {
dataType: FLOAT32
}
[Id: 0] Name: SemanticPredictions
Type: multiArrayType {
dataType: FLOAT32
}
出力が2種類ありますね。SemanticProbabilities
は実際には役立つこともあるかもしれませんが、XCodeのプレビューでこれを残していると頻繁にXCodeが落ちます。真剣にアプリを作る段階で、扱いを再検討することにして、これを削除し、SemanticPredictions
はint32の513x513のMultiArrayにし、名前をsemanticPredictions
に変更します。
spec.description.output.pop()
output = spec.description.output[0]
ct.utils.rename_feature(spec, "SemanticPredictions", "semanticPredictions")
output.type.multiArrayType.dataType = ct.proto.FeatureTypes_pb2.ArrayFeatureType.INT32
output.type.multiArrayType.shape.append(513)
output.type.multiArrayType.shape.append(513)
builder.inspect_output_features()
[Id: 0] Name: semanticPredictions
Type: multiArrayType {
shape: 513
shape: 513
dataType: INT32
}
出力についても修正できたので、メタ情報を充実させます。
spec.description.input[0].shortDescription = "Input image to be segmented"
spec.description.output[0].shortDescription = "Array of integers of the same size as the input image, where each value represents the class of the corresponding pixel."
spec.description.metadata.versionString = "--"
spec.description.metadata.shortDescription = "DeepLab is a state-of-art deep learning model for semantic image segmentation, where the goal is to assign semantic labels (e.g., person, dog, cat and so on) to every pixel in the input image."
spec.description.metadata.author = "Original Paper: Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, Liang-Chieh Chen. DeepLabV3+: Liang-Chieh Chen and Yukun Zhu and George Papandreou and Florian Schroff and Hartwig Adam."
spec.description.metadata.license = "Please see https://github.com/tensorflow/tensorflow for license information, and https://github.com/tensorflow/models/tree/master/research/deeplab for the original source of the model."
最後にプレビューが出るように設定します。
model = ct.models.MLModel(spec)
labels_json = {"labels": [
'background', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus',
'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike',
'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tv'
]}
model.user_defined_metadata["com.apple.coreml.model.preview.type"] = "imageSegmenter"
model.user_defined_metadata['com.apple.coreml.model.preview.params'] = json.dumps(labels_json)
保存します。
model.save("MyDeepLabV3.mlmodel")
最後にMacに移動させてXCodeで動作確認してみます。
写真は家族-ビーチ-人-海洋-6398107を使用しました。
Discussion