train_option_example.yaml

# DiOodMi 학습 설정 예시
# 사용법: hae train --config train.yaml
#
# 이 파일을 train.yaml로 복사한 후 경로를 수정하세요:
#   cp train.yaml.example train.yaml

version: "1.0"

# =============================================================================
# 데이터셋 설정
# =============================================================================
dataset: ./definitions/mvtec_ad_256_minimal/dataset.yaml

# =============================================================================
# 모델 설정
# =============================================================================
model:
  # 모델 유형: decodiff (권장), ddpm
  type: decodiff

  # 모델 크기: UNet_XS, UNet_S, UNet_M, UNet_L (권장), UNet_XL
  size: UNet_L

  # VAE 유형: ema (권장), ms
  vae_type: ema

  # 처음부터 학습 (true) 또는 사전 학습 가중치 사용 (false)
  from_scratch: true

# =============================================================================
# 학습 파라미터
# =============================================================================
training:
  # 총 에포크 수
  epochs: 5000000

  # 배치 크기 (GPU 메모리에 따라 조정)
  # RTX 3090: 64-128, RTX 4090: 96-128, A100: 128-256
  batch_size: 2

  # 학습률
  learning_rate: 0.0001

  # 웜업 에포크 수
  warmup_epochs: 10

  # Gradient Accumulation 스텝 (메모리 부족 시 증가)
  accumulation_steps: 2

  # DeCo-Diff 전용 옵션
  mask_ratio: 0.7
  patch_shuffle_ratio: 0.3

  # EMA (Exponential Moving Average)
  use_ema: true
  ema_decay: 0.999

  # Mixed Precision (메모리 절약 및 속도 향상)
  use_amp: true

  # 학습률 스케줄러 (선택 사항)
  # lr_scheduler: constant  # constant, cosine_annealing, cosine_warm_restarts, exponential, step, multi_step, plateau, onecycle
  # lr_T_max: 800           # [cosine_annealing] 최대 반복 횟수 (기본값: epochs)
  # lr_eta_min: 1e-8        # [cosine_annealing, plateau] 최소 학습률
  # lr_gamma: 0.1           # [exponential, step, multi_step] 감쇠 계수
  # lr_step_size: 30        # [step] 감쇠 주기
  # lr_milestones: "50,100,150"  # [multi_step] 감쇠 에포크 (쉼표 구분)
  # lr_patience: 2          # [plateau] 성능 정체 후 감소까지 대기 에포크

# =============================================================================
# 이미지 처리
# =============================================================================
image:
  # 입력 이미지 리사이즈 크기
  # 지원 형식: int (288), string ("WxH"), list [W, H]
  # Note: For VAE encoder, dimensions should be divisible by 8
  size: 256x256

