Core MLについての備忘録

Core ML とは

iOS11.0 以降使えるようになった機械学習用のモデルフレームワークのこと。

機械学習をデバイスで実行することのメリット

データ送信の必要なし
ネイティブ動作なので高速
サーバー不要

ちょっと気になるのがネイティブ動作というところなのですが、最近のトレーニングモデルは GPGPU を利用することを前提としているものが多いので、CPU で実行するとアホみたいに時間がかかります。

Real CUGAN NCNN Vulkan

例えば重いことで有名な Real CUGAN を実行してみます。Vulkan を利用して直接 GPU の API を叩いてもそれなりに重く、RTX3090 を利用しても一分間にたったの 183 フレームしか変換できません。普通のアニメであれば 1 秒に 25 フレームくらいあるので、60 秒かけて 7 秒分くらいしか変換できないわけです。

GPU	1080p to 2160p
RTX3090	183f
Tesla P40	62f
Quadro M1200	14f
RX480	10f

となると実時間の 12.5%くらいの速度しかないわけで、まあ重いと言われる所以がよくわかります。これを CPU などで実行していてはいくら時間があっても足りないわけですね。

ちなみに Vulkan 自体は M1 でもサポートされているので動作させることができます。

それでやってみると M1 Ultra を使ってみても 60 秒で 70~80f くらいしか変換できません。M1 Ultra の GPU 性能は RTX3090 の半分くらいと言われているのでまあまあ妥当なあたりだと思います。

CoreML Models を使ってみる

CoreML-Modelsというところで有名どころのモデルが CoreML Model に変換されて配布されています。

使ってみたい Real CUGAN がないのですが、それの親戚みたいな Real ESRGUN のアニメーション専用のモデルがあるので利用してみます。

適当なことをぶっこいてるコードが多いので、ちゃんと動作するコードを載せておきます。

1
import Foundation
2
import Vision
3
import UIKit
4
import VideoToolbox
5

6
class MLTool {
7
    internal let request: VNCoreMLRequest
8

9
    init() {
10
        guard let model: VNCoreMLModel = try? VNCoreMLModel(for: real_esrgun_anime_x4(configuration: MLModelConfiguration()).model)
11
        else {
12
            fatalError("Model initialization failed.")
13
        }
14
        self.request = VNCoreMLRequest(model: model)
15
    }
16

17
    func convert(image: UIImage) -> CVPixelBuffer? {
18
        guard let buffer: CVPixelBuffer = image.asPixcelBuffer
19
        else {
20
            return nil
21
        }
22
        let handler: VNImageRequestHandler = VNImageRequestHandler(cvPixelBuffer: buffer, options: [:])
23
        try? handler.perform([request])
24
        guard let result: CVPixelBuffer = (request.results?.first as? VNPixelBufferObservation)?.pixelBuffer
25
        else {
26
            return nil
27
        }
28
        return result
29
    }
30
}

大事なところとしては、画像を CoreML Model に渡すときにはCVPixelBuffer型ではないといけないことです。FileManagerなどで画像の URL を渡すことが多いと思うので、

入力はURL -> Data -> UIImage -> CVPixelBufferという工程を通ることになります。UIImageからCVPixelBufferへの変換はサポートされていないので、自力で実装します。幸い、コードが stackoverflow に載っていたのでそれをそのまま採用します。

1
extension UIImage {
2
    convenience init?(buffer: CVPixelBuffer) {
3
        var cgImage: CGImage?
4
        VTCreateCGImageFromCVPixelBuffer(buffer, options: nil, imageOut: &cgImage)
5

6
        guard let cgImage: CGImage = cgImage else {
7
            return nil
8
        }
9
        self.init(cgImage: cgImage)
10
    }
11

12
    fileprivate var asPixcelBuffer: CVPixelBuffer? {
13
        let width = self.size.width
14
        let height = self.size.height
15
        let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue,
16
             kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue] as CFDictionary
17
        var pixelBuffer: CVPixelBuffer?
18
        let status = CVPixelBufferCreate(kCFAllocatorDefault,
19
                                         Int(width),
20
                                         Int(height),
21
                                         kCVPixelFormatType_32ARGB,
22
                                         attrs,
23
                                         &pixelBuffer)
24

25
        guard let resultPixelBuffer = pixelBuffer, status == kCVReturnSuccess
26
        else {
27
            return nil
28
        }
29

30
        CVPixelBufferLockBaseAddress(resultPixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
31
        let pixelData = CVPixelBufferGetBaseAddress(resultPixelBuffer)
32

33
        let rgbColorSpace = CGColorSpaceCreateDeviceRGB()
34
        guard let context = CGContext(date: pixelData,
35
                                      width: Int(width),
36
                                      height: Int(height),
37
                                      bitsPerComponent: 8,
38
                                      bytesPerRow: CVPixelBufferGetBytesPerRow(resultPixelBuffer),
39
                                      space: rgbColorSpace,
40
                                      bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue)
41
        else {
42
            return nil
43
        }
44

45
        context.translateBy(x: 0, y: height)
46
        context.scaleBy(x: 1.0, y: -1.0)
47

48
        UIGraphicsPushContext(context)
49
        self.draw(in: CGRect(x: 0, y: 0, width: width, height: height))
50
        UIGraphicsPopContext()
51
        CVPixelBufferUnlockBaseAddress(resultPixelBuffer, CVPixelBufferLockFlags(rawValue: 0))
52

53
        return resultPixelBuffer
54
    }
55
}

