튜토리얼: 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 데이터프레임 형식)
    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)

    # 텐서로 변환하고 인덱싱할 수 있도록 레이블을 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)

    # 트레이닝 루프에서 비교할 초기 더미 loss 값을 설정합니다
    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()

        # loss가 개선되면 모델을 체크포인트로 저장합니다
        if (e % 100 == 0) and (loss <= prev_best_loss):
            print("epoch: ", e, "loss:", loss.item())

            # 새 최저 loss를 저장합니다
            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)