pytorch-patterns

PyTorch 개발 패턴

견고하고 효율적이며 재현 가능한 딥러닝 애플리케이션을 구축하기 위한 관용적인 PyTorch 패턴 및 모범 사례입니다.

활성화 시점

• 새로운 PyTorch 모델 또는 학습 스크립트를 작성할 때
• 딥러닝 코드를 리뷰할 때
• 학습 루프 또는 데이터 파이프라인을 디버깅할 때
• GPU 메모리 사용량 또는 학습 속도를 최적화할 때
• 재현 가능한 실험 설정을 구성할 때

핵심 원칙

1. 장치 독립적 코드 (Device-Agnostic Code)

장치를 하드코딩하지 않고 CPU와 GPU 모두에서 작동하는 코드를 작성하십시오.

# 좋음: 장치 독립적
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel().to(device)
data = data.to(device)

# 나쁨: 장치 하드코딩
model = MyModel().cuda()  # GPU가 없으면 크래시 발생
data = data.cuda()

2. 재현성 우선 (Reproducibility First)

재현 가능한 결과를 위해 모든 난수 시드(seed)를 설정하십시오.

# 좋음: 전체 재현성 설정
def set_seed(seed: int = 42) -> None:
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# 나쁨: 시드 제어 없음
model = MyModel()  # 실행할 때마다 가중치가 달라짐

3. 명시적인 Shape 관리

항상 텐서(tensor)의 형상(shape)을 문서화하고 검증하십시오.

# 좋음: Shape 어노테이션이 포함된 forward 패스
def forward(self, x: torch.Tensor) -> torch.Tensor:
    # x: (batch_size, channels, height, width)
    x = self.conv1(x)    # -> (batch_size, 32, H, W)
    x = self.pool(x)     # -> (batch_size, 32, H//2, W//2)
    x = x.view(x.size(0), -1)  # -> (batch_size, 32*H//2*W//2)
    return self.fc(x)    # -> (batch_size, num_classes)

# 나쁨: Shape 추적 없음
def forward(self, x):
    x = self.conv1(x)
    x = self.pool(x)
    x = x.view(x.size(0), -1)  # 크기가 어떻게 되나요?
    return self.fc(x)           # 작동은 할까요?

모델 아키텍처 패턴

깔끔한 nn.Module 구조

# 좋음: 잘 조직된 모듈
class ImageClassifier(nn.Module):
    def __init__(self, num_classes: int, dropout: float = 0.5) -> None:
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(64 * 16 * 16, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

# 나쁨: forward 안에 모든 것을 넣음
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = F.conv2d(x, weight=self.make_weight())  # 호출할 때마다 가중치를 생성!
        return x

올바른 가중치 초기화

# 좋음: 명시적 초기화
def _init_weights(self, module: nn.Module) -> None:
    if isinstance(module, nn.Linear):
        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
        if module.bias is not None:
            nn.init.zeros_(module.bias)
    elif isinstance(module, nn.Conv2d):
        nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
    elif isinstance(module, nn.BatchNorm2d):
        nn.init.ones_(module.weight)
        nn.init.zeros_(module.bias)

model = MyModel()
model.apply(model._init_weights)

학습 루프 패턴

표준 학습 루프

# 좋음: 모범 사례가 적용된 완전한 학습 루프
def train_one_epoch(
    model: nn.Module,
    dataloader: DataLoader,
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    device: torch.device,
    scaler: torch.amp.GradScaler | None = None,
) -> float:
    model.train()  # 항상 train 모드 설정
    total_loss = 0.0

    for batch_idx, (data, target) in enumerate(dataloader):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad(set_to_none=True)  # zero_grad()보다 효율적임

        # 혼합 정밀도(Mixed precision) 학습
        with torch.amp.autocast("cuda", enabled=scaler is not None):
            output = model(data)
            loss = criterion(output, target)

        if scaler is not None:
            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()

        total_loss += loss.item()

    return total_loss / len(dataloader)

검증 루프

# 좋음: 올바른 평가 방식
@torch.no_grad()  # torch.no_grad() 블록으로 감싸는 것보다 효율적임
def evaluate(
    model: nn.Module,
    dataloader: DataLoader,
    criterion: nn.Module,
    device: torch.device,
) -> tuple[float, float]:
    model.eval()  # 항상 eval 모드 설정 — 드롭아웃 비활성화, 배치 정규화 통계 사용
    total_loss = 0.0
    correct = 0
    total = 0

    for data, target in dataloader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        total_loss += criterion(output, target).item()
        correct += (output.argmax(1) == target).sum().item()
        total += target.size(0)

    return total_loss / len(dataloader), correct / total

데이터 파이프라인 패턴

커스텀 데이터셋 (Custom Dataset)

# 좋음: 타입 힌트가 포함된 깔끔한 Dataset
class ImageDataset(Dataset):
    def __init__(
        self,
        image_dir: str,
        labels: dict[str, int],
        transform: transforms.Compose | None = None,
    ) -> None:
        self.image_paths = list(Path(image_dir).glob("*.jpg"))
        self.labels = labels
        self.transform = transform

    def __len__(self) -> int:
        return len(self.image_paths)

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, int]:
        img = Image.open(self.image_paths[idx]).convert("RGB")
        label = self.labels[self.image_paths[idx].stem]

        if self.transform:
            img = self.transform(img)

        return img, label

효율적인 DataLoader 설정

