チュートリアル: PyTorch モデルをトレーニングする

torch
pandas
ucimlrepo
scikit-learn
pyyaml

learning_rate: 0.1
epochs: 1000
model_type: Multivariate_neural_network_classifier

import argparse
import torch
from torch import nn
import yaml
import pandas as pd
from ucimlrepo import fetch_ucirepo

from sklearn.model_selection import train_test_split

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear_stack = nn.Sequential(
            nn.Linear(in_features=16 , out_features=16),
            nn.Sigmoid(),
            nn.Linear(in_features=16, out_features=7)
        )

    def forward(self, x):
        logits = self.linear_stack(x)
        return logits

def main(args):
    # 指定された設定ファイルからハイパーパラメーターを読み込む
    with open(args.config, 'r') as f:
        hyperparameter_config = yaml.safe_load(f)

    # データセットを取得
    zoo = fetch_ucirepo(id=111)

    # データ（pandas dataframe として）
    X = zoo.data.features
    y = zoo.data.targets

    print("features: ", X.shape, "type: ", type(X))
    print("labels: ", y.shape, "type: ", type(y))

    ## データを処理
    # データの型はモデルのデータ型と一致している必要があります。nn.Linear のデフォルトの dtype は torch.float32 です
    dataset = torch.tensor(X.values).type(torch.float32)

    # tensor に変換し、インデックスで扱えるようにラベルを 0～6 に変換する
    labels = torch.tensor(y.values)  - 1

    print("dataset: ", dataset.shape, "dtype: ",dataset.dtype)
    print("labels: ", labels.shape, "dtype: ",labels.dtype)

    torch.save(dataset, "zoo_dataset.pt")
    torch.save(labels, "zoo_labels.pt")

    # 後で参照できるようにし、再現性を確保するために、トレーニング用データセットの分割方法を記述する
    config = {
        "random_state" : 42,
        "test_size" : 0.25,
        "shuffle" : True
    }

    # データセットをトレーニング用とテスト用に分割
    X_train, X_test, y_train, y_test = train_test_split(
        dataset,labels,
        random_state=config["random_state"],
        test_size=config["test_size"],
        shuffle=config["shuffle"]
    )

    # ファイルをローカルに保存
    torch.save(X_train, "zoo_dataset_X_train.pt")
    torch.save(y_train, "zoo_labels_y_train.pt")

    torch.save(X_test, "zoo_dataset_X_test.pt")
    torch.save(y_test, "zoo_labels_y_test.pt")


    ## モデルを定義
    model = NeuralNetwork()
    loss_fn = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=hyperparameter_config["learning_rate"])
    print(model)

    # トレーニングループ内で比較するため、初期のダミー損失値を設定
    prev_best_loss = 1e10

    # トレーニングループ
    for e in range(hyperparameter_config["epochs"] + 1):
        pred = model(X_train)
        loss = loss_fn(pred, y_train.squeeze(1))

        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        # 損失が改善した場合はモデルをチェックポイントとして保存
        if (e % 100 == 0) and (loss <= prev_best_loss):
            print("epoch: ", e, "loss:", loss.item())

            # 新しい最良の損失を保存
            prev_best_loss = loss

    print("Saving model...")
    PATH = 'zoo_wandb.pth'
    torch.save(model.state_dict(), PATH)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train a simple neural network on the zoo dataset.")
    parser.add_argument("--config", type=str, required=True, help="Path to the hyperparameter configuration file.")
    args = parser.parse_args()
    main(args)