dlshogiモデルをMacのCore MLで動かす

Macで高速にdlshogiモデルを動作させる話題の続きです。前回は、ONNX Runtimeを介してCore MLを呼び出してみましたが、ONNX Runtime単体(CPUのみの利用)と比べ高速化しなかったという結果でした。今回は、ONNX Runtimeを使わず直接Core MLを呼び出してみます。

※2022-06-02修正: メモリリークを回避

select766.hatenablog.com

ふかうら王の改造ではなく、dlshogiモデルを動作させるだけの最小限のコードを実装しました。

Core MLを直接呼び出すコードはこちら。以下のコード解説はこちらのリポジトリの内容を引用しています。

github.com

比較のため、ONNX Runtimeを介してCore MLを呼び出すコードも実装してあります。

github.com

Objective-CからCore MLを呼び出す

Core MLのAPIはSwiftかObjective-Cで呼び出す必要があります。型定義などがObjective-Cのヘッダでなされているため、C++から直接呼ぶことは難しいです。

今回のためにObjective-Cに入門して、まずはObjective-C単体でCore MLを呼び出すコードを実装しました。

dlshogiモデルは以前紹介した方法で、Core ML特有のDNNモデル形式であるmlmodelに変換しています。

select766.hatenablog.com

下準備として、mlmodelをコンパイルし、mlmodelcを作成する必要があります。以下のコマンドによって行うことができます。

/Applications/Xcode.app/Contents/Developer/usr/bin/coremlc compile DlShogiResnet15x224SwishBatch.mlmodel .

コンパイル結果は、 DlShogiResnet15x224SwishBatch.mlmodelc として出力されます。これはファイルではなくディレクトリで、中にメタデータや学習済みの重みなど複数のファイルが含まれていました。コンパイルオプションにターゲットとしてmacos, iosなどがあるため、プラットフォーム依存の内容と思われます。コマンドラインツールではなく、Core MLのAPIを用いてのコンパイルも可能なので、ユーザがモデルを差し替えられるような機構にするならそちらを使うほうがポータビリティが上がると思われます。

次に、DNNモデル固有の入出力インターフェースのコードを生成します。

/Applications/Xcode.app/Contents/Developer/usr/bin/coremlc generate DlShogiResnet15x224SwishBatch.mlmodel .

生成結果は、DlShogiResnet15x224SwishBatch.m (Objective-Cのソース)、DlShogiResnet15x224SwishBatch.h (Objective-Cのヘッダ)として出力されます。薄いラッパーに過ぎないので、モデルの差し替えを可能とするならこれも使用しないほうが良いかもしれません。

これらを用いて、モデルを実行するコードを実装します。

https://github.com/select766/mac-coreml-objc-cpp/blob/main/mainobjc.m

#import <CoreML/CoreML.h>
#import "DlShogiResnet15x224SwishBatch.h"

Core MLのヘッダをロードします。

MLModelConfiguration* config = [MLModelConfiguration new];
config.computeUnits = MLComputeUnitsAll;

モデルのロードに関する設定を行います。どのcompute unit (Neural Engine / GPU / CPU)で実行されるかを指定できます。

NSError *error = nil;
DlShogiResnet15x224SwishBatch* model = [[DlShogiResnet15x224SwishBatch alloc] initWithConfiguration:config error:&error];

モデルを読み込みます。カレントディレクトリの DlShogiResnet15x224SwishBatch.mlmodelc が読まれます。

float *input_data, *output_policy_expected, *output_value_expected;
read_test_case(&input_data, &output_policy_expected, &output_value_expected);

テスト用の入力データおよび、PyTorchで生成した正しい出力をfloat配列に読み込みます。

MLMultiArray *model_input = [[MLMultiArray alloc] initWithDataPointer:input_data shape:@[[NSNumber numberWithInt:batch_size], @119, @9, @9] dataType:MLMultiArrayDataTypeFloat32 strides:@[@(119*9*9), @(9*9), @9, @1] deallocator:NULL error:NULL];

