🎯 PyTorch Lightning, 왜 필요한 걸까?

딥러닝 모델을 만들다 보면 정말 골치 아픈 순간들이 있어. 학습 루프 작성하고, GPU 설정하고, 체크포인트 저장하고, 로깅 설정하고... 😫 코드가 점점 스파게티처럼 엉켜가는 걸 느낄 때가 있지 않아?

바로 이럴 때 PyTorch Lightning이 등장해! PyTorch의 유연성은 그대로 유지하면서, 반복적인 보일러플레이트 코드를 확 줄여주는 고마운 프레임워크야. 마치 요리할 때 미리 손질된 재료를 쓰는 것처럼, 핵심 로직에만 집중할 수 있게 해주지. 🎨

실제로 재능넷 같은 플랫폼에서 AI 관련 프로젝트를 진행하는 개발자들도 PyTorch Lightning을 활용해서 더 빠르고 효율적으로 작업하고 있어. 코드 품질이 올라가니까 협업도 훨씬 수월해지거든!

🏗️ LightningModule: 모델의 핵심 구조

PyTorch Lightning의 가장 중요한 개념이 바로 LightningModule이야. 이건 기존 PyTorch의 nn.Module을 상속받아서 확장한 거라서, 기존에 작성한 PyTorch 코드를 거의 그대로 사용할 수 있어. 😊

LightningModule은 크게 5가지 핵심 메서드로 구성돼:

실제 코드로 보면 훨씬 이해가 쉬워. 간단한 이미지 분류 모델을 만들어볼게! 🖼️

import torch
import torch.nn as nn
import torch.nn.functional as F
import pytorch_lightning as pl
from torchmetrics import Accuracy

class ImageClassifier(pl.LightningModule):
    def __init__(self, num_classes=10, learning_rate=1e-3):
        super().__init__()
        # 하이퍼파라미터 자동 저장
        self.save_hyperparameters()
        
        # 모델 아키텍처 정의
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(128 * 4 * 4, 512)
        self.fc2 = nn.Linear(512, num_classes)
        self.dropout = nn.Dropout(0.5)
        
        # 메트릭 정의
        self.train_accuracy = Accuracy(task='multiclass', num_classes=num_classes)
        self.val_accuracy = Accuracy(task='multiclass', num_classes=num_classes)
        
    def forward(self, x):
        # 순전파 로직
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = x.view(-1, 128 * 4 * 4)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x
    
    def training_step(self, batch, batch_idx):
        # 학습 단계
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        preds = torch.argmax(logits, dim=1)
        
        # 메트릭 계산
        self.train_accuracy(preds, y)
        
        # 로깅 (자동으로 TensorBoard 등에 기록됨)
        self.log('train_loss', loss, prog_bar=True)
        self.log('train_acc', self.train_accuracy, prog_bar=True)
        
        return loss
    
    def validation_step(self, batch, batch_idx):
        # 검증 단계
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        preds = torch.argmax(logits, dim=1)
        
        self.val_accuracy(preds, y)
        
        self.log('val_loss', loss, prog_bar=True)
        self.log('val_acc', self.val_accuracy, prog_bar=True)
        
        return loss
    
    def configure_optimizers(self):
        # 옵티마이저와 스케줄러 설정
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, 
            mode='min', 
            factor=0.5, 
            patience=5
        )
        return {
            'optimizer': optimizer,
            'lr_scheduler': {
                'scheduler': scheduler,
                'monitor': 'val_loss'
            }
        }

위 코드를 보면 알겠지만, 학습 루프를 직접 작성할 필요가 없어. training_step과 validation_step만 정의하면 Lightning이 알아서 에폭을 돌리고, 배치를 처리하고, GPU로 데이터를 옮기고, 그래디언트를 계산해줘. 마법 같지? ✨

🎮 Trainer: 학습의 지휘자

Trainer는 PyTorch Lightning의 핵심 엔진이야. 모든 학습 과정을 자동화하고 관리하는 역할을 하지. GPU 설정부터 분산 학습, 체크포인트 저장까지 모든 걸 처리해줘. 🎪

Trainer의 가장 큰 장점은 설정만 바꾸면 코드 수정 없이 다양한 환경에서 실행할 수 있다는 거야. 노트북에서 개발하다가 서버로 옮길 때도, 단일 GPU에서 멀티 GPU로 확장할 때도 코드를 거의 건드리지 않아도 돼!

from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping, LearningRateMonitor

# 콜백 설정
checkpoint_callback = ModelCheckpoint(
    dirpath='checkpoints/',
    filename='model-{epoch:02d}-{val_loss:.2f}',
    save_top_k=3,  # 상위 3개 모델만 저장
    monitor='val_loss',
    mode='min'
)

early_stop_callback = EarlyStopping(
    monitor='val_loss',
    patience=10,
    mode='min',
    verbose=True
)

lr_monitor = LearningRateMonitor(logging_interval='epoch')

