Docker로 첫 학습 실행

Docker 컨테이너를 사용하여 재현 가능한 학습 환경을 구성하고 실행합니다. 프로덕션 환경과 동일한 설정으로 모델을 학습할 수 있습니다.

사전 요구사항

이 가이드를 진행하기 전에:

• 로컬에서 첫 학습 실행을 완료했어야 합니다
• Docker가 설치되어 있어야 합니다 (docker --version)
• 로컬 실험 환경이 실행 중이어야 합니다 (docker compose up -d)
• GPU 사용 시 NVIDIA Docker가 설치되어 있어야 합니다

Docker를 사용하는 이유

환경 재현성

로컬 환경에서는 Python 버전, 라이브러리 버전 등이 달라 동일한 결과를 보장하기 어렵습니다. Docker는 완전히 동일한 실행 환경을 보장합니다.

프로덕션과 동일한 설정

실제 배포 환경(Kubernetes 등)에서도 Docker 이미지를 사용하므로, 개발 단계부터 동일한 환경에서 테스트할 수 있습니다.

팀 협업

팀원 모두가 동일한 Docker 이미지를 사용하여 "내 컴퓨터에서는 되는데" 문제를 방지합니다.

GPU 리소스 격리

여러 실험을 동시에 실행할 때 GPU 메모리를 격리하여 관리할 수 있습니다.

프로젝트 준비

작업 디렉토리 구조

# 새 프로젝트 디렉토리 생성
mkdir mnist-docker-example
cd mnist-docker-example

# 필요한 파일 구조
# mnist-docker-example/
# ├── train.py           # 학습 스크립트
# ├── requirements.txt   # Python 의존성
# └── Dockerfile         # Docker 이미지 정의

requirements.txt 작성

requirements.txt

keynet-train==0.8.5

train.py 작성

로컬에서 첫 학습 실행에서 사용한 동일한 코드를 사용합니다:

train.py

import argparse
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import mlflow
from keynet_train import trace_pytorch


# 간단한 CNN 모델 정의
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)
        self.dropout = nn.Dropout(0.25)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = self.dropout(x)
        x = torch.flatten(x, 1)
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return torch.log_softmax(x, dim=1)


def train_epoch(model, device, train_loader, optimizer, epoch):
    """한 epoch 학습"""
    model.train()
    total_loss = 0
    correct = 0
    total = 0

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = nn.functional.nll_loss(output, target)
        loss.backward()
        optimizer.step()

        # 통계 수집
        total_loss += loss.item()
        pred = output.argmax(dim=1, keepdim=True)
        correct += pred.eq(target.view_as(pred)).sum().item()
        total += len(data)

    avg_loss = total_loss / len(train_loader)
    accuracy = 100.0 * correct / total
    return avg_loss, accuracy


def validate(model, device, test_loader):
    """Validation 수행"""
    model.eval()
    test_loss = 0
    correct = 0

    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += nn.functional.nll_loss(output, target, reduction="sum").item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    accuracy = 100.0 * correct / len(test_loader.dataset)
    return test_loss, accuracy


