簡介本資源是面向深度學習初學者與計算機視覺實踐者的PyTorch多類別語義分割實戰(zhàn)項目聚焦Unet模型在自定義數(shù)據(jù)集上的完整實現(xiàn)適用于醫(yī)學影像、遙感解譯等需精細像素級分類的場景。壓縮包共46個文件含19個核心Python源碼如unet.py、train.py、datasets/模塊、24個編譯緩存pyc文件、2個配置說明txt及1個info.json元信息文件結構清晰體現(xiàn)數(shù)據(jù)加載、模型構建、損失計算、訓練調度與評估全流程整體僅69KB輕量易部署。已有15279人學習下載反映出其在入門級語義分割實踐中的高復用價值。讀者可直接復用該工程框架快速替換自有數(shù)據(jù)集路徑、調整類別數(shù)與標簽映射、調用預置的數(shù)據(jù)增強與權重計算邏輯并通過demo.py可視化預測結果目錄中custom_transforms、mypath、metrics等模塊封裝規(guī)范便于理解Unet各組件協(xié)作機制與PyTorch工程化組織方式。1. Pytorch下實現(xiàn)Unet對自己多類別數(shù)據(jù)集的語義分割不是調個庫就能跑通而是從標簽映射錯位、通道數(shù)爆炸到loss不降的完整排雷實錄你手頭有一批自己拍的工業(yè)零件圖像帶像素級標注——不是PASCAL VOC那種“背景人車”的三分類而是“螺栓A/螺栓B/墊片C/銹蝕區(qū)/劃痕D/正?；住惫?類你想用Unet做端到端語義分割但pip install torch torchvision后第一行model UNet(n_classes6)就報錯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.13CuDNN 8.6實打實跑通6類Unet的全流程——從原始label.png的灰度值校驗到DataLoader里transforms.Compose的順序陷阱再到驗證階段逐類計算IoU的腳本每一步都附可直接粘貼運行的代碼塊和血淚參數(shù)說明。2. 多類別語義分割的數(shù)據(jù)準備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]這種非連續(xù)整數(shù)——PyTorch的CrossEntropyLoss會直接報錯Target size must be the same as input size。正確做法是強制轉為單通道灰度圖并重映射為連續(xù)整數(shù)。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值或離散灰度值統(tǒng)一映射為0~len(class_map)-1的連續(xù)整數(shù) 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], dtypenp.uint8) for name, idx in class_map.items(): # 定義該類別的RGB閾值根據(jù)你實際導出的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: # 已是單通道但可能含非連續(xù)值如[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字典必須與你后續(xù)模型輸出的n_classes完全一致且鍵名順序無關值必須為0~N-1連續(xù)整數(shù)。執(zhí)行后檢查_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默認LANCZOScv2默認INTER_LINEAR同一縮放比例下像素偏移1~2像素。解決方案統(tǒng)一用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): # 統(tǒng)一使用F.resize且interpolation參數(shù)必須相同 # image用BILINEAR保留紋理細節(jié)target用NEAREST保持label整數(shù)性 image F.resize(image, self.size, interpolationF.InterpolationMode.BILINEAR) target F.resize(target, self.size, interpolationF.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), dtypetorch.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])參數(shù)說明F.InterpolationMode.BILINEAR用于圖像保證邊緣平滑F(xiàn).InterpolationMode.NEAREST用于label這是鐵律——任何其他插值BILINEAR/AREA都會把整數(shù)label變成浮點導致CrossEntropyLoss崩潰。torch.as_tensor(..., dtypetorch.long)顯式指定dtype避免默認float32引發(fā)的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讀取后未squeezeDataLoader會自動堆疊為[N, H, W, 1]觸發(fā)Target size must be the same as input size。必須在__getitem__末尾確保label是二維張量。class CustomSegmentationDataset(torch.utils.data.Dataset): def __init__(self, image_dir, label_dir, transformNone): 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, flabel must be 2D, got {label.ndim}D assert label.dtype torch.long, flabel dtype must be torch.long, got {label.dtype} return image, label def __len__(self): return len(self.image_paths) # 實例化DataLoader dataset CustomSegmentationDataset( image_dirdata/images, label_dirdata/labels, transformSegmentationTransform(size(512, 512)) ) dataloader torch.utils.data.DataLoader( dataset, batch_size4, shuffleTrue, num_workers4, pin_memoryTrue # 關鍵加速GPU傳輸 )邏輯說明squeeze()操作是防御性編程。即使transform已保證label為二維但在某些PIL版本或異常圖像下仍可能殘留單通道維度。assert語句在debug階段必加上線后可注釋。pin_memoryTrue能將CPU tensor預加載到page-locked memory使GPU數(shù)據(jù)搬運提速20%以上尤其在batch_size2時效果顯著。3. Unet模型構建與PyTorch適配通道數(shù)、初始化、輸出頭必須與多類別嚴格匹配3.1 自定義Unet為什么不能直接用timm或segmentation_models_pytorch的預訓練Unettimm.unet默認輸出通道數(shù)為1二值分割segmentation_models_pytorch.Unet雖支持classesn但其decoder最后一層是Conv2d(64, n, 1)而你的6類分割需要輸出[N,6,H,W]且loss函數(shù)要求logits未softmax。更重要的是這些庫的預訓練權重如encoder部分來自ImageNet在工業(yè)小樣本場景下易過擬合某實驗室實測直接加載imagenet權重的Unet在6類產線數(shù)據(jù)上val loss震蕩劇烈而從零初始化encoderdecoder配合強augmentation收斂更穩(wěn)。因此我們手寫輕量Unet控制總參數(shù)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, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue), nn.Conv2d(out_ch, out_ch, 3, padding1, biasFalse), nn.BatchNorm2d(out_ch), nn.ReLU(inplaceTrue) ) 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, bilinearTrue): super().__init__() if bilinear: self.up nn.Upsample(scale_factor2, modebilinear, align_cornersTrue) else: self.up nn.ConvTranspose2d(in_ch // 2, in_ch // 2, 2, stride2) 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], dim1) 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_channels3, n_classes6, bilinearTrue): 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_channels3, n_classes6, bilinearTrue).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})參數(shù)說明n_classes6直接決定OutConv的輸出通道數(shù)這是整個流程的錨點。bilinearTrue啟用雙線性上采樣比轉置卷積更穩(wěn)定align_cornersTrue解決PyTorch 1.13中upsample的坐標對齊bug避免分割邊界偏移。dummy_input測試必須在GPU上運行因為后續(xù)loss計算依賴CUDA kernel。3.2 權重初始化策略為什么He初始化比Xavier更適合Unet的ReLUUnet大量使用ReLU激活其輸出存在“死亡神經(jīng)元”風險。Xavier初始化適用于tanh/sigmoid會使ReLU前向傳播方差衰減而He初始化專為ReLU設計能保持各層輸出方差穩(wěn)定。我們在__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, modefan_in, nonlinearityrelu) 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_channels3, n_classes6).cuda() model.apply(init_weights) # 驗證檢查第一層卷積權重標準差 print(f? conv1 weight std: {model.inc.conv[0].weight.std().item():.4f}) # 應接近0.05邏輯說明nn.init.kaiming_normal_的modefan_in表示以輸入通道數(shù)fan_in為歸一化基準這符合Unet encoder中特征圖通道數(shù)遞增、空間尺寸遞減的特點。nonlinearityrelu告知初始化器激活函數(shù)類型自動計算最優(yōu)std。若跳過此步某實驗室實測未初始化的Unet在epoch5前l(fā)oss下降極慢且梯度norm波動劇烈。3.3 輸出頭與loss的耦合logits必須直接喂給CrossEntropyLoss禁止softmaxCrossEntropyLoss LogSoftmax NLLLoss它內部已包含softmax若你手動F.softmax(logits, dim1)再傳入會導致雙重softmax輸出概率全趨近于1loss爆炸。這是新手最高頻翻車點。# ? 錯誤示范手動softmax logits model(image_batch) # shape [N,6,H,W] pred_prob F.softmax(logits, dim1) # 錯dim1是channel維度 loss criterion(pred_prob, target_batch) # 報錯或lossnan # ? 正確示范logits直連 logits model(image_batch) # shape [N,6,H,W] loss criterion(logits, target_batch) # target_batch shape [N,H,W]dtype long關鍵參數(shù)criterion nn.CrossEntropyLoss(ignore_index255)。ignore_index用于跳過無效像素如標注中的ignore區(qū)域但你的數(shù)據(jù)若無ignore區(qū)域必須設為-100PyTorch默認值不可設為0背景類會被忽略。若你強行設ignore_index0則背景類像素不參與loss計算模型永遠學不會區(qū)分背景和其他類。4. 多類別分割的損失函數(shù)與評估DiceLossCE混合、逐類IoU計算不可少4.1 混合損失函數(shù)為什么單用CrossEntropyLoss在類別不平衡時失效你的6類數(shù)據(jù)中“background”占85%像素“rust”僅占0.3%。單用CrossEntropyLoss會導致模型偏向預測background驗證集上rust類IoU恒為0。解決方案DiceLoss懲罰前景類召回率CE Loss保證全局分類能力按0.5:0.5加權。class DiceLoss(nn.Module): def __init__(self, smooth1.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, dim1) # [N,C,H,W] # 將target轉為one-hot: [N,H,W] - [N,C,H,W] target_one_hot F.one_hot(target, num_classeslogits.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_weight0.5, dice_weight0.5, ignore_index-100): super().__init__() self.ce_weight ce_weight self.dice_weight dice_weight self.ce_loss nn.CrossEntropyLoss(ignore_indexignore_index) self.dice_loss DiceLoss(ignore_indexignore_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 # 實例化損失函數(shù) criterion CombinedLoss(ce_weight0.5, dice_weight0.5, ignore_index-100).cuda() # 測試構造模擬數(shù)據(jù) 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})參數(shù)說明smooth1.0防止分母為01e-8是額外保險ignore_index-100與CrossEntropyLoss默認值對齊確保一致性。ce_weight和dice_weight需根據(jù)數(shù)據(jù)平衡性調整若前景類極少1%可提高dice_weight至0.7若各類均衡0.5:0.5即可。F.one_hot是PyTorch 1.8特性低于此版本需用torch.eye手動構造。4.2 逐類IoU計算mIoU不是唯一指標必須看每一類的召回與精確率CrossEntropyLoss優(yōu)化的是像素級分類準確率但語義分割核心指標是IoU交并比。mIoU是各類IoU的平均值但若某類IoU0.95另一類0.01mIoU0.48看似尚可實則模型完全失效。必須輸出每類IoU表格。def compute_iou_per_class(pred_mask, target_mask, num_classes6, ignore_index-100): pred_mask: [N,H,W] int64, 每個像素的預測類別ID target_mask: [N,H,W] int64, 真實類別ID 返回: list of IoU per class, lengthnum_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 # 在驗證循環(huán)中調用 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(dim1) # [N,H,W] ious compute_iou_per_class(preds, targets, num_classes6) all_ious.append(ious) # 匯總對所有batch取平均忽略nan ious_array np.array(all_ious) # shape [num_batches, 6] mean_ious np.nanmean(ious_array, axis0) 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(fmIoU: {np.nanmean(mean_ious):.4f})邏輯說明preds logits.argmax(dim1)是獲取預測類別ID的標準操作dim1對應channel維度。compute_iou_per_class中union 0的判斷至關重要——當某類在當前batch中完全未出現(xiàn)如rust類樣本極少union為0此時IoU無意義返回nan最終np.nanmean會自動忽略。這樣得到的mIoU才真實反映模型能力。4.3 避坑多類別分割的5個高頻翻車現(xiàn)場與血淚解法現(xiàn)象1訓練loss下降但驗證mIoU始終為0且所有預測都是background原因label.png的灰度值未重映射為0~5而是[0,10,20,30,40,50]導致CrossEntropyLoss將所有非0像素視為ignore_index默認-100只學習background。解決嚴格執(zhí)行2.1節(jié)的convert_label_to_continuous()并用np.unique()打印clean label的值確認為[0 1 2 3 4 5]。現(xiàn)象2loss在0.3附近震蕩不下降梯度norm極小1e-5原因模型權重未初始化或init_weights()未應用到所有子模塊如忘記model.apply()。解決在模型實例化后立即調用model.apply(init_weights)并用print(model.inc.conv[0].weight.std())驗證首層權重標準差是否≈0.05?，F(xiàn)象3驗證時preds logits.argmax(dim1)報錯RuntimeError: argmax(): Expected dtype int64 or fp32原因logits是fp16混合精度訓練argmax不支持fp16輸入。解決preds logits.float().argmax(dim1)或在AMP上下文中用with torch.cuda.amp.autocast(enabledFalse):包裹預測代碼。現(xiàn)象4mIoU計算中某類IoU為nan且該類在驗證集出現(xiàn)頻率高原因ignore_index設為0而0是background類導致background被忽略union恒為0。解決ignore_index必須設為-100PyTorch默認或一個不在類別ID中的值如255絕不可設為任一有效類別ID。現(xiàn)象5訓練時GPU顯存爆滿batch_size1就OOM原因Unet的skip connection中torch.cat([x2, x1], dim1)將encoder特征與upsampled特征拼接若encoder輸出通道為512upsampled為256則拼接后為768通道顯存激增。解決在Up模塊中將self.up nn.ConvTranspose2d(in_ch // 2, in_ch // 2, 2, stride2)改為self.up nn.Upsample(scale_factor2, modebilinear, align_cornersTrue)并確保bilinearTrue可減少30%顯存占用。5. 訓練循環(huán)與早停策略動態(tài)學習率、梯度裁剪、模型保存的工業(yè)級實踐5.1 帶warmup的余弦退火學習率為什么固定lr1e-3在第10輪就過擬合Unet對學習率敏感。某實驗室對比實驗顯示固定lr1e-3時loss在epoch5驟降但val mIoU在epoch8達峰后持續(xù)下跌而采用warmupcosinelr從1e-5線性升至1e-3warmup 5 epoch再余弦退火至1e-6val mIoU穩(wěn)定提升至epoch30。代碼實現(xiàn)如下from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR def get_scheduler(optimizer, warmup_epochs5, total_epochs50, min_lr1e-6): warmup_scheduler LinearLR( optimizer, start_factor1e-3, # 從lr*1e-3開始 end_factor1.0, # 到lr*1.0結束 total_iterswarmup_epochs ) cosine_scheduler CosineAnnealingLR( optimizer, T_maxtotal_epochs - warmup_epochs, eta_minmin_lr ) scheduler SequentialLR( optimizer, schedulers[warmup_scheduler, cosine_scheduler], milestones[warmup_epochs] ) return scheduler # 使用示例 optimizer torch.optim.AdamW(model.parameters(), lr1e-3, weight_decay1e-4) scheduler get_scheduler(optimizer, warmup_epochs5, total_epochs50) # 在訓練循環(huán)中 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_norm1.0) # 梯度裁剪 optimizer.step() scheduler.step() # 更新學習率 print(fEpoch {epoch1}, LR: {scheduler.get_last_lr()[0]:.2e})參數(shù)說明LinearLR的start_factor1e-3表示warmup起始lr為1e-3 * 1e-3 1e-6end_factor1.0表示結束lr為1e-3 * 1.0 1e-3CosineAnnealingLR的eta_min1e-6是最終最小lr。torch.nn.utils.clip_grad_norm_將梯度norm限制在1.0內防止梯度爆炸這是訓練穩(wěn)定性的后悔藥。5.2 早停與最佳模型保存監(jiān)控val mIoU而非lossloss下降不代表分割質量提升必須以val mIoU為早停依據(jù)。保存最佳模型時不僅存state_dict還要存optimizer和scheduler狀態(tài)便于斷點續(xù)訓。class EarlyStopping: def __init__(self, patience7, delta0.001, save_pathbest_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 # 觸發(fā)早停 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}) # 在驗證循環(huán)后調用 early_stopping EarlyStopping(patience10, delta0.002, save_pathunet_best.pth) for epoch in range(50): # ... 訓 p a hrefhttps://download.csdn.net/download/brf_UCAS/14109453 stylecolor:#ec7500;font-size:14px; 本文還有配套的精品資源點擊獲取 /a img altmenu-r.4af5f7ec.gif srchttps://csdnimg.cn/release/wenkucmsfe/public/img/menu-r.4af5f7ec.gif stylewidth:16px;margin-left:4px;vertical-align:text-bottom;cursor:text; /p