PyTorch Lightning

PyTorch Lightning은 전문 AI 연구자와 머신러닝 엔지니어를 위한 딥러닝 프레임워크예요. PyTorch의 유연함은 그대로 두면서, 학습에 필요한 반복 보일러플레이트를 걷어내고 스케일에 맞춰 최적화해 줘요. 프로젝트가 아이디어 단계에서 논문·프로덕션까지 갈 때도 같은 코드를 계속 쓸 수 있는 게 핵심이에요.

"batteries included"라는 말처럼, 필요한 도구를 대부분 갖추고 있으면서도 PyTorch 코드를 숨기지 않아요. 그 덕분에 유연성과 성능을 동시에 잡으려는 사람들에게 특히 잘 맞아요.

출처: PyTorch Lightning 공식 문서

설치하기

pip과 conda 모두로 설치할 수 있어요.

pip install lightning

Conda를 쓴다면 conda-forge 채널에서 설치해요.

conda install lightning -c conda-forge

Lightning이 해결하는 문제

순수 PyTorch로 학습 루프를 직접 짜면 반복되는 코드가 많아져요. 에폭·배치 반복, optimizer.step(), loss.backward(), optimizer.zero_grad() 호출, 평가 중 model.eval()과 grad 켜고 끄기, TensorBoard 로깅, 멀티 GPU, 16비트 AMP까지 전부 직접 관리해야 하죠.

Lightning의 Trainer가 이 40여 가지 반복 작업을 자동으로 처리해 줘요. 그래서 연구 코드와 엔지니어링 코드를 분리할 수 있고, 재현성과 가독성도 좋아져요.

LightningModule로 코드 조직화하기

nn.Module을 LightningModule로 감싸면, 학습 루프의 핵심 로직을 훅(hook) 형태로 담을 수 있어요. training_step이 학습 루프를 정의하고, configure_optimizers가 옵티마이저를 돌려줘요.

import os
from torch import optim, nn
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
import lightning as L

# 어떤 nn.Module이든 그대로 사용 가능
encoder = nn.Sequential(nn.Linear(28 * 28, 64), nn.ReLU(), nn.Linear(64, 3))
decoder = nn.Sequential(nn.Linear(3, 64), nn.ReLU(), nn.Linear(64, 28 * 28))

# LightningModule 정의
class LitAutoEncoder(L.LightningModule):
    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def training_step(self, batch, batch_idx):
        # training_step이 학습 루프를 정의해요. forward와 분리돼 있어요.
        x, _ = batch
        x = x.view(x.size(0), -1)
        z = self.encoder(x)
        x_hat = self.decoder(z)
        loss = nn.functional.mse_loss(x_hat, x)
        # 기본적으로 TensorBoard에 로깅돼요
        self.log("train_loss", loss)
        return loss

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

autoencoder = LitAutoEncoder(encoder, decoder)

Trainer로 학습 실행하기

모델과 데이터셋을 준비한 뒤엔 Trainer가 나머지를 처리해요.

trainer = L.Trainer(limit_train_batches=100, max_epochs=1)
trainer.fit(model=autoencoder, train_dataloaders=train_loader)

데이터셋은 어떤 iterable이든 지원돼요. PyTorch의 DataLoader, numpy 배열 등 모두 가능해요.

dataset = MNIST(os.getcwd(), download=True, transform=ToTensor())
train_loader = utils.data.DataLoader(dataset)

학습 후 사용하기

학습이 끝나면 체크포인트를 로드해 추론에 쓰거나, ONNX·TorchScript로 내보낼 수 있어요.

checkpoint = "./lightning_logs/version_0/checkpoints/epoch=0-step=100.ckpt"
autoencoder = LitAutoEncoder.load_from_checkpoint(checkpoint, encoder=encoder, decoder=decoder)

encoder = autoencoder.encoder
encoder.eval()

fake_image_batch = torch.rand(4, 28 * 28, device=autoencoder.device)
embeddings = encoder(fake_image_batch)

스케일 확장하기

Trainer의 인자만 바꾸면 멀티 GPU나 1TB 이상의 거대 모델 학습까지 코드 수정 없이 가능해요.

# 4 GPU 학습
trainer = L.Trainer(devices=4, accelerator="gpu")

# DeepSpeed/FSDP로 1TB+ 파라미터 모델 학습
trainer = L.Trainer(devices=4, accelerator="gpu",
                    strategy="deepspeed_stage_2", precision=16)

PyTorch Lightning