import logging, argparse, shutil from datetime import datetime from pathlib import Path from typing import Dict, Tuple import gc import numpy as np import segmentation_models_pytorch as smp import torch import torch.optim as optim from torch.optim import lr_scheduler # 学习率调度器 from torch.amp import GradScaler from torch.utils.data import DataLoader # 本地应用/库的导入 import config, os import utils from dataset import SegmentationDataset from loss import MultiClassLoss, UNetPlusPlusLoss # --- 日志设置 --- logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", handlers=[ logging.StreamHandler() # 输出日志到控制台 ] ) # --- 辅助函数：移动训练结果文件夹从源输出目录移动到配置的硬盘目录 --- def move_results_to_hardisk(project_folder: str): """ 将指定的训练结果文件夹从源输出目录移动到配置的硬盘目录。 移动操作包括复制整个文件夹树，然后在复制成功后删除原始文件夹。 参数: project_folder (str): 要移动的项目文件夹的名称。 该名称通常由模型架构和时间戳构成。 """ source_dir = Path(config.OUTPUTS_DIR) / project_folder # 确保目标硬盘目录存在 outputs_folder_name = Path(config.OUTPUTS_DIR).name destination_dir = Path(config.HARDISK_DIR) / outputs_folder_name / project_folder destination_dir.parent.mkdir(parents=True, exist_ok=True) logging.info(f"准备将结果从 {source_dir} 移动到 {destination_dir}...") try: # 步骤 1: 复制文件夹 shutil.copytree(source_dir, destination_dir) logging.info(f"成功复制结果到: {destination_dir}") # 步骤 2: 复制成功后，删除原文件夹 logging.info(f"正在删除原文件夹: {source_dir}") shutil.rmtree(source_dir) logging.info("成功删除原文件夹。") logging.info(f"任务完成，最终结果已保存至: {destination_dir}") except Exception as e: logging.error(f"移动文件夹时发生错误: {e}") logging.warning(f"原始训练结果仍保留在: {source_dir}") # 1. 根据配置文件初始化模型、损失函数和优化器。 def initialize_components(model_architecture:str, num_classes: int, seg_mode: str) -> Tuple: """ 根据配置文件动态初始化模型、损失函数和优化器。 """ # --- 初始化模型 --- logging.info(f"正在初始化模型: '{model_architecture}'...") # 使用传入的 model_architecture 作为 key 来获取参数 try: model_params = config.ALL_MODEL_CONFIGS[model_architecture] model_class = getattr(smp, model_architecture) except KeyError: logging.error(f"模型 '{model_architecture}' 的配置未在 config.py 的 ALL_MODEL_CONFIGS 中定义！") raise except AttributeError: logging.error(f"模型 '{model_architecture}' 在 segmentation_models_pytorch 库中不存在！") raise # 2. 准备参数字典 # 首先复制 config 中的参数，然后添加固定的 `in_channels` 和 `classes` params = model_params.copy() params['in_channels'] = 3 params['classes'] = num_classes # ======================== 【新增代码段开始】 ======================== # # 自动检测 encoder_name 是否包含 'vit' encoder_name = model_params.get('encoder_name', '') if 'vit' in encoder_name.lower(): params['dynamic_img_size'] = True logging.info(f"检测到 ViT 编码器 ('{encoder_name}')。自动设置 dynamic_img_size=True。") # ======================== 【新增代码段结束】 ======================== # # 3. 使用字典解包 (**) 将所有参数传递给模型构造函数 # 这种方式非常灵活，无论你在 config 中定义了多少参数，都能正确传递 model = model_class(**params).to(config.DEVICE) # --- 设置损失函数 (保持不变) --- logging.info(f"为 '{seg_mode}' 模式设置损失函数。") if seg_mode == 'multiclass': loss_fn = MultiClassLoss(mode=seg_mode) elif seg_mode == 'multilabel': loss_fn = UNetPlusPlusLoss(mode=seg_mode) else: raise ValueError(f"无效的 SEG_MODE: '{seg_mode}'。必须是 'multiclass' 或 'multilabel'。") # --- 优化器和梯度缩放器 (保持不变) --- optimizer = optim.Adam(model.parameters(), lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY) # --- TODO 添加学习率调度器 TODO --- # T_max 是调度器周期的最大迭代次数，通常设置为总的 epoch 数量 scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.EPOCHS, eta_min=1e-6) # ----------------------- scaler = GradScaler('cuda') return model, loss_fn, optimizer, scaler def main(model_architecture: str) -> None: """ 编排图像分割模型的主要训练流程。 该函数执行以下步骤: 1. 设置输出目录。 2. 准备数据增强，并为训练集和验证集创建 DataLoader 实例。 3. 初始化 U-Net++ 模型、损失函数、优化器和梯度缩放器。 4. 运行主训练循环，其中包括: - 训练一个 epoch。 - 在验证集上评估模型。 - 将各项指标记录到 CSV 文件。 - 基于验证损失实现早停机制。 - 保存性能最佳的模型和定期的断点。 - 绘制训练进度曲线图。 返回: str: 本次训练运行的文件夹名称 (run_name)。 """ logging.info(f"使用设备: {config.DEVICE}") # --- 1.1. 使用 pathlib 进行现代化的路径管理 --- timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") run_name = f"{model_architecture}_{timestamp}" output_dir = Path(config.OUTPUTS_DIR) / run_name # <-- 使用 / 操作符 metrics_csv_path = output_dir / config.METRICS_CSV_NAME best_model_path = output_dir / config.BEST_MODEL_SAVE_NAME utils.setup_directories(output_dir) # --- 1.2. 探测模型以确定最佳图像尺寸 --- logging.info("正在探测模型以确定输入尺寸要求...") # num_classes = len(config.CLASSES) # # 创建一个临时模型实例，仅用于检查其属性 # 注意：这里的调用现在更简单，不需要传递 dynamic_img_size probe_model, _, _, _ = initialize_components( model_architecture, num_classes, config.SEG_MODE # ) target_height, target_width = config.IMAGE_HEIGHT, config.IMAGE_WIDTH # try: # 检查编码器是否需要固定输入尺寸 if probe_model.encoder.is_fixed_input_size: required_size = probe_model.encoder.input_size target_height = required_size[1] target_width = required_size[2] logging.warning( f"模型 '{model_architecture}' 的编码器需要固定的输入尺寸 " f"({target_height}x{target_width})。" f"将忽略 config.py 中的尺寸设置。" ) else: logging.info(f"模型 '{model_architecture}' 支持动态输入尺寸。将使用 config.py 中定义的尺寸 ({target_height}x{target_width})。") except AttributeError: logging.info("无法自动检测模型尺寸要求。将使用 config.py 中定义的尺寸。") del probe_model # 删除临时模型，释放内存 # --- 1.3. 数据加载 --- logging.info(f"正在设置数据增强和加载器，目标尺寸为 {target_height}x{target_width}...") # # 使用探测到的或配置中指定的目标尺寸 train_transform = utils.get_training_augmentation(target_height, target_width) # val_transform = utils.get_validation_augmentation(target_height, target_width) # train_dataset = SegmentationDataset( image_dir=Path(config.TRAIN_IMAGE_DIR), mask_dir=Path(config.TRAIN_MASK_DIR), classes=config.CLASSES, class_rgb_values=config.CLASS_RGB_VALUES, augmentation=train_transform, seg_mode=config.SEG_MODE ) val_dataset = SegmentationDataset( # 已修正路径，使用配置文件中的验证数据目录 image_dir=Path(config.VAL_IMAGE_DIR), mask_dir=Path(config.VAL_MASK_DIR), classes=config.CLASSES, class_rgb_values=config.CLASS_RGB_VALUES, augmentation=val_transform, seg_mode=config.SEG_MODE ) train_loader = DataLoader(train_dataset, batch_size=config.BATCH_SIZE, shuffle=True, pin_memory=config.PIN_MEMORY, num_workers=config.NUM_WORKERS) val_loader = DataLoader(val_dataset, batch_size=config.BATCH_SIZE, shuffle=False, pin_memory=config.PIN_MEMORY, num_workers=config.NUM_WORKERS) # --- 2. 模型、损失函数、优化器、梯度缩放器 --- num_classes = len(config.CLASSES) model, loss_fn, optimizer, scaler = initialize_components(model_architecture, num_classes, config.SEG_MODE) # --- 启用多 GPU 训练 --- if torch.cuda.device_count() >= 1: logging.info(f"正在使用 {torch.cuda.device_count()} 块 GPU 进行并行训练。") model = torch.nn.DataParallel(model) # --- 3. 训练循环 --- best_val_loss = float('inf') epochs_no_improve = 0 for epoch in range(config.EPOCHS): logging.info(f"\n--- 第 {epoch+1}/{config.EPOCHS} 轮 ---") # --- 3.1. 训练集训练数据 --- train_loss = utils.train_fn(train_loader, model, optimizer, loss_fn, scaler, config.DEVICE, model_architecture) # --- 3.2. 验证集测试数据【获得指标】 --- val_metrics_dict = utils.evaluate_fn(val_loader, model, loss_fn, config.DEVICE, num_classes, model_architecture, beta=config.FBETA_BETA) # --- 3.3. 日志记录 及 显示 --- log_data = {'epoch': epoch + 1, 'train_loss': train_loss} log_data.update(val_metrics_dict) utils.log_metrics_to_csv(log_data, metrics_csv_path) val_loss = val_metrics_dict.get('val_loss') iou_score = val_metrics_dict.get('iou_score', 0) fbeta_score = val_metrics_dict.get('fbeta_score', 0) logging.info(f"训练集Loss: {train_loss:.4f} | 验证集Loss: {val_loss:.4f} | IoU: {iou_score:.4f} | F-beta: {fbeta_score:.4f}") # --- 3.4. 早停与模型检查点 --- if val_loss < best_val_loss: best_val_loss = val_loss epochs_no_improve = 0 if isinstance(model, torch.nn.DataParallel): torch.save(model.module.state_dict(), best_model_path) else: torch.save(model.state_dict(), best_model_path) logging.info(f"验证损失改善至 {best_val_loss:.4f}。正在保存最佳模型至 '{best_model_path}'。") else: epochs_no_improve += 1 logging.info(f"验证损失已连续 {epochs_no_improve} 轮未改善。") if isinstance(config.EARLY_STOPPING_PATIENCE, int) and epochs_no_improve >= config.EARLY_STOPPING_PATIENCE: logging.info(f"\n训练在 {epoch + 1} 轮后触发早停。") logging.warning(f"验证损失已连续 {config.EARLY_STOPPING_PATIENCE} 轮未改善。") break # --- 3.5. 检查点及曲线图保存 --- checkpoint_path = output_dir / f"epoch_{epoch+1}.pth" if isinstance(model, torch.nn.DataParallel): torch.save(model.module.state_dict(), checkpoint_path) else: torch.save(model.state_dict(), checkpoint_path) logging.info(f"第 {epoch+1} 轮的检查点已保存至 '{checkpoint_path}'。") # ---------- 清理旧检查点 ---------- # 只在 epoch > 0 时才开始删，避免第一轮就误删 if epoch > 0: # 要尝试删除的 epoch 编号 = 上一轮 last_epoch = epoch # 如果上一轮不是“10 的整数倍”就删掉 if last_epoch % config.CKPT_SAVE_INTERVAL != 0: old_ckpt = output_dir / f"epoch_{last_epoch}.pth" if old_ckpt.exists(): old_ckpt.unlink() # 文件删除 utils.plot_training_progress(metrics_csv_path, output_dir) logging.info("所有训练曲线图已更新并保存。") return run_name if __name__ == "__main__": # 创建一个参数解析器 parser = argparse.ArgumentParser(description="训练图像分割模型。") # 添加 --architecture 参数 # - `choices` 会自动从你的 config 文件中获取所有可用的模型名称 # - `default` 设置了在不提供参数时的默认模型 parser.add_argument( "-a", "--architecture", type=str, choices=list(config.ALL_MODEL_CONFIGS.keys()), required=True, help="选择要训练的模型架构。" ) # 解析命令行传入的参数 args = parser.parse_args() # 将解析出的架构名称传入 main 函数并执行 project_folder = None try: project_folder = main(model_architecture=args.architecture) finally: # 强制进行垃圾回收，释放 Python 对象占用的系统内存 (RAM) gc.collect() logging.info("已执行垃圾回收 (gc.collect())。") # 清理 PyTorch 在 CUDA 上缓存的显存 torch.cuda.empty_cache() logging.info("训练结束，已清理 CUDA 缓存。") # 训练和清理完成后，移动结果到硬盘 if project_folder: move_results_to_hardisk(project_folder) else: logging.error("由于训练失败或未生成项目文件夹，未执行结果移动操作。")