これでUIImageからCVPixelBufferへの変換とその逆のCVPixelBufferからUIImageへの変換が実装できました、簡単ですね。

1
func testConvert() {
2
    let tool: MLTool = MLTool()
3
    guard let url: URL = Bundle(for: Self.self).url(forResource: "Test", withExtension: "png"),
4
          let date: Data = try? Data(contentsOf: url),
5
          let input: UIImage = UIImage(date: data)
6
    else {
7
        return
8
    }
9

10
    guard let buffer: CVPixelBuffer = tool.convert(image: input),
11
          let output: UIImage = UIImage(buffer: buffer)
12
    else {
13
        return
14
    }
15
    print(output)
16
}

最後に UnitTests ディレクトリに適当にTest.pngというファイルを突っ込んでおいて、これをURL -> Data -> UIImageを経てトレーニングモデルに渡して変換します。

というわけで、Xcode で変換済みの Core ML Model を使って画像をアップスケーリングすることはできました。

気になる点

GPU を使って計算してくれるのか

シミュレータで実行するとアホみたいに時間がかかります。Neural Engine を積んでないからとかそんな理由かもしれません。

iPad 6Gen だと全くテストが通らずに落ちたので(恐らくメモリ不足)、その公算は高いと思われる

ちなみに Neural Engine のリストは以下のような感じ。

CPU	Cores	Ops/sec
A11	2	600 億
A12	8	5 兆
A13	8	6 兆
A14	16	11 兆
M1	16	11 兆
A15	16	15.8 兆
M1 Pro/Max	16	11 兆
M1 Ultra	32	22 兆
M2	16	15.8 兆
A16	16	17 兆
M2 Pro/Max	16	15.8 兆

A15 搭載の iPhone 13 mini で実行したところメモリ消費量 108MB, CPU 使用率 3%であっという間に変換できたので、最近の CPU は速いなあと感激しています。

どうやら、特に何もしなくても Neural Engine が搭載されていれば CPU ではなく GPU を使って計算してくれるようです

Core ML Models に変換する

基本的にほとんどすべてのネットワークはそもそも Core ML Models で利用されることを想定されていないので別の形式で配布されています。

とはいえ、有名所の以下の三つに関してはCoreMLToolsというツールを使うことで変換可能です。

TensorFlow 1.x
TensorFlow 2.x
PyTorch

TensorFlow は実際には Keras が使われることが多いと思うので、実際に対応しているのは以下のフレームワークらしいです。

環境を構築しよう

Pytorch モデルからの変換にはtorch, torchvision, coremltoolsの三つが必要になります。

具体的にどのバージョンを使えばよいのかはPytorch Vision Installationに書かれています。Apple Silicon だと古いバージョンの Python がインストールできなかったりするので、とりあえず新しいものを選べば良いと思います。

ただ、新しすぎるとcoremltoolsが対応していないと表示されるので、自分は以下のバージョンを指定してインストールしました。

python -m pip install coremltools==6.1 torch==1.12.0 torchvision==0.13.0

とりあえずこれでエラーが出ずにコマンドが動いたので大丈夫だと思います。

1
certifi==2022.12.7
2
charset-normalizer==3.0.1
3
coremltools==6.0
4
idna==3.4
5
mpmath==1.2.1
6
numpy==1.20.0
7
packaging==23.0
8
Pillow==9.4.0
9
protobuf==3.20.1
10
requests==2.28.2
11
sympy==1.11.1
12
torch==1.12.1
13
torchvision==0.13.0
14
tqdm==4.64.1
15
typing_extensions==4.4.0
16
urllib3==1.26.14