企业官网建设流程全解析

简介：面向PyTorch自定义多类别数据集语义分割的实战代码包，适合具备一定深度学习基础、希望跑通U-Net训练与评估流程的开发者，也可用于医学影像、遥感图像等像素级分类场景的学习参考。压缩包共46个文件，约69KB，以19个Python脚本为核心，辅以24个编译缓存文件、2个txt说明和1个json配置，涵盖数据加载与增强、网络结构、损失函数、评价指标、学习率调度及训练日志保存等模块。已有15284人学习，说明其在同类教程中具有较高关注度。读者可从中获得多类别分割的关键实现思路，包括输出通道与类别数对应、交叉熵损失使用、IoU等指标计算，以及训练过程监控和结果记录方式，便于在自己的数据集上复用和改造完整流程。

1. 多类别分割不是把类别数改大：一份 PyTorch Unet 工程能替你省掉哪些返工

把 Unet 的输出通道从 1 改成 8，loss 照样往下掉，mIoU 却卡在 0.12，预测图还是一片黑——这是自建多类别语义分割最常见的第一课。问题往往不在网络结构，而在标签像素值、忽略像素、损失函数和指标统计没对齐。这份 PyTorch 下实现 Unet 多类别语义分割的工程，把目录约定、Dataset、模型输出、训练验证、滑窗推理串成一条可复现链路，适合刚接手自建数据集、课程设计或原型验证的开发者。它解决的是“能跑但结果不可信”的返工，而不是只给一个模型文件。先按数据链路跑通，再谈改 backbone 和加注意力。

2. 从标签图到 DataLoader：多类别数据集的目录约定与通道处理

多类别语义分割的第一道分水岭不在模型，而在标签图怎么读。单通道标签里每个像素值就是类别 ID，0 是背景，1 到 N-1 是目标类，255 常被拿来做忽略像素；一旦把标签图先转成 RGB，或者用双线性插值去缩放掩码，类别 ID 就会被搅成小数和新颜色，后面损失函数不报错，但模型学到的边界全是糊的。我一般先把目录、通道、忽略值三件事固定死，再动网络。目录约定用下面这种最朴素的拆分，训练集和验证集各自带images与masks，文件名一一对应，不靠文件名前缀猜类别。

data/ ├── train/ │ ├── images/ │ │ ├── 0001.png │ │ └── 0002.png │ └── masks/ │ ├── 0001.png │ └── 0002.png └── val/ ├── images/ └── masks/

这段目录结构的作用是把图像和掩码的对应关系收敛到同名文件，sorted(os.listdir(...))之后两边顺序天然一致，省掉一张 CSV 清单。参数上，train与val只是目录名，换成training也可以，但 Dataset 里的split要同步；images和masks不建议混在一起，否则做同步增强时很容易把掩码当图像读进来。验证集不要做随机裁剪和翻转，目录可以先按原图放，推理阶段再统一处理尺寸。

2.1 标签图的像素值就是类别 ID，不要先转 RGB

标签图用 PIL 打开时，convert("L")得到单通道 8 位数组，像素值 0、1、2、3 直接对应类别编号，这是多类别 Unet 最稳的输入形态。转成 RGB 之后，一个像素变成三个通道，np.array(mask).shape从(H, W)变成(H, W, 3)，CrossEntropyLoss的 target 维度对不上还是小事，更常见的是有人把 RGB 再转灰度，类别值被加权平均成 0 到 255 之间的连续值，训练时target出现 37、128 这种值，直接触发 CUDA 越界。我的习惯是读进来先打印np.unique(mask)，确认最大值小于num_classes，并且忽略值只有 255。

