694 views

この記事は最終更新から 397日 が経過しています。

【1】やりたいこと

前回の投稿 (33)【PyTorchでMNIST #4】プログラムの保守性を向上させる。 で残件となっていた 学習済みパラメータファイルのコマンドライン指定 を実装したい。

これにより、複数の実行結果を保持、再現することができる。
技術的には特筆すべきことはないが…

学習済みパラメータについて

・拡張子 pth は PyTorch で学習済みパラメータを保存する際の標準・慣習的な形式だ。
・このファイルには、Pythonオブジェクト (state_dictなど) を torch.save() で保存している。
・通常は dict（例：model.state_dict()）がシリアライズされている。
・実体は Pythonの pickle 形式（バイナリデータ）で保存されている。

【2】やってみる

今回変更するファイルは、ハイライト表示した 2ファイルのみだ。
他のファイルは、そのまま使用できる。

1) プログラムソースコード

(1) control_train.py

コマンドライオプション -p でファイル名を指定する。
指定がなかった場合、固定名＋日時分秒でファイル名を自動生成する。

from torch.utils.data import DataLoader
from net_model import Net
from trainer import Trainer
from dataset_MNIST import MyDataset
import torch
import argparse
from datetime import datetime

#//////////////////////////////////////////////////////////////////////////////
def execTrain( save_path ):
    # データ
    dataset = MyDataset()
    train_loader = DataLoader(dataset.get_train_dataset(), batch_size=64, shuffle=True)
    test_loader = DataLoader(dataset.get_test_dataset(), batch_size=1000, shuffle=False)
    # 環境
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = Net()
    trainer = Trainer(model, device)
    # 実行
    trainer.train(train_loader, epochs=3)
    trainer.test(test_loader)
    trainer.save(save_path)

#//////////////////////////////////////////////////////////////////////////////
def main():
    parser = argparse.ArgumentParser(description="Train and save MNIST model")
    parser.add_argument('-p', type=str, help='Filename to save trained model')
    args = parser.parse_args()
    # ファイル名の決定
    if args.p:
        save_path = args.p
    else:
        now = datetime.now().strftime("%Y%m%d_%H%M%S")
        save_path = f"learned_model_{now}.pth"
    execTrain(save_path)

#//////////////////////////////////////////////////////////////////////////////
if __name__ == "__main__":
    main()

(2) control_test.py

こちらは -p オプション指定が必須だ。

import torch
from torch.utils.data import DataLoader
from net_model import Net
from trainer import Trainer
from dataset_MNIST import MyDataset
import argparse

#//////////////////////////////////////////////////////////////////////////////
def execTest( prmfile ):
    # データ
    dataset = MyDataset()
    test_loader = DataLoader(dataset.get_test_dataset(), batch_size=1000, shuffle=False)
    # 環境
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = Net()
    trainer = Trainer(model, device)
    # 実行
    trainer.load(prmfile)
    trainer.test(test_loader)

#//////////////////////////////////////////////////////////////////////////////
def main():
    parser = argparse.ArgumentParser(description="Load and test trained MNIST model.")
    parser.add_argument('-p', type=str, required=True, help='Filename of trained model to load')
    args = parser.parse_args()
    execTest(args.p)

#//////////////////////////////////////////////////////////////////////////////
if __name__ == "__main__":
    main()

2) 実行結果

2回学習を実行する。
テスト結果は、1回目が 97.19%, 2回目が 97.03%だ。

$ python control_train.py -p AAA.pth
[1/3] Epoch complete
[2/3] Epoch complete
[3/3] Epoch complete
Test Accuracy: 97.19%
Model saved to AAA.pth
$ python control_train.py -p BBB.pth
[1/3] Epoch complete
[2/3] Epoch complete
[3/3] Epoch complete
Test Accuracy: 97.03%
Model saved to BBB.pth

1回目、2回目それぞれの学習済みパラメータファイルを指定し、テストを単独実行する。
学習実行時と同じ結果が得られたので OK！

$ python control_test.py -p AAA.pth
Model loaded from AAA.pth
Test Accuracy: 97.19%
$ python control_test.py -p BBB.pth
Model loaded from BBB.pth
Test Accuracy: 97.03%