2. 機械学習チュートリアル
発展的な PyTorch プログラムを題材に、MLSDK の機械学習用の仕様について紹介します。
2.1. 事前準備
はじめに の内容を一通り進め、既に必要な環境構築は完了しているものとします。これに加え、 MNIST のデータセットをサンプルプログラム内部でダウンロードするため、インターネット接続があることを前提にしています。DevKit などインターネット接続がない環境の場合、以下のようにプログラムを修正することで任意のパスに置いたデータセットを参照できます。
train_dataset = datasets.MNIST(
- "/tmp", train=True, transform=transform, download=True
+ "<path-to-MNIST>/", train=True, transform=transform, download=False
)
...
eval_dataset = datasets.MNIST(
- "/tmp",
+ "<path-to-MNIST>/",
train=False,
transform=transform,
- download=True,
+ download=False,
)
2.2. サンプルプログラムの実行
まず、 はじめに で取り上げた add.py と同様に mnist.py を動かしてみましょう。動かす手順は Example: MNIST on MN-Core 2 を参考にしてください。もし add.py と別の環境や Notebook で走らせる場合、 codegen_preload.sh と codegen_pythonpath.sh の読み込みを忘れないようにしてください。
出力の最後に以下のような結果が出力されていれば、実行に成功しています。MLSDK のバージョンによっては Accuracy の値が一致しない場合もありますが、値が0.95を上回っていれば問題ありません。
Correct: 9609 / 10000. Accuracy: 0.9609
また、出力の途中に以下のような学習中のログが出ているはずです。
epoch 0, iter 0, loss 2.3125
epoch 0, iter 100, loss 0.6226431969368814
...
epoch 9, iter 900, loss 0.10909322893182918
epoch 9, loss 0.11064393848594248
こちらは loss の値が次第に下がっていれば問題ありません。
mnist.py プログラムではモデルの学習と推論を連続して行うため、ログ出力の順番はおおまかに以下の順番になっています。
MNIST データセットのダウンロード (初回のみ)
学習 1 iter 分の処理 (
train_step) のコンパイル10 epoch 分の学習
推論 1 iter 分の処理 (
eval_step) のコンパイル10000 のケースに対して分類 (推論)
2.3. サンプルプログラムの解説
2.3.1. drop_lastの指定
22 train_loader = torch.utils.data.DataLoader(
23 train_dataset,
24 batch_size=batch_size,
25 shuffle=True,
26 drop_last=True,
27 collate_fn=list_to_dict,
28 )
drop_last フラグを指定した場合、全データセットを batch_size 毎に分割した際の端数分のイテレーションをスキップします。そして MN-Core 2 を使用する場合は、 train_loader の作成時に drop_last フラグを指定してください。
Context.compile は処理全体を静的な計算グラフとして扱うため、サンプル入力のバッチサイズを前提にコンパイルを行います。そのため batch_size に満たない入力を CompiledFunction に指定した場合、次元が異なるとしてエラーとなる場合があり、これを回避するため drop_last を指定します。
注釈
バッチ軸のサイズなど一部が可変になっているような次元を Dynamic shape と呼びます。現在 drop_last の指定が必要なのは MLSDK が Dynamic shape をサポートしていないためで、将来的にサポートが完了すれば指定は不要になります。
2.3.2. MNCoreClassifier
MNCoreClassifier という名前で多層パーセプトロンによる学習モデルを実装しています。後述するパラメータは torch.nn.Linear に対応する重みやバイアスのテンソルを指します。
2.3.3. パラメータをContextへ登録
36 set_tensor_name_in_module(model_with_loss_fn, "model_with_loss_fn")
37 for p in model_with_loss_fn.parameters():
38 context.register_param(p)
mlsdk.set_tensor_name_in_module() はモデル中の各テンソルに Context が識別するための名前を付与します。この名前は mlsdk.Context.register_param() もしくは mlsdk.Context.register_buffer() の内部で参照されるため、これらの API を呼ぶ前に set_tensor_name_in_module を呼ぶ必要があります。
この例ではパラメータにセットされる名前は以下のようになります。これらの名前は mlsdk.get_tensor_name() で取得できます。
model_with_loss_fn@linear1/weight
model_with_loss_fn@linear1/bias
model_with_loss_fn@linear2/weight
model_with_loss_fn@linear2/bias
注釈
名前は各パラメータのテンソルに FX2ONNX_EXPORTER_TENSOR_NAME_ATTR を setattr することで付与しています。
36 set_tensor_name_in_module(model_with_loss_fn, "model_with_loss_fn")
37 for p in model_with_loss_fn.parameters():
38 context.register_param(p)
register_param は学習によって更新されるパラメータに対して適用します。もし register_param が行われなかった場合、一見学習は進みますがそれはデバイス上のパラメータが更新されているためであり、 Context.synchronize のタイミングでホスト上の実パラメータへの反映が行われません。そのためパラメータの Context への登録は、モデルの学習プログラムにおいて必須です。パラメータと同様に、バッファーにも register_buffer を行います。
また、同一の Context で扱うモデルが複数ある場合、それぞれに対して set_tensor_name_in_module と register_param 及び register_buffer が必要になります。実際に複数のモデルを扱う例が Example: Inference With Multiple Models にあります。
2.3.4. オプティマイザの内部バッファをContextへ登録
40 optimizer = MNCoreSGD(model_with_loss_fn.parameters(), 0.1, 0.9, 0.0)
41 set_buffer_name_in_optimizer(optimizer, "optimizer")
42 context.register_optimizer_buffers(optimizer)
mlsdk.MNCoreSGD は torch.optim.SGD を MLSDK 用に再実装したものです。Learning rage (lr) など基本的なオプションについては共通していますが、完全な互換性は保証していません。この例では SGD を使用していますが、他にも mlsdk.MNCoreAdam や mlsdk.MNCoreAdamW が使用可能です。
40 optimizer = MNCoreSGD(model_with_loss_fn.parameters(), 0.1, 0.9, 0.0)
41 set_buffer_name_in_optimizer(optimizer, "optimizer")
42 context.register_optimizer_buffers(optimizer)
mlsdk.set_buffer_name_in_optimizer() はオプティマイザの持つ内部状態を表現する各テンソルに Context が識別するための名前を付与します。この名前は mlsdk.Context.register_optimizer_buffers() の内部で参照されます。登録することにより、オプティマイザ内部の state_dict が Context.synchronize のタイミングで更新されるため、学習のチェックポイントにそれを含めることができます。
2.3.5. 学習・推論処理の関数化
44 def train_step(inp: Mapping[str, torch.Tensor]) -> Mapping[str, torch.Tensor]:
45 x = inp["x"]
46 t = inp["t"]
47 optimizer.zero_grad()
48 output = model_with_loss_fn(x, t)
49 loss = output["loss"]
50 loss.backward()
51 optimizer.step()
52 return {"loss": loss}
学習ループの中心である Forward → Backward → Optimize の処理 (4-8行目) を Context.compile() で受け取れるように関数化しています。引数と返り値を Mapping[str, torch.Tensor] にする手間はありますが、関数としてまとめる際に処理本体を修正する必要は基本的にありません。
87 def eval_step(inp: Mapping[str, torch.Tensor]) -> Mapping[str, torch.Tensor]:
88 x = inp["x"]
89 t = inp["t"]
90 output = model_with_loss_fn(x, t)
91 y = output["y"]
92 _, predicted = torch.max(y, 1)
93 correct = (predicted == t).sum()
94 return {"correct": correct}
推論についても、 Forward → Max+Sum の処理を train_step と同様に関数化しています。モデルやオプティマイザ本体の処理だけでなく、その結果の後処理もまとめて MN-Core 2 で行っています。
2.3.6. コンパイルオプションの指定
54 compile_options = {}
55 if option_json_path is not None:
56 compile_options["option_json"] = str(option_json_path)
MLSDK のバックエンドである codegen にはコンパイル時に指定可能な環境変数やコマンドラインオプションが多数ありますが、その中でもベストな組み合わせをまとめているのが Preset Options です。各 Preset Option は /opt/pfn/pfcomp/codegen/preset_options/ 以下に JSON 形式で保存されており、デバッグ用の debug.json から高度な最適化をする O4.json まであります。注意点として、最適化オプションによっては演算順序が変わることにより、最終的な計算結果が変わる可能性があります。
mnist.py では以下のように Preset Option を指定可能になっています。
$ cd /opt/pfn/pfcomp/codegen/examples/
$ ./exec_with_env.sh python3 mnist.py --option_json /opt/pfn/pfcomp/codegen/preset_options/O1.json
実際にどれほど性能が変わるのかを調べる場合、 Codegen Dashboard を参照してください。また、コンパイルオプションとして指定可能なものについては、 コンパイルオプション を参照してください。
2.3.7. 学習結果の同期
66 for epoch in range(10):
67 loss = 0.0
68 for i, sample in enumerate(train_loader):
69 curr_loss = compiled_train_step(sample)["loss"].item()
70 loss += (curr_loss - loss) / (i + 1)
71 if i % 100 == 0:
72 print(f"epoch {epoch}, iter {i:4}, loss {loss}")
73 print(f"epoch {epoch}, loss {loss}")
74
75 context.synchronize()
この例では 10 epoch 分の学習ループを抜けたあとに Context.synchronize を呼んでいます。Context に登録されたパラメータはこのタイミングで同期されるため、 state_dict を保存する場合は同期後に行います。
2.3.8. 学習結果の保存
77 torch.save(
78 {
79 "model_state_dict": model_with_loss_fn.state_dict(),
80 "optim_state_dict": optimizer.state_dict(),
81 },
82 storage.path(outdir) / "checkpoint.pt",
83 )
学習したモデルやオプティマイザの state_dict をここの torch.save で保存しています。保存先は --outdir を特に指定しない限り /tmp/mlsdk_mnist/checkpoint.pt です。
このように保存した内容は、 torch.load を使うことで復元することが出来、デバイスに依存せず学習結果の活用ができます。
2.4. 発展的なトピック
移行作業チュートリアル でプログラムを MN-Core シリーズで動かす際の流れを説明しています
発展的な機能 で
add.pyとmnist.pyでは紹介できていない機能を説明していますギャラリー から MLSDK を利用した例を参照できます