import os, random import numpy as np import torch from torch.utils.data import Dataset from PIL import Image class SegDataset(Dataset): def __init__(self, root, split="train", image_size=(512, 512), ignore_index=255, aug=True): self.img_dir = os.path.join(root, split, "images") self.mask_dir = os.path.join(root, split, "masks") self.names = sorted(os.listdir(self.img_dir)) self.image_size = image_size self.ignore_index = ignore_index self.aug = aug def __len__(self): return len(self.names) def _sync_aug(self, img, mask): # 图像和掩码共用同一组随机参数，避免空间错位 if random.random() < 0.5: img = np.ascontiguousarray(img[:, ::-1, :]) mask = np.ascontiguousarray(mask[:, ::-1]) if self.aug: h, w = mask.shape th, tw = self.image_size if h < th or w < tw: pad_h, pad_w = max(0, th - h), max(0, tw - w) img = np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)), mode="reflect") mask = np.pad(mask, ((0, pad_h), (0, pad_w)), mode="constant", constant_values=self.ignore_index) h, w = mask.shape y0 = random.randint(0, h - th) x0 = random.randint(0, w - tw) img = img[y0:y0 + th, x0:x0 + tw] mask = mask[y0:y0 + th, x0:x0 + tw] return img, mask def __getitem__(self, idx): name = self.names[idx] img = np.array(Image.open(os.path.join(self.img_dir, name)).convert("RGB"), dtype=np.float32) / 255.0 # 单通道读取：像素值就是类别 ID，0 为背景，255 为忽略 mask = np.array(Image.open(os.path.join(self.mask_dir, name)).convert("L"), dtype=np.int64) img, mask = self._sync_aug(img, mask) img = torch.from_numpy(img).permute(2, 0, 1).float() mask = torch.from_numpy(mask).long() return img, mask

这段 Dataset 的关键点不在代码长短，而在三个参数：convert("L")保证掩码是单通道，permute(2, 0, 1)把图像从 HWC 转成 CHW，ignore_index在填充和裁剪时只作用到掩码上。image_size控制训练输入尺寸，太小会丢小目标，太大显存吃紧；我一般先按 512 跑通，再根据显存和最小目标尺寸调整。aug=False用在验证集和推理集上，训练集才开翻转和随机裁剪。裁剪前如果原图小于目标尺寸，用反射填充图像、用忽略值填充掩码，避免randint上界为负。掩码不要做归一化，也不要转 float，CrossEntropyLoss需要 long 类型的类别索引。

2.2 自定义 Dataset 的读取、同步增强与忽略像素

同步增强是多类别分割最容易翻车的位置之一。图像翻转了、掩码没翻转，模型会学到左右相反的边界；图像随机旋转 15 度、掩码旋转 16 度，边缘类别会整体错位。上面_sync_aug用同一个随机数决定是否翻转，裁剪也用同一组y0、x0，这是最小可用的同步策略。更复杂的颜色抖动只作用在图像上，几何变换必须同时作用到图像和掩码；掩码的缩放只能用最近邻，不能用双线性，否则两个类别之间会插值出第三个 ID，听起来玄学，实际排查时np.unique一跑就能看到多出来的像素值。

忽略像素要贯穿整条链路。数据集里有些区域无法标注，比如遮挡、边界模糊、无关背景，标成 255 之后，损失函数用ignore_index=255跳过，指标统计也要跳过。如果 Dataset 把 255 当普通类别送进模型，num_classes又只设了 8，训练到第一个 batch 就会报Target 255 is out of bounds。我的做法是让ignore_index在 Dataset、损失函数、混淆矩阵三处用同一个常量，配置文件里写一次，别在三个文件里各写各的。

2.3 用配置表锁定类别数、输入尺寸和归一化参数

多类别工程跑崩，很多时候是配置散落。类别数在模型里写 8，损失函数里写 7，指标统计里写 9，训练能跑几个 batch，验证时维度对不上。我一般把和数据集强相关的参数收在一张表里，训练脚本只读表，不读魔法数。下面这张表里的值按自建数据集填，类别数包含背景，忽略值要和标签图的实际值一致。

归一化参数不要照搬 ImageNet 的均值方差，自建数据的成像条件不同，照搬会让模型前几个 epoch 收敛很慢。统计训练集 RGB 三通道均值和标准差，写进配置表，验证和推理沿用同一组值。base_channels控制模型宽度，32 在 512 输入下通常能在一张消费级显卡上跑起来，64 更稳但显存翻倍；先小后大，别一上来就追高。

3. Unet 结构与多类别输出：从通道数、上采样到损失函数选择

