增加图像分割相应的代码
This commit is contained in:
@@ -0,0 +1,679 @@
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user