coreML学習済みネットワークをplaygroundで試してみた

coreMLを使えば簡単にiOSで学習済みネットワークが実行できる、という話を聞いたので勉強し始めてみました。

最終的にはiOSにデプロイするつもりなんですが、今回はcoreMLの挙動を検証するため、デバッグしやすいplaygroundで環境を組んでみました。

この記事では学習済みモデルをplaygroundからロードする方法と、coreMLで学習済みモデルを動かすための最小実装をメモとして残しておきます。

今回作成したplaygroundはこちらのリポジトリに置いています。

※ 容量と著作権の関係で学習済みモデルは置いていないので、READMEに従って各自ダウンロードしてください

0. 開発環境

Xcode 10.3
Swift 5.0.1

1. 学習済みモデルと画像を用意する

coreMLでニューラルネットを動かすときは、TensorFlowやPyTorchなどのライブラリでネットワークを学習させてからcoreML用のファイルに変換するのが一般的な方法だと思いますが、今回は学習を回すのが面倒だったのでとりあえずAppleが配布している学習済みモデルを動かしてみました。

配布元：https://developer.apple.com/jp/machine-learning/models/

*.mlmodelという拡張子のファイルに諸々の情報が入っているらしいので、これをそのまま利用します。

動作検証用の画像はGoogleが配布しているOpen Images Dataset v5から拝借しました。

配布元：https://storage.googleapis.com/openimages/web/index.html

2. playgroundを作成する

普段どおりにplaygroundを作成します。 Xcodeを起動してGet started with a playground、もしくは起動中のXcodeのメニューバーからFile > New > Playgroundです。

次に作成したplaygroundの中にResourcesフォルダを作成します。 Finderから見るとplaygroundは1つのファイルであるかのように見えますが、実際にはパッケージとして扱われているフォルダなので、中にファイルを追加することができます。

playgroundの上で右クリック > Show Package Contentsで中のファイル構造が表示されます。

f:id:kamino-dev:20190822152636p:plain:w400

ここにResourcesというフォルダを作成します。このフォルダの中にファイルを置いておけば、playgroundのSwiftコードからそれらのファイルにアクセスできるようになります。

今回はResourcesの中に学習済みモデル*.mlmodelと動作検証用の画像を置いておきます。

f:id:kamino-dev:20190822153055p:plain:w450

3. coreMLを動かすための最小実装

あとはplaygroundでcoreMLを動かすためのコードを書くだけです。

import UIKit
import Vision

// 学習済みモデルを読み込む
// ビルド時にmlmodelをコンパイルしたmlmodelcが作られるのでそれをロードする
let modelURL = Bundle.main.url(forResource: "YOLOv3", withExtension: "mlmodelc")!
guard let model = try? VNCoreMLModel(for: MLModel(contentsOf: modelURL)) else {
    fatalError()
}

// ネットワーク実行のリクエストを発行するための関数
func createRequest(model: VNCoreMLModel) -> VNCoreMLRequest {
    return VNCoreMLRequest(model: model, completionHandler: { (request, error) in
        // ネットワーク実行完了後の処理を定義する
        DispatchQueue.main.async(execute: {
            // ネットワークによってrequest.resultsの型は違う
            // YOLOv3の場合は[VNRecognizedObjectObservation]型
            guard let results = request.results as? [VNRecognizedObjectObservation] else {
                fatalError("Error results")
            }
            for result in results {
                print("\(result.confidence) : \(result.boundingBox)")
                let len = result.labels.count > 5 ? 5 : result.labels.count
                for i in 0..<len{
                    print("\(result.labels[i].identifier), ", terminator: "")
                }
                print()
            }
            print()
        })
    })
}

// 検証用画像の読み込み
let img1 = UIImage(named: "sample1.jpg")!
let img2 = UIImage(named: "sample2.jpg")!

// 1枚目の画像の分析をリクエスト
let handler1 = VNImageRequestHandler(cgImage: img1.cgImage!, options: [:])
let request1 = createRequest(model: model)
try? handler1.perform([request1])

// 2枚目の画像の分析をリクエスト
let handler2 = VNImageRequestHandler(cgImage: img2.cgImage!, options: [:])
let request2 = createRequest(model: model)
try? handler2.perform([request2])