Unet 在多类别任务里不是把最后一层卷积分成 N 个通道就完事。编码器下采样四次，解码器上采样四次，skip connection 把浅层纹理和深层语义拼在一起，拼接时通道数要对，空间尺寸也要对。输入尺寸不是 16 的倍数时，下采样取整会让torch.cat报尺寸不匹配，这个问题在自建数据集里比结构选型更常见。输出层返回 logits，不接 softmax，损失函数用CrossEntropyLoss，它内部会做 log-softmax；如果模型末尾已经 softmax，再送进 CE，梯度会被压扁，训练表现为 loss 下降很慢、预测概率全挤在 1/N 附近。

3.1 编码器-解码器的最小闭环与 skip connection 对齐

下面这套 Unet 是够用的最小闭环：四次下采样、三次上采样拼接、最后一层 1×1 卷积输出num_classes个通道。base控制第一层通道数，后面每下采样一次翻倍；上采样用转置卷积，把深层特征放大到和对应浅层特征一样的空间尺寸，再按通道维拼接。输入尺寸建议取 16 的倍数，512、256、128 都行；如果原图不是 16 的倍数，训练时先 pad 到最近的 16 倍数，推理完再裁回原尺寸。

import torch import torch.nn as nn class DoubleConv(nn.Module): def __init__(self, in_ch, out_ch): super().__init__() self.block = 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.block(x) class UNet(nn.Module): def __init__(self, in_ch=3, num_classes=8, base=32): super().__init__() self.d1 = DoubleConv(in_ch, base) self.d2 = DoubleConv(base, base * 2) self.d3 = DoubleConv(base * 2, base * 4) self.d4 = DoubleConv(base * 4, base * 8) self.pool = nn.MaxPool2d(2) self.u3 = nn.ConvTranspose2d(base * 8, base * 4, 2, stride=2) self.c3 = DoubleConv(base * 8, base * 4) self.u2 = nn.ConvTranspose2d(base * 4, base * 2, 2, stride=2) self.c2 = DoubleConv(base * 4, base * 2) self.u1 = nn.ConvTranspose2d(base * 2, base, 2, stride=2) self.c1 = DoubleConv(base * 2, base) self.out = nn.Conv2d(base, num_classes, 1) def forward(self, x): x1 = self.d1(x) x2 = self.d2(self.pool(x1)) x3 = self.d3(self.pool(x2)) x4 = self.d4(self.pool(x3)) y = self.c3(torch.cat([self.u3(x4), x3], dim=1)) y = self.c2(torch.cat([self.u2(y), x2], dim=1)) y = self.c1(torch.cat([self.u1(y), x1], dim=1)) return self.out(y) # 返回 logits，不接 softmax

参数上，in_ch=3对应 RGB 输入，如果数据是灰度图就改 1；num_classes必须包含背景；base=32时第一层输出 32 通道，最后一次拼接后c1输入是 64 通道。forward里没有对浅层做裁剪，所以输入尺寸要能被 16 整除，或者你在拼接前手动F.interpolate对齐。self.out是 1×1 卷积，把特征压到类别数，后面直接送损失函数。这个结构不追求 SOTA，但胜在通道和尺寸关系清楚，排查时能一眼看出是哪一层对不上。

3.2 输出层与 CrossEntropyLoss 的配合：logits、ignore_index、类别权重

多类别分割的主损失我一般用CrossEntropyLoss，它接收N, C, H, W的 logits 和N, H, W的 long 掩码，内部做 log-softmax 和 NLL，不需要模型末尾再 softmax。ignore_index要和 Dataset 里的忽略值一致，这样 255 的像素不参与梯度。类别不平衡明显时，给损失函数传weight，按训练集每个类别的像素频率取倒数或平方根倒数，背景类权重压低一点，小目标类权重抬高一点；权重不要拍脑袋，先统计像素占比，再跑一版不加权的基线做对比。

import numpy as np import torch import torch.nn as nn # counts 是训练集每个类别的像素数，顺序对应 0..num_classes-1 counts = np.array([...], dtype=np.float64) freq = counts / counts.sum() class_weight = torch.tensor((1.0 / np.sqrt(freq + 1e-8)), dtype=torch.float32) class_weight = class_weight / class_weight.mean() criterion = nn.CrossEntropyLoss(ignore_index=255, weight=class_weight.cuda())