入力テンソルを表すMLMultiArrayを生成します。入力データへのポインタ、テンソルの形状を指定します。

DlShogiResnet15x224SwishBatchOutput *model_output = [model predictionFromInput:model_input error:&error];

モデルを実行します。実行は同期的に行われます。

check_result(output_policy_expected, (float*)output_policy.dataPointer, batch_size * policy_size);

結果に対して処理を行います。結果は、Core ML内部で動的に確保されたバッファ内に保存されており、floatポインタとして取り出すことが可能です。

gcc -o mainobjc -fobjc-arc -O3 mainobjc.m DlShogiResnet15x224SwishBatch.m -framework Foundation -framework CoreML

コンパイルはgccコマンドでできます。-frameworkという特有の指定があることに注意してください。-fobjc-arcをつけることで、明示的なメモリ開放なくObjective-Cのオブジェクトに対して参照カウント式のGC (ARC: Automatic Reference Counting) が動きます。

さらに、モデルを差し替えることを考慮し、mlmodelを動的にコンパイルするバージョンも作成しました。詳細は mainobjcdynamic.m を参照ください。

C++からCore MLを呼び出す

ふかうら王はC++で実装されており、DNN評価を行うクラスをC++で実装する必要があります。C++で定義されたクラスの中で、Objective-CのAPIを呼び出すことができますのでこれを用いてCore MLをラップしました。

ラッパークラスのヘッダ nnwrapper.hpp を作成します。これは、Objective-Cの要素がない純粋なC++です。

// pure C++ code

class NNWrapper {
public:
    NNWrapper(const char* computeUnits);
    bool run(int batch_size, float* input, float* output_policy, float* output_value);
    static const int input_size = 119 * 9 * 9;
    static const int policy_size = 2187;
    static const int value_size = 1;
private:
    void* model;
};

メンバ変数void* modelに、Objective-Cのオブジェクトを格納するようになっています。

ラッパークラスの実装 nnwrapper.mm を作成します。これは、Objective-CとC++が混在したコード(Objective-C++と呼ばれる)です。

// Objective-C, C++混在コード
#import <Foundation/Foundation.h>
#import <CoreML/CoreML.h>
#import <stdlib.h>
#import "nnwrapper.hpp"
#import "DlShogiResnet15x224SwishBatch.h"

NNWrapper::NNWrapper(const char* computeUnits) {
    MLModelConfiguration* config = [MLModelConfiguration new];
    // 使用デバイス
    // MLComputeUnitsCPUOnly = 0,
    // MLComputeUnitsCPUAndGPU = 1,
    // MLComputeUnitsAll = 2
    MLComputeUnits cu;
    if (strcmp(computeUnits, "cpuonly") == 0) {
        cu = MLComputeUnitsCPUOnly;
    } else if (strcmp(computeUnits, "cpuandgpu") == 0) {
        cu = MLComputeUnitsCPUAndGPU;
    } else if (strcmp(computeUnits, "all") == 0) {
        cu = MLComputeUnitsAll;
    } else {
        fprintf(stderr, "Unknown computeUnits %s\n", computeUnits);
        exit(1);
    }
    config.computeUnits = cu;
    NSError *error = nil;
    DlShogiResnet15x224SwishBatch* model = [[DlShogiResnet15x224SwishBatch alloc] initWithConfiguration:config error:&error];
    NSLog(@"%@", model);
    NSLog(@"%@", error);

    if (!model) {
        NSLog(@"Failed to load model, %@", error);
        exit(1);
    }
    // 所有権をARCからプログラマに移す
    this->model = (void*)CFBridgingRetain(model);
}

