로컬에서 첫 학습 실행

로컬 환경에서 keynet-train을 사용하여 간단한 model을 학습하고, 자동으로 ONNX 변환 및 배포까지 완료하는 전체 workflow를 경험합니다.

사전 요구사항

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

• 로컬 실험 환경이 실행 중이어야 합니다 (docker compose up -d)
• Python 가상환경이 활성화되어 있어야 합니다

프로젝트 준비

작업 디렉토리 생성

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

# Python 3.12 가상환경 생성
uv venv --python 3.12
source .venv/bin/activate

필요한 패키지 설치

방법 1: 직접 설치

# keynet-train 설치 (MLflow 자동 포함)
uv pip install keynet-train==0.8.5

# PyTorch 설치 (CPU 버전)
uv pip install torch==2.5.1 torchvision==0.20.1

방법 2: requirements.txt 사용

requirements.txt 파일 생성:

requirements.txt

keynet-train
torch==2.5.1
torchvision==0.20.1

설치:

uv pip install -r requirements.txt

GPU 사용하기

GPU를 사용하려면 먼저 시스템의 CUDA 버전을 확인하세요:

nvidia-smi

출력에서 "CUDA Version"을 확인한 후, 같거나 낮은 버전의 PyTorch를 설치하세요:

CUDA 13.0:

uv pip install torch==2.9.0 torchvision==0.21.0 --index-url https://download.pytorch.org/whl/cu130

CUDA 12.4:

uv pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cu124

CUDA 11.8:

uv pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/cu118

더 많은 버전: PyTorch 공식 사이트

호환성 주의

시스템 CUDA 버전보다 높은 PyTorch CUDA 버전을 설치하면 GPU를 인식하지 못하거나 런타임 에러가 발생할 수 있습니다.

:::

설치 검증

PyTorch가 올바르게 설치되었는지 확인:

python -c "import torch; print(f'PyTorch: {torch.__version__}'); print(f'CUDA available: {torch.cuda.is_available()}')"

예상 출력:

PyTorch: 2.5.1+cpu
CUDA available: False

GPU 버전을 설치했다면:

PyTorch: 2.5.1+cu124
CUDA available: True

학습 코드 작성

MNIST 분류 모델

keynet-train의 @trace_pytorch 데코레이터를 사용하여 간단한 MNIST 분류 model을 작성합니다. 다음 코드를 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 데코레이터로 자동화된 학습 함수
# - MLflow 실험/런 자동 생성 및 관리
# - 학습된 모델 자동으로 ONNX 변환 및 배포
# - MinIO/S3 업로드 자동 처리
@trace_pytorch(
    model_name="mnist-cnn",
    sample_input=torch.randn(1, 1, 28, 28),  # 샘플 입력으로 ONNX 스키마 자동 추론
)
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}%")

    # ⚠️ 중요: 반드시 torch.nn.Module 객체만 반환
    # @trace_pytorch 데코레이터가 자동으로 처리:
    # 1. MLflow에 모델 로깅 (autolog)
    # 2. ONNX 변환
    # 3. MinIO/S3 업로드
    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,
    )

학습 실행

기본 실행

기본 hyperparameter로 학습:

python train.py

Hyperparameter 변경

원하는 hyperparameter를 지정하여 실행:

# Epoch 수 변경
python train.py --epochs 10

# Batch size와 learning rate 변경
python train.py --batch-size 128 --learning-rate 0.0001

# 모든 hyperparameter 지정
python train.py --batch-size 32 --epochs 10 --learning-rate 0.0005

예상 출력

Using device: cpu
Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
Failed to download (trying next):
HTTP Error 404: Not Found

Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz
Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz to ./data/MNIST/raw/train-images-idx3-ubyte.gz
Extracting ./data/MNIST/raw/train-images-idx3-ubyte.gz to ./data/MNIST/raw
...

🚀 학습 시작...

Epoch 1/5: Train Loss=0.1501, Train Acc=95.42% | Val Loss=0.0453, Val Acc=98.56%
Epoch 2/5: Train Loss=0.0553, Train Acc=98.33% | Val Loss=0.0416, Val Acc=98.58%
Epoch 3/5: Train Loss=0.0423, Train Acc=98.67% | Val Loss=0.0389, Val Acc=98.71%
Epoch 4/5: Train Loss=0.0356, Train Acc=98.89% | Val Loss=0.0367, Val Acc=98.82%
Epoch 5/5: Train Loss=0.0312, Train Acc=99.01% | Val Loss=0.0351, Val Acc=98.91%