这段代码里counts要按类别统计，不要包含忽略像素；1.0 / np.sqrt(freq)比纯倒数温和，避免极小类权重过大导致训练震荡；归一化到均值 1 是为了让学习率不用跟着权重重调。ignore_index=255和 Dataset 对齐，weight放在 GPU 上要和 logits 同设备。如果小目标还是被吞，可以加 Dice 或 Tversky 辅助损失，但主损失仍建议保留 CE，因为 CE 对类别索引的监督更直接，Dice 单独用时早期梯度不稳。

3.3 训练循环里的 AMP、调度器和指标累计

训练循环里最容易忽略的是 mask 的类型和设备。图像送进 GPU 后是 float32，mask 送进去必须是 long，且和 logits 同设备；AMP 只对前向和损失做 autocast，mask 不要转 half。优化器用 AdamW 或 SGD 都能跑，学习率先按 1e-3 到 1e-4 试，调度器用余弦退火或按验证集 mIoU 触发的 ReduceLROnPlateau。下面这段循环把 AMP、梯度清零、反向传播和调度器步进串起来，验证循环单独写，不要混在训练里。

import torch scaler = torch.cuda.amp.GradScaler() for epoch in range(epochs): model.train() for img, mask in train_loader: img = img.cuda(non_blocking=True) mask = mask.cuda(non_blocking=True).long() optimizer.zero_grad(set_to_none=True) with torch.cuda.amp.autocast(): logits = model(img) loss = criterion(logits, mask) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() scheduler.step()

GradScaler负责损失缩放，避免 half 下小梯度下溢；set_to_none=True比清零更省显存；mask.long()不能省，AMP 不会自动把 mask 变成 long。scheduler.step()放在每个 batch 后是按步调度，如果用的是按验证集指标调度的策略，就要挪到 epoch 末尾。训练 loss 降但 mIoU 不涨时，先看验证集有没有做同样的归一化和尺寸处理，再看类别权重是不是把背景压得太狠，最后才怀疑模型容量。

4. 训练、验证与推理：把 IoU 算对，把可视化做成人能看懂的图

多类别分割的指标不是“像素准确率”一个数就够。背景占比高时，像素准确率能到 0.9，但小目标类一个都没分出来，mIoU 只有 0.1。IoU 要按类别算，再对类别取平均；混淆矩阵要跳过忽略像素，预测值要先 argmax 再统计。验证阶段不要开随机增强，也不要随机裁剪，否则同一张图每次验证的输入都不一样，mIoU 曲线会抖得没法判断。推理阶段大图直接送进模型容易爆显存，用滑窗加概率平均，边缘回退窗口，拼接后再 argmax。

4.1 混淆矩阵累计与 mIoU 计算公式

混淆矩阵的横轴是预测类别，纵轴是真实类别，conf[i, j]表示真实为 i、预测为 j 的像素数。对每个 batch，先对 logits 做argmax(dim=1)得到预测类别，再用target != ignore_index筛掉忽略像素，把target * num_classes + pred拉平后做bincount，最后 reshape 成num_classes × num_classes累加。下面这段函数只做统计，不参与反向传播，conf放在 CPU 上累加，避免显存越跑越少。

import torch def update_confusion(conf, pred_logits, target, num_classes, ignore_index=255): pred = pred_logits.argmax(dim=1) valid = target != ignore_index idx = target[valid] * num_classes + pred[valid] hist = torch.bincount(idx, minlength=num_classes * num_classes) conf += hist.reshape(num_classes, num_classes).cpu() return conf def miou_from_confusion(conf): inter = torch.diag(conf) union = conf.sum(0) + conf.sum(1) - inter iou = inter / union.clamp_min(1) return iou.mean().item(), iou

pred_logits是模型原始输出，不要先 softmax；target是 long 类型；num_classes要和模型输出通道一致。union.clamp_min(1)防止某个类别在验证集里没出现时除零，但这类别的 IoU 会变成 0，看结果时要结合conf判断是“真没分出来”还是“验证集里没有”。iou.mean()是所有类别的平均，背景类通常很高，会把 mIoU 拉上去，所以我一般额外打印每个类别的 IoU，小目标类单独看。

