[파이토치] 05 - Training step 내부 동작 알아보기

Lightning 모듈 내부 해부¶

일반적인 pytorch 에서 학습 sequence 는 다음과 같다. (traning loop)

In [10]:

def train_loop( dataloader, model, lossfn, optimizer):
    size = len(dataloader.dataset)
    for batch, (x, y) in enumerate(dataloader):

        pred = model(x) # model class 가 callable 한 경우! 
                        # forward를 직접 불러도 결과는 같겠지만, 권장하지 않음 -> 모델을 직접 call해서 부르지 않으면  hook 이 동장하지 않음!
        loss = lossfn( pred, y) 

         
        if batch%100 == 0 :
            loss, current = loss.item(), batch*len(x)
            print(f"loss : { loss }, [{current}/{size}")
        
        

반면 torch lightning 에서 학습을 정의하는 부분은 다음과 같다. (traning_step)

In [ ]:

#in the class 
    def traning_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x) ## callable 한 경우  ->  forward 를 포함한 진행 
        loss = loss_function(y_hat, y) ## loss_function 은 미리 주어져야함
        self.log('traninig loss', loss)
        return loss

구조는 매우 간단하다. 먼저 데이터를 받고 (x, y를 batch로부터, 이 부분은 pl.traniner.fit 이 알아서 해준다.) 이를 모델에 넣어서 나온 값(예측 y값)과, 실제 y값 의 차이를 loss_function 함수에 넣어서 loss 를 계산하고 이를 return 한다. 별도로 back propagation 이나 optimizer 설정은 필요 없음에 유의하자!

이와 완벽하게 똑같은 구조로 validation 또는 test 를 위한 내부 메쏘드 설정도 가능하다. 이름은 똑같이 validation_step, test_step 이라고 정의하면 되고 또한 traninig step 에서의 return 은 loss 이지만 만약 역전파가 필요없는 validation 이나 test 에서는 return 을 정의하지 않고, 결과만 log 에 표현해도 된다.

In [ ]:

History Log¶

Log 의 경우에는 다음의 문서에 잘 나와 있는데 (https://lightning.ai/docs/pytorch/stable/api_references.html#loggers), 각각 Step, Epoch level 에서 로깅이 가능하다.

LightningModule.log(name, value,
    prog_bar=False, logger=True, on_step=None, on_epoch=None, reduce_fx='mean',
    enable_graph=False, sync_dist=False, sync_dist_group=None, add_dataloader_idx=True,
    batch_size=None, metric_attribute=None, rank_zero_only=False)

예를들어 다음의 3개의 차이를 보면

In [ ]:

from IPython.core.display import ProgressBar

## 1) 
def traninig_step( self, batch, batch_idx ) :
    x, y = batch
    y_hat = self(x)
    loss = loss_ftn(y_hat, y)
    self.log("loss", loss, on_step = True, on_epoch = False ) ## Step level 에서 
    return loss

## 2)
def training_step( self, batch, batch_idx ):
    x, y = batch
    y_hat = self(x)
    loss = loss_ftn(y_hat, y)
    self.log("loss", loss, on_step=False, on_epoch = True ) ## Epoch level 마다 log
    return loss

## 3)
def traning_step( self, batch, batch_idx ):
    x, y = batch
    y_hat = self(x)
    loss = loss_ftn(y_hat, y)
    acc = FM.accuracy(y_hat, y, task="multiclass", num_classes = 10)
    metrics = {'loss' : loss, 'acc' : acc }
    self.log_dict( metrics, prog_bar = True ) # by default, on_step = True, on_epoch = False 
    return loss
    

log 를 활용한 1) , 2) 의 경우의 차이는 step 마다 logging을 할 건지, 또는 epoch 마다 logging 을 할 건지의 차이고, 기본적으로는 on_step 이 default 이다. 만약 log 에 보다 자세한 정보를 기록하고 싶다면 먼저 log 에 담길 정보를 dict 형식으로 만든 다음에 (3번처럼), 이를 log_dict 매소드를 사용해서 기록하면 된다. 또한 prog_bar 를 활성화하면 진행정도를 볼 수 있는데 (마치 Keras 처럼, 만약 이 옵션을 False 로 하면 iteration 진행도만 나오고, 실제로 acc 라던가 loss 가 어떻게 변하는지는 보이지 않는다.) log 를 사용하면 자주 활성화 시키는 옵션이다. 이를 이용해서 실제 MNIST 모델을 만들어서 모델이 어떻게 돌아가는지 확인해보자.