bool NNWrapper::run(int batch_size, float* input, float* output_policy, float* output_value) {
    // 所有権を移さない(プログラマのまま)
    DlShogiResnet15x224SwishBatch* model = (__bridge DlShogiResnet15x224SwishBatch*)(this->model);

    MLMultiArray *model_input = [[MLMultiArray alloc] initWithDataPointer:input shape:@[[NSNumber numberWithInt:batch_size], @119, @9, @9] dataType:MLMultiArrayDataTypeFloat32 strides:@[@(119*9*9), @(9*9), @9, @1] deallocator:NULL error:NULL];

    NSError *error = nil;
    @autoreleasepool { // Core ML内部で確保されたメモリを解放するのに必要
        DlShogiResnet15x224SwishBatchOutput *model_output = [model predictionFromInput:model_input error:&error];
        if (error) {
            NSLog(@"%@", error);
            return false;
        }

        // 出力は動的確保された領域に書き出されるため、これを自前のバッファにコピー
        memcpy(output_policy, model_output.output_policy.dataPointer, batch_size * NNWrapper::policy_size * sizeof(float));
        memcpy(output_value, model_output.output_value.dataPointer, batch_size * NNWrapper::value_size * sizeof(float));
    }

    return true;
}

NNWrapper::~NNWrapper() {
    // 所有権をARCに返す
    DlShogiResnet15x224SwishBatch* model = CFBridgingRelease(this->model);
    // スコープを外れるので解放される
}

C++のメンバ関数内で、Objective-CのAPIを呼び出すことができます。呼び出しの手順はObjective-Cのみで実装したコードと同じです。ただし、メモリ管理について注意点があります。C++オブジェクトのメンバ変数(this->model)にObjective-Cで生成したモデルオブジェクトへのポインタを格納しています。voidポインタにキャストしてしまうとARCが機能しないため、CFBridgingRetainという関数を用いてオブジェクトの所有権をプログラマに渡します（プログラマの責任で明示的に開放する必要がある）。(__bridge DlShogiResnet15x224SwishBatch*)(this->model);は所有権を移さないキャスト、CFBridgingReleaseはARCに所有権を返すキャストです。これによりデストラクタでモデルが解放されます。本来これだけでよいはずですが、[model predictionFromInput:...]の箇所を@autoreleasepoolで囲う必要があります。これがないとメモリリークすることが観測されました。C++と連携しない、Objective-C単独のプログラムではリークしませんでした。以前Swiftから呼び出したときも@autoreleasepoolがないとメモリリークする事例があったため、Core ML特有の問題と思われます。以上の処理で、推論を数分間にわたり繰り返してもメモリリークは見られませんでした。ただし、モデル自体のロード・解放を1000回繰り返すとプロセスのメモリ使用量が50MB程度となりました。このような極端な使い方はなかなかしないと思われるので、許容することとします。

このように定義したラッパークラスは、C++のアプリのコードから次のように呼び出します。特にObjective-Cの要素はありません。

#include "nnwrapper.hpp"

NNWrapper nnwrapper(backend);
nnwrapper.run(batch_size, input.data(), output_policy.data(), output_value.data()); // vector<float>.data()

ビルドに用いるMakefileの抜粋を示します。Objective-Cソースのコンパイルには-fobjc-arcをつけています。

maincpp: maincpp.o nnwrapper.o DlShogiResnet15x224SwishBatch.o DlShogiResnet15x224SwishBatch.mlmodelc
    g++ -o maincpp maincpp.o nnwrapper.o DlShogiResnet15x224SwishBatch.o -framework Foundation -framework CoreML

maincpp.o: maincpp.cpp
    g++ -c --std=c++11 -O3 $<

nnwrapper.o: nnwrapper.mm
    g++ -c -fobjc-arc -O3 $<

DlShogiResnet15x224SwishBatch.o: DlShogiResnet15x224SwishBatch.m
    g++ -c -fobjc-arc -O3 $<

それぞれのソースを別個にコンパイルした後、リンクしています。非常に単純なコードであればgccコマンド1回でビルドも可能でしたが、C++固有のオプションをつけた場合はエラーとなりました。

gcc -o maincpp --std=c++11 maincpp.cpp nnwrapper.mm DlShogiResnet15x224SwishBatch.m -framework Foundation -framework CoreML
# error: invalid argument '--std=c++11' not allowed with 'Objective-C'

このやり方で、Core MLを動作させることができました。次回、ベンチマーク結果を示します。