318 lines
9.3 KiB
Python
318 lines
9.3 KiB
Python
"""
|
||
训练配置文件
|
||
|
||
功能说明：
|
||
1. 集中管理所有训练超参数
|
||
2. 区分CPU和GPU训练配置
|
||
3. 便于实验管理和超参数调优
|
||
|
||
使用方式：
|
||
from config import TrainConfig
|
||
config = TrainConfig()
|
||
# 根据需要修改配置
|
||
config.batch_size = 2
|
||
"""
|
||
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
import torch
|
||
|
||
|
||
@dataclass
|
||
class TrainConfig:
|
||
"""
|
||
训练配置类
|
||
|
||
使用dataclass装饰器，自动生成__init__等方法
|
||
所有配置都有类型注解和默认值
|
||
"""
|
||
|
||
# ==================== 路径配置 ====================
|
||
# 数据集根目录
|
||
data_root: str = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/data/segformer_format"
|
||
|
||
# 模型保存目录
|
||
output_dir: str = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/outputs"
|
||
|
||
# 预训练模型名称
|
||
pretrained_model: str = "nvidia/segformer-b0-finetuned-ade-512-512"
|
||
|
||
# ==================== 模型配置 ====================
|
||
# 类别数（包括背景）
|
||
# 0: 背景, 1: 食物
|
||
num_classes: int = 2
|
||
|
||
# 输入图像尺寸
|
||
image_size: int = 512
|
||
|
||
# SegFormer模型版本
|
||
# B0: 最轻量 (3.7M参数)
|
||
# B1: 轻量 (13.7M)
|
||
# B2: 中等 (24.7M)
|
||
# B3: 较大 (44.6M)
|
||
# B4: 大型 (61.4M)
|
||
# B5: 超大 (81.9M)
|
||
model_variant: str = "b0"
|
||
|
||
# ==================== 训练配置 ====================
|
||
# 批大小
|
||
# CPU训练建议: 1-2
|
||
# GPU训练建议: 4-8 (4090可以用8-16)
|
||
batch_size: int = 1
|
||
|
||
# 总训练轮数
|
||
num_epochs: int = 50
|
||
|
||
# 学习率
|
||
learning_rate: float = 6e-5
|
||
|
||
# 权重衰减（L2正则化）
|
||
weight_decay: float = 0.01
|
||
|
||
# 学习率调度器类型
|
||
# 'cosine': 余弦退火
|
||
# 'linear': 线性衰减
|
||
# 'polynomial': 多项式衰减
|
||
lr_scheduler_type: str = "cosine"
|
||
|
||
# Warmup轮数（学习率逐步增加的轮数）
|
||
warmup_epochs: int = 5
|
||
|
||
# 梯度裁剪（防止梯度爆炸）
|
||
max_grad_norm: float = 1.0
|
||
|
||
# ==================== 优化器配置 ====================
|
||
# 优化器类型: 'adamw' 或 'sgd'
|
||
optimizer_type: str = "adamw"
|
||
|
||
# AdamW的beta参数
|
||
adam_betas: tuple = (0.9, 0.999)
|
||
|
||
# SGD动量
|
||
sgd_momentum: float = 0.9
|
||
|
||
# ==================== 损失函数配置 ====================
|
||
# 是否使用混合损失 (CrossEntropy + Dice)
|
||
use_mixed_loss: bool = True
|
||
|
||
# Dice Loss的权重
|
||
dice_loss_weight: float = 0.5
|
||
|
||
# 类别权重（用于处理类别不平衡）
|
||
# None: 自动计算
|
||
# List[float]: 手动指定 [背景权重, 食物权重]
|
||
class_weights: list = None # 例如: [0.3, 0.7]
|
||
|
||
# ==================== 数据加载配置 ====================
|
||
# 数据加载进程数
|
||
# CPU训练: 0 (避免进程间通信开销)
|
||
# GPU训练: 4-8
|
||
num_workers: int = 0
|
||
|
||
# 是否使用数据增强
|
||
use_augmentation: bool = True
|
||
|
||
# ==================== 训练策略 ====================
|
||
# 是否使用两阶段训练
|
||
# Stage 1: 冻结Encoder，只训练Decoder
|
||
# Stage 2: Fine-tune整个模型
|
||
use_two_stage_training: bool = True
|
||
|
||
# Stage 1的训练轮数（冻结Encoder）
|
||
stage1_epochs: int = 10
|
||
|
||
# ==================== 验证和保存 ====================
|
||
# 验证频率（每N个epoch验证一次）
|
||
eval_every_n_epochs: int = 5
|
||
|
||
# 保存checkpoint频率（每N个epoch保存一次）
|
||
save_every_n_epochs: int = 10
|
||
|
||
# 是否只保存最佳模型
|
||
save_best_only: bool = True
|
||
|
||
# 最佳模型的评估指标: 'miou', 'loss', 'pixel_acc'
|
||
best_metric: str = "miou"
|
||
|
||
# ==================== 设备配置 ====================
|
||
# 设备: 'cpu', 'cuda', 'auto'
|
||
device: str = "auto"
|
||
|
||
# 是否使用混合精度训练（仅GPU）
|
||
use_amp: bool = False
|
||
|
||
# ==================== 日志配置 ====================
|
||
# 打印频率（每N个batch打印一次）
|
||
print_every_n_batches: int = 5
|
||
|
||
# 是否保存训练日志
|
||
save_logs: bool = True
|
||
|
||
# 是否启用详细日志（包括每个batch的详细信息）
|
||
verbose: bool = True
|
||
|
||
# ==================== 随机种子 ====================
|
||
# 随机种子（确保可复现）
|
||
random_seed: int = 42
|
||
|
||
def __post_init__(self):
|
||
"""
|
||
初始化后的处理
|
||
|
||
在所有参数赋值后自动调用，用于：
|
||
1. 自动推断设备
|
||
2. 创建输出目录
|
||
3. 验证配置的合理性
|
||
"""
|
||
# 1. 自动推断设备
|
||
if self.device == "auto":
|
||
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
|
||
# 2. CPU训练时自动调整配置
|
||
if self.device == "cpu":
|
||
self.use_amp = False # CPU不支持混合精度
|
||
self.num_workers = 0 # CPU训练避免多进程开销
|
||
if self.batch_size > 2:
|
||
print(f"⚠️ CPU训练建议batch_size<=2，当前值 {self.batch_size} 可能很慢")
|
||
|
||
# 3. 创建输出目录
|
||
Path(self.output_dir).mkdir(parents=True, exist_ok=True)
|
||
|
||
# 4. 验证配置
|
||
assert self.num_classes >= 2, "类别数必须>=2"
|
||
assert self.batch_size > 0, "batch_size必须>0"
|
||
assert self.num_epochs > 0, "num_epochs必须>0"
|
||
assert self.learning_rate > 0, "learning_rate必须>0"
|
||
|
||
if self.use_two_stage_training:
|
||
assert self.stage1_epochs < self.num_epochs, \
|
||
"stage1_epochs必须小于num_epochs"
|
||
|
||
def to_dict(self) -> dict:
|
||
"""
|
||
将配置转换为字典（用于保存）
|
||
|
||
Returns:
|
||
config_dict: 配置字典
|
||
"""
|
||
return {
|
||
k: v for k, v in self.__dict__.items()
|
||
if not k.startswith('_')
|
||
}
|
||
|
||
def print_config(self):
|
||
"""打印所有配置信息"""
|
||
print("\n" + "="*60)
|
||
print("训练配置")
|
||
print("="*60)
|
||
|
||
print("\n【路径配置】")
|
||
print(f" 数据集: {self.data_root}")
|
||
print(f" 输出目录: {self.output_dir}")
|
||
print(f" 预训练模型: {self.pretrained_model}")
|
||
|
||
print("\n【模型配置】")
|
||
print(f" 类别数: {self.num_classes}")
|
||
print(f" 图像尺寸: {self.image_size}×{self.image_size}")
|
||
print(f" 模型版本: SegFormer-{self.model_variant.upper()}")
|
||
|
||
print("\n【训练配置】")
|
||
print(f" 批大小: {self.batch_size}")
|
||
print(f" 训练轮数: {self.num_epochs}")
|
||
print(f" 学习率: {self.learning_rate}")
|
||
print(f" 权重衰减: {self.weight_decay}")
|
||
print(f" 学习率调度: {self.lr_scheduler_type}")
|
||
print(f" Warmup轮数: {self.warmup_epochs}")
|
||
|
||
print("\n【训练策略】")
|
||
if self.use_two_stage_training:
|
||
print(f" 两阶段训练: 是")
|
||
print(f" Stage 1 (冻结Encoder): {self.stage1_epochs} epochs")
|
||
print(f" Stage 2 (全模型Fine-tune): {self.num_epochs - self.stage1_epochs} epochs")
|
||
else:
|
||
print(f" 两阶段训练: 否")
|
||
|
||
print("\n【损失函数】")
|
||
if self.use_mixed_loss:
|
||
print(f" 混合损失: CrossEntropy + {self.dice_loss_weight}×Dice")
|
||
else:
|
||
print(f" 损失函数: CrossEntropy")
|
||
if self.class_weights:
|
||
print(f" 类别权重: {self.class_weights}")
|
||
|
||
print("\n【设备配置】")
|
||
print(f" 设备: {self.device}")
|
||
print(f" 混合精度: {'是' if self.use_amp else '否'}")
|
||
print(f" 数据加载进程: {self.num_workers}")
|
||
|
||
print("\n【验证和保存】")
|
||
print(f" 验证频率: 每{self.eval_every_n_epochs}轮")
|
||
print(f" 保存频率: 每{self.save_every_n_epochs}轮")
|
||
print(f" 最佳指标: {self.best_metric}")
|
||
|
||
print("="*60 + "\n")
|
||
|
||
|
||
def get_cpu_config() -> TrainConfig:
|
||
"""
|
||
获取CPU训练的推荐配置
|
||
|
||
适用于：
|
||
- 本地开发调试
|
||
- 快速验证代码正确性
|
||
- 小数据集实验
|
||
"""
|
||
config = TrainConfig()
|
||
|
||
# CPU优化配置
|
||
config.device = "cpu"
|
||
config.batch_size = 1
|
||
config.num_workers = 0
|
||
config.use_amp = False
|
||
config.image_size = 256 # 降低分辨率加快训练
|
||
|
||
# 快速验证配置
|
||
# config.num_epochs = 20
|
||
config.num_epochs = 50
|
||
config.eval_every_n_epochs = 5
|
||
config.save_every_n_epochs = 10
|
||
|
||
return config
|
||
|
||
|
||
def get_gpu_config() -> TrainConfig:
|
||
"""
|
||
获取GPU训练的推荐配置
|
||
|
||
适用于：
|
||
- 4090等高性能GPU
|
||
- 正式训练
|
||
- 追求最佳性能
|
||
"""
|
||
config = TrainConfig()
|
||
|
||
# GPU优化配置
|
||
config.device = "cuda"
|
||
config.batch_size = 8
|
||
config.num_workers = 4
|
||
config.use_amp = True
|
||
config.image_size = 512
|
||
|
||
# 完整训练配置
|
||
config.num_epochs = 100
|
||
config.eval_every_n_epochs = 5
|
||
config.save_every_n_epochs = 10
|
||
|
||
return config
|
||
|
||
|
||
if __name__ == "__main__":
|
||
"""测试配置"""
|
||
print("CPU配置:")
|
||
cpu_config = get_cpu_config()
|
||
cpu_config.print_config()
|
||
|
||
print("\n\nGPU配置:")
|
||
gpu_config = get_gpu_config()
|
||
gpu_config.print_config()
|