增加图像分割相应的代码
This commit is contained in:
@@ -0,0 +1,316 @@
|
||||
"""
|
||||
训练配置文件
|
||||
|
||||
功能说明：
|
||||
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.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()
|
||||
@@ -0,0 +1,344 @@
|
||||
"""
|
||||
SegFormer训练数据集类
|
||||
|
||||
功能说明：
|
||||
1. 加载图像和对应的分割mask
|
||||
2. 数据增强（翻转、旋转、颜色变换等）
|
||||
3. 预处理和标准化
|
||||
4. 批量数据加载
|
||||
"""
|
||||
|
||||
import os
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from pathlib import Path
|
||||
from typing import Optional, Tuple, List
|
||||
import albumentations as A
|
||||
from albumentations.pytorch import ToTensorV2
|
||||
|
||||
|
||||
class FoodSegmentationDataset(Dataset):
|
||||
"""
|
||||
食物分割数据集
|
||||
|
||||
数据格式：
|
||||
- 图像：RGB图像 (.jpg, .png等)
|
||||
- 标注：PNG格式mask，像素值为类别ID (0=背景, 1=食物)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
data_root: str,
|
||||
split: str = 'train',
|
||||
image_size: int = 512,
|
||||
num_classes: int = 2,
|
||||
augmentation: bool = True
|
||||
):
|
||||
"""
|
||||
初始化数据集
|
||||
|
||||
Args:
|
||||
data_root: 数据集根目录（包含images和annotations文件夹）
|
||||
split: 'train' 或 'val'
|
||||
image_size: 输入图像尺寸（将resize到此尺寸）
|
||||
num_classes: 类别数（包括背景）
|
||||
augmentation: 是否使用数据增强（仅训练集）
|
||||
"""
|
||||
self.data_root = Path(data_root)
|
||||
self.split = split
|
||||
self.image_size = image_size
|
||||
self.num_classes = num_classes
|
||||
self.augmentation = augmentation and (split == 'train')
|
||||
|
||||
# 构建图像和mask的路径
|
||||
self.images_dir = self.data_root / "images" / split
|
||||
self.masks_dir = self.data_root / "annotations" / split
|
||||
|
||||
# 获取所有图像文件
|
||||
self.image_files = sorted(list(self.images_dir.glob("*.*")))
|
||||
|
||||
# 过滤：只保留有对应mask的图像
|
||||
self.valid_samples = []
|
||||
for img_path in self.image_files:
|
||||
mask_path = self.masks_dir / (img_path.stem + '.png')
|
||||
if mask_path.exists():
|
||||
self.valid_samples.append((img_path, mask_path))
|
||||
|
||||
print(f"✓ {split}集加载完成: {len(self.valid_samples)} 个样本")
|
||||
|
||||
# 构建数据增强pipeline
|
||||
self.transform = self._build_transforms()
|
||||
|
||||
def _build_transforms(self):
|
||||
"""
|
||||
构建数据增强和预处理pipeline
|
||||
|
||||
使用albumentations库进行高效的数据增强
|
||||
注意：对于分割任务，增强操作需要同时应用到图像和mask
|
||||
"""
|
||||
if self.augmentation:
|
||||
# 训练集：激进的数据增强（因为数据量小）
|
||||
transform = A.Compose([
|
||||
# 1. 尺寸调整
|
||||
A.Resize(self.image_size, self.image_size),
|
||||
|
||||
# 2. 几何变换（同时作用于图像和mask）
|
||||
A.HorizontalFlip(p=0.5), # 50%概率水平翻转
|
||||
A.VerticalFlip(p=0.3), # 30%概率垂直翻转
|
||||
A.Rotate(limit=30, p=0.5), # ±30度旋转
|
||||
A.ShiftScaleRotate(
|
||||
shift_limit=0.1, # 平移±10%
|
||||
scale_limit=0.2, # 缩放±20%
|
||||
rotate_limit=20, # 旋转±20度
|
||||
p=0.5
|
||||
),
|
||||
|
||||
# 3. 颜色增强（仅作用于图像）
|
||||
A.RandomBrightnessContrast(
|
||||
brightness_limit=0.2,
|
||||
contrast_limit=0.2,
|
||||
p=0.5
|
||||
),
|
||||
A.HueSaturationValue(
|
||||
hue_shift_limit=20,
|
||||
sat_shift_limit=30,
|
||||
val_shift_limit=20,
|
||||
p=0.5
|
||||
),
|
||||
|
||||
# 4. 模糊和噪声
|
||||
A.OneOf([
|
||||
A.GaussianBlur(blur_limit=(3, 5), p=1.0),
|
||||
A.MedianBlur(blur_limit=5, p=1.0),
|
||||
], p=0.3),
|
||||
|
||||
# 5. 标准化（使用ImageNet均值和标准差）
|
||||
A.Normalize(
|
||||
mean=[0.485, 0.456, 0.406],
|
||||
std=[0.229, 0.224, 0.225],
|
||||
),
|
||||
|
||||
# 6. 转换为Tensor
|
||||
ToTensorV2(),
|
||||
])
|
||||
else:
|
||||
# 验证集：仅resize和标准化
|
||||
transform = A.Compose([
|
||||
A.Resize(self.image_size, self.image_size),
|
||||
A.Normalize(
|
||||
mean=[0.485, 0.456, 0.406],
|
||||
std=[0.229, 0.224, 0.225],
|
||||
),
|
||||
ToTensorV2(),
|
||||
])
|
||||
|
||||
return transform
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""返回数据集大小"""
|
||||
return len(self.valid_samples)
|
||||
|
||||
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
获取一个样本
|
||||
|
||||
Args:
|
||||
idx: 样本索引
|
||||
|
||||
Returns:
|
||||
image: 图像Tensor (3, H, W)
|
||||
mask: 分割mask Tensor (H, W)，像素值为类别ID
|
||||
"""
|
||||
img_path, mask_path = self.valid_samples[idx]
|
||||
|
||||
# 1. 读取图像和mask
|
||||
image = np.array(Image.open(img_path).convert('RGB'))
|
||||
mask = np.array(Image.open(mask_path))
|
||||
|
||||
# 2. 应用数据增强
|
||||
# albumentations会自动将增强同时应用到image和mask
|
||||
transformed = self.transform(image=image, mask=mask)
|
||||
image = transformed['image'] # Tensor (3, H, W)
|
||||
mask = transformed['mask'] # ndarray (H, W)
|
||||
|
||||
# 3. 将mask转换为Tensor
|
||||
mask = torch.from_numpy(mask).long()
|
||||
|
||||
# 4. 检查mask的有效性
|
||||
# 确保mask的值在[0, num_classes-1]范围内
|
||||
if mask.max() >= self.num_classes:
|
||||
print(f"警告: {mask_path.name} 包含无效类别ID {mask.max()}")
|
||||
mask = torch.clamp(mask, 0, self.num_classes - 1)
|
||||
|
||||
return image, mask
|
||||
|
||||
def get_sample_info(self, idx: int) -> dict:
|
||||
"""
|
||||
获取样本的元信息（用于调试和可视化）
|
||||
|
||||
Args:
|
||||
idx: 样本索引
|
||||
|
||||
Returns:
|
||||
info: 包含文件名、路径等信息的字典
|
||||
"""
|
||||
img_path, mask_path = self.valid_samples[idx]
|
||||
return {
|
||||
'image_name': img_path.name,
|
||||
'mask_name': mask_path.name,
|
||||
'image_path': str(img_path),
|
||||
'mask_path': str(mask_path),
|
||||
}
|
||||
|
||||
|
||||
def get_dataloaders(
|
||||
data_root: str,
|
||||
batch_size: int = 4,
|
||||
image_size: int = 512,
|
||||
num_workers: int = 0,
|
||||
num_classes: int = 2
|
||||
) -> Tuple[torch.utils.data.DataLoader, torch.utils.data.DataLoader]:
|
||||
"""
|
||||
创建训练集和验证集的DataLoader
|
||||
|
||||
Args:
|
||||
data_root: 数据集根目录
|
||||
batch_size: 批大小
|
||||
image_size: 图像尺寸
|
||||
num_workers: 数据加载进程数（CPU训练时设为0）
|
||||
num_classes: 类别数
|
||||
|
||||
Returns:
|
||||
train_loader: 训练集DataLoader
|
||||
val_loader: 验证集DataLoader
|
||||
"""
|
||||
# 创建训练集
|
||||
train_dataset = FoodSegmentationDataset(
|
||||
data_root=data_root,
|
||||
split='train',
|
||||
image_size=image_size,
|
||||
num_classes=num_classes,
|
||||
augmentation=True # 训练集使用数据增强
|
||||
)
|
||||
|
||||
# 创建验证集
|
||||
val_dataset = FoodSegmentationDataset(
|
||||
data_root=data_root,
|
||||
split='val',
|
||||
image_size=image_size,
|
||||
num_classes=num_classes,
|
||||
augmentation=False # 验证集不使用数据增强
|
||||
)
|
||||
|
||||
# 创建DataLoader
|
||||
train_loader = torch.utils.data.DataLoader(
|
||||
train_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=True, # 训练集打乱顺序
|
||||
num_workers=num_workers,
|
||||
pin_memory=False, # CPU训练时设为False
|
||||
drop_last=True if len(train_dataset) > batch_size else False
|
||||
)
|
||||
|
||||
val_loader = torch.utils.data.DataLoader(
|
||||
val_dataset,
|
||||
batch_size=batch_size,
|
||||
shuffle=False, # 验证集不打乱
|
||||
num_workers=num_workers,
|
||||
pin_memory=False,
|
||||
drop_last=False
|
||||
)
|
||||
|
||||
print(f"\n✓ DataLoader创建完成")
|
||||
print(f" 训练集: {len(train_dataset)} 样本, {len(train_loader)} 批次")
|
||||
print(f" 验证集: {len(val_dataset)} 样本, {len(val_loader)} 批次")
|
||||
|
||||
return train_loader, val_loader
|
||||
|
||||
|
||||
def test_dataset():
|
||||
"""
|
||||
测试数据集加载是否正常
|
||||
|
||||
用于开发调试，验证：
|
||||
1. 数据能否正确加载
|
||||
2. 数据增强是否正常工作
|
||||
3. 数据的shape和类型是否正确
|
||||
"""
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
print("="*60)
|
||||
print("数据集测试")
|
||||
print("="*60)
|
||||
|
||||
# 配置
|
||||
DATA_ROOT = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/data/segformer_format"
|
||||
|
||||
# 创建数据集
|
||||
dataset = FoodSegmentationDataset(
|
||||
data_root=DATA_ROOT,
|
||||
split='train',
|
||||
image_size=512,
|
||||
augmentation=True
|
||||
)
|
||||
|
||||
# 测试读取第一个样本
|
||||
image, mask = dataset[0]
|
||||
info = dataset.get_sample_info(0)
|
||||
|
||||
print(f"\n样本信息:")
|
||||
print(f" 文件名: {info['image_name']}")
|
||||
print(f" 图像shape: {image.shape}") # 应该是 (3, 512, 512)
|
||||
print(f" Mask shape: {mask.shape}") # 应该是 (512, 512)
|
||||
print(f" Mask唯一值: {torch.unique(mask).numpy()}") # 应该是 [0, 1]
|
||||
print(f" 图像数据范围: [{image.min():.3f}, {image.max():.3f}]")
|
||||
|
||||
# 可视化前3个样本（含数据增强效果）
|
||||
fig, axes = plt.subplots(3, 3, figsize=(12, 12))
|
||||
|
||||
for i in range(3):
|
||||
image, mask = dataset[i]
|
||||
|
||||
# 反标准化用于显示
|
||||
mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
|
||||
std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
|
||||
image_denorm = image * std + mean
|
||||
image_denorm = torch.clamp(image_denorm, 0, 1)
|
||||
|
||||
# 转换为numpy用于显示
|
||||
image_np = image_denorm.permute(1, 2, 0).numpy()
|
||||
mask_np = mask.numpy()
|
||||
|
||||
# 显示图像
|
||||
axes[i, 0].imshow(image_np)
|
||||
axes[i, 0].set_title(f"样本{i+1}: 图像")
|
||||
axes[i, 0].axis('off')
|
||||
|
||||
# 显示mask
|
||||
axes[i, 1].imshow(mask_np, cmap='gray', vmin=0, vmax=1)
|
||||
axes[i, 1].set_title(f"样本{i+1}: Mask")
|
||||
axes[i, 1].axis('off')
|
||||
|
||||
# 显示叠加
|
||||
overlay = image_np.copy()
|
||||
red_mask = np.zeros_like(overlay)
|
||||
red_mask[mask_np == 1] = [1, 0, 0]
|
||||
overlay = overlay * 0.6 + red_mask * 0.4
|
||||
axes[i, 2].imshow(overlay)
|
||||
axes[i, 2].set_title(f"样本{i+1}: 叠加")
|
||||
axes[i, 2].axis('off')
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig("dataset_test_result.png", dpi=120)
|
||||
print(f"\n✓ 测试结果已保存: dataset_test_result.png")
|
||||
plt.show()
|
||||
|
||||
print("\n" + "="*60)
|
||||
print("✓ 数据集测试通过!")
|
||||
print("="*60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_dataset()
|
||||
@@ -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