# Trainer 초기화
trainer = Trainer(
    max_epochs=100,
    accelerator='gpu',  # 'cpu', 'gpu', 'tpu' 등 선택 가능
    devices=1,  # 사용할 GPU 개수
    precision=16,  # Mixed Precision Training (16-bit)
    callbacks=[checkpoint_callback, early_stop_callback, lr_monitor],
    log_every_n_steps=10,
    gradient_clip_val=1.0,  # 그래디언트 클리핑
    deterministic=True,  # 재현 가능한 결과
    enable_progress_bar=True
)

# 모델 초기화
model = ImageClassifier(num_classes=10, learning_rate=1e-3)

# 학습 시작 (이게 전부!)
trainer.fit(model, train_dataloader, val_dataloader)

특히 멀티 GPU 학습으로 확장하고 싶을 때는 정말 간단해:

# 단일 GPU
trainer = Trainer(accelerator='gpu', devices=1)

# 멀티 GPU (DDP)
trainer = Trainer(accelerator='gpu', devices=4, strategy='ddp')

# 특정 GPU 선택
trainer = Trainer(accelerator='gpu', devices=[0, 2, 3])

# TPU 사용
trainer = Trainer(accelerator='tpu', devices=8)

코드 한 줄만 바꾸면 되는 거야! 이게 바로 PyTorch Lightning의 진정한 파워지. 💪

📊 DataModule: 데이터 관리의 정석

데이터 로딩 코드가 여기저기 흩어져 있으면 정말 관리하기 힘들어. LightningDataModule은 데이터 준비, 전처리, 로딩을 한 곳에 깔끔하게 정리할 수 있게 해줘. 🗂️

DataModule은 5가지 주요 단계로 구성돼:

실제 예제를 보면 훨씬 명확해져:

import pytorch_lightning as pl
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms
import os

class CIFAR10DataModule(pl.LightningDataModule):
    def __init__(self, data_dir='./data', batch_size=32, num_workers=4):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.num_workers = num_workers
        
        # 데이터 증강 및 정규화
        self.transform_train = transforms.Compose([
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), 
                               (0.2023, 0.1994, 0.2010))
        ])
        
        self.transform_test = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.4914, 0.4822, 0.4465), 
                               (0.2023, 0.1994, 0.2010))
        ])
    
    def prepare_data(self):
        # 데이터 다운로드 (단일 프로세스에서만 실행)
        datasets.CIFAR10(self.data_dir, train=True, download=True)
        datasets.CIFAR10(self.data_dir, train=False, download=True)
    
    def setup(self, stage=None):
        # 데이터셋 분할 (각 프로세스에서 실행)
        if stage == 'fit' or stage is None:
            cifar_full = datasets.CIFAR10(
                self.data_dir, 
                train=True, 
                transform=self.transform_train
            )
            # 학습/검증 분할 (45000/5000)
            self.cifar_train, self.cifar_val = random_split(
                cifar_full, [45000, 5000]
            )
        
        if stage == 'test' or stage is None:
            self.cifar_test = datasets.CIFAR10(
                self.data_dir, 
                train=False, 
                transform=self.transform_test
            )
    
    def train_dataloader(self):
        return DataLoader(
            self.cifar_train,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.num_workers,
            pin_memory=True,
            persistent_workers=True
        )
    
    def val_dataloader(self):
        return DataLoader(
            self.cifar_val,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers,
            pin_memory=True,
            persistent_workers=True
        )
    
    def test_dataloader(self):
        return DataLoader(
            self.cifar_test,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers,
            pin_memory=True
        )

# 사용법
data_module = CIFAR10DataModule(batch_size=128, num_workers=8)
model = ImageClassifier(num_classes=10)

trainer = Trainer(max_epochs=50, accelerator='gpu', devices=1)
trainer.fit(model, data_module)
trainer.test(model, data_module)

pin_memory=True와 persistent_workers=True 옵션은 GPU 학습 시 데이터 로딩 속도를 크게 향상시켜줘. 특히 persistent_workers는 에폭마다 워커를 재시작하지 않아서 오버헤드가 줄어들어! 🚄

🎨 Callbacks: 학습 과정 커스터마이징

Callbacks는 학습 과정의 특정 시점에 자동으로 실행되는 함수들이야. 체크포인트 저장, 조기 종료, 학습률 조정 등 다양한 기능을 모듈화해서 쉽게 추가할 수 있어. 🎯

PyTorch Lightning은 기본적으로 많은 유용한 콜백을 제공하고, 필요하면 직접 만들 수도 있어!

from pytorch_lightning.callbacks import (
    ModelCheckpoint, 
    EarlyStopping, 
    LearningRateMonitor,
    StochasticWeightAveraging,
    RichProgressBar,
    DeviceStatsMonitor
)

