인공지능/딥러닝

Pytorch - Tutorials(파이토치 모델 학습 코드 설명)

BangPro 2024. 4. 9. 15:33
728x90

데이터셋

Fashion MNIST

모델 학습 순서

  1. 데이터 로드
  2. 네트워크 구조 만들기
  3. Loss
  4. Optimizer
  5. Train
  6. Test
  7. Save

import할 라이브러리들

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

Dataset과 Dataloader

Dataset 저장

# 공개 데이터셋에서 학습 데이터를 내려받습니다.
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

# 공개 데이터셋에서 테스트 데이터를 내려받습니다.
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)
  • Train, Test 데이터를 따로 내려받는다
  • transform = ToTensor()를 통해서 tensor로 변환해서 학습을 한다. 여기는 필수다.

DataLoader

batch_size = 64

# 데이터로더를 생성합니다.
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

for X, y in test_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

DataLoader를 생성하는데 Train, Test 따로 만들고 데이터를 각각 넣어준다.

batch_size를 정의해서 학습 및 검증시에 batch_size 만큼 한번에 계산을 수행한다.

데이터가 너무 크면 OOM(Out Of Memory) 문제가 생길수도 있으니 데이터 크기를 고려해서 batch_size를 조절하자.

근데 또 너무 작으면 학습이 너무 느리다.

DataLoader의 역할

  • 데이터를 shuffle해준다 (매번 같은 순서로 학습하면 NN이 순서도 학습해버려서 데이터 순서를 섞어야한다.)
    • 학습 데이터는 무조건 shuffle 해줘야한다.
  • 총 데이터셋을 batch_size만큼 잘라서 모델 학습시키는 역할 -> SGD 방식

Model 정의

# 모델을 정의합니다.
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )

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

Nerual Network 함수를 정의할때는 무조건 두 함수를 override해야한다.
init은 NN layer를 정의해주고 flatten()을 통해서 데이터를 평탄화(1차원)한다.
forward의 경우 init에서 정의한 layer에 input값을 넣어서 결과(예측값)을 얻는다.
이때 출력값은 logits라 하여 반환한다.

device = 'cuda'  if torch.cuda.is_available() else 'cpu'
print(device)

만약 gpu를 사용할 수 있는 환경이라면 device='cuda'가 되고 그렇지 않다면 device = 'cpu'가 된다.

model = NeuralNetwork().to(device)
print(model)

우리가 정의한 모델을 cuda 즉, gpu에 넣고 model 변수에 저장한다.

Loss Function과 Optimizer

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

Loss functionr과 Optimizer를 정의해준다.

  • Loss Function에는 MseLoss(), CrossEntropyLoss()등 다양한 함수가 있다.
  • Optimizer에는 SGD를 비롯해서 ADAM, adaGrad 등 다양한 Optimizer가 존재한다. ADAM이 가장 일반적으로 사용된다.

Train Test 함수 정의

Train

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)

        # 예측 오류 계산
        pred = model(X)
        loss = loss_fn(pred, y)

        # 역전파
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")
  • model() 이 Neural network 즉, 신경망이다. 다른말로 f(x)
  • dataloader에는 데이터가 있어서 for문에 넣으면 데이터 (X,y)가 튀어나온다. 물론 batch_size만큼!
  • X를 model()에 넣고 정답값 y와의 차이를 구하면 그게 바로 loss값
    • 여기서는 CE loss
    • loss 계산을 하면 forward propagation
  • backpropagation은 optimizer가 미분값을 계산한다.
  • 우선 optimizer.zero_grad()로 기존의 미분값 즉, 쓰레기 값을 비워준다.
  • loss.backward()를 하면 미분값을 계산을 한다.
  • optimizer.step()에서는 미분값을 빼서 미분값을 update한다.
  • if문은 batch 100개마다 loss가 주는지 체크하는 것

Test

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")
  • 테스트시에는 torch.no_grad()해서 학습못하도록 막는다
  • dataloader에서 데이터를 뽑아서 model에 넣고 예측을 한다.
    • test_loss에다 loss값들을 더해준다.
    • .item()은 값만 빼주는 함수로 loss_fn(pred,y)에는 tensor(숫자) 이런식으로 저장이 되어있는데 .item()을 해서 숫자만 뽑는다.
    • 계산 편할려고 .item()을 쓴다
  • 모델 출력(pred)은 일반적으로 각 데이터 포인트에 대한 클래스 확률을 포함하는 2차원 텐서.
    예를 들어, 분류 문제에서 10개의 클래스를 가진다면, 모델 출력은 각 데이터 포인트에 대해 10개의 확률 값을 포함
    출력 텐서의 모양은 [batch_size, n_classes]가 된다.
  • argmax() 함수는 input tensor에 있는 모든 element들 중에서 가장 큰 값을 가지는 항목의 인덱스 번호를 반환하는 함수이다.
    • 인자로 차원을 전달할 수 있다. pred에서 정답 레이블은 2번째 차원에 있으므로 인자로 1(0,1 중에서)을 전달한다.
  • pred.argmax(1) == y을 하면 예측 레이블과 정답 레이블이 일치하면 True, 아니면 False를 반환한다.
  • 그 값에 .type(torch.float)를 추가해서 boolean값을 1.0 혹은 0.0으로 변환을 해준다.
  • .sum으로 해당 값들의 총 합을 구하고 .item()으로 숫자만 변수에 저장한다.

학습

epochs = 20
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
print("Done!")

에폭마다 학습하고 테스트한다.

save와 load

torch.save(model.state_dict(), '가중치파일이름.pth')

위 코드로 모델의 구조는 버리고 가중치만 저장할 수 있다

torch.save(model,'가중치파일이름.pth')

위 코드로 모델 자체를 저장할 수 있다.

Reference

https://tutorials.pytorch.kr/beginner/basics/quickstart_tutorial.html