In [14]:

import torch
from torch import nn
from torch.nn import functional as F
import torch.optim as optim

import pytorch_lightning as pl
from pytorch_lightning.accelerators import accelerator
from torchmetrics import functional as FM
from torchinfo import summary

from torchvision.datasets import MNIST
import torchvision.transforms as transforms
import torch.utils.data as data
from torch.utils.data import DataLoader

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

In [15]:

loss_ftn = nn.CrossEntropyLoss()

In [65]:

## 4강에서 가져온 모델 사용
class Model(pl.LightningModule):

    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.dense1 = nn.Linear(28*28, 32)
        self.dense2 = nn.Linear(28*28, 32)
        self.dense2_2= nn.Linear(32, 16)
        self.dense3 = nn.Linear(32+16, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.flatten(x)

        x1 = self.dense1(x)
        x1 = self.relu(x1)

        x2_1 = self.dense2(x)
        x2_1 = self.relu(x2_1)
        x2_2 = self.dense2_2(x2_1)
        x2_2 = self.relu(x2_2)

        x = torch.cat([x1, x2_2], dim=1)
        x = self.dense3(x)

        return(x)
        

In [66]:

loss_function = nn.CrossEntropyLoss()
class SimpleModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layers = Model()  ## 위에서 정의된 모델을 사용 (어차피 in-out 만 맞으면 된다)

    def forward(self, x):
        out = self.layers(x)
        return out

    def training_step( self, batch, batch_idx):
        x, y = batch
        y_pred = self(x)
        loss = loss_function(y_pred, y)
        acc = FM.accuracy( y_pred, y, task='multiclass', num_classes=10)
        
        metrics = {'loss':loss, 'acc' : acc}

        self.log_dict(metrics, prog_bar=True, on_step=True, on_epoch=True)

        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.001)
     

In [73]:

model = SimpleModel()

In [68]:

summary(model, input_size=(8, 1, 28, 28))

Out[68]:

==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
SimpleModel                              [8, 10]                   --
├─Model: 1-1                             [8, 10]                   --
│    └─Flatten: 2-1                      [8, 784]                  --
│    └─Linear: 2-2                       [8, 32]                   25,120
│    └─ReLU: 2-3                         [8, 32]                   --
│    └─Linear: 2-4                       [8, 32]                   25,120
│    └─ReLU: 2-5                         [8, 32]                   --
│    └─Linear: 2-6                       [8, 16]                   528
│    └─ReLU: 2-7                         [8, 16]                   --
│    └─Linear: 2-8                       [8, 10]                   490
==========================================================================================
Total params: 51,258
Trainable params: 51,258
Non-trainable params: 0
Total mult-adds (Units.MEGABYTES): 0.41
==========================================================================================
Input size (MB): 0.03
Forward/backward pass size (MB): 0.01
Params size (MB): 0.21
Estimated Total Size (MB): 0.24
==========================================================================================

모델 생성은 원만하게 완료하였다. 이제 실제로 학습을 돌려보자

In [69]:

def dataLoader(batch_size=128):
    train_dataset = MNIST('', transform=transforms.ToTensor(), train=True, download=True) ## 한 번 인터넷으로 가져온걸 매번 가져올 필요가 없기 때문에 가져올때 download True 로 하면 다음 부터는 다운로드 된 데이터를 사용한다.
    test_dataset = MNIST('', transform=transforms.ToTensor(), train=False, download=True)
    trainDataLoader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    valDataLoader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    return (trainDataLoader,valDataLoader)

trainDataLoader,valDataLoader = dataLoader()

In [74]:

from pytorch_lightning.accelerators import accelerator
epoch = 5
logger = pl.loggers.CSVLogger("logs", name = "train_history_log")
traniner = pl.Trainer( max_epochs= epoch, logger= logger, accelerator = 'auto', 
                      log_every_n_steps=10) # default 는 50 

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs

In [75]:

traniner.fit(model, trainDataLoader)

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

  | Name   | Type  | Params
