torchvisionの画像分類モデルをCoreMLに変換する

事前トレーニング済みモデルをCoreMLの画像分類モデルに変換してiOSで使う方法です

画像分類モデルを変換したい

そのままTorchVisionの画像分類モデルを変換すると、出力は１０００の数値になります。
これを、信頼度のパーセンテージに直し、画像分類モデルとしてプレビューできるようにしたい。

クラス名とソフトマックス層を追加する

クラス名とソフトマックス層を追加すれば、１０００クラス内のクラスごとの信頼度のパーセンテージが得られ、xcodeのプレビューも使えます。

方法

torchvisionからモデルをダウンロード。

import torchvision.models as models

regnet_y_400mf = models.regnet_y_400mf(pretrained=True)

%表示のためにソフトマックス層を追加したラッパーモデルクラスを作る。

import torch.nn as nn 

class RegNetClassificationModel(nn.Module):
    def __init__(self):
        super(RegNetClassificationModel, self).__init__()
        self.layers = nn.Sequential(
            regnet_y_400mf,
            nn.Softmax(dim=1)
        )
    def forward(self, x):
        return self.layers(x)

model = RegNetClassificationModel().eval()

ImageNetのクラスラベルをダウンロードして、リストにし、リストからCoreMLのClassifierConfigを作る。

import urllib
import coremltools as ct

label_url = 'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt'
class_labels = urllib.request.urlopen(label_url).read().splitlines()
class_labels = class_labels[1:] # remove the first class which is background
assert len(class_labels) == 1000

for i, label in enumerate(class_labels):
  if isinstance(label, bytes):
    class_labels[i] = label.decode("utf8")
classifier_config = ct.ClassifierConfig(class_labels)

変換。

ex_input=torch.randn((1,3,224,224))
traced = torch.jit.trace(model,ex_input)

mlmodel = ct.convert(traced,inputs=[ct.ImageType(shape=ex_input.shape,bias=[-0.485/0.229,-0.456/0.224,-0.406/0.225],scale=1.0/255.0/0.226)],classifier_config=classifier_config)
mlmodel.save("regnet_y_400mf.mlmodel")

これでXcodeでプレビューできます。

iOSでの使い方

let model = try! regnet_y_400mf(configuration: MLModelConfiguration()).model
let vnModel = try! VNCoreMLModel(for: model)
let request = VNCoreMLRequest(model: vnModel)
let handler = VNImageRequestHandler(ciImage: ci1, options: [:])
try! handler.perform([request])
let result = request.results?.first as! VNClassificationObservation
print(result.identifier,result.confidence)

king penguin 0.61...

🐣

---

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

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

Twitter
Medium