# 다양한 콜백 설정
callbacks = [
    # 최고 성능 모델 저장
    ModelCheckpoint(
        dirpath='checkpoints/',
        filename='best-{epoch:02d}-{val_acc:.3f}',
        monitor='val_acc',
        mode='max',
        save_top_k=1,
        save_last=True
    ),
    
    # 여러 메트릭 기준으로 저장
    ModelCheckpoint(
        dirpath='checkpoints/loss/',
        filename='loss-{epoch:02d}-{val_loss:.3f}',
        monitor='val_loss',
        mode='min',
        save_top_k=3
    ),
    
    # 조기 종료
    EarlyStopping(
        monitor='val_loss',
        patience=15,
        mode='min',
        verbose=True,
        min_delta=0.001  # 최소 개선 폭
    ),
    
    # 학습률 모니터링
    LearningRateMonitor(logging_interval='step'),
    
    # Stochastic Weight Averaging
    StochasticWeightAveraging(swa_lrs=1e-2),
    
    # 예쁜 진행 표시줄
    RichProgressBar(),
    
    # 디바이스 통계
    DeviceStatsMonitor()
]

trainer = Trainer(
    max_epochs=100,
    callbacks=callbacks,
    accelerator='gpu',
    devices=1
)

이제 커스텀 콜백을 만드는 방법을 알아볼게. 예를 들어, 특정 에폭마다 샘플 예측 결과를 이미지로 저장하는 콜백을 만들어보자! 📸

import pytorch_lightning as pl
from pytorch_lightning.callbacks import Callback
import matplotlib.pyplot as plt
import torch
import os

class VisualizationCallback(Callback):
    def __init__(self, save_dir='visualizations', num_samples=8):
        super().__init__()
        self.save_dir = save_dir
        self.num_samples = num_samples
        os.makedirs(save_dir, exist_ok=True)
    
    def on_validation_epoch_end(self, trainer, pl_module):
        # 검증 에폭이 끝날 때마다 실행
        if trainer.current_epoch % 5 != 0:  # 5 에폭마다만 실행
            return
        
        # 검증 데이터에서 샘플 가져오기
        val_dataloader = trainer.val_dataloaders[0]
        batch = next(iter(val_dataloader))
        images, labels = batch
        images = images[:self.num_samples].to(pl_module.device)
        labels = labels[:self.num_samples]
        
        # 예측
        pl_module.eval()
        with torch.no_grad():
            logits = pl_module(images)
            preds = torch.argmax(logits, dim=1)
        
        # 시각화
        fig, axes = plt.subplots(2, 4, figsize=(12, 6))
        for idx, ax in enumerate(axes.flat):
            img = images[idx].cpu().permute(1, 2, 0)
            # 정규화 해제
            img = img * torch.tensor([0.2023, 0.1994, 0.2010])
            img = img + torch.tensor([0.4914, 0.4822, 0.4465])
            img = torch.clamp(img, 0, 1)
            
            ax.imshow(img)
            ax.set_title(f'True: {labels[idx]}, Pred: {preds[idx].item()}')
            ax.axis('off')
        
        plt.tight_layout()
        plt.savefig(f'{self.save_dir}/epoch_{trainer.current_epoch}.png')
        plt.close()
        
        pl_module.train()

# 사용법
vis_callback = VisualizationCallback(save_dir='predictions', num_samples=8)
trainer = Trainer(callbacks=[vis_callback, ...])

실무에서 자주 사용하는 또 다른 커스텀 콜백 예제를 하나 더 보여줄게. 학습 중 메모리 사용량을 추적하는 콜백이야:

import torch
import psutil
from pytorch_lightning.callbacks import Callback

class MemoryMonitorCallback(Callback):
    def __init__(self):
        super().__init__()
        self.process = psutil.Process()
    
    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        # GPU 메모리
        if torch.cuda.is_available():
            gpu_memory = torch.cuda.memory_allocated() / 1024**3  # GB
            gpu_memory_cached = torch.cuda.memory_reserved() / 1024**3
            pl_module.log('gpu_memory_allocated_gb', gpu_memory)
            pl_module.log('gpu_memory_cached_gb', gpu_memory_cached)
        
        # CPU 메모리
        cpu_memory = self.process.memory_info().rss / 1024**3  # GB
        pl_module.log('cpu_memory_gb', cpu_memory)
    
    def on_train_epoch_end(self, trainer, pl_module):
        # 에폭 종료 시 GPU 캐시 정리
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

# 사용
memory_callback = MemoryMonitorCallback()
trainer = Trainer(callbacks=[memory_callback])

이런 식으로 콜백을 활용하면 학습 과정을 모니터링하고 제어하는 게 정말 쉬워져. 재능넷에서 AI 프로젝트를 진행하는 개발자들도 이런 콜백들을 활용해서 더 안정적이고 효율적인 학습 파이프라인을 구축하고 있어! 🎓

⚡ Mixed Precision Training: 속도와 메모리 최적화

Mixed Precision Training은 16-bit 부동소수점 연산을 활용해서 학습 속도를 높이고 메모리 사용량을 줄이는 기법이야. PyTorch Lightning에서는 정말 간단하게 적용할 수 있어! 🚀

# 16-bit Mixed Precision
trainer = Trainer(
    precision=16,  # 또는 '16-mixed'
    accelerator='gpu',
    devices=1
)

# bfloat16 (더 안정적, Ampere 이상 GPU에서 지원)
trainer = Trainer(
    precision='bf16',
    accelerator='gpu',
    devices=1
)

# 32-bit (기본값)
trainer = Trainer(
    precision=32,
    accelerator='gpu',
    devices=1
)

