Core MLの仕組みを利用した.mlmodelの配信・暗号化の実情

# 概要

※この記事の結論は2021年4月時点での調査結果に基づいています。

AppleはWWDC 2020で、Core MLモデルの「配信」「暗号化」の仕組みを提供すると発表しました。

https://developer.apple.com/videos/play/wwdc2020/10152/


それぞれのコンセプト・仕様は上記動画を見ていただくとして(日本語字幕もあります)、この記事では

- 実際に調査・検証してわかった利用する際の問題点
- その問題点を解消するためにとり得る選択肢

を書きます。

結論を先にまとめると、残念ながら「配信」の仕組みは2021年4月時点で実用可能な状態ではありませんでした。
「暗号化」の方は実用可能でしたが、「配信」の問題に引きずられる形で使用に制限があります。


# モデルの配信

## 配信するまでの簡単な流れ

利用手順や諸々のスクショなど記載されている公式ドキュメントは以下です。

- [Creating and Deploying a Model Collection](https://developer.apple.com/documentation/coreml/core_ml_api/generating_a_model_archive)

Xcode 12上で```.mlmodel```を選択すると、Utilitiesのタブに「Create Model Archive」のボタンが存在しているのがわかると思います。それをタップすると```.mlarchive```というファイルが生成されます。このファイルを[Core ML Model Deployment](https://ml.developer.apple.com/)というダッシュボードにアップロードすると、モデルの配信は完了です。

モデルはアップロードしたダッシュボードと同一team idのアプリのみでDL可能です。また、配信する地域やデバイスなどある程度セグメントを分けてモデルを配信することもできます。

利用する場合は、アプリ内で[MLModelCollection](https://developer.apple.com/documentation/coreml/mlmodelcollection)のstatic methodを一つ叩くだけで自動DLされます。基本的にはいつどのようにDLされているのかはブラックボックスで、iOSが適切なタイミングを判断してくれるようです。

- [beginAccessing(identifier:completionHandler:)](https://developer.apple.com/documentation/coreml/mlmodelcollection/3649655-beginaccessing)

```swift
// identifierにはダッシュボードのModel Collectionの名前を入れる。
// ダッシュボード上で配信のマークがついているバージョンが自動でDLされる。
let progress = MLModelCollection.beginAccessing(identifier: "HogeMLModel") { result in
   switch result {
   case .success(modelCollection):
       let fileURL = modelCollection.entries["HogeModel"] // ★ アップロードしたときのモデルの名前を指定するとDL先のファイルパスが取得できる
   case .failure(error):
       break
   }
}
```

また、一度アクセスすれば次からはディスクに保存したmlmodelが利用されるようになるため通信は発生しません。

## まだ仕様通りに動作していないという課題

とても便利な仕組みなのですが、実際に使ってみるとうまくDLが走らない不可解な挙動が多くありました。例えば、ありえるシチュエーションとして次の３つのパターンを考えます。

1. アプリをDL直後の起動で初めて.mlmodelをDLする
1. .mlmodelを一度DLした後に、ダッシュボードで新しい.mlmodelへ差し替える
1. .mlmodelを一度DLした後に、アプリを一度削除してさらに再度アプリのDLを行って起動

1のケースでは正常にモデルのDLは成功で返ってきて、★のコードが実行されたときにファイルパスが取得できました。
しかし、2と3のケースではモデルのDLは成功で返ってくるもののファイルパスが取得できず、モデル自体の利用ができなくなりました。

この問題に関しては、Developer Forum上で同様の問題についていくつも指摘がされています。ただしApple側からの反応はなく、問題が解消したという開発者の報告もありません。

https://developer.apple.com/forums/thread/676873

https://developer.apple.com/forums/thread/676429

https://developer.apple.com/forums/thread/660333

https://developer.apple.com/forums/thread/653259

https://developer.apple.com/forums/thread/669969

https://developer.apple.com/forums/thread/675847

WWDCでの話を聞く限り、この機能は「基本的には予めアプリの中にあるモデルを使いつつ、場合によって補助的にアップデートできる仕組み」というコンセプトのようにも読み取れます。そこまで正確な制御を期待できない仕様なのかもしれません。

DLの仕組みはブラックボックスのため調査も難しく、現状では配信を前提とした設計では利用可能な状態になさそうです。

# モデルの暗号化

## 暗号化できた前提での鍵生成と復号化の簡単な流れ

公式ドキュメントは以下です。

- [Generating a Model Encryption Key](https://developer.apple.com/documentation/coreml/core_ml_api/generating_a_model_encryption_key)

「配信」の時と流れは同様です。先程のUtilitiesのタブに「Create Encryption Key」のボタンが存在しています。これをクリックすると```.mlmodelkey```というファイルが生成されます。この鍵を使ってモデルを暗号化します。

鍵を生成する際にteam idを選択するステップがあり、このとき選択したteam idのアプリでのみモデルの復号化ができます。復号化に使われる鍵は鍵生成時に自動的にAppleのサーバーにアップロードされるようです。

暗号化のステップを一旦飛ばします。その次の復号化では実装時に意識しなければならないことは殆どありません。一点だけ注意が必要なのはiOS14以降で使えるloadメソッドを経由させないといけない点です。


```swift
HogeHogeModel.load(contentOf: fileUrl) { result in
    switch result {
        case .success(let model):
            print(model)
        case .failure(let error):
            break
    }
}
```

- [load(contentsOf:configuration:completionHandler:)](https://developer.apple.com/documentation/coreml/mlmodel/3600218-load)


loadメソッドを呼んだときに、まだ鍵が端末に存在していなければ鍵のDLが走ります。そのため初回の利用時にはユーザーは必ず通信できる環境である必要があります。それ以降は端末に保存された鍵が利用されて、モデルの利用タイミングで都度復号化されます。

鍵の生成・復号化については以上になります。
もう一つモデルの暗号化が必要になるのですが、ここに関して現状では課題があります。

## 暗号化の選択可能な方法

暗号化のステップには現状3つの選択肢があります。

1. アプリをコンパイルする際に一緒にモデルも暗号化する
1. .mlarchive生成時に暗号化して、Core ML Model Deploymentで配信する
1. コマンドラインでモデルをコンパイルして、一緒に暗号化を行う

### 1. アプリをコンパイルする際に一緒にモデルも暗号化する

この方法が一番考えることが少なく簡単です。
Xcode内に```.mlmodel```を入れると、「Build Phases」の「Compile Sources」にmlmodelも追加されます。モデルもコンパイルされた上で利用されるためここに表示されるのですが、Compiler Flagとして鍵を指定可能です。次のオプションを追加します。

```
--encrypt $SRCROOT/HogeHogeModel.mlmodelkey
```

こうすると、暗号化済み```.mlmodelc```(コンパイルされた```.mlmodel```)が生成されて、アプリ内に組み込まれます。

あとは復号化の説明に記載したとおり```load```でモデルを読み込むと自動的に復号化されます。

### 2. .mlarchive生成時に暗号化して、Core ML Model Deploymentで配信する

この方法は「配信」と「暗号化」を組み合わせて使いたいケースを想定しています。やり方は簡単で、「配信」の時に説明した「Create Model Archive」のボタンをクリックすると鍵を指定するチェックボックスがあるためこれにチェックをつけるだけです。これで暗号化した状態で.mlarchiveが生成されます。 

ただし、この方法は使えません。

「配信」のところで説明したとおり、Core ML Model Deploymentは2021/04時点では実用的な状態でないためです。暗号化したモデルをアプリへ組み込まずオンデマンドで取得・復号化するためには他の方法で暗号化する必要があります。

### 3. コマンドラインでモデルをコンパイルして、一緒に暗号化を行う

残念ながら、1と2の選択肢以外の方法は正式なドキュメントとしての記載がありません。
しかし、mlmodelには実はコマンドでコンパイルする方法があります。

```swift
xcrun coremlcompiler compile HogeHogeModel.mlmodel . // 対象.mlmodelとコンパイルしたファイルの保存先を指定する
```

Appleがこのコマンドの使い方を公式に言及しているのはDeveloper Forumになります。

https://developer.apple.com/forums/thread/84401

https://developer.apple.com/forums/thread/127769

このコマンドはオプションも受け取ることができます。どこにも記載がないものの、先程のCompiler Flagも受け取ることができます。
したがって、コマンドとCompiler Flagを組み合わせて、

```
xcrun coremlcompiler compile HogeHogeModel.mlmodel . --encrypt $SRCROOT/HogeHogeModel.mlmodelkey
```

と書くと、暗号化済みの```.mlmodelc```が生成されます。
実際にそのコンパイル済みモデルを独自の配信サーバーに置いて```URLSession```でDLし、```load```メソッドで読み込んだところ、無事復号化されモデルを使うことができました。
また、鍵と異なるteam idでビルドしたアプリでも試したところ復号化に失敗していました。
モデルはちゃんと特定のteam idのアプリでのみ利用可能になっているようです。

## 自前の暗号化ツールは利用可能か？

そもそも```.mlmodelkey```を利用せず、自分で鍵を用意してそれでモデルを暗号化するのはどうでしょうか。例えばAppleが提供するフレームワーク、CryptoKitでは「AES-GCM」方式の暗号化をサポートしています。 手元で共通鍵を生成してモデルを暗号化し、同じ鍵をアプリのバイナリへハードコードした上でCryptoKitで復号化することが可能です。

暗号化をSwiftで行う場合はこのようになります。

```swift
import CryptoKit

let destinationURL = // 暗号化したデータの保存先
let rawData = // .mlmodelをzip化したものなど暗号化したいデータ

let symmetricKey = SymmetricKey(size: .bits256) // 鍵の生成
let sealedBox = try AES.GCM.seal(data, using: symmetricKey)

// 鍵をBase64文字列にする
let keyString = symmetricKey.withUnsafeBytes { pointer in
    Data(Array(pointer)).base64EncodedString()
}

let encryptoData = sealedBox.combined!

encryptoData.write(to: destinationURL) // モデルを暗号化してファイルに保存

print(keyString) // 鍵をBase64形式で保存する
```

アプリ内での復号化のコードはこのようになります。

```swift
import CryptoKit

let keyString = // 保存した文字列をソルトを使うなど何らかのルールを組み合わせてハードコードする
let encryptoData = // DLしてきた暗号化されたモデルデータ

let keyData = Data(base64Encoded: keyString)!
let symmetricKey = SymmetricKey(data: keyData)
let sealedBox = try AES.GCM.SealedBox(combined: encryptoData)

let decryptoData = try AES.CGM.open(sealedBox, using: symmetricKey) // 共通鍵を使って復号化の実行
```

上記コードで通信経路上のデータそのものと保存したファイルに関しては暗号化された状態を作ることができます。
しかしこの方法には問題があります。MLModelのロードには.mlmodel(もしくは.mlmodelc)のファイルパスを指定する方法しかなく、メモリ上のデータをロードする方法がありません。そのため、復号化した素のモデルを一旦ファイルに書き込まなくてなりません。

MLModelのAPIリファレンスを見る限りそれを回避する方法は今の所ありませんでした。

そのためこの方法は避けたほうが良さそうです。

## その他暗号化を使わないで似たような効果を得る方法
機械学習モデルの一部のレイヤーをアプリ内の実装に移すことで、.mlmodelを不完全版にしてしまうという方法も考えられます。
その場合は.mlmodel単体では機能が利用できないという状態にできます。


# 結論と現状のベストプラクティス

以上をまとめると、現時点(2021年4月)では.mlmodelの「配信」「暗号化」を実現するためには次のやり方が最も適切です。

1. Core ML Model Deploymentは使わず、自前の配信サーバーを用意する
1. xcrunで.mlmodelのコンパイルと暗号化を予めやっておく
1. 配信サーバー経由で暗号化済み.mlmodelcを配信し、URLSessionによるDLを自分で実装する
1. 生成した鍵と同じteam idでアプリをビルドする

この方法で最新の.mlmodelをDLして、自分のアプリでのみ復号化できます。