# @trace_pytorch 데코레이터로 자동화된 학습 함수
@trace_pytorch(
    model_name="mnist-docker",
    sample_input=torch.randn(1, 1, 28, 28),
)
def train_mnist(batch_size, epochs, learning_rate):
    # Device 설정
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    print(f"Using device: {device}")

    # Dataset 준비
    data_dir = "./data"  # MNIST 데이터셋 저장 경로
    transform = transforms.Compose(
        [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
    )

    train_dataset = datasets.MNIST(
        data_dir, train=True, download=True, transform=transform
    )
    test_dataset = datasets.MNIST(data_dir, train=False, transform=transform)

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    # Model 초기화
    model = SimpleCNN().to(device)
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)

    # Hyperparameters logging
    mlflow.log_params({
        "batch_size": batch_size,
        "epochs": epochs,
        "learning_rate": learning_rate,
        "optimizer": "Adam",
        "device": str(device),
    })

    print("\n🚀 학습 시작...\n")

    # 학습 loop
    for epoch in range(1, epochs + 1):
        # Training
        train_loss, train_acc = train_epoch(
            model, device, train_loader, optimizer, epoch
        )

        # Validation
        val_loss, val_acc = validate(model, device, test_loader)

        # Metric logging
        mlflow.log_metrics({
            "train_loss": train_loss,
            "train_accuracy": train_acc,
            "val_loss": val_loss,
            "val_accuracy": val_acc,
        }, step=epoch)

        print(
            f"Epoch {epoch}/{epochs}: "
            f"Train Loss={train_loss:.4f}, Train Acc={train_acc:.2f}% | "
            f"Val Loss={val_loss:.4f}, Val Acc={val_acc:.2f}%"
        )

    # 최종 결과
    print(f"\n✅ 학습 완료! 최종 Validation Accuracy: {val_acc:.2f}%")

    return model


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="MNIST Classification Training")
    parser.add_argument("--batch-size", type=int, default=64, help="Training batch size (default: 64)")
    parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs (default: 5)")
    parser.add_argument("--learning-rate", type=float, default=0.001, help="Learning rate (default: 0.001)")
    args = parser.parse_args()

    train_mnist(
        batch_size=args.batch_size,
        epochs=args.epochs,
        learning_rate=args.learning_rate,
    )

GPU 환경 확인

Docker에서 GPU를 사용하려면 시스템의 CUDA 버전과 호환되는 PyTorch 베이스 이미지를 선택해야 합니다.

CUDA 버전 확인

먼저 시스템의 NVIDIA Driver와 CUDA 버전을 확인하세요:

nvidia-smi

출력 예시:

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 580.82.07              Driver Version: 580.82.07      CUDA Version: 13.0     |
+-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
|   0  NVIDIA RTX PRO 6000 Blac...    Off |   00000000:61:00.0 Off |                    0 |
+-----------------------------------------+------------------------+----------------------+

확인할 정보:

• Driver Version: NVIDIA GPU 드라이버 버전
• CUDA Version: 지원하는 최대 CUDA 버전 (예: 13.0)
• GPU Name: GPU 모델 (최신 아키텍처는 최신 PyTorch 필요)

베이스 이미지 선택

keynet-train 제약: PyTorch 2.8.0 이하만 지원

선택 방법:

1. nvidia-smi로 지원하는 CUDA 버전 확인
2. PyTorch Docker Hub에서 검색:
   • PyTorch 버전: 2.8.0 이하
   • CUDA 버전: nvidia-smi와 같거나 낮은 버전

호환성 제약

PyTorch 2.8.0 이하만 사용 가능

keynet-train 의 onnxruntime 1.17.x는 ONNX IR version 9까지만 지원합니다. PyTorch 2.9.0부터 IR v10을 생성하므로 사용 불가합니다.

발생 가능한 에러:

# PyTorch 2.9.0 사용 시
ERROR: Unsupported model IR version: 10, max supported IR version: 9

# CUDA 버전 부족 시
RuntimeError: CUDA error: no kernel image is available for execution on the device

Dockerfile 작성

GPU 버전

GPU를 사용하는 경우:

Dockerfile

# PyTorch 공식 이미지 사용
# ⚠️ 중요: GPU 아키텍처에 맞는 이미지를 선택하세요
FROM pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime

# 작업 디렉토리 설정
WORKDIR /app

# requirements 먼저 복사 (캐싱 최적화)
COPY requirements.txt .

# 의존성 설치
RUN pip install --no-cache-dir -r requirements.txt

# 학습 코드 복사
COPY train.py .

# 컨테이너 실행 시 학습 시작
ENTRYPOINT ["python", "train.py"]

Docker 이미지 빌드

# GPU Dockerfile 사용
docker build -f Dockerfile -t mnist-training:gpu .

# 빌드 확인
docker images | grep mnist-training