하지만 주의할 점도 있어:

만약 수치적 불안정성 문제가 발생하면, 특정 레이어만 32-bit로 유지할 수 있어:

class MixedPrecisionModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(100, 50)
        self.layer2 = nn.Linear(50, 10)
        
        # 특정 레이어를 32-bit로 유지
        self.layer2.float()
    
    def forward(self, x):
        x = self.layer1(x)
        # 명시적으로 32-bit로 변환
        x = x.float()
        x = self.layer2(x)
        return x

🔄 Gradient Accumulation: 큰 배치 효과 내기

GPU 메모리가 부족해서 큰 배치 크기를 사용할 수 없을 때, Gradient Accumulation을 사용하면 작은 배치를 여러 번 누적해서 큰 배치와 동일한 효과를 낼 수 있어! 📊

# 배치 크기 32를 4번 누적 = 실질적으로 배치 크기 128
trainer = Trainer(
    accumulate_grad_batches=4,
    max_epochs=100
)

# 에폭에 따라 동적으로 변경
trainer = Trainer(
    accumulate_grad_batches={
        0: 1,   # 처음 5 에폭은 누적 없음
        5: 2,   # 5~10 에폭은 2배 누적
        10: 4   # 10 에폭 이후는 4배 누적
    }
)

Gradient Accumulation을 사용할 때는 학습률 조정이 필요할 수 있어. 실질적인 배치 크기가 커지니까 학습률도 그에 맞게 조정해야 해:

class MyModel(pl.LightningModule):
    def __init__(self, base_lr=1e-3, accumulate_batches=4):
        super().__init__()
        self.base_lr = base_lr
        self.accumulate_batches = accumulate_batches
        # ... 모델 정의 ...
    
    def configure_optimizers(self):
        # 누적 배치에 맞춰 학습률 조정
        adjusted_lr = self.base_lr * self.accumulate_batches
        optimizer = torch.optim.Adam(self.parameters(), lr=adjusted_lr)
        return optimizer

🌐 분산 학습: 멀티 GPU/멀티 노드

PyTorch Lightning의 진정한 파워는 분산 학습에서 빛을 발해. 코드 한 줄만 바꾸면 단일 GPU에서 멀티 GPU, 심지어 멀티 노드 클러스터로 확장할 수 있어! 🌟

# 단일 GPU
trainer = Trainer(accelerator='gpu', devices=1)

# 멀티 GPU - DDP (권장)
trainer = Trainer(
    accelerator='gpu',
    devices=4,
    strategy='ddp'
)

# 멀티 GPU - DDP Spawn (디버깅용)
trainer = Trainer(
    accelerator='gpu',
    devices=4,
    strategy='ddp_spawn'
)

# FSDP (매우 큰 모델)
trainer = Trainer(
    accelerator='gpu',
    devices=8,
    strategy='fsdp',
    precision=16
)

# DeepSpeed Stage 2
from pytorch_lightning.strategies import DeepSpeedStrategy

trainer = Trainer(
    accelerator='gpu',
    devices=8,
    strategy=DeepSpeedStrategy(stage=2),
    precision=16
)

# 멀티 노드
trainer = Trainer(
    accelerator='gpu',
    devices=8,
    num_nodes=4,  # 4개 노드
    strategy='ddp'
)

분산 학습을 사용할 때 주의해야 할 점들이 있어:

import pytorch_lightning as pl

class DistributedModel(pl.LightningModule):
    def __init__(self, base_lr=1e-3):
        super().__init__()
        self.base_lr = base_lr
        # ... 모델 정의 ...
    
    def configure_optimizers(self):
        # GPU 개수에 맞춰 학습률 조정
        num_gpus = self.trainer.num_devices
        adjusted_lr = self.base_lr * num_gpus
        
        optimizer = torch.optim.SGD(
            self.parameters(),
            lr=adjusted_lr,
            momentum=0.9,
            weight_decay=1e-4
        )
        
        # Warmup + Cosine Annealing
        scheduler = torch.optim.lr_scheduler.OneCycleLR(
            optimizer,
            max_lr=adjusted_lr,
            total_steps=self.trainer.estimated_stepping_batches,
            pct_start=0.1  # 10% warmup
        )
        
        return {
            'optimizer': optimizer,
            'lr_scheduler': {
                'scheduler': scheduler,
                'interval': 'step'
            }
        }
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        
        # 모든 GPU에서 동기화된 메트릭
        self.log('train_loss', loss, sync_dist=True)
        
        return loss

# 사용
pl.seed_everything(42)  # 재현성을 위한 시드 고정

trainer = Trainer(
    accelerator='gpu',
    devices=4,
    strategy='ddp',
    precision=16,
    max_epochs=100
)

📈 로깅과 실험 추적

모델 학습 과정을 제대로 추적하고 분석하는 건 정말 중요해. PyTorch Lightning은 다양한 로깅 프레임워크와 쉽게 통합돼! 📊

from pytorch_lightning.loggers import TensorBoardLogger, WandbLogger, CSVLogger

# TensorBoard
tensorboard_logger = TensorBoardLogger(
    save_dir='logs/',
    name='my_model',
    version='v1'
)