# =============================================================================
# 증강 설정 (스플릿별) - 모든 옵션 명시적 설정
# =============================================================================
# 참조: mldosa unified augmentation spec v15
# 파이프라인 순서 (7단계): Resize → Geometric → Crop → Color → Noise → Blur → Erasing
augmentation:
  # 학습 데이터 증강 (standard preset 기반, 모든 옵션 명시)
  train:
    # 파이프라인 순서 (커스터마이징 가능)
    augmentation_order:
      - resize      # Stage 1: 리사이즈
      - geometric   # Stage 2: 기하학적 변환 (Scale, Rotation, Shift, Flip)
      - crop        # Stage 3: 크롭 (RandomCrop / CenterCrop)
      - color       # Stage 4: 색상 변환 (ColorJitter, HistogramEqualization)
      - noise       # Stage 5: 노이즈 (Gaussian, Salt&Pepper)
      - blur        # Stage 6: 블러 (Gaussian blur)
      - erasing     # Stage 7: 지우기 (Random erasing, Cutout)

    # --- Stage 1: RESIZE ---
    resize_enabled: false           # 리사이즈 활성화 여부 (true 시 resize_size 필수)
    resize_mode: stretch            # stretch | fit | fill
    # resize_size: 256              # int (정사각형) 또는 [width, height] (직사각형)
    pad_color: [0, 0, 0]            # fit 모드 패딩 색상 (RGB)
    pad_position: center            # fit 모드 패딩 위치: center | top-left | bottom-right

    # --- Stage 2: GEOMETRIC ---
    # 수평 뒤집기
    horizontal_flip_enabled: true   # 수평 뒤집기 활성화
    horizontal_flip: 0.5            # 확률 [0.0, 1.0]

    # 수직 뒤집기
    vertical_flip_enabled: true    # 수직 뒤집기 활성화
    vertical_flip: 0.5              # 확률 [0.0, 1.0]

    # 회전
    rotation_enabled: true          # 회전 활성화
    rotation_limit: 15              # 최대 회전 각도 (±도)
    rotation_probability: 0.5       # 확률 [0.0, 1.0]

    # 스케일 (확대/축소)
    scale_enabled: false            # 스케일 활성화
    scale_min: 0.8                  # 최소 스케일 (0.8 = 80%로 축소)
    scale_max: 1.2                  # 최대 스케일 (1.2 = 120%로 확대)
    scale_probability: 0.5          # 확률 [0.0, 1.0]

    # 이동 (Shift)
    shift_enabled: true             # 이동 활성화
    shift_limit: null               # 이동 범위 (비율) [0.0, 1.0] - 0.1 = 이미지 크기의 10%
    shift_limit_px: 2               # 이동 범위 (픽셀) - 예: 32 = 32픽셀 (설정 시 shift_limit보다 우선)
    shift_probability: 0.5          # 확률 [0.0, 1.0]

    # --- Stage 3: CROP ---
    crop_enabled: true              # 크롭 활성화 (true 시 crop_size 필수)
    crop_mode: random               # random | center (학습: random 권장)
    crop_size: 128x128              # int (정사각형) 또는 [width, height] 또는 "WxH"

    # --- Stage 4: COLOR ---
    # ColorJitter (밝기, 대비, 채도, 색조 동시 조정)
    color_jitter_enabled: true      # ColorJitter 활성화
    color_jitter:
      brightness: 0.2               # 밝기 변동 범위 [0.0, 1.0]
      contrast: 0.2                 # 대비 변동 범위 [0.0, 1.0]
      saturation: 0.1               # 채도 변동 범위 [0.0, 1.0]
      hue: 0.05                     # 색조 변동 범위 [0.0, 0.5]
      probability: 0.5              # 확률 [0.0, 1.0]

    # NOTE: 밝기/대비 조정은 color_jitter를 통해 설정하세요 (brightness, contrast 필드)

    # 히스토그램 평활화 (v15 신규)
    histogram_equalization_enabled: true  # 히스토그램 평활화 활성화
    histogram_equalization:
      mode: clahe                   # standard | clahe (CLAHE 권장)
      clip_limit: 10.0               # CLAHE 클립 제한 (1.0-10.0)
      tile_grid_size: 8             # CLAHE 타일 그리드 크기
      probability: 1.0              # 확률 [0.0, 1.0]

    # --- Stage 5: NOISE ---
    # 가우시안 노이즈
    gaussian_noise_enabled: false   # 가우시안 노이즈 활성화
    gaussian_noise: 0.0             # 노이즈 표준편차 (0.0 = 비활성화)

    # 소금&후추 노이즈
    salt_pepper_enabled: true      # 소금&후추 노이즈 활성화
    salt_pepper_noise:
      amount: 0.01                  # 노이즈 양 [0.0, 1.0]
      size: 9                   # 노이즈 입자 크기 (픽셀)
      salt_ratio: 0.0               # 소금(흰색) 비율: 0.0=모두 후추, 0.5=균형, 1.0=모두 소금
      roundness: 1.0                # 입자 모양: 0.0=사각형, 1.0=원형
      antialiasing: true            # 안티앨리어싱 (부드러운 가장자리)
      probability: 0.05             # 확률 [0.0, 1.0]

    # --- Stage 6: BLUR ---
    gaussian_blur_enabled: false    # 가우시안 블러 활성화
    # gaussian_blur:
    #   kernel_size: 3              # 커널 크기 (홀수)
    #   sigma: [0.1, 2.0]           # 시그마 범위 [min, max]

    # --- Stage 7: ERASING ---
    # Random Erasing
    erasing_enabled: false          # Random Erasing 활성화
    # erasing:
    #   probability: 0.5            # 확률 [0.0, 1.0]
    #   scale: [0.02, 0.33]         # 지우기 영역 비율 범위

    # Cutout
    cutout_enabled: false           # Cutout 활성화
    # cutout:
    #   num_holes: 1                # 구멍 개수
    #   hole_size: [32, 32]         # 구멍 크기 [height, width]

  # 검증 데이터 증강 (결정론적 - 일관된 검증 메트릭을 위해)
  val:
    # 파이프라인 순서
    augmentation_order:
      - resize
      - geometric
      - crop
      - color
      - noise
      - blur
      - erasing

    # --- Stage 1: RESIZE ---
    resize_enabled: false
    resize_mode: stretch

    # --- Stage 2: GEOMETRIC ---
    horizontal_flip_enabled: false
    horizontal_flip: 0.0
    vertical_flip_enabled: false
    vertical_flip: 0.0
    rotation_enabled: false
    rotation_limit: 0
    rotation_probability: 0.0
    scale_enabled: false
    shift_enabled: false
    shift_limit: 0.0
    shift_limit_px: null

    # --- Stage 3: CROP ---
    crop_enabled: true
    crop_mode: center               # 검증: center crop 권장
    crop_size: 256x256

    # --- Stage 4: COLOR ---
    # NOTE: 밝기/대비 조정은 color_jitter를 통해 설정하세요
    color_jitter_enabled: false
    histogram_equalization_enabled: false

    # --- Stage 5: NOISE ---
    gaussian_noise_enabled: false
    salt_pepper_enabled: false

    # --- Stage 6: BLUR ---
    gaussian_blur_enabled: false

    # --- Stage 7: ERASING ---
    erasing_enabled: false
    cutout_enabled: false

