vit/config.yaml

# ViT Model Configuration
model:
  name: "vit_base"  # vit_tiny, vit_small, vit_base
  num_classes: 3
  pretrained: false
  image_size: 224
  dropout: 0 # 增加dropout减少过拟合
  drop_path: 0  # Drop path regularization

data:
  root_dir: "data"
  train_split: 0.8
  val_split: 0.2
  test_split: 0
  batch_size: 64
  num_workers: 4
  augmentation:
    enabled: true
    use_vit_transforms: true  # 使用ViT专用的数据增强
    # resize策略: 'squash'(强制拉伸，不推荐), 'crop'(保持比例+crop，推荐), 'pad'(保持比例+pad)
    resize_strategy: "crop"  # 对于高分辨率图片，推荐使用crop保持长宽比
    resize_scale: 1.1  # 训练时随机crop的放大倍数（仅crop策略有效）
    horizontal_flip: 0.5
    color_jitter: true
    color_jitter_params:
      brightness: 0.3  # 增加亮度变化，减少对红色的依赖
      contrast: 0.3
      saturation: 0.4  # 增加饱和度变化，让模型学习不同饱和度的红色
      hue: 0.2  # 增加色相变化，让红色可以变成橙色、黄色等
      p: 0.8  # 提高颜色增强的概率
    random_brightness_contrast: true
    brightness_limit: 0.3  # 增加亮度变化范围
    contrast_limit: 0.3
    brightness_contrast_p: 0.7
    # 添加RGB通道独立调整，可以降低红色通道的影响
    channel_shuffle: false  # 可选：通道打乱，但可能破坏语义
    # 添加颜色空间转换增强
    rgb_shift: true  # 随机调整RGB通道
    rgb_shift_limit: 20  # RGB通道偏移范围
    rgb_shift_p: 0.5
    rotate: true
    rotate_limit: 15
    rotate_p: 0.5
    shift_scale_rotate: true
    shift_limit: 0.1
    scale_limit: 0.1
    shift_scale_rotate_p: 0.5
    gaussian_noise: true
    noise_var_limit: [10.0, 50.0]
    noise_p: 0.3
    gaussian_blur: true
    blur_limit: [3, 7]
    blur_p: 0.3
    cutout: true
    max_holes: 8
    max_height: 32
    max_width: 32
    cutout_p: 0.3
    normalize:
      mean: [0.485, 0.456, 0.406]
      std: [0.229, 0.224, 0.225]

training:
  epochs: 50
  learning_rate: 5e-5
  weight_decay: 1e-4
  optimizer: "adamw"
  scheduler: "cosine"
  warmup_epochs: 10
  save_interval: 5
  # 使用Focal Loss处理类别不平衡和困难样本
  use_focal_loss: False
  focal_loss_alpha: 1.0  # Focal loss alpha参数
  focal_loss_gamma: 2.0  # Focal loss gamma参数，gamma越大，对困难样本关注越多
  # 类别权重（如果不用focal loss）
  use_class_weights: false  # 与focal_loss二选一
  # Label smoothing减少对单一特征的过度依赖
  label_smoothing: 0.1

paths:
  checkpoint_dir: "checkpoints/vit"
  log_dir: "logs/vit"
  result_dir: "results/vit"

wandb:
  enabled: true  # 是否启用wandb记录
  project: "fire_detection_vit"  # wandb项目名称
  name: null  # 运行名称，null则自动生成