[PyTorch-Lightning: v2.5.1] LightningModule class 파악

필자는 요즘 고전적인 Deep Learning 모델을 개발하기 위해 PyTorch-Lightning 을 활용한 모델 학습 및 평가 모듈을 개발하고 있다.
개발을 하며 가장 최신 버전으로 공부중인 PyTorch-Lightning 의 기능에 대해 포스팅하고자 한다.
PyTorch-Lightning 기능에 대한 포스팅은 대략 3개 정도 올릴 예정이다.
그 중 본 포스팅에서는 LightningModule 에 대해 다뤄보도록 하겠다.

목차

1. 왜 PyTorch-Lightning 을 사용하는가?

2. LightningModule 이란

3. v2.5.1 변경 사항 및 특징

출처: https://github.com/Lightning-AI/pytorch-lightning

---

1. 왜 PyTorch-Lightning 을 사용하는가?

요즘 인공지능 쪽에서 가장 핫한 분야는 AI-Agent 이다.
AI-Agent 에서 가장 중요한 Deep Learning 분야는 당연 LLM 또는 VLM 일거다.
이를 위해 Transformers 프레임워크가 가장 많이 활용되고 분석되고 있다.
Transformers 는 NLP-Transformer 모델 및 Multi-Modal 모델 대상으로 굉장히 잘 짜여졌다.

하지만 해당 모델이 아닌 다른 모델을 개발해야 한다면 어떻게 될까?

 

다른 모델을 개발해야 한다면 필자는 PyTorch-Lightning 을 활용하여 모델 학습 및 평가, 추론하는 걸 추천한다.

그 이유는 다음과 같다.

1. Transformers 프레임워크 대비 폭 넓은 Model Architecture 에 적용 가능하다.
   • CNN, RNN 과 같은 고전적인 모델 뿐만 아니라 Transformer 구조의 모델에도 적용할 수 있다.
   • 즉, Model Architecture 로 폭 넓게, 유연하게 지원해줄 수 있다. 
2. 학습 루프를 직접 수정하거나 커스텀 로직을 추가할 수 있다.
   • 모델, 학습 로직, 데이터 처리를 명확히 분리해 유지보수가 편리하다.
   • 새로 합류하는 팀원이 F/U 하기 편하다. 
3. 다중 GPU/TPU 학습을 위한 고급 설정을 지원한다.
   • 분산 학습 셋팅이 간단하다.
   • 필자 생각에 이 부분은 Transformers 프레임워크가 더 많은 기능을 제공하는 것 같다. 

 

필자는 개인적으로 2번의 이유로 많은 연구원들이 PyTorch-Lightning 을 사용했으면 좋겠다.

실무자 입장에서 유지보수와 새로운 팀원의 F/U 이 용이하다는 이점이 굉장히 크기 때문이다.

 

자 그러면 PyTorch-Lightning 프레임워크 중 중요한 class 인 LightningModule 에 대해 다뤄보도록 하겠다.

---

2. LightningModule 이란

PyTorch-Lightning 프레임워크에서 모델의 구조, 학습/검증/테스트 루프, 최적화 방법 등을 하나의 클래스로 통합해 관리할 수 있도록 해주는 핵심 클래스이다.

PyTorch에서는 모델, 데이터 로딩, 학습 루프, 최적화, 로깅 등을 모두 별도의 코드로 구현해야 했지만, LightningModule을 사용하면 이 모든 과정을 하나의 클래스 안에 체계적으로 정의할 수 있다.(이를 보일러플레이트 코드라 한다.)

 

LightningModule은 다음과 같은 주요 메서드(함수)로 구성된다.

1. __init__: 모델의 레이어, 손실함수, 메트릭 등 필요한 구성요소를 초기화
2. forward: 입력 데이터를 받아 모델의 순전파(forward pass) 연산을 정의
3. training_step: 학습 데이터의 한 배치(batch)에 대해 손실(loss) 계산 및 학습 과정의 핵심 로직을 구현
4. validation_step: 검증 데이터의 한 배치에 대해 성능 평가 로직을 구현
5. test_step: 테스트 데이터의 한 배치에 대해 평가 로직을 구현
6. predict_step: 추론(inference) 시 사용할 로직을 정의(선택 사항)
7. configure_optimizers: 사용할 옵티마이저(optimizer)와 학습률 스케줄러 등을 정의