빌드 최적화

requirements.txt를 먼저 복사하고 설치하면, 학습 코드가 변경되어도 의존성은 캐시에서 재사용됩니다.

# ✅ 좋은 예: requirements 먼저 복사
COPY requirements.txt .
RUN pip install -r requirements.txt
COPY train.py .

# ❌ 나쁜 예: 모든 파일 한 번에 복사
COPY . .
RUN pip install -r requirements.txt

Docker 컨테이너 실행

GPU 사용

# 기본 hyperparameter로 실행
docker run --rm \
  --network host \
  --gpus all \
  mnist-training:gpu

# Hyperparameter 변경
docker run --rm \
  --network host \
  --gpus all \
  mnist-training:gpu \
  --epochs 10 --learning-rate 0.0001

특정 GPU 지정

# GPU 0번만 사용
docker run --rm \
  --network host \
  --gpus '"device=0"' \
  mnist-training:gpu

# GPU 0번, 1번 사용 + hyperparameter 변경
docker run --rm \
  --network host \
  --gpus '"device=0,1"' \
  mnist-training:gpu \
  --batch-size 256 --epochs 20

--network host 사용 이유

MLflow tracking server와 MinIO가 localhost에서 실행 중이므로 --network host를 사용합니다. keynet-train이 자동으로 localhost 환경을 인식하여 연결합니다.

실행 결과 확인

예상 출력

Using device: cuda:0
Downloading MNIST dataset...

🚀 학습 시작...

Epoch 1/5: Train Loss=0.3421, Train Acc=89.23% | Val Loss=0.1234, Val Acc=96.12%
Epoch 2/5: Train Loss=0.1123, Train Acc=96.54% | Val Loss=0.0891, Val Acc=97.23%
Epoch 3/5: Train Loss=0.0834, Train Acc=97.45% | Val Loss=0.0723, Val Acc=97.89%
Epoch 4/5: Train Loss=0.0687, Train Acc=97.89% | Val Loss=0.0654, Val Acc=98.12%
Epoch 5/5: Train Loss=0.0598, Train Acc=98.23% | Val Loss=0.0612, Val Acc=98.34%

✅ 학습 완료! 최종 Validation Accuracy: 98.34%

🔄 Converting PyTorch model to ONNX...
✅ ONNX conversion successful: model.onnx
📦 Uploading model to MinIO...
✅ Model uploaded: s3://mlflow/mnist-docker/abc123xyz/artifacts/onnx_model/
🚀 Triton deployment config generated: config.pbtxt

🔍 MLflow UI에서 확인: http://localhost:5000/#/experiments/0/runs/abc123xyz

MLflow UI 확인

웹 브라우저에서 MLflow UI 접속:

http://localhost:5000

mnist-docker experiment에서 실행된 run을 확인할 수 있습니다.

문제 해결

GPU 인식 실패

증상:

NVIDIA RTX PRO 6000 Blackwell with CUDA capability sm_120 is not compatible
The current PyTorch install supports CUDA capabilities sm_50 sm_60 ... sm_90
RuntimeError: CUDA error: no kernel image is available for execution on the device

원인: PyTorch 베이스 이미지의 CUDA 버전이 GPU를 지원하지 않습니다.

해결:

1. GPU 아키텍처 확인: nvidia-smi에서 GPU 모델 확인
2. Blackwell GPU → CUDA 12.9 이상 필요
3. Dockerfile의 베이스 이미지를 pytorch/pytorch:2.8.0-cuda12.9-cudnn9-runtime으로 변경

ONNX 변환 실패 (IR version 불일치)

증상:

ERROR: Unsupported model IR version: 10, max supported IR version: 9

원인: PyTorch 2.9.0이 생성한 ONNX IR version 10을 keynet-train의 onnxruntime이 지원하지 않습니다.

해결: PyTorch 2.8.0 이하로 다운그레이드 (권장)

FROM pytorch/pytorch:2.8.0-cuda12.9-cudnn9-runtime