4.2 验证阶段不增强、不裁剪的边界处理

验证集用aug=False只是第一步，尺寸处理也要统一。训练时随机裁剪到 512×512，验证时如果直接 resize 到 512×512，长宽比变了，边界类别会被拉变形；如果原图比 512 大，验证时又直接送整图，显存和训练不一致，BN 统计也会偏。我的做法是验证集保持原图尺寸，按 16 的倍数 pad 到可下采样，推理完再裁回原图大小；如果显存不够，用滑窗，不要用 resize 硬压。验证循环里关掉model.train()，开torch.no_grad()，不要漏掉model.eval()，否则 BN 用 batch 统计，mIoU 会随着 batch 内容波动。

验证阶段还要统一归一化。训练集统计的mean、std要原样用在验证集上，不要重新统计验证集均值，否则训练和验证的输入分布对不上。图像 resize 用双线性，掩码 resize 用最近邻；如果验证集不 resize，就只做 pad，不要对掩码做插值。边界处理上，pad 的图像用反射或复制边缘，pad 的掩码用ignore_index，这样指标不会把填充区算进去。

4.3 滑窗推理与大图拼接的尺寸对齐

大图推理我一般用滑窗，窗口 512×512，步长 256，重叠区域把 softmax 概率累加，最后除以覆盖次数再 argmax。窗口走到图像边缘时，如果直接取y:y+512会越界，所以先算y1 = min(y+win, h)，再回退y0 = max(0, y1-win)，保证窗口尺寸固定。概率累加比直接拼 argmax 稳，因为重叠区的类别概率会被平均，边缘不会出现硬切缝。

import torch @torch.no_grad() def sliding_infer(model, image, num_classes, window=(512, 512), stride=(256, 256), device="cuda"): model.eval() n, c, h, w = image.shape prob = torch.zeros((n, num_classes, h, w), device=device) count = torch.zeros((1, 1, h, w), device=device) for y in range(0, h, stride[0]): for x in range(0, w, stride[1]): y1, x1 = min(y + window[0], h), min(x + window[1], w) y0, x0 = max(0, y1 - window[0]), max(0, x1 - window[1]) patch = image[:, :, y0:y1, x0:x1].to(device) logits = model(patch) prob[:, :, y0:y1, x0:x1] += logits.softmax(dim=1) count[:, :, y0:y1, x0:x1] += 1 return (prob / count).argmax(dim=1)

window和stride控制重叠率，512 窗口配 256 步长是 50% 重叠，边缘更稳但计算量翻倍；显存紧就把窗口降到 256，步长 128。count用单通道累计覆盖次数，prob按类别累计概率，最后逐像素除，避免重叠区被重复计权。返回值是类别索引图，可视化时用固定颜色表映射，不要用随机颜色，否则同一类别两次推理颜色不一样，验收时容易误判。

5. 避坑与排查：多类别 Unet 最容易翻车的五个位置

多类别 Unet 的坑大多不在反向传播，而在数据进入模型之前。下面五个位置是我在自建数据集上反复踩过的：全黑预测、标签值越界、尺寸不整除、增强错位、显存波动。每个坑都按“现象、原因、解决”拆开，排查顺序从数据到模型，别一上来就调学习率。先跑一个 batch 过拟合，能过拟合说明模型和损失没大问题，问题在数据分布或验证链路；过拟合都过不了，先查标签值和输出通道。

5.1 全黑预测与类别不平衡

现象是训练 loss 从 2.0 降到 0.3，验证集预测图几乎全是一个颜色，mIoU 0.1 左右。原因通常不是模型没学会，而是背景像素占 90% 以上，CE 被背景主导，小目标类的梯度被淹没；如果再加了过大的类别权重，又可能反过来让背景崩掉。解决分三步：先统计每个类别的像素占比，确认不平衡程度；再给 CE 加温和的类别权重，权重按平方根倒数算，不要直接按倒数；最后拿 20 张图做小样本过拟合，关掉增强，学习率调大，看能不能把训练图分割出来。小样本都分不出，先查标签和输入是否对齐。