# Weights & Biases
wandb_logger = WandbLogger(
    project='image-classification',
    name='resnet50-experiment',
    log_model=True  # 모델 체크포인트도 업로드
)

# CSV Logger
csv_logger = CSVLogger(
    save_dir='logs/',
    name='csv_logs'
)

# 여러 로거 동시 사용
trainer = Trainer(
    logger=[tensorboard_logger, wandb_logger, csv_logger],
    max_epochs=100
)

모델 내에서 다양한 방식으로 로깅할 수 있어:

class MyModel(pl.LightningModule):
    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        
        # 스칼라 로깅
        self.log('train_loss', loss)
        
        # 여러 메트릭 한 번에
        self.log_dict({
            'train_loss': loss,
            'train_acc': accuracy,
            'learning_rate': self.optimizers().param_groups[0]['lr']
        })
        
        # 진행 표시줄에 표시
        self.log('train_loss', loss, prog_bar=True)
        
        # 에폭 단위로 집계
        self.log('train_loss', loss, on_step=False, on_epoch=True)
        
        # 분산 학습 시 동기화
        self.log('train_loss', loss, sync_dist=True)
        
        return loss
    
    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        
        # 이미지 로깅 (TensorBoard/Wandb)
        if batch_idx == 0:
            # 첫 배치의 이미지만 로깅
            self.logger.experiment.add_images(
                'val_images',
                x[:8],
                self.current_epoch
            )
        
        # 히스토그램 로깅
        for name, param in self.named_parameters():
            self.logger.experiment.add_histogram(
                name,
                param,
                self.current_epoch
            )
        
        return loss

Weights & Biases를 사용하면 더 풍부한 로깅이 가능해:

import wandb
from pytorch_lightning.loggers import WandbLogger

class AdvancedModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        # ... 모델 정의 ...
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        
        # 기본 로깅
        self.log('train_loss', loss)
        
        # Confusion Matrix 로깅 (검증 시)
        if batch_idx % 100 == 0:
            preds = torch.argmax(logits, dim=1)
            self.logger.experiment.log({
                'confusion_matrix': wandb.plot.confusion_matrix(
                    probs=None,
                    y_true=y.cpu().numpy(),
                    preds=preds.cpu().numpy(),
                    class_names=['class_0', 'class_1', ...]
                )
            })
        
        return loss
    
    def on_train_epoch_end(self):
        # 에폭 종료 시 커스텀 차트
        self.logger.experiment.log({
            'custom_metric': wandb.plot.line_series(
                xs=range(len(self.train_losses)),
                ys=[self.train_losses, self.val_losses],
                keys=['train', 'val'],
                title='Loss Curves',
                xname='Step'
            )
        })

# 사용
wandb_logger = WandbLogger(
    project='my-project',
    config={
        'learning_rate': 1e-3,
        'batch_size': 32,
        'architecture': 'ResNet50'
    }
)

trainer = Trainer(logger=wandb_logger)

💾 체크포인트와 모델 저장

학습 중간에 모델을 저장하고, 나중에 이어서 학습하거나 추론에 사용하는 건 필수야. PyTorch Lightning은 이 과정을 정말 쉽게 만들어줘! 💪

from pytorch_lightning.callbacks import ModelCheckpoint

# 최고 성능 모델 저장
checkpoint_callback = ModelCheckpoint(
    dirpath='checkpoints/',
    filename='best-{epoch:02d}-{val_loss:.2f}',
    monitor='val_loss',
    mode='min',
    save_top_k=3,  # 상위 3개만 저장
    save_last=True,  # 마지막 체크포인트도 저장
    save_weights_only=False,  # 전체 상태 저장
    every_n_epochs=1  # 매 에폭마다 체크
)

trainer = Trainer(callbacks=[checkpoint_callback])
trainer.fit(model, datamodule)

# 저장된 체크포인트에서 모델 로드
model = MyModel.load_from_checkpoint('checkpoints/best-epoch=10-val_loss=0.25.ckpt')

# 학습 재개
trainer = Trainer(resume_from_checkpoint='checkpoints/last.ckpt')
trainer.fit(model, datamodule)

체크포인트에는 다음 정보들이 모두 저장돼:

커스텀 체크포인트 로직을 추가할 수도 있어:

class MyModel(pl.LightningModule):
    def on_save_checkpoint(self, checkpoint):
        # 체크포인트 저장 시 추가 정보 저장
        checkpoint['custom_data'] = {
            'best_metric': self.best_metric,
            'training_history': self.history,
            'special_config': self.config
        }
    
    def on_load_checkpoint(self, checkpoint):
        # 체크포인트 로드 시 추가 정보 복원
        if 'custom_data' in checkpoint:
            self.best_metric = checkpoint['custom_data']['best_metric']
            self.history = checkpoint['custom_data']['training_history']
            self.config = checkpoint['custom_data']['special_config']

프로덕션 환경에서는 모델을 더 가볍게 저장하고 싶을 수 있어:

# 가중치만 저장 (용량 절약)
checkpoint_callback = ModelCheckpoint(
    save_weights_only=True,
    dirpath='weights/',
    filename='model-weights'
)