위와 같은 메서드가 hook 에 의해 정의만 하면 모듈 내부적으로 동작되기 때문에 메서드 별 동작 순서와 방식만 이해하면 코드가 유연성과 확장성이 좋아지고 재현성과 협업하기 쉬운 이점이 있다.

이러한 구조 덕분에 모델 코드와 학습 로직이 명확하게 분리되고, 코드의 재사용성과 유지보수성이 크게 향상되는 거다.

 

간단한 예시 코드는 아래를 참고하기 바란다.

import torch
import torch.nn as nn
from lightning.pytorch import LightningModule

class LitClassifier(LightningModule):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(28*28, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
        self.loss_fn = nn.CrossEntropyLoss()

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = self.loss_fn(logits, y)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

 

이번에는 PyTorch-Lightning 의 가장 최신 버전인 v2.5.1 의 특징에 대해 알아보도록 하자.

---

 

3. v2.5.1 변경 사항 및 특징

v2.x 대로 넘어오면서 가장 크게 변경된 특징은 다음과 같다.

training_epoch_end 및 validation_epoch_end method 가 사라지고 on_train_epoch_end 및 on_validation_epoch_end 로 변경됐다는 점이다.

단순히 method name 만 변경된게 아니다. 기능도 변경됐다.

 

1.x 버전에서 training_epoch_end method 는 인자로 받은 outputs 값을 모든 배치의 반환값 리스트(모든 배치의 training_step 출력 리스트)가 자동으로 전달되었으나 2.x 부터는 아무 인자도 받지 않게 됐다. 이는 validation_epoch_end 도 마찬가지다.

이게 무슨 말인지 예시와 함께 확인해보자. ( validation_epoch_end 도 마찬가지니 동일하게 변경됐다고 이해하면 된다. )

# 모든 배치의 출력값이 필요하다면, 직접 인스턴스 변수(리스트 등)에 저장해야 됨
class MyLightningModule(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.training_step_outputs = []
        
    def training_step(self, batch, batch_idx):
        loss = ...
        self.training_step_outputs.append(loss)
        return loss
        
    def on_train_epoch_end(self):
        epoch_mean = torch.stack(self.training_step_outputs).mean()
        self.log("training_epoch_mean", epoch_mean)
        self.training_step_outputs.clear()

v1.x 에서는 training_step method 에서 매 batch 마다 연산한 값을 자동으로 리스트 형태로 training_epoch_end method 에 전달해줬다. 때문에 training_epoch_end 는 인자로 output 을 사용했다. 이게 바로 training_step method 를 통해 한 epoch 내에서 연산된 모든 batch 리스트가 종합된 output 인거다. 

 

하지만 v2.x 에서는 자동으로 리스트 형태로 전달해주지 않는다. 그래서인  on_train_epoch_end 에 인자로 아무 값도 사용하지 않는다.

직접 batch 마다 연산된 값을 list 의 형태로 append 해준 뒤 on_train_epoch_end 에 전달해줘야 된다.

그리고 이 부분이 중요한데 on_train_epoch_end 에 전달을 해준 뒤 모든 처리가 끝나면 .clear() 를 통해 list 값을 초기화해줘야 한다. 초기화해주지 않으면 다음 train epoch 때에 이전 train epoch 에서 계산된 값들이 유지되기 때문이다.

 

필자가 생각에는 이러한 변경이 유지보수와 새로운 팀원의 F/U 이 용이하다는 점에서 더 좋아진 것 같다.

해당 부분 외에는 크게 달라진 점이 없는걸 보면 이 목적으로 변경한게 아닌가...? 라는 생각도 든다.

---

 

마무리,,

 AI-Agent 에는 LLM, VLM 과 같은 초거대 모델만 필요하다고 생각하지는 않는다.
반드시 간단한 모델도 Agent 의 한 부분으로 사용되어야 할 거다.
그렇기 때문에 당신이 데이터 사이언스 실무자라면 Transformers 프레임워크 뿐만 아니라 PyTorch-Lightning 프레임워크도 다룰 줄 알아야 한다고 생각한다.

다음 포스팅으로는 LightningDataModule 에 대해서 다루도록 하겠다.

To Be Continued ~~