5.2 标签值不连续导致的 CUDA 越界

现象是训练到某个 batch 报CUDA error: device-side assert triggered，或者Target 5 is out of bounds。原因一般是标签图里有不连续的类别值，比如背景 0、目标 1、目标 2，但某张图里出现了 3，而num_classes=3；也可能是把 255 当普通类别送进去了。解决是训练前跑一遍全量标签检查，统计每张图的np.unique，把非 0 到 N-1 且非忽略值的像素列出来，再做重映射。

import numpy as np from PIL import Image def check_mask_values(mask_path, num_classes=8, ignore_index=255): m = np.array(Image.open(mask_path).convert("L")) vals = np.unique(m) bad = [int(v) for v in vals if v != ignore_index and not (0 <= v < num_classes)] return vals.tolist(), bad

num_classes要包含背景，ignore_index和损失函数一致；bad非空就说明标签值越界，需要写映射表把旧值改成连续 ID，或者把num_classes调到最大 ID 加一。检查脚本不要只跑训练集，验证集和测试集也要跑，验证集里混进一个越界值，训练能跑，验证一开就崩。

5.3 输入尺寸非 2 的幂与 skip 拼接报错

现象是torch.cat报Sizes of tensors must match except in dimension 1，或者上采样后尺寸差 1 像素。原因是输入长宽不是 16 的倍数，四次下采样每次都向下取整，解码器上采样回来对不上浅层特征。解决是训练和推理前把图像 pad 到 16 的倍数，pad 的图像用反射填充，pad 的掩码用忽略值；推理完再裁回原图尺寸。不要用 resize 硬拉到 2 的幂，长宽比变了，小目标会被压扁。

5.4 数据增强同步错位与掩码插值

现象是预测边缘整体偏移，或者掩码里出现训练集没有的类别值。原因是图像和掩码用了不同的随机数做翻转、旋转、裁剪；或者掩码缩放用了双线性插值，两个类别之间插出了新 ID。解决是几何变换用同一组随机参数，掩码缩放统一最近邻，颜色变换只作用在图像上。增强库如果有image和mask双输入接口，优先用它的同步变换；自己写也不难，随机数只生成一次，图像和掩码共用。

5.5 显存波动与 batch size 调参

现象是训练几个 epoch 后突然 OOM，或者验证时 batch size 明明比训练小却爆显存。原因是训练图像尺寸不一致，DataLoader 拼 batch 时按最大图分配；或者验证阶段没开torch.no_grad()，计算图一直留着；也可能是混淆矩阵放在 GPU 上越累越大。解决是训练前统一尺寸，验证循环加torch.no_grad()，混淆矩阵放 CPU，显存吃紧时用梯度累积代替大 batch。batch_size不要一次拉满，先跑 2 个 batch 看显存占用，再按 80% 显存上限调。

6. 从能跑到能交付：用滑动窗口+后处理把多类别分割结果稳住

模型能跑出 mIoU 只是中间态，交付前还要看单张图的边界和小目标。我一般做两件事：固定随机种子，把同一张图连续推理三次，确认输出完全一致；再对预测图做一次小连通域过滤，把面积小于阈值的孤立类别块并回背景或邻类。下面这段后处理只处理非背景类，面积阈值按最小目标像素数设，别把真实小目标误删。

import cv2 import numpy as np def remove_small_components(pred, num_classes, min_area=64): out = pred.copy() for cls in range(1, num_classes): mask = (pred == cls).astype(np.uint8) n, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8) for i in range(1, n): if stats[i, cv2.CC_STAT_AREA] < min_area: out[labels == i] = 0 return out

pred是argmax后的类别索引图，min_area按数据集里最小目标的像素面积设，太小起不到过滤作用，太大会吃掉真实小目标。后处理不要放在验证指标里，先看原始 mIoU，再看后处理后的可视化，避免用后处理掩盖模型问题。验收时我会按类别看 IoU，再看边界图，最后固定种子复现一轮。从那以后我每次接自建多类别数据集，都强制先跑一遍标签值检查和固定种子复现，再谈调模型。希望帮到你。

本文还有配套的精品资源，点击获取