파이토치 4 - 쿠버네티스에서 학습하기

파이토치 모델을 학습을 하는데 여러 사람이 동시에 작업을 하거나 또는 학습이 멀티 GPU, 멀티 노드를 사용할 경우에는 한정된 GPU 자원을 어떻게 배치할지에 대한 어려움이 생긴다. 또한 파이토치 작업이 학습도중에 중단되거나 했을때 자동으로 학습을 재시작하는 것과 같은 자동화된 관리가 필요한데, 이런 문제를 해결하기 위해서 가장 보편적으로 사용되는 접근 방법이 쿠버네티스나 Slurm등이 될 수 있다. 

 

향후에 멀티 머신 기반의 학습이나 Fault Torelent 아키텍처를 기반으로한 분산 모델 학습 방법에도 설명하겠지만, 이 글에서는 간단하게 모델을 쿠버네티스에서 학습할 수 있는 방법에 대해서 알아본다. 

 

아래 코드는 파이토치 라이트닝을 이용하여 Fashion Mnist 데이터를 학습하는 간단한 DNN 모델이다. 테스트용 모델이기 때문에 모델의 정확도등은 고려하지 않고 작성하였다. 참고로 GPU 사용을 많이 유도하기 위해서 배치 사이즈를 늘려서 GPU 메모리 사용량을 늘렸으며, 또한 불필요한 레이어를 많이 추가하여 GPU의 사용량이 올라가도록 유도하였다. 

 

구글 쿠버네티스 환경을 사용할 경우, 클러스터 생성 방법은 다음과 같다

https://cloud.google.com/deep-learning-containers/docs/kubernetes-container?hl=ko

클러스터 생성시 몇가지 주의해야할 점은 다음과 같다.

1. 쿠버네티스 노트풀을 생성할때 자동으로 GPU 드라이버를 설치하도록 한다. 디폴트는 수동 설치이다.
2. 컨테이너를 ArtifactRegistry에서 다운 받을때, Artifact Registry는 서비스 어카운트 기반으로 인증을 한다. 쿠버네티스는 디폴트로 GCE SA를 사용하기 때문에, 레지스트리 접근 권한이 없음으로, ImagePullOff 에러가 난다. 그래서 GKE 클러스터의 노드풀을 생성할때 Artifact Registry의 Read 권한이 있는 SA를 생성해서 지정하거나 또는 GCE SA에 Artifact Registry의 Read 권한을 부여해야 한다.

#pip install torch torchvision pytorch-lightning GPUtil

import torch
import pytorch_lightning as pl
from pytorch_lightning import Trainer
from pytorch_lightning.loggers import TensorBoardLogger
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch import nn
import torch.nn.functional as F
import GPUtil  # GPU 정보를 얻기 위한 라이브러리

# XLA 컴파일러 사용 CUDA 사용하도록 설정
import os
os.environ['PJRT_DEVICE'] = 'CUDA'

# GPU 사용 여부 확인
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 데이터 변환 (정규화 등)
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 데이터셋 로드
trainset = datasets.FashionMNIST(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = DataLoader(trainset, batch_size=256, shuffle=True)

class FashionMNISTModel(pl.LightningModule):

    
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
                nn.Conv2d(1, 64, 3, padding=1),
                nn.BatchNorm2d(64),  # 배치 정규화 추가
                nn.ReLU(),
                nn.Conv2d(64, 128, 3, padding=1),
                nn.BatchNorm2d(128),
                nn.ReLU(),
                nn.MaxPool2d(2),
                nn.Conv2d(128, 256, 3, padding=1),
                nn.BatchNorm2d(256),
                nn.ReLU(),
                nn.Conv2d(256, 512, 3, padding=1),
                nn.BatchNorm2d(512),
                nn.ReLU(),
                nn.MaxPool2d(2),
                # ... (더 많은 레이어 추가 가능)
                nn.Flatten(),
                nn.Linear(25088, 8192),  # Added linear layer with input 25088
                nn.ReLU(),
                nn.Dropout(0.5),
                nn.Linear(8192, 4096),
                nn.ReLU(),
                nn.Dropout(0.5),
                nn.Linear(4096, 10)
        )

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

    def training_step(self, batch, batch_idx):
        x, y = batch
        logits = self(x)
        loss = F.cross_entropy(logits, y)
        self.log("train_loss", loss)
        return loss

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

    def on_train_epoch_end(self):
        # 50 epoch마다 loss, GPU 사용률, 메모리 사용률 출력
        if self.current_epoch % 2 == 0:
            print(f"Epoch {self.current_epoch}: Loss = {self.trainer.callback_metrics['train_loss']:.4f}")
            if torch.cuda.is_available():
                gpu_stats = GPUtil.getGPUs()[0]
                print(f"GPU Utilization: {gpu_stats.load*100}%")
                print(f"GPU Memory Used: {gpu_stats.memoryUsed}MB")

# 모델 인스턴스 생성
model = FashionMNISTModel()

# 체크포인트 로드 (만약 존재하면)
checkpoint_callback = pl.callbacks.ModelCheckpoint(
    dirpath='checkpoints',
    filename='fashion-mnist-{epoch:02d}-{train_loss:.2f}',
    save_top_k=1,
    mode='min'
)

# 로거 설정
logger = TensorBoardLogger("logs", name="fashion_mnist")

# 트레이너 설정
if torch.cuda.is_available():
  accelerator = 'gpu'
  devices = 1
else:
  accelerator = 'cpu'
  devices = 1

trainer = Trainer(
    accelerator=accelerator,  # Specify hardware type (gpu or cpu)
    devices=devices,        # Specify number of devices to use
    max_epochs=100,
    logger=logger,
    callbacks=[checkpoint_callback]
)

# 학습 시작
trainer.fit(model, trainloader)

 

다음은, 이 모델을 쿠버네티스에서 학습하기 위해서 컨테이너로 패킹하기 위한 도커  파일이다.

# Base image with PyTorch and other dependencies
#FROM nvidia/cuda:12.2.0-base-ubuntu20.04 AS base
FROM gcr.io/deeplearning-platform-release/pytorch-gpu

# Install PyTorch Lightning and other requirements
RUN pip3 install torch torchvision pytorch-lightning tensorboard GPUtil

# Create working directory
WORKDIR /app

# Copy application code
COPY . /app

# Command to run training script
CMD python fashion-mnist.py

 

마지막으로 쿠버네티스 Job 파일이다. 

싱글 머신을 사용하는 학습으로 Job을 사용하였다. (나중에 분산 머신에서 학습하는 경우에는 Indexed Job을 사용하여 Master Node 를 식별할 수 있도록 변경할것이다.)

apiVersion: batch/v1
kind: Job
metadata:
  name: fashion-mnist-trainer
spec:
  template:
    metadata:
      labels:
        app: fashion-mnist-trainer
    spec:
      containers:
      - name: fashion-mnist-trainer
        image: [도커 이미지 경로] :latest
        env:
          - name: DATA_ROOT
            value: /app/data
        resources:
          limits:
            memory: "8Gi"
            cpu: "2"
            nvidia.com/gpu: 1  # 1개의 GPU 요청
          requests:
            memory: "4Gi"
            cpu: "1"
            nvidia.com/gpu: 1  # 1개의 GPU 요청
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-accelerator: nvidia-l4  # GPU 유형 (예시)