# ONNX 형식으로 내보내기
model = MyModel.load_from_checkpoint('checkpoints/best.ckpt')
input_sample = torch.randn(1, 3, 224, 224)
model.to_onnx('model.onnx', input_sample, export_params=True)

# TorchScript로 변환
script = model.to_torchscript()
torch.jit.save(script, 'model.pt')

🧪 하이퍼파라미터 튜닝

최적의 하이퍼파라미터를 찾는 건 정말 중요하지만 시간이 많이 걸려. PyTorch Lightning은 다양한 하이퍼파라미터 튜닝 도구와 통합돼 있어! 🎯

from pytorch_lightning.tuner import Tuner

# 자동 배치 크기 찾기
trainer = Trainer(auto_scale_batch_size='binsearch')
tuner = Tuner(trainer)

# 최적 배치 크기 찾기
tuner.scale_batch_size(model, datamodule=datamodule)

# 자동 학습률 찾기
trainer = Trainer(auto_lr_find=True)
tuner.lr_find(model, datamodule=datamodule)

# 결과 시각화
fig = tuner.lr_find(model, datamodule=datamodule).plot(suggest=True)
fig.show()

Optuna를 사용한 고급 하이퍼파라미터 튜닝:

import optuna
from optuna.integration import PyTorchLightningPruningCallback

def objective(trial):
    # 하이퍼파라미터 샘플링
    lr = trial.suggest_loguniform('lr', 1e-5, 1e-1)
    batch_size = trial.suggest_categorical('batch_size', [16, 32, 64, 128])
    dropout = trial.suggest_uniform('dropout', 0.1, 0.5)
    hidden_dim = trial.suggest_categorical('hidden_dim', [128, 256, 512])
    
    # 모델 생성
    model = MyModel(
        learning_rate=lr,
        dropout=dropout,
        hidden_dim=hidden_dim
    )
    
    # 데이터 모듈
    datamodule = MyDataModule(batch_size=batch_size)
    
    # Trainer 설정
    trainer = Trainer(
        max_epochs=30,
        accelerator='gpu',
        devices=1,
        callbacks=[
            PyTorchLightningPruningCallback(trial, monitor='val_loss')
        ],
        enable_progress_bar=False,
        logger=False
    )
    
    # 학습
    trainer.fit(model, datamodule)
    
    return trainer.callback_metrics['val_loss'].item()

# Optuna 스터디 실행
study = optuna.create_study(direction='minimize')
study.optimize(objective, n_trials=100, timeout=3600)

print('Best trial:')
print(f'  Value: {study.best_trial.value}')
print('  Params: ')
for key, value in study.best_trial.params.items():
    print(f'    {key}: {value}')

Ray Tune을 사용한 분산 하이퍼파라미터 튜닝:

from ray import tune
from ray.tune.integration.pytorch_lightning import TuneReportCallback

def train_model(config):
    model = MyModel(
        learning_rate=config['lr'],
        dropout=config['dropout']
    )
    
    datamodule = MyDataModule(batch_size=config['batch_size'])
    
    trainer = Trainer(
        max_epochs=30,
        callbacks=[
            TuneReportCallback(
                metrics=['val_loss', 'val_acc'],
                on='validation_end'
            )
        ]
    )
    
    trainer.fit(model, datamodule)

# 검색 공간 정의
config = {
    'lr': tune.loguniform(1e-5, 1e-1),
    'batch_size': tune.choice([16, 32, 64, 128]),
    'dropout': tune.uniform(0.1, 0.5)
}

# Ray Tune 실행
analysis = tune.run(
    train_model,
    config=config,
    num_samples=50,
    resources_per_trial={'gpu': 1},
    metric='val_loss',
    mode='min'
)

print('Best config:', analysis.best_config)

🎓 실전 예제: 완전한 학습 파이프라인

지금까지 배운 모든 내용을 종합해서 실전에서 바로 사용할 수 있는 완전한 파이프라인을 만들어볼게! 🚀

import torch
import torch.nn as nn
import torch.nn.functional as F
import pytorch_lightning as pl
from torch.utils.data import DataLoader, random_split
from torchvision import datasets, transforms, models
from torchmetrics import Accuracy, F1Score, ConfusionMatrix
from pytorch_lightning.callbacks import (
    ModelCheckpoint, EarlyStopping, LearningRateMonitor,
    RichProgressBar, DeviceStatsMonitor
)
from pytorch_lightning.loggers import WandbLogger
import wandb