# 좋음: 최적화된 DataLoader
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,            # 학습용 셔플
    num_workers=4,           # 병렬 데이터 로딩
    pin_memory=True,         # CPU->GPU 전송 속도 향상
    persistent_workers=True, # 에포크 간에 워커 유지
    drop_last=True,          # BatchNorm을 위한 일관된 배치 크기
)

# 나쁨: 느린 기본값
dataloader = DataLoader(dataset, batch_size=32)  # num_workers=0, pin_memory 미사용

가변 길이 데이터를 위한 커스텀 Collate

# 좋음: collate_fn에서 시퀀스 패딩 처리
def collate_fn(batch: list[tuple[torch.Tensor, int]]) -> tuple[torch.Tensor, torch.Tensor]:
    sequences, labels = zip(*batch)
    # 배치의 최대 길이에 맞춰 패딩
    padded = nn.utils.rnn.pad_sequence(sequences, batch_first=True, padding_value=0)
    return padded, torch.tensor(labels)

dataloader = DataLoader(dataset, batch_size=32, collate_fn=collate_fn)

체크포인트 패턴

체크포인트 저장 및 로드

# 좋음: 모든 학습 상태를 포함하는 완전한 체크포인트
def save_checkpoint(
    model: nn.Module,
    optimizer: torch.optim.Optimizer,
    epoch: int,
    loss: float,
    path: str,
) -> None:
    torch.save({
        "epoch": epoch,
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "loss": loss,
    }, path)

def load_checkpoint(
    path: str,
    model: nn.Module,
    optimizer: torch.optim.Optimizer | None = None,
) -> dict:
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    model.load_state_dict(checkpoint["model_state_dict"])
    if optimizer:
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    return checkpoint

# 나쁨: 모델 가중치만 저장 (학습 재개 불가)
torch.save(model.state_dict(), "model.pt")

성능 최적화

혼합 정밀도(Mixed Precision) 학습

# 좋음: GradScaler를 사용한 AMP
scaler = torch.amp.GradScaler("cuda")
for data, target in dataloader:
    with torch.amp.autocast("cuda"):
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad(set_to_none=True)

대형 모델을 위한 그래디언트 체크포인팅 (Gradient Checkpointing)

# 좋음: 메모리를 위해 연산량을 사용함
from torch.utils.checkpoint import checkpoint

class LargeModel(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 메모리 절약을 위해 backward 시에 활성화를 재계산함
        x = checkpoint(self.block1, x, use_reentrant=False)
        x = checkpoint(self.block2, x, use_reentrant=False)
        return self.head(x)

속도를 위한 torch.compile

# 좋음: 더 빠른 실행을 위해 모델 컴파일 (PyTorch 2.0+)
model = MyModel().to(device)
model = torch.compile(model, mode="reduce-overhead")

# 모드: "default" (안전), "reduce-overhead" (더 빠름), "max-autotune" (가장 빠름)

빠른 참조: PyTorch 관용구

관용구	설명
model.train() / model.eval()	학습/평가 전에 항상 모드 설정
torch.no_grad()	추론 시 그래디언트 계산 비활성화
optimizer.zero_grad(set_to_none=True)	더 효율적인 그래디언트 초기화
.to(device)	장치 독립적인 텐서/모델 배치
torch.amp.autocast	2배 속도 향상을 위한 혼합 정밀도
pin_memory=True	더 빠른 CPU→GPU 데이터 전송
torch.compile	속도 향상을 위한 JIT 컴파일 (2.0+)
weights_only=True	보안이 강화된 모델 로드
torch.manual_seed	재현 가능한 실험
gradient_checkpointing	메모리 절약을 위해 연산량 사용

피해야 할 안티 패턴

# 나쁨: 검증 중에 model.eval()을 잊음
model.train()
with torch.no_grad():
    output = model(val_data)  # 드롭아웃이 여전히 활성화됨! BatchNorm이 배치 통계를 사용함!

# 좋음: 항상 eval 모드 설정
model.eval()
with torch.no_grad():
    output = model(val_data)

# 나쁨: autograd를 깨뜨리는 In-place 연산
x = F.relu(x, inplace=True)  # 그래디언트 계산을 방해할 수 있음
x += residual                  # In-place 더하기는 autograd 그래프를 깨뜨림

# 좋음: Out-of-place 연산
x = F.relu(x)
x = x + residual

# 나쁨: 학습 루프 안에서 반복적으로 데이터를 GPU로 이동
for data, target in dataloader:
    model = model.cuda()  # 매 반복마다 모델을 이동!

# 좋음: 루프 시작 전에 한 번만 이동
model = model.to(device)
for data, target in dataloader:
    data, target = data.to(device), target.to(device)

# 나쁨: backward 전에 .item() 사용
loss = criterion(output, target).item()  # 그래프에서 분리됨!
loss.backward()  # 오류: .item()을 통해서는 역전파가 불가능함

# 좋음: 로깅 시에만 .item() 호출
loss = criterion(output, target)
loss.backward()
print(f"Loss: {loss.item():.4f}")  # backward 이후 .item() 호출은 괜찮음

# 나쁨: torch.save를 부적절하게 사용
torch.save(model, "model.pt")  # 모델 전체 저장 (취약하고 이식성이 낮음)

# 좋음: state_dict 저장
torch.save(model.state_dict(), "model.pt")

기억하십시오: PyTorch 코드는 장치 독립적이고, 재현 가능하며, 메모리를 고려해야 합니다. 의심스러울 때는 torch.profiler로 프로파일링하고 torch.cuda.memory_summary()로 GPU 메모리를 확인하십시오.