8.1.11. Example: Training MNIST

mnist.py から MLSDK API を取り除き、PyTorch で学習を行うサンプルプログラム (移行作業チュートリアル で使用します)

Example: MNIST on MN-Core 2 と同様ですが、単に checkpoint.pt を --output で指定された場所に保存するだけです。 (デフォルトでは /tmp/mlsdk_mnist_train/checkpoint.pt)

実行方法

$ cd /opt/pfn/pfcomp/codegen/MLSDK/examples/
$ ./exec_with_env.sh python3 mnist_train.py

想定出力

• 学習中のログ

  • Loss curve が Example: MNIST on MN-Core 2 のものと異なることがありますが、これは異なるバックエンドが使用されているためです。

epoch 0, iter    0, loss 2.29758358001709
epoch 0, iter  100, loss 0.6065061688423157
...
epoch 9, iter  900, loss 0.12388602644205093
epoch 9, loss 0.12544165551662445

• チェックポイント (checkpoint.pt)

  • 学習が正常に完了したかは mnist_infer.py を使ってチェックします。

  • Accuracy 指標が 0.94 よりも大きければ良いです。

関連リンク

• 移行作業チュートリアル

  • MLSDK API を段階的に導入する際の参考資料です。

サンプルプログラム

リスト 8.11 /opt/pfn/pfcomp/codegen/MLSDK/examples/mnist_train.py

import argparse
import os
import random
from pathlib import Path
from typing import Mapping, Optional

import numpy as np
import torch
from mlsdk import storage
from mnist_common import MNCoreClassifier, mnist_loaders

torch.manual_seed(0)
random.seed(0)
np.random.seed(0)


def main(outdir: str, option_json_path: Optional[Path], device_str: str) -> None:
    batch_size = 64
    eval_batch_size = 125

    train_loader, _ = mnist_loaders(batch_size, eval_batch_size)

    model_with_loss_fn = MNCoreClassifier()
    model_with_loss_fn.train()

    optimizer = torch.optim.SGD(model_with_loss_fn.parameters(), 0.1, 0.9, 0.0)

    def train_step(inp: Mapping[str, torch.Tensor]) -> Mapping[str, torch.Tensor]:
        x = inp["x"]
        t = inp["t"]
        optimizer.zero_grad()
        output = model_with_loss_fn(x, t)
        loss = output["loss"]
        loss.backward()
        optimizer.step()
        return {"loss": loss}

    for epoch in range(10):
        loss = 0.0
        for i, sample in enumerate(train_loader):
            curr_loss = train_step(sample)["loss"]
            loss += (curr_loss - loss) / (i + 1)
            if i % 100 == 0:
                print(f"epoch {epoch}, iter {i:4}, loss {loss}")
        print(f"epoch {epoch}, loss {loss}")

    os.makedirs(outdir, exist_ok=True)
    torch.save(
        {
            "model_state_dict": model_with_loss_fn.state_dict(),
            "optim_state_dict": optimizer.state_dict(),
        },
        storage.path(outdir) / "checkpoint.pt",
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="""
        A script designed to be used with mnist_infer.py,
        specifically for running MNIST training operations.
        """
    )
    parser.add_argument("--outdir", type=str, default="/tmp/mlsdk_mnist_train")
    parser.add_argument("--option_json", type=Path, default=None)
    parser.add_argument("--device", type=str, default="mncore2:auto")
    args = parser.parse_args()
    main(args.outdir, args.option_json, args.device)