# =============================================================================
# 이미지 로딩
# =============================================================================
loading:
  # 로딩 전략: resize_first (권장), keep_original, adaptive
  strategy: keep_original

  # 데이터 로더 워커 수 (CPU 코어 수에 따라 조정)
  num_workers: 4

# =============================================================================
# 출력 설정
# =============================================================================
output:
  # 결과 저장 디렉토리
  dir: ./results

  # 체크포인트 저장 주기 (에포크)
  checkpoint_every: 5000

# =============================================================================
# 에포크 샘플 이미지 저장
# =============================================================================
# 학습 중 학습/검증 이미지 재구성 결과 저장
epoch_samples:
  enabled: true
  count: 1        # 저장할 이미지 수
  frequency: 1    # N 에포크마다 저장

# =============================================================================
# 분할 전략 (검증 데이터 생성) - 선택 사항
# =============================================================================
# validation/ 폴더가 없는 경우, 명시적으로 설정 필요
# (설정하지 않으면 검증 없이 학습 진행)
# split_strategy:
#   type: random_split    # static, leave_k_out, k_fold, random_split
#   ratios: [0.8, 0.1, 0.1]  # train 80%, val 10%, test 10%
#   seed: 42

# =============================================================================
# 기타 설정
# =============================================================================
# 디버그 모드
debug: false

# 재현성을 위한 시드
seed: 1000

# 데이터 셔플
shuffle: true

visualization:
  live_plot:
    enabled: true
    update_on: validation # validation | epoch | steps
    output: "{checkpoint.dir}/loss_plot.png"
    format: png # png | html
    metrics: [train_loss]

# =============================================================================
# 학습 재개 (Resume Training)
# =============================================================================
# checkpoint 경로는 두 가지 형식을 지원합니다:
#   - 디렉토리: ./results/model/ → checkpoints/ 하위 폴더에서 우선순위에 따라 체크포인트 선택
#   - 파일: ./results/model/checkpoints/epoch_100.pt → 지정된 특정 체크포인트 파일 사용
#
# 디렉토리 모드 체크포인트 선택 우선순위:
#   1. val_best.pt   - 최저 검증 손실 (프로덕션/일반화에 권장)
#   2. train_best.pt - 최저 학습 손실 (검증 데이터 없을 때 유용)
#   3. last.pt       - 가장 최근 학습 상태 (중단된 학습 재개에 적합)
#   4. epoch_*.pt    - 위 파일들이 없으면 가장 최신 에포크 체크포인트로 폴백
#
# Checkpoint selection priority (directory mode):
#   1. val_best.pt   - Best validation loss (preferred for production/generalization)
#   2. train_best.pt - Best training loss (useful when no validation data available)
#   3. last.pt       - Most recent training state (good for continuing interrupted training)
#   4. epoch_*.pt    - Falls back to latest epoch checkpoint if none of the above exist
resume:
  enabled: true
  # 예시 1: 디렉토리 경로 - 우선순위에 따라 체크포인트 자동 선택
  # checkpoint: ./train_result_1/decodiff_custom_wood_UNet_L/
  # 예시 2: 파일 경로 - 특정 체크포인트 지정
  #checkpoint: ./train_result_1/decodiff_custom_wood_UNet_L/checkpoints/val_best.pt
  #additional_epochs: 10  # 체크포인트 이후 추가로 학습할 에포크 수