---------------------------------
0 | layers | Model | 51.3 K
---------------------------------
51.3 K    Trainable params
0         Non-trainable params
51.3 K    Total params
0.205     Total estimated model params size (MB)

Training: 0it [00:00, ?it/s]

`Trainer.fit` stopped: `max_epochs=5` reached.

Log 에 단순히 matrics 라는 이름의 dict 만 출력했으며, 이 dict 에는 loss 와 acc 만 들어있기 때문에 CSV 파일로 output 을 남겼다. log 옵션은 위와 같이

self.log_dict(metrics, prog_bar=True, on_step=True, on_epoch=True)

으로 남겼기 때문에 1 epoch 안에서도 step (1개 batch가 다 돌 때마다) 로그가 계속 남는다. 이 파일은 기본적으로 'CSVLogger' 를 생성할 때 실행 위치의 하위 폴더 (여기서는 logs)의 name (여기서는 train_history_log)이라는 폴더에 파일을 남기는데, 같은 모델을 반복해서 돌리거나 학습할 때, 또는 같은 이름의 폴더가 존재할 구분을 하기 위해서 version_x 이런 식으로 버전이 하나씩 높아진다.

가장 최근에 돈 버전을 확인하고 싶다면

In [78]:

v_num = logger.version

history = pd.read_csv(f'./logs/train_history_log/version_{v_num}/metrics.csv')
history

Out[78]:

	loss_step	acc_step	epoch	step	loss_epoch	acc_epoch
0	2.093160	0.406250	0	9	NaN	NaN
1	1.744891	0.617188	0	19	NaN	NaN
2	1.477360	0.718750	0	29	NaN	NaN
3	1.163913	0.734375	0	39	NaN	NaN
4	0.953659	0.835938	0	49	NaN	NaN
...	...	...	...	...	...	...
234	0.101236	0.968750	4	2309	NaN	NaN
235	0.179500	0.960938	4	2319	NaN	NaN
236	0.154053	0.953125	4	2329	NaN	NaN
237	0.107376	0.960938	4	2339	NaN	NaN
238	NaN	NaN	4	2344	0.128552	0.962667

239 rows × 6 columns

와 같이 확인할 수 있다.

step 로그가 남는 동안에는 epoch 정보가 없기 때문에 epoch 이 Nan 으로 나오고, epoch log 가 남을 때에는 step 정보가 없다. 만약 로그를 epoch 단위로만 남긴다면 아래와 같은 모양으로 남게 된다.

In [85]:

class SimpleModel2 (SimpleModel):
 
    def training_step( self, batch, batch_idx):
        x, y = batch
        y_pred = self(x)
        loss = loss_function(y_pred, y)
        acc = FM.accuracy( y_pred, y, task='multiclass', num_classes=10)
        
        metrics = {'loss':loss, 'acc' : acc}

        self.log_dict(metrics, prog_bar=False , on_step=False, on_epoch=True) ### << 이부분 수정 

        return loss

In [86]:

model2 = SimpleModel2()

In [87]:

epoch = 5
logger = pl.loggers.CSVLogger("logs", name = "train_history_log_epoch_only")
traniner = pl.Trainer( max_epochs= epoch, logger= logger, accelerator = 'auto', 
                      log_every_n_steps=10) # default 는 50 

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs

In [88]:

traniner.fit(model2, trainDataLoader)

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

  | Name   | Type  | Params
---------------------------------
0 | layers | Model | 51.3 K
---------------------------------
51.3 K    Trainable params
0         Non-trainable params
51.3 K    Total params
0.205     Total estimated model params size (MB)

Training: 0it [00:00, ?it/s]

`Trainer.fit` stopped: `max_epochs=3` reached.

on_step=False, on_epoch=True 로 했을 경우 (log_every_n_steps가 설정되어 있더라도 기록이 되지 않음에 주의!!)

In [95]:

df = history.groupby('epoch').last().drop('step', axis=1)

In [110]:

import matplotlib.pyplot as plt
%matplotlib inline

In [117]:

df["loss_epoch"].plot(kind='line', title = "train loss")

Out[117]:

<Axes: title={'center': 'train loss'}, xlabel='epoch'>

In [118]:

df["acc_epoch"].plot(kind='line',  title = "train accuracy")

Out[118]:

<Axes: title={'center': 'train accuracy'}, xlabel='epoch'>

In [ ]: