Google の Magika を Python から C#に移植するまでの過程 (Day 5 / 7)

Google のファイル判定プログラム Magika を Python から C# に移植する過程を共有する記事の第5回目です。

前回は移植作業の効率を上げるために GitHub Copilot を使ってみました。今回はクラスライブラリ完成までを目指して作りこみを進めていきます。

## 目次
- [Day 1 : まずは Magika の中身を見てみよう](https://qiita.com/mkht/items/64b02b8648266d40380a)
- [Day 2 : C# で 概念実証コードを書いてみる](https://qiita.com/mkht/items/fa76b0238ed976d20ca9)
- [Day 3 : C# クラスライブラリとして Magika を移植していく](https://qiita.com/mkht/items/a7fa36c5446b5fad5092)
- [Day 4 : GitHub Copilot を使って作業効率アップ](https://qiita.com/mkht/items/2c2b23ceb0bc1df66599)
- [Day 5 : クラスライブラリとしての Magika を完成させる](https://qiita.com/mkht/items/9d7a8a16ed1d60e353fa)
- [Day 6 : コンソールアプリを作成する](https://qiita.com/mkht/items/0b2d1546cf8a25ab4cce)
- [Day 7 : 移植したMagikaをビルドし、動作確認する](https://qiita.com/mkht/items/f36d87afe50eae2be5f1)

## クラスライブラリとしての Magika を完成させる

GitHub Copilot の助けも借りつつ、根気よく移植を続けていきます。

### 機械学習モデルによる推論処理

周辺の処理の移植が進み、いよいよ Magika のキモである機械学習モデルによる推論処理の移植に取りかかります。概念実証コードで ONNX Runtime を使って推論処理を行う基本的な流れは分かっていますが、実際の Python 版 Magika の推論処理はもう少し複雑になっているので、それを C# でどう表現するかが課題です。

まずは ONNX 形式の機械学習モデルを読み込んで、推論セッションを初期化する部分。ここは簡単です。

```csharp
private InferenceSession InitOnnxSession()
{
    var modelPath = $".\models\standard_v1\model.onnx";
    var onnxSession = new InferenceSession(modelPath);
    return onnxSession;
}
```

頭を悩ませたのが実際に推論を行う部分。Python 版 Magika の`_get_raw_predictions`メソッドの移植です。

概念実証では単一の入力から単一の出力を得るだけの簡単なコードで済ませていましたが、オリジナルの Magika では高速化のためにバッチ処理が組み込まれており、複数ファイルの情報を一度にまとめて入力し、一回の推論実行で多数のファイルの判定処理を一気に行う形になっています。

```python
X = np.array(X_bytes).astype(np.float32)

raw_predictions_list = []
samples_num = X.shape[0]
max_internal_batch_size = 1000
batches_num = samples_num // max_internal_batch_size

for batch_idx in range(batches_num):
    start_idx = batch_idx * max_internal_batch_size
    end_idx = min((batch_idx + 1) * max_internal_batch_size, samples_num)
    batch_raw_predictions = self._onnx_session.run(
        ["target_label"], {"bytes": X[start_idx:end_idx, :]}
    )[0]
    raw_predictions_list.append(batch_raw_predictions)
return np.concatenate(raw_predictions_list)
```

変数`X`が多次元配列(`NDArray`)になっていて、バッチサイズにあわせて`X[start_idx:end_idx, :]`でスライスして推論処理の入力としているようです。

機械学習モデルでバッチ推論をするという概念が良くわかっておらず「え？そんなことできるの？」という感じでしたし、`NDArray`という NumPy に特有の型、多次元配列という直感的ではない構造にも惑わされ、この処理の内容を理解するのにかなり苦戦しました。

さらにこれを C# でどう表現するかも難しい問題でした。C# の ONNX Runtime では入力は単純な配列ではなく`Tensor<T>`や`NamedOnnxValue`という独自のデータ型で扱う必要があるため、これらを使って多次元配列をどう表現するのか、調べてもほとんど情報が見つかりません。

さらに`InferenceSession.Run()`メソッドも Python 版では一種類しかないのに対し、C# 版では複数のオーバーロード（17種類！）が用意されており、どれを使えばいいのかもわかりませんでした。 

https://onnxruntime.ai/docs/api/csharp/api/Microsoft.ML.OnnxRuntime.InferenceSession.html

最終的に、いろいろ試しながら以下のようなコードでバッチ推論を実現することができました。

```csharp
Memory<float> X = XBytes.Select(b => (float)b).ToArray().AsMemory();

List<float[]> rawPredictionsList = [];
int _input_column_size = 512 * 3;
int _output_column_size = 113;
int samplesNum = features.Count;
int maxInternalBatchSize = 1000;
int batchesNum = samplesNum / maxInternalBatchSize;

for (int batchIdx = 0; batchIdx < batchesNum; batchIdx++)
{
    int startIdx = _input_column_size * batchIdx * maxInternalBatchSize;
    int endIdx = Math.Min((batchIdx + 1) * maxInternalBatchSize, samplesNum);
    int count = endIdx - startIdx;

    List<NamedOnnxValue> input = [];
    Tensor<float> dt = new DenseTensor<float>(
            X.Slice(startIdx, _input_column_size * count),
            [count, _input_column_size],
            false
        );
    input.Add(NamedOnnxValue.CreateFromTensor("bytes", dt));

    using var batchRawPredictions = this._onnx_session.Run(
            outputNames: ["target_label"],
            inputs: input
        );

    foreach (var rawPreds in batchRawPredictions[0].AsEnumerable<float>().Chunk(_output_column_size))
    {
        rawPredictionsList.Add(rawPreds);
    }
}
return rawPredictionsList;
```

200 個のファイルを 100 個ごとに 2 回に分けてバッチ処理するケースを考えてみます。

まず、判定する 200 個のファイルそれぞれから抽出した 1536 バイトの特徴量(`Features`)を全て単一の float 型の配列`X`にまとめておきます。Python 版とは違い多次元配列ではなく普通の一次元配列です。この時点で`X`の要素数は`1536 * 200 = 307200`あります。

次に、バッチごとに一回の推論で使うだけの`Tensor<float>`型のデータを作成します。`DenseTensor(Memory<T>, ReadOnlySpan<Int32>, Boolean)`コンストラクタを使用します。

コンストラクタの第一引数には`X`から一回のバッチで使用する分のデータをスライスで取り出して渡します（`X`の最初から`100 * 1536 ＝ 153600`個のデータ）。

第二引数には各次元のサイズを Span 指定します。第一引数のデータは全部で`153600`個あるけど、これは`10 * 15360`でも`30 * 5120`でもなく、`100 * 1536`ですよと教えてあげているようなイメージです。

あとは`Tensor<float>`を`NamedOnnxValue`型に変換して`List<>`に格納し、`InferenceSession.Run()`メソッドで推論処理を行います。どのオーバーロードを使うかは、どういう方式で推論結果を受け取るかの違いくらいで、どれを使ってもあまり問題にはならないようです。Python 版と最も近い`Run(IReadOnlyCollection<NamedOnnxValue> inputs, IReadOnlyCollection<string> outputNames)`を使っています。

単一入力の場合は 113 要素数の float 配列が出力されましたが、バッチ処理の場合は 100 個のファイルを一度に推論させる例で`100 * 113`要素数の float 配列が出力されます。これを`Chunk()`を使って 113 要素ずつ分割すれば、それぞれのファイルに対する 113 個の要素の float 配列を取り出すことができます。


### Logger の実装

ロギングのための Logger の実装方法について検討します。

オリジナルの Magika では、ログ情報を`stderr`に出力するだけの非常にシンプルな Logger を独自実装しています（その名も`SimpleLogger`）

```python
class SimpleLogger:
    """
    We implement a simple logger to not rely on additional python packages,
    e.g., rich. This is written in way that, by default, log messages (e.g.,
    debug/info/...) are sent to stderr.
    """
```

C# ではどういう実装が一般的なのか調べてみると、同様に独自 Logger を実装している例も多いようですが、お作法としては`ILogger`インターフェースを使用するのが良いようです。

https://learn.microsoft.com/ja-jp/dotnet/core/extensions/logging?tabs=command-line

`ILogger`インターフェースを使うと、ログの出力先を柔軟に変更できるなどのメリットがあるようです。（ふわっとした理解）

コンソールにログを出力するだけならコードもわりとシンプルです。

```csharp
using ILoggerFactory factory = LoggerFactory.Create(builder => builder.AddConsole());
ILogger logger = factory.CreateLogger("Category");
logger.LogInformation("Hello World! Logging is {Description}.", "fun");
```

ということで最初は`ILogger`を使って実装する方針で書き始めたのですが、いざ実装してみるといくつか問題が見えてきました。

#### 1. コンソールログの表示形式がPython版と異なる

`ILogger`でコンソールにログ出力すると、出力が2行になります。

![day4-6-img-01.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/177762/67452219-ea37-7eb4-7257-ca41aadec817.png)


Python 版 Magika のログ出力をなるべくそのまま再現したいと考えていたので、この表示は気になりました。

#### 2. デフォルトのログレベル問題

`ILogger`を使うと、デフォルトで`Information`レベル以上のログが表示されるようになっています。一方でオリジナルの Magika では`verbose`や`debug`を明示的に指定しないと`Information`ログは表示されません。この差異をスマートに吸収する方法が思いつきませんでした。

結局これらの問題をうまいこと何とかする方法がわからず、そこまで`ILogger`にこだわる理由もなかったので、Python 版 Magika の`SimpleLogger`をほぼそのまま C# に移植することにしました。

ただ、一つオリジナルから変更したのがログの出力先で、エラー出力`stderr`ではなく標準出力`stdout`に変えています。これは PowerShell から呼び出す際のことを考えてそうしています。

というのも、PowerShell のエラーハンドリングはやや特殊な形になっていて、実行する外部プログラムが`stderr`に何らかのメッセージを出力すると実行に失敗したとみなし、例外を発生させるのです。そのため、本当にエラーが発生した場合はともかく、そうではないログ出力にまで`stderr`を使うのは避けた方が良いと考えています。


```python
def warning(self, msg: str) -> None:
    if logging.WARNING >= self.level:
        if self.use_colors:
            self.raw_print(f"{colors.YELLOW}WARNING: {msg}{colors.RESET}")
        else:
            self.raw_print(f"WARNING: {msg}")
```

```csharp
internal void Warning(string message)
{
    if (this.LogLevel <= LogLevel.Warning)
    {
        if (this.useColors)
        {
            Console.WriteLine(WARN_MSG_COLOR, message);
        }
        else
        {
            Console.WriteLine(WARN_MSG, message);
        }
    }
}
```

### 続く
クラスライブラリとしての Magika の移植が完了しました。実際には細かいバグ修正や各種プラットフォームでのテストなどがまだまだ残っている状態ですが、とりあえず開発環境では動いているので一段落とします。

[次回](https://qiita.com/mkht/items/0b2d1546cf8a25ab4cce)はこのクラスライブラリ版 Magika を使って、コマンドラインから実行できる Magika コンソールアプリを作成していきます。