✅ 학습 완료! 최종 Validation Accuracy: 98.91%
🏃 View run bright-mink-77 at: http://localhost:5000/#/experiments/1/runs/abc123xyz
🧪 View experiment at: http://localhost:5000/#/experiments/1

INFO:keynet_train.decorators.pytorch:✅ MLflow PyTorch autolog 활성화 완료
INFO:keynet_train.decorators.pytorch:새 실험 생성: mnist-classification
INFO:keynet_train.decorators.pytorch:MLflow 실행 시작 (run_id: abc123xyz)
INFO:keynet_train.decorators.pytorch:ONNX 변환 시작 - 입력: ['input_0'], 출력: ['output_0']
INFO:keynet_train.decorators.pytorch:동적 크기 ONNX 모델 변환 완료
INFO:keynet_train.decorators.pytorch:ONNX 변환 완료: /tmp/tmpXXXXXX.onnx (4.58MB)
INFO:keynet_train.clients.onnx:✅ ONNX 모델 유효성 검사 완료
INFO:keynet_train.clients.onnx:MLflow에 ONNX 모델 저장 완료
INFO:keynet_train.decorators.pytorch:✅ ONNX 모델 업로드 및 RabbitMQ 발행 완료
INFO:keynet_train.decorators.pytorch:🚀 ONNX 모델 서비스 업로드 완료
INFO:keynet_train.decorators.pytorch:🎉 모델 추적 완료 (실행시간: 87.46초)

자동화된 작업 확인

@trace_pytorch 데코레이터가 자동으로 처리한 작업:

1. MLflow 실험/런 관리

• mnist-classification experiment 자동 생성
• 새로운 run 시작 및 종료 처리

2. ONNX 변환

• PyTorch 모델을 ONNX 형식으로 자동 변환
• 샘플 입력 (1, 1, 28, 28)로 스키마 자동 추론
• 변환된 파일: model.onnx

3. 모델 배포

• MinIO/S3에 ONNX 모델 자동 업로드
• Triton Inference Server용 config.pbtxt 자동 생성
• MLflow artifact로 모델 등록

4. 메타데이터 저장

• 모델 입출력 스키마
• ONNX 버전 정보
• 변환 시간 등

수동으로 처리한 것

다음 항목은 사용자가 직접 mlflow.log_* 함수로 로깅했습니다:

• Hyperparameters (mlflow.log_params)
• Training/Validation metrics (mlflow.log_metrics)

@trace_pytorch는 모델 배포 자동화에 집중하며, 실험 추적은 MLflow API를 직접 사용합니다.

MLflow UI에서 확인

Experiment 확인

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

http://localhost:5000

확인할 수 있는 내용:

1. Experiments

• mnist-classification experiment 자동 생성
• 모든 run 목록과 상태 확인

2. Parameters

로깅한 hyperparameters:

• batch_size: 64
• learning_rate: 0.001
• epochs: 5
• optimizer: Adam
• device: cpu 또는 cuda

3. Metrics

각 epoch별 변화 추이를 그래프로 확인:

• train_loss: Training loss (감소하는 패턴)
• train_accuracy: Training accuracy (증가하는 패턴)
• val_loss: Validation loss
• val_accuracy: Validation accuracy

4. Artifacts

@trace_pytorch가 자동으로 저장한 파일들:

📁 onnx_model/

• model.onnx: 변환된 ONNX 모델
• config.pbtxt: Triton Inference Server 설정 파일
• MLmodel: MLflow 모델 메타데이터
• conda.yaml, requirements.txt: 의존성 정보

확인 방법:

1. Run 상세 페이지 → Artifacts 탭 클릭
2. onnx_model/ 폴더 확인
3. 각 파일 다운로드 또는 내용 미리보기 가능

MinIO에 자동 업로드

Model artifact는 MLflow가 자동으로 MinIO에 저장합니다. MLflow UI에서 바로 확인하고 다운로드할 수 있으므로, MinIO Console에 직접 접속할 필요는 없습니다.

다음 단계

축하합니다! keynet-train을 사용한 첫 번째 학습을 완료했습니다.

학습한 내용:

• @trace_pytorch 데코레이터로 자동화된 모델 배포
• MLflow를 활용한 실험 추적 (params, metrics)
• PyTorch → ONNX 자동 변환
• MinIO/S3 자동 업로드 및 Triton 배포 설정

Docker 환경에서 실행하기

프로덕션 환경과 동일한 설정으로 학습을 실행하려면 Docker로 첫 학습 실행을 참고하세요.