""" SegFormer Fine-tuning训练脚本（CPU版本） 功能说明： 1. 加载预训练的SegFormer模型 2. 在自定义数据集上Fine-tune 3. 支持两阶段训练（冻结Encoder + 全模型Fine-tune） 4. 混合损失函数（CrossEntropy + Dice Loss） 5. 完整的训练监控和模型保存 适用场景： - 本地CPU开发调试 - 验证代码正确性 - 小数据集快速实验 pip install numpy==1.26.4 目前是2.2.6 """ import os import sys import time import json import numpy as np from pathlib import Path from tqdm import tqdm import torch import torch.nn as nn import torch.nn.functional as F from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR from transformers import SegformerForSemanticSegmentation import matplotlib.pyplot as plt # 导入自定义模块 from config import TrainConfig, get_cpu_config from dataset import get_dataloaders class DiceLoss(nn.Module): """ Dice Loss for semantic segmentation Dice系数衡量预测和真实mask的重叠程度： Dice = 2 * |A ∩ B| / (|A| + |B|) Dice Loss = 1 - Dice 优点： - 对类别不平衡不敏感 - 关注重叠区域，适合分割任务 """ def __init__(self, smooth=1.0): """ Args: smooth: 平滑因子，防止除零 """ super(DiceLoss, self).__init__() self.smooth = smooth def forward(self, logits, targets): """ Args: logits: 模型输出 (B, C, H, W) targets: 真实标签 (B, H, W) Returns: loss: Dice Loss标量 """ # 1. 将logits转换为概率 probs = F.softmax(logits, dim=1) # (B, C, H, W) # 2. 将targets转换为one-hot编码 num_classes = logits.shape[1] targets_one_hot = F.one_hot(targets, num_classes) # (B, H, W, C) targets_one_hot = targets_one_hot.permute(0, 3, 1, 2).float() # (B, C, H, W) # 3. 展平空间维度 probs = probs.view(probs.shape[0], probs.shape[1], -1) # (B, C, H*W) targets_one_hot = targets_one_hot.view(targets_one_hot.shape[0], targets_one_hot.shape[1], -1) # 4. 计算Dice系数 intersection = (probs * targets_one_hot).sum(dim=2) # (B, C) union = probs.sum(dim=2) + targets_one_hot.sum(dim=2) # (B, C) dice = (2.0 * intersection + self.smooth) / (union + self.smooth) # (B, C) # 5. 对所有类别和batch取平均 dice_loss = 1.0 - dice.mean() return dice_loss class MixedLoss(nn.Module): """ 混合损失：CrossEntropy + Dice Loss 结合两种损失的优点： - CrossEntropy: 像素级分类，梯度稳定 - Dice Loss: 关注重叠区域，处理类别不平衡 """ def __init__(self, num_classes, class_weights=None, dice_weight=0.5): """ Args: num_classes: 类别数 class_weights: 类别权重（用于CrossEntropy） dice_weight: Dice Loss的权重 """ super(MixedLoss, self).__init__() # CrossEntropy Loss if class_weights is not None: class_weights = torch.tensor(class_weights, dtype=torch.float32) self.ce_loss = nn.CrossEntropyLoss(weight=class_weights) # Dice Loss self.dice_loss = DiceLoss() self.dice_weight = dice_weight def forward(self, logits, targets): """ Args: logits: 模型输出 (B, C, H, W) targets: 真实标签 (B, H, W) Returns: loss: 混合损失 loss_dict: 各损失分量的字典 """ ce = self.ce_loss(logits, targets) dice = self.dice_loss(logits, targets) total_loss = ce + self.dice_weight * dice return total_loss, {'ce_loss': ce.item(), 'dice_loss': dice.item()} class SegmentationMetrics: """ 分割任务的评估指标 包括： - mIoU (Mean Intersection over Union) - Pixel Accuracy - Per-class IoU """ def __init__(self, num_classes): """ Args: num_classes: 类别数 """ self.num_classes = num_classes self.reset() def reset(self): """重置所有统计""" self.confusion_matrix = np.zeros((self.num_classes, self.num_classes)) def update(self, preds, targets): """ 更新混淆矩阵 Args: preds: 预测标签 (B, H, W) 或 (N,) targets: 真实标签 (B, H, W) 或 (N,) """ preds = preds.flatten() targets = targets.flatten() # 过滤无效值 mask = (targets >= 0) & (targets < self.num_classes) preds = preds[mask] targets = targets[mask] # 更新混淆矩阵 for pred, target in zip(preds, targets): self.confusion_matrix[target, pred] += 1 def compute_miou(self): """ 计算mIoU (Mean Intersection over Union) IoU = TP / (TP + FP + FN) mIoU = mean(IoU_per_class) Returns: miou: 平均IoU iou_per_class: 每个类别的IoU """ # 对角线是TP（预测和真实都是该类别） tp = np.diag(self.confusion_matrix) # 每行之和是该类别的真实数量（TP + FN） # 每列之和是预测为该类别的数量（TP + FP） fp_fn = self.confusion_matrix.sum(axis=0) + self.confusion_matrix.sum(axis=1) - tp # IoU = TP / (TP + FP + FN) iou_per_class = tp / (fp_fn + 1e-10) # 忽略无效类别（分母为0） valid_classes = fp_fn > 0 miou = iou_per_class[valid_classes].mean() return miou, iou_per_class def compute_pixel_accuracy(self): """ 计算像素准确率 Pixel Accuracy = (TP + TN) / Total Returns: pixel_acc: 像素准确率 """ total_correct = np.diag(self.confusion_matrix).sum() total_pixels = self.confusion_matrix.sum() pixel_acc = total_correct / (total_pixels + 1e-10) return pixel_acc def get_results(self): """ 获取所有评估指标 Returns: results: 包含所有指标的字典 """ miou, iou_per_class = self.compute_miou() pixel_acc = self.compute_pixel_accuracy() return { 'miou': miou, 'pixel_accuracy': pixel_acc, 'iou_per_class': iou_per_class.tolist(), } class SegFormerTrainer: """SegFormer训练器""" def __init__(self, config: TrainConfig): """ 初始化训练器 Args: config: 训练配置 """ self.config = config self.device = torch.device(config.device) # 设置随机种子 self._set_seed(config.random_seed) # 创建输出目录 self.output_dir = Path(config.output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) # 加载数据 print("\n" + "="*60) print("加载数据集") print("="*60) self.train_loader, self.val_loader = get_dataloaders( data_root=config.data_root, batch_size=config.batch_size, image_size=config.image_size, num_workers=config.num_workers, num_classes=config.num_classes ) # 加载模型 print("\n" + "="*60) print("加载模型") print("="*60) self.model = self._build_model() # 构建损失函数 self.criterion = self._build_criterion() # 构建优化器和调度器 self.optimizer = self._build_optimizer() self.scheduler = self._build_scheduler() # 训练状态 self.current_epoch = 0 self.best_metric = 0.0 self.train_losses = [] self.val_metrics = [] print("\n" + "="*60) print("✓ 训练器初始化完成") print("="*60) def _set_seed(self, seed): """设置随机种子""" torch.manual_seed(seed) np.random.seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def _build_model(self): """ 构建SegFormer模型 Returns: model: SegFormer模型 """ print(f"加载预训练模型: {self.config.pretrained_model}") # 加载预训练模型 model = SegformerForSemanticSegmentation.from_pretrained( self.config.pretrained_model, num_labels=self.config.num_classes, ignore_mismatched_sizes=True # 允许分类头尺寸不匹配 ) model.to(self.device) # 统计参数量 total_params = sum(p.numel() for p in model.parameters()) trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"✓ 模型加载完成") print(f" 总参数量: {total_params:,}") print(f" 可训练参数: {trainable_params:,}") return model def _build_criterion(self): """构建损失函数""" if self.config.use_mixed_loss: criterion = MixedLoss( num_classes=self.config.num_classes, class_weights=self.config.class_weights, dice_weight=self.config.dice_loss_weight ) print(f"✓ 损失函数: Mixed Loss (CE + {self.config.dice_loss_weight}×Dice)") else: weights = None if self.config.class_weights: weights = torch.tensor(self.config.class_weights, dtype=torch.float32) criterion = nn.CrossEntropyLoss(weight=weights) print(f"✓ 损失函数: CrossEntropy") return criterion.to(self.device) def _build_optimizer(self): """构建优化器""" if self.config.optimizer_type == 'adamw': optimizer = AdamW( self.model.parameters(), lr=self.config.learning_rate, weight_decay=self.config.weight_decay, betas=self.config.adam_betas ) print(f"✓ 优化器: AdamW (lr={self.config.learning_rate})") else: optimizer = torch.optim.SGD( self.model.parameters(), lr=self.config.learning_rate, weight_decay=self.config.weight_decay, momentum=self.config.sgd_momentum ) print(f"✓ 优化器: SGD (lr={self.config.learning_rate})") return optimizer def _build_scheduler(self): """构建学习率调度器""" if self.config.lr_scheduler_type == 'cosine': scheduler = CosineAnnealingLR( self.optimizer, T_max=self.config.num_epochs, eta_min=1e-7 ) print(f"✓ 学习率调度: Cosine Annealing") else: scheduler = None print(f"✓ 学习率调度: None") return scheduler def _freeze_encoder(self): """冻结Encoder（仅训练Decoder）""" for name, param in self.model.named_parameters(): if 'decode_head' not in name: param.requires_grad = False trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad) print(f" 冻结Encoder，可训练参数: {trainable:,}") def _unfreeze_all(self): """解冻所有参数""" for param in self.model.parameters(): param.requires_grad = True trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad) print(f" 解冻全模型，可训练参数: {trainable:,}") def train_epoch(self): """ 训练一个epoch Returns: avg_loss: 平均损失 """ self.model.train() epoch_loss = 0.0 epoch_ce_loss = 0.0 epoch_dice_loss = 0.0 pbar = tqdm(self.train_loader, desc=f"Epoch {self.current_epoch}/{self.config.num_epochs}") for batch_idx, (images, masks) in enumerate(pbar): # 1. 数据迁移到设备 images = images.to(self.device) masks = masks.to(self.device) # 2. 前向传播 outputs = self.model(pixel_values=images) logits = outputs.logits # 3. 上采样logits到mask尺寸 logits = F.interpolate( logits, size=masks.shape[-2:], mode='bilinear', align_corners=False ) # 4. 计算损失 if isinstance(self.criterion, MixedLoss): loss, loss_dict = self.criterion(logits, masks) epoch_ce_loss += loss_dict['ce_loss'] epoch_dice_loss += loss_dict['dice_loss'] else: loss = self.criterion(logits, masks) # 5. 反向传播 self.optimizer.zero_grad() loss.backward() # 6. 梯度裁剪 if self.config.max_grad_norm > 0: torch.nn.utils.clip_grad_norm_( self.model.parameters(), self.config.max_grad_norm ) # 7. 更新参数 self.optimizer.step() # 8. 累计损失 epoch_loss += loss.item() # 9. 更新进度条 pbar.set_postfix({'loss': f'{loss.item():.4f}'}) # 计算平均损失 avg_loss = epoch_loss / len(self.train_loader) if isinstance(self.criterion, MixedLoss): avg_ce = epoch_ce_loss / len(self.train_loader) avg_dice = epoch_dice_loss / len(self.train_loader) print(f" Train Loss: {avg_loss:.4f} (CE: {avg_ce:.4f}, Dice: {avg_dice:.4f})") else: print(f" Train Loss: {avg_loss:.4f}") return avg_loss @torch.no_grad() def validate(self): """ 验证模型 Returns: metrics: 评估指标字典 """ self.model.eval() metrics = SegmentationMetrics(self.config.num_classes) val_loss = 0.0 pbar = tqdm(self.val_loader, desc="Validation") for images, masks in pbar: # 1. 数据迁移 images = images.to(self.device) masks = masks.to(self.device) # 2. 前向传播 outputs = self.model(pixel_values=images) logits = outputs.logits # 3. 上采样 logits = F.interpolate( logits, size=masks.shape[-2:], mode='bilinear', align_corners=False ) # 4. 计算损失 if isinstance(self.criterion, MixedLoss): loss, _ = self.criterion(logits, masks) else: loss = self.criterion(logits, masks) val_loss += loss.item() # 5. 计算预测 preds = logits.argmax(dim=1) # 6. 更新指标 metrics.update(preds.cpu().numpy(), masks.cpu().numpy()) # 计算最终指标 results = metrics.get_results() results['loss'] = val_loss / len(self.val_loader) print(f"\n Validation Results:") print(f" Loss: {results['loss']:.4f}") print(f" mIoU: {results['miou']:.4f}") print(f" Pixel Acc: {results['pixel_accuracy']:.4f}") print(f" IoU per class: {[f'{x:.3f}' for x in results['iou_per_class']]}") return results def save_checkpoint(self, metrics, is_best=False): """ 保存checkpoint Args: metrics: 评估指标 is_best: 是否是最佳模型 """ checkpoint = { 'epoch': self.current_epoch, 'model_state_dict': self.model.state_dict(), 'optimizer_state_dict': self.optimizer.state_dict(), 'metrics': metrics, 'config': self.config.to_dict(), } # 保存最新checkpoint checkpoint_path = self.output_dir / f"checkpoint_epoch_{self.current_epoch}.pt" torch.save(checkpoint, checkpoint_path) print(f" ✓ Checkpoint已保存: {checkpoint_path}") # 保存最佳模型 if is_best: best_path = self.output_dir / "best_model.pt" torch.save(checkpoint, best_path) print(f" ✓ 最佳模型已保存: {best_path}") def train(self): """完整训练流程""" print("\n" + "="*60) print("开始训练") print("="*60) self.config.print_config() # 两阶段训练 if self.config.use_two_stage_training: print("\n【Stage 1: 冻结Encoder，训练Decoder】") self._freeze_encoder() stage1_end = self.config.stage1_epochs else: stage1_end = 0 # 主训练循环 for epoch in range(1, self.config.num_epochs + 1): self.current_epoch = epoch # 切换到Stage 2 if self.config.use_two_stage_training and epoch == stage1_end + 1: print("\n" + "="*60) print("【Stage 2: Fine-tune全模型】") print("="*60) self._unfreeze_all() # 重建优化器（因为参数变了） self.optimizer = self._build_optimizer() self.scheduler = self._build_scheduler() print(f"\n{'='*60}") print(f"Epoch {epoch}/{self.config.num_epochs}") print(f"{'='*60}") # 训练 train_loss = self.train_epoch() self.train_losses.append(train_loss) # 学习率调度 if self.scheduler: self.scheduler.step() print(f" Learning Rate: {self.optimizer.param_groups[0]['lr']:.2e}") # 验证 if epoch % self.config.eval_every_n_epochs == 0: print(f"\n执行验证...") val_metrics = self.validate() self.val_metrics.append((epoch, val_metrics)) # 判断是否是最佳模型 current_metric = val_metrics[self.config.best_metric] is_best = current_metric > self.best_metric if is_best: self.best_metric = current_metric print(f" 🎉 新的最佳{self.config.best_metric}: {current_metric:.4f}") # 保存checkpoint if epoch % self.config.save_every_n_epochs == 0: if self.config.save_best_only and not is_best: print(f" 跳过保存（非最佳模型）") else: self.save_checkpoint(val_metrics, is_best=is_best) print("\n" + "="*60) print("✓ 训练完成!") print("="*60) print(f"最佳{self.config.best_metric}: {self.best_metric:.4f}") # 保存训练曲线 self.plot_training_curves() def plot_training_curves(self): """绘制训练曲线""" fig, axes = plt.subplots(1, 2, figsize=(15, 5)) # 损失曲线 axes[0].plot(range(1, len(self.train_losses) + 1), self.train_losses) axes[0].set_xlabel('Epoch') axes[0].set_ylabel('Training Loss') axes[0].set_title('Training Loss Curve') axes[0].grid(True) # mIoU曲线 if len(self.val_metrics) > 0: epochs = [x[0] for x in self.val_metrics] mious = [x[1]['miou'] for x in self.val_metrics] axes[1].plot(epochs, mious, marker='o') axes[1].set_xlabel('Epoch') axes[1].set_ylabel('mIoU') axes[1].set_title('Validation mIoU Curve') axes[1].grid(True) plt.tight_layout() save_path = self.output_dir / "training_curves.png" plt.savefig(save_path, dpi=120) print(f"\n✓ 训练曲线已保存: {save_path}") plt.close() def main(): """主函数""" print("="*60) print("SegFormer Fine-tuning训练脚本") print("="*60) # 获取配置 config = get_cpu_config() # 根据需要修改配置 # config.num_epochs = 20 # config.image_size = 256 # config.eval_every_n_epochs = 5 # 创建训练器 trainer = SegFormerTrainer(config) # 开始训练 trainer.train() print("\n训练完成! 接下来可以:") print(" 1. 运行 5_test_model.py 测试模型") print(" 2. 查看输出目录中的训练曲线和checkpoint") if __name__ == "__main__": main()