# ==================== 데이터 모듈 ====================
class ImageDataModule(pl.LightningDataModule):
    def __init__(
        self,
        data_dir='./data',
        batch_size=32,
        num_workers=4,
        image_size=224
    ):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.image_size = image_size
        
        # 데이터 증강
        self.train_transform = transforms.Compose([
            transforms.RandomResizedCrop(image_size),
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(15),
            transforms.ColorJitter(brightness=0.2, contrast=0.2),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])
        
        self.val_transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(image_size),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])
    
    def prepare_data(self):
        # 데이터 다운로드
        datasets.CIFAR10(self.data_dir, train=True, download=True)
        datasets.CIFAR10(self.data_dir, train=False, download=True)
    
    def setup(self, stage=None):
        if stage == 'fit' or stage is None:
            full_dataset = datasets.CIFAR10(
                self.data_dir,
                train=True,
                transform=self.train_transform
            )
            self.train_dataset, self.val_dataset = random_split(
                full_dataset, [45000, 5000]
            )
        
        if stage == 'test' or stage is None:
            self.test_dataset = datasets.CIFAR10(
                self.data_dir,
                train=False,
                transform=self.val_transform
            )
    
    def train_dataloader(self):
        return DataLoader(
            self.train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.num_workers,
            pin_memory=True,
            persistent_workers=True
        )
    
    def val_dataloader(self):
        return DataLoader(
            self.val_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers,
            pin_memory=True,
            persistent_workers=True
        )
    
    def test_dataloader(self):
        return DataLoader(
            self.test_dataset,
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.num_workers,
            pin_memory=True
        )

# ==================== 모델 ====================
class ImageClassifier(pl.LightningModule):
    def __init__(
        self,
        num_classes=10,
        learning_rate=1e-3,
        weight_decay=1e-4,
        pretrained=True
    ):
        super().__init__()
        self.save_hyperparameters()
        
        # ResNet50 백본
        self.backbone = models.resnet50(pretrained=pretrained)
        num_features = self.backbone.fc.in_features
        self.backbone.fc = nn.Linear(num_features, num_classes)
        
        # 메트릭
        self.train_acc = Accuracy(task='multiclass', num_classes=num_classes)
        self.val_acc = Accuracy(task='multiclass', num_classes=num_classes)
        self.test_acc = Accuracy(task='multiclass', num_classes=num_classes)
        self.val_f1 = F1Score(task='multiclass', num_classes=num_classes)
        
        # 학습 기록
        self.training_step_outputs = []
        self.validation_step_outputs = []
    
    def forward(self, x):
        return self.backbone(x)
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        preds = torch.argmax(logits, dim=1)
        
        # 메트릭 계산
        self.train_acc(preds, y)
        
        # 로깅
        self.log('train_loss', loss, prog_bar=True, on_step=True, on_epoch=True)
        self.log('train_acc', self.train_acc, prog_bar=True, on_epoch=True)
        
        self.training_step_outputs.append(loss)
        return loss
    
    def validation_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        preds = torch.argmax(logits, dim=1)
        
        # 메트릭 계산
        self.val_acc(preds, y)
        self.val_f1(preds, y)
        
        # 로깅
        self.log('val_loss', loss, prog_bar=True, sync_dist=True)
        self.log('val_acc', self.val_acc, prog_bar=True, sync_dist=True)
        self.log('val_f1', self.val_f1, sync_dist=True)
        
        self.validation_step_outputs.append({
            'loss': loss,
            'preds': preds,
            'targets': y
        })
        
        return loss
    
    def test_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        preds = torch.argmax(logits, dim=1)
        
        self.test_acc(preds, y)
        self.log('test_loss', loss)
        self.log('test_acc', self.test_acc)
        
        return loss
    
    def on_train_epoch_end(self):
        # 에폭 종료 시 처리
        avg_loss = torch.stack(self.training_step_outputs).mean()
        self.log('train_loss_epoch', avg_loss)
        self.training_step_outputs.clear()
    
    def on_validation_epoch_end(self):
        # 검증 에폭 종료 시 처리
        avg_loss = torch.stack([x['loss'] for x in self.validation_step_outputs]).mean()
        self.log('val_loss_epoch', avg_loss)
        self.validation_step_outputs.clear()
    
    def configure_optimizers(self):
        # 옵티마이저
        optimizer = torch.optim.AdamW(
            self.parameters(),
            lr=self.hparams.learning_rate,
            weight_decay=self.hparams.weight_decay
        )
        
        # 스케줄러
        scheduler = torch.optim.lr_scheduler.OneCycleLR(
            optimizer,
            max_lr=self.hparams.learning_rate,
            total_steps=self.trainer.estimated_stepping_batches,
            pct_start=0.1,
            anneal_strategy='cos'
        )
        
        return {
            'optimizer': optimizer,
            'lr_scheduler': {
                'scheduler': scheduler,
                'interval': 'step'
            }
        }

