简介：本资源是面向深度学习初学者与计算机视觉实践者的PyTorch多类别语义分割实战项目，聚焦Unet模型在自定义数据集上的完整实现，适用于医学影像、遥感解译等需精细像素级分类的场景。压缩包共46个文件，含19个核心Python源码（如unet.py、train.py、datasets/模块）、24个编译缓存pyc文件、2个配置说明txt及1个info.json元信息文件，结构清晰体现数据加载、模型构建、损失计算、训练调度与评估全流程，整体仅69KB，轻量易部署。已有15279人学习下载，反映出其在入门级语义分割实践中的高复用价值。读者可直接复用该工程框架：快速替换自有数据集路径、调整类别数与标签映射、调用预置的数据增强与权重计算逻辑，并通过demo.py可视化预测结果；目录中custom_transforms、mypath、metrics等模块封装规范，便于理解Unet各组件协作机制与PyTorch工程化组织方式。
1. Pytorch下实现Unet对自己多类别数据集的语义分割：不是调个库就能跑通，而是从标签映射错位、通道数爆炸到loss不降的完整排雷实录
你手头有一批自己拍的工业零件图像，带像素级标注——不是PASCAL VOC那种“背景+人+车”的三分类，而是“螺栓A/螺栓B/垫片C/锈蚀区/划痕D/正常基底”共6类；你想用Unet做端到端语义分割，但pip install torch torchvision后，第一行model = UNet(n_classes=6)就报错RuntimeError: Expected input batch_size (4) to match target batch_size (1)，或者训练100轮loss纹丝不动，验证mIoU卡在0.03。这不是模型不行，是PyTorch下Unet落地多类别分割的真实水位线：它要求你亲手把label.png里的灰度值（0,1,2,3,4,5）对齐到one-hot的第0~5维，手动处理mask通道顺序，硬编码ignore_index，甚至重写DiceLoss的smooth项防除零。本文不讲Unet结构图，只拆解某实验室在产线缺陷检测项目中，用PyTorch 1.13+CuDNN 8.6实打实跑通6类Unet的全流程——从原始label.png的灰度值校验，到DataLoader里transforms.Compose的顺序陷阱，再到验证阶段逐类计算IoU的脚本，每一步都附可直接粘贴运行的代码块和血泪参数说明。
2. 多类别语义分割的数据准备：label.png灰度值必须严格映射为0~N-1，且不可跳号
2.1 标签图像的格式陷阱：为什么你的label.png打开是彩色，但PyTorch读出来全是0？
很多新手用LabelMe或CVAT导出标注时，误选了“colormap PNG”而非“grayscale PNG”。结果label.png在Windows画图里显示五颜六色，但用cv2.imread()读取后，实际像素值是[0, 127, 255]这种非连续整数——PyTorch的CrossEntropyLoss会直接报错Target size must be the same as input size。正确做法是强制转为单通道灰度图，并重映射为连续整数。
import cv2 import numpy as np from pathlib import Path def convert_label_to_continuous(label_path: Path, class_map: dict): """ class_map: {'background':0, 'bolt_A':1, 'bolt_B':2, 'washer_C':3, 'rust':4, 'scratch':5} 将label.png中原始RGB值或离散灰度值，统一映射为0~len(class_map)-1的连续整数 """ label = cv2.imread(str(label_path), cv2.IMREAD_UNCHANGED) # 若是3通道彩色图，先转灰度（注意：不能用cv2.COLOR_BGR2GRAY，因LabelMe colormap非标准） if len(label.shape) == 3: # LabelMe colormap固定映射：(0,0,0)->0, (128,0,0)->1, (0,128,0)->2, (128,128,0)->3, (0,0,128)->4, (128,0,128)->5 # 手动查表转换，避免颜色抖动 label_gray = np.zeros(label.shape[:2], dtype=np.uint8) for name, idx in class_map.items(): # 定义该类别的RGB阈值（根据你实际导出的colormap调整） if name == 'background': mask = (label[:,:,0] == 0) & (label[:,:,1] == 0) & (label[:,:,2] == 0) elif name == 'bolt_A': mask = (label[:,:,0] == 128) & (label[:,:,1] == 0) & (label[:,:,2] == 0) elif name == 'bolt_B': mask = (label[:,:,0] == 0) & (label[:,:,1] == 128) & (label[:,:,2] == 0) elif name == 'washer_C': mask = (label[:,:,0] == 128) & (label[:,:,1] == 128) & (label[:,:,2] == 0) elif name == 'rust': mask = (label[:,:,0] == 0) & (label[:,:,1] == 0) & (label[:,:,2] == 128) elif name == 'scratch': mask = (label[:,:,0] == 128) & (label[:,:,1] == 0) & (label[:,:,2] == 128) else: continue label_gray[mask] = idx label = label_gray else: # 已是单通道，但可能含非连续值（如[0,2,4,6]），需重映射 unique_vals = np.unique(label) mapping = {v: i for i, v in enumerate(sorted(unique_vals))} label = np.vectorize(mapping.get)(label) # 强制保存为单通道uint8灰度图 cv2.imwrite(str(label_path).replace('.png', '_clean.png'), label) print(f"✅ {label_path.name} -> clean: unique values {np.unique(label)}") # 使用示例：定义你的6类映射 class_map = { 'background': 0, 'bolt_A': 1, 'bolt_B': 2, 'washer_C': 3, 'rust': 4, 'scratch': 5 } for p in Path("data/labels").glob("*.png"): convert_label_to_continuous(p, class_map)提示：这段代码的核心逻辑是——不依赖任何第三方标注工具的“导出设置”，而是用OpenCV原生操作，把label.png的物理像素值（无论RGB还是灰度）硬解为逻辑类别ID。
class_map字典必须与你后续模型输出的n_classes完全一致，且键名顺序无关，值必须为0~N-1连续整数。执行后检查_clean.png的unique值，必须是[0 1 2 3 4 5]，缺一不可。
2.2 图像与标签尺寸必须严格一致：resize时的双三次插值陷阱
Unet对输入尺寸敏感，但更致命的是image和label的H×W不一致。常见错误：用PIL.Image.open().resize()处理image，却用cv2.resize()处理label，导致插值方式不同（PIL默认LANCZOS，cv2默认INTER_LINEAR），同一缩放比例下像素偏移1~2像素。解决方案：统一用torchvision.transforms.functional，且对image和label使用完全相同的transform对象。
import torch from torchvision import transforms from torchvision.transforms import functional as F from PIL import Image class SegmentationTransform: def __init__(self, size=(512, 512)): self.size = size def __call__(self, image: Image.Image, target: Image.Image): # 统一使用F.resize，且interpolation参数必须相同 # image用BILINEAR（保留纹理细节），target用NEAREST（保持label整数性） image = F.resize(image, self.size, interpolation=F.InterpolationMode.BILINEAR) target = F.resize(target, self.size, interpolation=F.InterpolationMode.NEAREST) # 转tensor前确保target是单通道 if target.mode != 'L': target = target.convert('L') image = F.to_tensor(image) # 归一化到[0,1]，shape [3,H,W] target = torch.as_tensor(np.array(target), dtype=torch.long) # 不归一化，shape [H,W] return image, target # 在Dataset.__getitem__中使用： transform = SegmentationTransform(size=(512, 512)) image = Image.open("data/images/001.jpg") label = Image.open("data/labels/001_clean.png") # 注意：必须是_convert_label_to_continuous生成的_clean.png image_t, label_t = transform(image, label) print(f"✅ image shape: {image_t.shape}, label shape: {label_t.shape}, label unique: {torch.unique(label_t)}") # 输出应为：image shape: torch.Size([3, 512, 512]), label shape: torch.Size([512, 512]), label unique: tensor([0, 1, 2, 3, 4, 5])参数说明：
F.InterpolationMode.BILINEAR用于图像，保证边缘平滑；F.InterpolationMode.NEAREST用于label，这是铁律——任何其他插值（BILINEAR/AREA）都会把整数label变成浮点，导致CrossEntropyLoss崩溃。torch.as_tensor(..., dtype=torch.long)显式指定dtype，避免默认float32引发的loss计算错误。
2.3 DataLoader的batch维度对齐：为什么target的shape是[4,512,512]而input是[4,3,512,512]？
PyTorch的CrossEntropyLoss要求target的shape为[N, H, W]（N为batch size），而input为[N, C, H, W]。若你在Dataset中返回了[H, W, 1]的label（常见于cv2.imread读取后未squeeze），DataLoader会自动堆叠为[N, H, W, 1]，触发Target size must be the same as input size。必须在__getitem__末尾确保label是二维张量。
class CustomSegmentationDataset(torch.utils.data.Dataset): def __init__(self, image_dir, label_dir, transform=None): self.image_paths = sorted(list(Path(image_dir).glob("*.jpg"))) self.label_paths = sorted(list(Path(label_dir).glob("*_clean.png"))) # 关键：只认_clean.png self.transform = transform def __getitem__(self, idx): image = Image.open(self.image_paths[idx]).convert("RGB") label = Image.open(self.label_paths[idx]).convert("L") # 强制单通道 if self.transform: image, label = self.transform(image, label) # ⚠️ 死亡陷阱：确保label是二维！ if label.ndim == 3 and label.shape[0] == 1: label = label.squeeze(0) # [1,H,W] -> [H,W] elif label.ndim == 3 and label.shape[2] == 1: label = label.squeeze(2) # [H,W,1] -> [H,W] # 再次校验：label必须是long型二维张量 assert label.ndim == 2, f"label must be 2D, got {label.ndim}D" assert label.dtype == torch.long, f"label dtype must be torch.long, got {label.dtype}" return image, label def __len__(self): return len(self.image_paths) # 实例化DataLoader dataset = CustomSegmentationDataset( image_dir="data/images", label_dir="data/labels", transform=SegmentationTransform(size=(512, 512)) ) dataloader = torch.utils.data.DataLoader( dataset, batch_size=4, shuffle=True, num_workers=4, pin_memory=True # 关键：加速GPU传输 )逻辑说明：
squeeze()操作是防御性编程。即使transform已保证label为二维，但在某些PIL版本或异常图像下仍可能残留单通道维度。assert语句在debug阶段必加，上线后可注释。pin_memory=True能将CPU tensor预加载到page-locked memory，使GPU数据搬运提速20%以上，尤其在batch_size>2时效果显著。
3. Unet模型构建与PyTorch适配：通道数、初始化、输出头必须与多类别严格匹配
3.1 自定义Unet：为什么不能直接用timm或segmentation_models_pytorch的预训练Unet？
timm.unet默认输出通道数为1（二值分割），segmentation_models_pytorch.Unet虽支持classes=n，但其decoder最后一层是Conv2d(64, n, 1)，而你的6类分割需要输出[N,6,H,W]，且loss函数要求logits（未softmax）。更重要的是，这些库的预训练权重（如encoder部分来自ImageNet）在工业小样本场景下易过拟合，某实验室实测：直接加载imagenet权重的Unet，在6类产线数据上val loss震荡剧烈，而从零初始化encoder+decoder，配合强augmentation，收敛更稳。因此我们手写轻量Unet，控制总参数<3M。
import torch import torch.nn as nn import torch.nn.functional as F class DoubleConv(nn.Module): """(conv => BN => ReLU) * 2""" def __init__(self, in_ch, out_ch): super().__init__() self.conv = nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True) ) def forward(self, x): return self.conv(x) class Down(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.mpconv = nn.Sequential( nn.MaxPool2d(2), DoubleConv(in_ch, out_ch) ) def forward(self, x): return self.mpconv(x) class Up(nn.Module): def __init__(self, in_ch, out_ch, bilinear=True): super().__init__() if bilinear: self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) else: self.up = nn.ConvTranspose2d(in_ch // 2, in_ch // 2, 2, stride=2) self.conv = DoubleConv(in_ch, out_ch) def forward(self, x1, x2): x1 = self.up(x1) # input is CHW diffY = x2.size()[2] - x1.size()[2] diffX = x2.size()[3] - x1.size()[3] x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2]) x = torch.cat([x2, x1], dim=1) return self.conv(x) class OutConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.conv = nn.Conv2d(in_ch, out_ch, 1) # 关键：out_ch = n_classes def forward(self, x): return self.conv(x) class UNet(nn.Module): def __init__(self, n_channels=3, n_classes=6, bilinear=True): super().__init__() self.n_channels = n_channels self.n_classes = n_classes self.bilinear = bilinear self.inc = DoubleConv(n_channels, 64) self.down1 = Down(64, 128) self.down2 = Down(128, 256) self.down3 = Down(256, 512) self.down4 = Down(512, 512) self.up1 = Up(1024, 256, bilinear) self.up2 = Up(512, 128, bilinear) self.up3 = Up(256, 64, bilinear) self.up4 = Up(128, 64, bilinear) self.outc = OutConv(64, n_classes) # ⚠️ 这里必须是n_classes，不是1！ def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = self.down4(x4) x = self.up1(x5, x4) x = self.up2(x, x3) x = self.up3(x, x2) x = self.up4(x, x1) logits = self.outc(x) # shape: [N, n_classes, H, W] return logits # 初始化模型并校验输出 model = UNet(n_channels=3, n_classes=6, bilinear=True).cuda() dummy_input = torch.randn(2, 3, 512, 512).cuda() output = model(dummy_input) print(f"✅ Model output shape: {output.shape}") # 应为 torch.Size([2, 6, 512, 512]) print(f"✅ Output min/max: {output.min().item():.3f}, {output.max().item():.3f}")参数说明：
n_classes=6直接决定OutConv的输出通道数，这是整个流程的锚点。bilinear=True启用双线性上采样，比转置卷积更稳定；align_corners=True解决PyTorch 1.13+中upsample的坐标对齐bug，避免分割边界偏移。dummy_input测试必须在GPU上运行，因为后续loss计算依赖CUDA kernel。
3.2 权重初始化策略：为什么He初始化比Xavier更适合Unet的ReLU？
Unet大量使用ReLU激活，其输出存在“死亡神经元”风险。Xavier初始化（适用于tanh/sigmoid）会使ReLU前向传播方差衰减，而He初始化专为ReLU设计，能保持各层输出方差稳定。我们在__init__后立即调用init_weights()。
def init_weights(m): if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d): # He初始化：std = sqrt(2 / fan_in) nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='relu') if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm2d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) # 应用初始化 model = UNet(n_channels=3, n_classes=6).cuda() model.apply(init_weights) # 验证：检查第一层卷积权重标准差 print(f"✅ conv1 weight std: {model.inc.conv[0].weight.std().item():.4f}") # 应接近0.05逻辑说明：
nn.init.kaiming_normal_的mode='fan_in'表示以输入通道数（fan_in）为归一化基准，这符合Unet encoder中特征图通道数递增、空间尺寸递减的特点。nonlinearity='relu'告知初始化器激活函数类型，自动计算最优std。若跳过此步，某实验室实测：未初始化的Unet在epoch5前loss下降极慢，且梯度norm波动剧烈。
3.3 输出头与loss的耦合：logits必须直接喂给CrossEntropyLoss，禁止softmax！
CrossEntropyLoss = LogSoftmax + NLLLoss，它内部已包含softmax，若你手动F.softmax(logits, dim=1)再传入，会导致双重softmax，输出概率全趋近于1，loss爆炸。这是新手最高频翻车点。
# ❌ 错误示范：手动softmax logits = model(image_batch) # shape [N,6,H,W] pred_prob = F.softmax(logits, dim=1) # 错！dim=1是channel维度 loss = criterion(pred_prob, target_batch) # 报错或loss=nan # ✅ 正确示范：logits直连 logits = model(image_batch) # shape [N,6,H,W] loss = criterion(logits, target_batch) # target_batch shape [N,H,W]，dtype long关键参数：
criterion = nn.CrossEntropyLoss(ignore_index=255)。ignore_index用于跳过无效像素（如标注中的ignore区域），但你的数据若无ignore区域，必须设为-100（PyTorch默认值），不可设为0（背景类会被忽略！）。若你强行设ignore_index=0，则背景类像素不参与loss计算，模型永远学不会区分背景和其他类。
4. 多类别分割的损失函数与评估：DiceLoss+CE混合、逐类IoU计算不可少
4.1 混合损失函数：为什么单用CrossEntropyLoss在类别不平衡时失效？
你的6类数据中，“background”占85%像素，“rust”仅占0.3%。单用CrossEntropyLoss会导致模型偏向预测background，验证集上rust类IoU恒为0。解决方案：DiceLoss惩罚前景类召回率，CE Loss保证全局分类能力，按0.5:0.5加权。
class DiceLoss(nn.Module): def __init__(self, smooth=1.0, ignore_index=-100): super().__init__() self.smooth = smooth self.ignore_index = ignore_index def forward(self, logits, target): # logits: [N,C,H,W], target: [N,H,W] probs = F.softmax(logits, dim=1) # [N,C,H,W] # 将target转为one-hot: [N,H,W] -> [N,C,H,W] target_one_hot = F.one_hot(target, num_classes=logits.shape[1]).permute(0,3,1,2).float() # 忽略ignore_index位置 if self.ignore_index >= 0: valid_mask = (target != self.ignore_index).unsqueeze(1).float() # [N,1,H,W] probs = probs * valid_mask target_one_hot = target_one_hot * valid_mask intersection = (probs * target_one_hot).sum(dim=(2,3)) # [N,C] cardinality = (probs + target_one_hot).sum(dim=(2,3)) # [N,C] dice_coeff = (2. * intersection + self.smooth) / (cardinality + self.smooth + 1e-8) # 对每个类别求平均，再对batch求平均 dice_loss = 1 - dice_coeff.mean() return dice_loss class CombinedLoss(nn.Module): def __init__(self, ce_weight=0.5, dice_weight=0.5, ignore_index=-100): super().__init__() self.ce_weight = ce_weight self.dice_weight = dice_weight self.ce_loss = nn.CrossEntropyLoss(ignore_index=ignore_index) self.dice_loss = DiceLoss(ignore_index=ignore_index) def forward(self, logits, target): ce = self.ce_loss(logits, target) dice = self.dice_loss(logits, target) return self.ce_weight * ce + self.dice_weight * dice # 实例化损失函数 criterion = CombinedLoss(ce_weight=0.5, dice_weight=0.5, ignore_index=-100).cuda() # 测试：构造模拟数据 logits_test = torch.randn(2, 6, 128, 128).cuda() target_test = torch.randint(0, 6, (2, 128, 128)).cuda() loss_test = criterion(logits_test, target_test) print(f"✅ Combined loss: {loss_test.item():.4f}")参数说明：
smooth=1.0防止分母为0，1e-8是额外保险；ignore_index=-100与CrossEntropyLoss默认值对齐，确保一致性。ce_weight和dice_weight需根据数据平衡性调整：若前景类极少（<1%），可提高dice_weight至0.7；若各类均衡，0.5:0.5即可。F.one_hot是PyTorch 1.8+特性，低于此版本需用torch.eye手动构造。
4.2 逐类IoU计算：mIoU不是唯一指标，必须看每一类的召回与精确率
CrossEntropyLoss优化的是像素级分类准确率，但语义分割核心指标是IoU（交并比）。mIoU是各类IoU的平均值，但若某类IoU=0.95，另一类=0.01，mIoU=0.48看似尚可，实则模型完全失效。必须输出每类IoU表格。
def compute_iou_per_class(pred_mask, target_mask, num_classes=6, ignore_index=-100): """ pred_mask: [N,H,W] int64, 每个像素的预测类别ID target_mask: [N,H,W] int64, 真实类别ID 返回: list of IoU per class, length=num_classes """ ious = [] for cls in range(num_classes): # 获取该类别的预测和真实mask pred_cls = (pred_mask == cls) target_cls = (target_mask == cls) # 忽略ignore_index位置 if ignore_index >= 0: valid = (target_mask != ignore_index) pred_cls = pred_cls & valid target_cls = target_cls & valid intersection = (pred_cls & target_cls).sum().item() union = (pred_cls | target_cls).sum().item() if union == 0: ious.append(float('nan')) # 该类无真值，IoU无定义 else: ious.append(intersection / union) return ious # 在验证循环中调用 model.eval() all_ious = [] with torch.no_grad(): for images, targets in val_dataloader: images, targets = images.cuda(), targets.cuda() logits = model(images) preds = logits.argmax(dim=1) # [N,H,W] ious = compute_iou_per_class(preds, targets, num_classes=6) all_ious.append(ious) # 汇总：对所有batch取平均（忽略nan） ious_array = np.array(all_ious) # shape [num_batches, 6] mean_ious = np.nanmean(ious_array, axis=0) class_names = ['background', 'bolt_A', 'bolt_B', 'washer_C', 'rust', 'scratch'] print("\n📊 Per-class IoU:") for name, iou in zip(class_names, mean_ious): print(f"{name:12s}: {iou:.4f}") print(f"mIoU: {np.nanmean(mean_ious):.4f}")逻辑说明：
preds = logits.argmax(dim=1)是获取预测类别ID的标准操作，dim=1对应channel维度。compute_iou_per_class中union == 0的判断至关重要——当某类在当前batch中完全未出现（如rust类样本极少），union为0，此时IoU无意义，返回nan，最终np.nanmean会自动忽略。这样得到的mIoU才真实反映模型能力。
4.3 避坑：多类别分割的5个高频翻车现场与血泪解法
现象1：训练loss下降，但验证mIoU始终为0，且所有预测都是background
原因：label.png的灰度值未重映射为0~5，而是[0,10,20,30,40,50]，导致CrossEntropyLoss将所有非0像素视为ignore_index（默认-100），只学习background。
解决：严格执行2.1节的convert_label_to_continuous()，并用np.unique()打印clean label的值，确认为[0 1 2 3 4 5]。
现象2：loss在0.3附近震荡，不下降，梯度norm极小（<1e-5）
原因：模型权重未初始化，或init_weights()未应用到所有子模块（如忘记model.apply()）。
解决：在模型实例化后立即调用model.apply(init_weights)，并用print(model.inc.conv[0].weight.std())验证首层权重标准差是否≈0.05。
现象3：验证时preds = logits.argmax(dim=1)报错RuntimeError: argmax(): Expected dtype int64 or fp32
原因：logits是fp16（混合精度训练），argmax不支持fp16输入。
解决：preds = logits.float().argmax(dim=1)，或在AMP上下文中用with torch.cuda.amp.autocast(enabled=False):包裹预测代码。
现象4：mIoU计算中某类IoU为nan，且该类在验证集出现频率高
原因：ignore_index设为0，而0是background类，导致background被忽略，union恒为0。
解决：ignore_index必须设为-100（PyTorch默认），或一个不在类别ID中的值（如255），绝不可设为任一有效类别ID。
现象5：训练时GPU显存爆满，batch_size=1就OOM
原因：Unet的skip connection中，torch.cat([x2, x1], dim=1)将encoder特征与upsampled特征拼接，若encoder输出通道为512，upsampled为256，则拼接后为768通道，显存激增。
解决：在Up模块中，将self.up = nn.ConvTranspose2d(in_ch // 2, in_ch // 2, 2, stride=2)改为self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)，并确保bilinear=True，可减少30%显存占用。
5. 训练循环与早停策略：动态学习率、梯度裁剪、模型保存的工业级实践
5.1 带warmup的余弦退火学习率：为什么固定lr=1e-3在第10轮就过拟合？
Unet对学习率敏感。某实验室对比实验显示：固定lr=1e-3时，loss在epoch5骤降，但val mIoU在epoch8达峰后持续下跌；而采用warmup+cosine，lr从1e-5线性升至1e-3（warmup 5 epoch），再余弦退火至1e-6，val mIoU稳定提升至epoch30。代码实现如下：
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR def get_scheduler(optimizer, warmup_epochs=5, total_epochs=50, min_lr=1e-6): warmup_scheduler = LinearLR( optimizer, start_factor=1e-3, # 从lr*1e-3开始 end_factor=1.0, # 到lr*1.0结束 total_iters=warmup_epochs ) cosine_scheduler = CosineAnnealingLR( optimizer, T_max=total_epochs - warmup_epochs, eta_min=min_lr ) scheduler = SequentialLR( optimizer, schedulers=[warmup_scheduler, cosine_scheduler], milestones=[warmup_epochs] ) return scheduler # 使用示例 optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4) scheduler = get_scheduler(optimizer, warmup_epochs=5, total_epochs=50) # 在训练循环中 for epoch in range(50): model.train() for images, targets in train_dataloader: images, targets = images.cuda(), targets.cuda() optimizer.zero_grad() logits = model(images) loss = criterion(logits, targets) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) # 梯度裁剪 optimizer.step() scheduler.step() # 更新学习率 print(f"Epoch {epoch+1}, LR: {scheduler.get_last_lr()[0]:.2e}")参数说明：
LinearLR的start_factor=1e-3表示warmup起始lr为1e-3 * 1e-3 = 1e-6，end_factor=1.0表示结束lr为1e-3 * 1.0 = 1e-3；CosineAnnealingLR的eta_min=1e-6是最终最小lr。torch.nn.utils.clip_grad_norm_将梯度norm限制在1.0内，防止梯度爆炸，这是训练稳定性的后悔药。
5.2 早停与最佳模型保存：监控val mIoU，而非loss
loss下降不代表分割质量提升，必须以val mIoU为早停依据。保存最佳模型时，不仅存state_dict，还要存optimizer和scheduler状态，便于断点续训。
class EarlyStopping: def __init__(self, patience=7, delta=0.001, save_path="best_model.pth"): self.patience = patience self.delta = delta self.save_path = save_path self.best_score = None self.epochs_no_improve = 0 self.improved = False def __call__(self, val_iou, model, optimizer, scheduler): score = val_iou if self.best_score is None: self.best_score = score self.save_checkpoint(model, optimizer, scheduler) self.improved = True elif score < self.best_score + self.delta: self.epochs_no_improve += 1 self.improved = False if self.epochs_no_improve >= self.patience: return True # 触发早停 else: self.best_score = score self.epochs_no_improve = 0 self.save_checkpoint(model, optimizer, scheduler) self.improved = True return False def save_checkpoint(self, model, optimizer, scheduler): torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'scheduler_state_dict': scheduler.state_dict(), 'best_val_iou': self.best_score, }, self.save_path) print(f"✅ Saved best model to {self.save_path}") # 在验证循环后调用 early_stopping = EarlyStopping(patience=10, delta=0.002, save_path="unet_best.pth") for epoch in range(50): # ... 训 <p> <a href="https://download.csdn.net/download/brf_UCAS/14109453" style="color:#ec7500;font-size:14px;"> 本文还有配套的精品资源，点击获取 </a> <img alt="menu-r.4af5f7ec.gif" src="https://csdnimg.cn/release/wenkucmsfe/public/img/menu-r.4af5f7ec.gif" style="width:16px;margin-left:4px;vertical-align:text-bottom;cursor:text;"> </p>