# ==================== 학습 실행 ====================
def main():
    # 시드 고정
    pl.seed_everything(42, workers=True)
    
    # 데이터 모듈
    datamodule = ImageDataModule(
        batch_size=128,
        num_workers=8,
        image_size=224
    )
    
    # 모델
    model = ImageClassifier(
        num_classes=10,
        learning_rate=1e-3,
        pretrained=True
    )
    
    # 로거
    wandb_logger = WandbLogger(
        project='cifar10-classification',
        name='resnet50-experiment',
        log_model=True
    )
    
    # 콜백
    callbacks = [
        ModelCheckpoint(
            dirpath='checkpoints/',
            filename='best-{epoch:02d}-{val_acc:.3f}',
            monitor='val_acc',
            mode='max',
            save_top_k=3,
            save_last=True
        ),
        EarlyStopping(
            monitor='val_loss',
            patience=10,
            mode='min',
            verbose=True
        ),
        LearningRateMonitor(logging_interval='step'),
        RichProgressBar(),
        DeviceStatsMonitor()
    ]
    
    # Trainer
    trainer = pl.Trainer(
        max_epochs=100,
        accelerator='gpu',
        devices=2,  # 멀티 GPU
        strategy='ddp',
        precision=16,  # Mixed Precision
        callbacks=callbacks,
        logger=wandb_logger,
        gradient_clip_val=1.0,
        accumulate_grad_batches=2,
        deterministic=True,
        log_every_n_steps=10
    )
    
    # 학습
    trainer.fit(model, datamodule)
    
    # 테스트
    trainer.test(model, datamodule, ckpt_path='best')
    
    # 최종 모델 저장
    trainer.save_checkpoint('final_model.ckpt')
    
    # Wandb 종료
    wandb.finish()

if __name__ == '__main__':
    main()

이 코드는 실전에서 바로 사용할 수 있는 완전한 파이프라인이야. 데이터 로딩, 모델 정의, 학습, 검증, 테스트, 로깅, 체크포인트 저장까지 모든 게 포함되어 있어! 😊

🔧 디버깅과 문제 해결

실전에서 모델을 학습하다 보면 다양한 문제에 부딪히게 돼. PyTorch Lightning은 디버깅을 쉽게 만들어주는 여러 기능을 제공해! 🐛

# 빠른 개발 모드 (소량 데이터로 빠르게 테스트)
trainer = Trainer(
    fast_dev_run=True  # 각 단계를 1배치씩만 실행
)

# 오버피팅 테스트 (단일 배치로 오버피팅 확인)
trainer = Trainer(
    overfit_batches=1  # 1개 배치만 반복 학습
)

# 일부 데이터만 사용
trainer = Trainer(
    limit_train_batches=0.1,  # 학습 데이터의 10%만
    limit_val_batches=0.1,    # 검증 데이터의 10%만
    limit_test_batches=0.1    # 테스트 데이터의 10%만
)

# 검증 주기 조정
trainer = Trainer(
    val_check_interval=0.25,  # 에폭의 25%마다 검증
    # 또는
    val_check_interval=100    # 100 스텝마다 검증
)

# 그래디언트 추적
trainer = Trainer(
    track_grad_norm=2,  # L2 노름 추적
    log_every_n_steps=1  # 매 스텝마다 로깅
)

# 프로파일링
from pytorch_lightning.profilers import SimpleProfiler, AdvancedProfiler

trainer = Trainer(
    profiler=AdvancedProfiler(dirpath='profiler_logs', filename='profile')
)

그래디언트 문제를 진단하는 유용한 콜백:

from pytorch_lightning.callbacks import Callback
import torch

class GradientMonitorCallback(Callback):
    def on_after_backward(self, trainer, pl_module):
        # 그래디언트 통계
        grad_norms = []
        for name, param in pl_module.named_parameters():
            if param.grad is not None:
                grad_norm = param.grad.norm().item()
                grad_norms.append(grad_norm)
                
                # 이상한 그래디언트 감지
                if torch.isnan(param.grad).any():
                    print(f'NaN gradient in {name}')
                if torch.isinf(param.grad).any():
                    print(f'Inf gradient in {name}')
        
        # 평균 그래디언트 노름
        if grad_norms:
            avg_grad_norm = sum(grad_norms) / len(grad_norms)
            pl_module.log('avg_grad_norm', avg_grad_norm)

# 사용
trainer = Trainer(callbacks=[GradientMonitorCallback()])

🎬 마무리하며

PyTorch Lightning으로 모델 학습 파이프라인을 구성하는 방법을 자세히 알아봤어! 처음에는 복잡해 보일 수 있지만, 한 번 익숙해지면 정말 강력한 도구가 돼. 🚀

핵심을 정리하자면:

PyTorch Lightning을 사용하면 연구 코드와 프로덕션 코드 사이의 간극을 크게 줄일 수 있어. 실험 단계에서 작성한 코드를 거의 그대로 프로덕션에 배포할 수 있거든! 💪

특히 팀으로 작업할 때 진가를 발휘해. 코드가 표준화되어 있으니까 다른 사람의 코드를 이해하기도 쉽고, 협업도 훨씬 수월해져. 실제로 많은 AI 스타트업과 연구팀들이 PyTorch Lightning을 표준으로 채택하고 있어.

재능넷 같은 플랫폼에서 AI 프로젝트를 진행할 때도, PyTorch Lightning을 활용하면 더 전문적이고 체계적인 결과물을 만들 수 있어. 클라이언트에게 깔끔한 코드와 재현 가능한 결과를 제공할 수 있으니까 신뢰도도 높아지지! 🎓

이제 여러분도 PyTorch Lightning으로 멋진 딥러닝 프로젝트를 만들어보길 바라! 처음에는 조금 낯설 수 있지만, 계속 사용하다 보면 "이거 없이 어떻게 개발했지?"라는 생각이 들 거야. 화이팅! 🔥