unet模型学习笔记

一、核心文件夹 (The Folders)1.nets/(神经网络定义库)职责存放 U-Net 的核心代码。unet.py定义了 U-Net 的整体架构编码器、解码器、跳跃连接。vgg.py/resnet.pyU-Net 通常需要一个“主干网络”来提取特征这里存放的就是常用的主干网络实现。nets/unet.py是整个仓库的核心它像一个“拼装车间”把特征提取网络VGG或ResNet和上采样网络组合在一起。1. 初始化部分 (__init__)定义网络组件这部分代码的作用是准备好所有需要的“积木”。import torch.nn as nn from nets.vgg import VGG16 from nets.resnet import resnet50 class Unet(nn.Module): def __init__(self, num_classes 21, pretrained False, backbone vgg): super(Unet, self).__init__() # num_classes: 分类数量如背景猫狗3类。 # pretrained: 是否加载预训练权重即别人练好的“经验”。 # backbone: 选择主干网络默认用 vgg。 if backbone vgg: self.vgg VGG16(pretrained pretrained) # 加载 VGG16 模型作为编码器左侧下采样部分。 in_filters [192, 384, 768, 1024] # 定义解码器右侧每一层接收到的通道数总和。 elif backbone resnet50: self.resnet resnet50(pretrained pretrained) # 或者选择 ResNet50 作为编码器。 in_filters [192, 512, 1024, 3072] # ResNet 通道多所以这里的数字比 VGG 大。 out_filters [64, 128, 256, 512] # 定义解码器每一层输出的通道数。 # 下面定义 4 个上采样模块U-Net右侧的四个上升台阶 # unetUp 是作者在下面自定义的一个类负责“上采样 特征融合 两次卷积” self.up_concat4 unetUp(in_filters[3], out_filters[3]) # 处理最小的特征图 self.up_concat3 unetUp(in_filters[2], out_filters[2]) # 处理倒数第二层 self.up_concat2 unetUp(in_filters[1], out_filters[1]) # 处理中间层 self.up_concat1 unetUp(in_filters[0], out_filters[0]) # 恢复到较大的特征图 # 最后的一层卷积将通道数调整为最终的类别数 self.final nn.Conv2d(out_filters[0], num_classes, 1) # kernel_size1 代表 1x1 卷积作用是把特征图映射到类别概率上。2. 前向传播部分 (forward)定义数据流动逻辑这是代码运行时的逻辑路线图决定了图片进来到结果出去的过程。def forward(self, inputs): # inputs: 输入的原始图片比如 [1, 3, 512, 512] (1张图, 3通道RGB, 512x512分辨率) # 第一步利用主干网络提取 5 组不同尺度的特征图 if self.backbone vgg: feat1, feat2, feat3, feat4, feat5 self.vgg.forward(inputs) # feat1 是最大的特征图细节多feat5 是最小的语义强。 elif self.backbone resnet50: feat1, feat2, feat3, feat4, feat5 self.resnet.forward(inputs) # 第二步开始由小变大进行上采样和特征融合跳跃连接 # 这一步是 U-Net 的灵魂把深层信息feat5和浅层信息feat4拼在一起 up4 self.up_concat4(feat4, feat5) # 把 feat5 放大然后和 feat4 左右拼接。 up3 self.up_concat3(feat3, up4) # 把刚才的结果 up4 再次放大并和 feat3 拼接。 up2 self.up_concat2(feat2, up3) # 继续放大和 feat2 拼接。 up1 self.up_concat1(feat1, up2) # 最后一次放大和最初始的浅层特征 feat1 拼接找回丢失的边缘细节。 # 第三步输出最终结果 final self.final(up1) # 把拼接完的所有信息变成类别预测图。 return final3. 解码器核心类 (unetUp) 的逻辑在unet.py文件的上方你会看到一个unetUp类它是实现“跳跃连接Skip Connection”的具体工具class unetUp(nn.Module): def __init__(self, in_size, out_size): super(unetUp, self).__init__() # nn.Upsample: 上采样。将图片尺寸放大 2 倍比如从 16x16 变 32x32。 self.upsample nn.Upsample(scale_factor 2, mode bilinear, align_corners True) # 拼接后通道数会增加这里用两次卷积来融合这些信息 self.conv1 nn.Conv2d(in_size, out_size, kernel_size 3, padding 1) self.conv2 nn.Conv2d(out_size, out_size, kernel_size 3, padding 1) def forward(self, inputs1, inputs2): # inputs1: 来自左边的特征浅层inputs2: 来自下方的特征深层 outputs2 self.upsample(inputs2) # 先把深层特征图放大 # torch.cat: 【最关键】在通道维度(dim1)把两张特征图“粘”在一起 outputs torch.cat([inputs1, outputs2], dim 1) # 经过两次卷积让网络自己学习如何融合这两种不同的特征 outputs self.conv1(outputs) outputs self.conv2(outputs) return outputs直观总结backbone是挖掘机负责把图片里的特征轮廓、颜色、物体含义挖出来。torch.cat是缝合线负责把刚才挖出来的深层信息知道这是猫和浅层信息知道猫的胡须在哪缝在一起。nn.Upsample是放大镜把被挖出来的、已经变小的特征图变回原来图片的大小。final层是调色板把模型对每个像素的理解涂上颜色比如红色代表细胞背景是黑色。nets/resnet.py通常被用作 U-Net 的主干特征提取网络Backbone。相比于原始的 U-Net使用简单的卷积堆叠使用 ResNet50 作为主干可以提取更深层的特征且更容易训练。由于代码较长我为你拆解最核心的Bottleneck瓶颈结构和ResNet主类这是理解 ResNet 的关键。1. 核心组件Bottleneck类这是 ResNet50/101/152 的基本单元。它之所以叫“瓶颈”是因为中间用了一个 3x3 的小卷积两头用 1x1 卷积加宽。class Bottleneck(nn.Module): # expansion 4 意味着输出通道数是输入通道数的 4 倍 expansion 4 def __init__(self, inplanes, planes, stride1, downsampleNone): super(Bottleneck, self).__init__() # 1x1 卷积用来压缩通道数减少计算量 self.conv1 nn.Conv2d(inplanes, planes, kernel_size1, biasFalse) # BatchNorm对卷积后的结果做标准化让数据分布更稳定 self.bn1 nn.BatchNorm2d(planes) # 3x3 卷积核心特征提取stride2 时会缩小图片尺寸下采样 self.conv2 nn.Conv2d(planes, planes, kernel_size3, stridestride, padding1, biasFalse) self.bn2 nn.BatchNorm2d(planes) # 1x1 卷积用来放大通道数planes * 4 self.conv3 nn.Conv2d(planes, planes * self.expansion, kernel_size1, biasFalse) self.bn3 nn.BatchNorm2d(planes * self.expansion) # ReLU激活函数给网络增加非线性能力 self.relu nn.ReLU(inplaceTrue) # downsample捷径分支Shortcut。如果输入输出尺寸对不上用它来调整 self.downsample downsample self.stride stride def forward(self, x): # identity保存原始输入为了后面的“残差连接” identity x # 第一层卷积 标准化 激活 out self.conv1(x) out self.bn1(out) out self.relu(out) # 第二层卷积 标准化 激活 out self.conv2(out) out self.bn2(out) out self.relu(out) # 第三层卷积 标准化注意这里最后不激活要加完残差再激活 out self.conv3(out) out self.bn3(out) # 如果输入和输出的尺寸/通道数不同就对输入进行调整 if self.downsample is not None: identity self.downsample(x) # 【核心残差连接】将处理后的特征和原始输入相加 out identity # 最后再统一做一次激活 out self.relu(out) return out2. 主类ResNet结构这个类负责把上面的Bottleneck积木搭成一个高楼。class ResNet(nn.Module): def __init__(self, block, layers, num_classes1000): # block: 使用哪种积木比如 Bottleneck # layers: 每一层有多少个积木ResNet50 是 [3, 4, 6, 3] self.inplanes 64 super(ResNet, self).__init__() # 第一层大的 7x7 卷积直接把图片缩小一半 self.conv1 nn.Conv2d(3, 64, kernel_size7, stride2, padding3, biasFalse) self.bn1 nn.BatchNorm2d(64) self.relu nn.ReLU(inplaceTrue) # 最大池化再次缩小图片尺寸 self.maxpool nn.MaxPool2d(kernel_size3, stride2, padding1) # 下面是 ResNet 的四个大阶段Stage # 每个 stage 都会提取不同层次的特征U-Net 会在这些地方“取走”特征做跳跃连接 self.layer1 self._make_layer(block, 64, layers[0]) self.layer2 self._make_layer(block, 128, layers[1], stride2) self.layer3 self._make_layer(block, 256, layers[2], stride2) self.layer4 self._make_layer(block, 512, layers[3], stride2)3. 为什么 U-Net 要看这个文件在nets/unet.py中你会看到类似这样的逻辑# 在 unet 的 forward 函数里 feat1 self.resnet.layer1(x) # 拿到浅层特征边缘、颜色 feat2 self.resnet.layer2(feat1) # 拿到中层特征形状 feat3 self.resnet.layer3(feat2) # 拿到深层特征语义代码中关键参数的意思inplanes输入到当前层的通道数水的流量。planes这一层基础的输出通道数。stride步长如果等于 2图片的长宽就会缩小一半。padding填充在图片外围补一圈 0保证卷积后图片尺寸不会因为边缘损耗而莫名缩小。biasFalse因为后面接了BatchNorm偏置项bias会被抵消掉所以为了省计算量直接关掉。直观总结卷积层 (Conv)是过滤器负责找特征。标准化层 (BN)是调节器防止模型练着练着“走火入魔”梯度爆炸。残差连接 (out identity)是 ResNet 的神来之笔它允许信息直接跨层传递解决了深层网络难以训练的问题。_make_layer是一个自动化工厂你告诉它你需要多少个Bottleneck它就自动帮你串联起来。nets/vgg.py在这个项目中被用作 U-Net 的特征提取器Backbone。VGG16 的结构非常规律由一系列“卷积激活池化”的积木块组成。为了让 U-Net 实现“跳跃连接”这个 VGG 的实现与标准的分类 VGG 不同它会在中间阶段把特征图“截断”并输出。1. 构造卷积层序列的工具函数 (make_layers)这个函数是按照预定义的配置表自动像搭积木一样叠放卷积层。import torch.nn as nn def make_layers(cfg, batch_normFalse): # cfg: 一个列表数字代表卷积核数量M 代表最大池化层。 # batch_norm: 是否使用标准归一化通常选 False 以节省内存。 layers [] # 创建一个空列表用来存放层级。 in_channels 3 # 输入通道数初始是 RGB 3 通道。 for v in cfg: if v M: # 如果是 M就加一个 2x2 的最大池化层图片长宽缩小一半。 layers [nn.MaxPool2d(kernel_size2, stride2)] else: # 否则就是一个卷积层3x3 卷积核保持图片尺寸padding1。 conv2d nn.Conv2d(in_channels, v, kernel_size3, padding1) if batch_norm: # 如果开启 BN则按顺序加入卷积 - 标准化 - 激活。 layers [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplaceTrue)] else: # 否则只加入卷积 - 激活。 layers [conv2d, nn.ReLU(inplaceTrue)] in_channels v # 更新输入通道数下一层的输入等于这一层的输出。 return nn.Sequential(*layers) # 将列表里的层组合成一个连续的神经网络。2. VGG16 主类结构 (VGG类)这里定义了 VGG16 的具体层数配置。# 这是 VGG16 的经典配置数字代表通道数M 代表下采样 cfgs [64, 64, M, 128, 128, M, 256, 256, 256, M, 512, 512, 512, M, 512, 512, 512, M] class VGG(nn.Module): def __init__(self, features, num_classes1000, image_setvoc): super(VGG, self).__init__() # features: 就是上面 make_layers 生成的那一长串卷积层。 self.features features # 对于 U-Net 来说我们通常不需要下面这些全连接层分类头 # 但为了代码完整性作者保留了这部分结构。 self.avgpool nn.AdaptiveAvgPool2d((7, 7)) # 自适应平均池化。 self.classifier nn.Sequential( nn.Linear(512 * 7 * 7, 4096), # 全连接层。 nn.ReLU(True), nn.Dropout(), # 随机失活防止过拟合。 nn.Linear(4096, 4096), nn.ReLU(True), nn.Dropout(), nn.Linear(4096, num_classes), # 映射到最终的分类数量。 )3. 前向传播为 U-Net 量身定制 (forward)这是初学者最需要关注的地方。标准的 VGG 只返回最后的结果但 U-Net 需要五个不同阶段的特征图。def forward(self, x): # x: 输入的原始图片 [batch, 3, 512, 512] # 阶段 1经过前两个卷积层图片还是 512x512 feat1 self.features[:4](x) # 阶段 2经过第一次下采样变为 256x256 feat2 self.features[4:9](feat1) # 阶段 3经过第二次下采样变为 128x128 feat3 self.features[9:16](feat2) # 阶段 4经过第三次下采样变为 64x64 feat4 self.features[16:23](feat3) # 阶段 5经过第四次下采样变为 32x32 feat5 self.features[23:30](feat4) # 最后把这 5 个不同大小的特征图全部返回。 # U-Net 的右侧上采样部分会像接力赛一样挨个把它们接过去。 return [feat1, feat2, feat3, feat4, feat5]核心函数意思总结nn.Conv2d(..., padding1): 卷积。由于设置了padding1它在提取特征的同时不会改变图片的长宽尺寸。nn.MaxPool2d: 最大池化。它的作用是“降维”就像把 512 像素的图拍扁成 256 像素保留最明显的特征比如最亮的点。nn.ReLU(inplaceTrue): 激活函数。inplaceTrue表示直接在原始内存上修改这样可以节省显存防止在训练大模型时内存溢出。self.features[:4]: 这是 Python 的切片语法。作者根据 VGG 的固定结构算好了前 4 层是第一个特征阶段16 到 23 层是第四个阶段。为什么 VGG 适合做 BackboneVGG 的结构非常直观它通过不断的池化M让特征图越来越小同时通过卷积让通道数64 - 128 - 512越来越多。这符合人类视觉系统的逻辑看的范围越来越广感受野变大但看到的内容越来越抽象从线条变成物体概念。2.utils/(工具箱)职责存放所有辅助训练和测试的脚本。dataloader.py负责把硬盘里的图片加载到内存并进行格式转换变成张量。loss.py定义“损失函数”即如何判断模型画得准不准。放在了nets/unet_training.py这个文件里。utils_metrics.py计算评价指标如 mIoU、准确率等。callbacks.py在训练过程中自动保存模型、记录日志。以下是nets/unet_training.py中关键代码的逐行解释1. 损失函数类CE_Loss(交叉熵损失)这是语义分割最基础的损失函数用来判断每个像素点分类是否正确。def CE_Loss(inputs, target, cls_weights, num_classes21): # inputs: 模型预测的结果 # target: 真实的标签图Ground Truth # cls_weights: 类别权重如果某些物体很小可以给它更高的权重 n, c, h, w inputs.size() # 获取预测图的数量、类别数、高度、宽度 nt, ht, wt target.size() # 获取标签图的尺寸 # 如果尺寸对不上进行插值缩放保证能对比 if h ! ht or w ! wt: inputs F.interpolate(inputs, size(ht, wt), modebilinear, align_cornersTrue) # 重点转置并改变形状为了符合 PyTorch 交叉熵函数的输入要求 inputs inputs.transpose(1, 2).transpose(2, 3).contiguous().view(-1, num_classes) target target.view(-1) # 调用 PyTorch 内置的交叉熵计算误差 loss nn.CrossEntropyLoss(weightcls_weights, ignore_indexnum_classes)(inputs, target) return loss2. 损失函数类Dice_Loss(Dice 损失)这是 U-Net 的灵魂损失函数。它专门用于处理“目标很小、背景很大”的情况比如医学图像里的小肿瘤。def Dice_Loss(inputs, target, beta1, smooth1e-5): # inputs: 预测结果target: 真实标签 # smooth: 一个很小的数防止分母为 0 导致报错 n, c, h, w inputs.size() nt, ht, wt target.size() # 同样进行尺寸对齐 if h ! ht or w ! wt: inputs F.interpolate(inputs, size(ht, wt), modebilinear, align_cornersTrue) # 将预测值转化为 0-1 之间的概率值 inputs torch.softmax(inputs, dim1) # 【核心逻辑】计算预测结果和真实结果的交集Intersection # 就像两个圆的重合面积越大Dice 损失就越小 tp torch.sum(target * inputs, axis[0, 2, 3]) fp torch.sum(inputs, axis[0, 2, 3]) - tp fn torch.sum(target, axis[0, 2, 3]) - tp # 根据公式计算 Dice 系数 score ((1 beta ** 2) * tp smooth) / ((1 beta ** 2) * tp beta ** 2 * fn fp smooth) dice_loss 1 - torch.mean(score) # 1 减去相似度就是误差值 return dice_loss3. 其他重要函数意思除了损失函数这个文件里还有几个初学者必须知道的函数weights_init(权重初始化)意思神经网络刚出生时是一张白纸这个函数负责给模型里的每一个“神经元”分配一个初始的随机数字。作用合理的初始值如 Xavier 或 Kaiming 初始化能让模型练得更快不至于在一开始就跑偏。get_lr_scheduler(学习率调度器)意思学习率就像模型走路的“步子”。作用刚开始练的时候步子大一点快速接近目标快练好时步子小一点精细打磨这个函数就是控制步子大小的。set_optimizer_lr(设置优化器学习率)意思手动把计算好的步长学习率告诉给“修理工”优化器。总结为什么放在nets而不是utils在bubbliiiing的架构逻辑里utils放置的是通用的、不依赖于具体模型的工具比如图片缩放、计算 mIoU 指标。nets/unet_training.py放置的是专门针对 U-Net 训练定制的代码。因为不同的网络如 YOLO, DeepLab往往需要不同的 Loss 组合方式。utils/dataloader.py文件在深度学习项目中相当于“物流部门”。它的任务是从硬盘读取原始图片和标签Mask把它们处理成统一的大小进行各种“变花样”的数据增强最后打包成机器能听懂的张量Tensor送入 GPU。1. 导入库与预处理函数import cv2 import numpy as np import torch from PIL import Image from torch.utils.data.dataset import Dataset def preprocess_input(image): image / 255.0 # 将像素值从 0-255 压缩到 0-1 之间这叫“归一化”有助于模型收敛。 return image2.UnetDataset类数据的“加工厂”这个类继承了 PyTorch 的Dataset。它是告诉程序给出一个索引 i你应该去哪里找第 i 张图并把它加工成什么样。(1) 初始化部分__init__class UnetDataset(Dataset): def __init__(self, annotation_lines, input_shape, num_classes, train, dataset_path): super(UnetDataset, self).__init__() self.annotation_lines annotation_lines # 存储所有图片路径的列表通常来自 train.txt。 self.length len(self.annotation_lines) # 数据集的总长度图片总数。 self.input_shape input_shape # 模型要求的图片尺寸如 [512, 512]。 self.num_classes num_classes # 分类的数量。 self.train train # 是否为训练模式训练模式下会开启随机翻转等增强。 self.dataset_path dataset_path # 数据集所在的根目录。(2) 获取长度__len__def __len__(self): return self.length # 告诉 PyTorch 这个数据集里一共有多少张照片。(3) 核心获取函数__getitem__这是整个文件最重要的部分。每当训练需要一张新图时就会调用这个函数。def __getitem__(self, index): # 1. 根据索引找到对应的图片名称 line self.annotation_lines[index].split() # 2. 从硬盘打开原始图片并转换成 RGB 模式 image Image.open(os.path.join(os.path.join(self.dataset_path, VOC2007/JPEGImages), line[0] .jpg)) # 3. 从硬盘打开标签图即 Mask转换成灰度模式L # 标签图中0 通常代表背景1 代表第一个目标以此类推。 png Image.open(os.path.join(os.path.join(self.dataset_path, VOC2007/SegmentationClass), line[0] .png)) # 4. 数据增强Data Augmentation # 如果是训练模式会随机对图片和标签同时进行缩放、翻转、颜色抖动等。 # 目的是让模型见多识广不至于只认识这一张图。 if self.train: image, png self.get_random_data(image, png, self.input_shape, random self.train) else: # 如果是验证模式只进行不失真的大小调整。 image, png self.get_random_data(image, png, self.input_shape, random False) # 5. 格式转换从图片格式转为数组 image np.transpose(preprocess_input(np.array(image, np.float32)), (2, 0, 1)) # 为什么要 transpose因为图片是 [H, W, C]而 PyTorch 要求 [C, H, W]通道数在前。 png np.array(png) # 将标签图里大于类别数的像素点强行设为“背景”防止数据标注错误。 png[png self.num_classes] self.num_classes # 6. 返回处理好的图片和标签 return image, png3.unet_dataset_collate打包函数如果说__getitem__是加工一件商品那么collate就是把加工好的商品装进集装箱Batch。def unet_dataset_collate(batch): images [] pngs [] for img, png in batch: images.append(img) pngs.append(png) # 将一组图片和标签堆叠在一起形成一个批次Batch。 # 比如 batch_size 是 4那这里就会返回 [4, 3, 512, 512] 的张量。 images torch.from_numpy(np.array(images)).type(torch.FloatTensor) pngs torch.from_numpy(np.array(pngs)).long() return images, pngs直观总结Dataset就像是一个“菜单”它知道厨房里有什么菜图片并且知道怎么把每一道菜洗干净、切好预处理。input_shape这是厨房的“统一盘子尺寸”。不管原始图片多大最后都要缩放到这个尺寸否则模型吃不下。get_random_data这是“加佐料”。通过旋转、裁剪把一张图变成“多张图”让模型学习得更扎实。np.transpose这是“换个摆盘姿势”。模型这个“客人”很挑剔它习惯先看通道RGB再看像素坐标。3.model_data/(模型资产)职责存放静态资源。.pth文件这是权重文件保存了模型学到的“经验”。cls_classes.txt定义了你的分类目标比如第一类是背景第二类是细胞第三类是血管。4.VOCdevkit/(数据集仓库)职责存放你的训练素材。VOC2007/JPEGImages存放原始彩色图片。VOC2007/SegmentationClass存放对应的标签图Mask即告诉机器哪里是什么。img/职责仅用于展示项目 README 文档里的示意图对代码运行没有实际功能影响。5.img/二、根目录下的关键脚本 (The Scripts)1.train.py(训练总控制台)这是最核心的文件。它把nets的模型、utils的数据加载器、model_data的权重全部组合起来开始训练。初学者注意你大部分的学习率、训练次数Epoch、批次大小Batch Size都在这里设置。1. 参数配置区 (Configuration)这部分代码主要是设置“游戏规则”比如用什么显卡、分多少类。# 是否使用 GPU 进行训练。如果显卡支持设为 True 能极大地提高速度。 Cuda True # 分类个数 11是背景。比如你要识别猫和狗这里就是 213。 num_classes 21 # 主干网络选择。可选 vgg 或 resnet50。 backbone vgg # 是否使用预训练权重。建议初学者开启这样模型就像是“读过高中”后再来学专业课比从零开始快。 pretrained False # 预训练权重文件的路径。 model_path model_data/unet_vgg_voc.pth # 输入图片的大小。必须是 32 的倍数因为 U-Net 有多次下采样每次缩小一半。 input_shape [512, 512]2. 训练阶段设置 (Training Phases)这个仓库采用了**“冻结训练”**的策略。这是一种非常聪明的训练方法先冻结主干特征提取网络Backbone只练后面的解码器等稳定了再全线开启训练。# 冻结阶段只训练 U-Net 的加强特征提取部分。 Init_Epoch 0 # 起始世代 Freeze_Epoch 50 # 冻结训练到第几个世代结束 Freeze_batch_size 2 # 冻结阶段每批次喂给模型多少张图显存小就调小点 # 解冻阶段全网络一起训练精细打磨参数。 UnFreeze_Epoch 100 # 总训练世代 Unfreeze_batch_size 2 # 解冻后每批次的数量3. 模型与初始化 (Model Weights)这部分代码负责把模型“造”出来并给它装载记忆权重。# 实例化 U-Net 模型 model Unet(num_classesnum_classes, pretrainedpretrained, backbonebackbone).train() # 如果没有预训练权重就进行随机初始化给神经元分配随机的初始数字 if not pretrained: weights_init(model) # 如果指定了 model_path就加载已经练过的权重 if model_path ! : device torch.device(cuda if torch.cuda.is_available() else cpu) model_dict model.state_dict() pretrained_dict torch.load(model_path, map_locationdevice) # 只加载匹配的部分防止报错 model_dict.update(pretrained_dict) model.load_state_dict(model_dict)4. 数据加载 (Data Loading)这里负责把voc_annotation.py生成的名单读进来并交给“搬运工”。# 读取训练集和验证集的图片路径和标签信息 with open(train_annotation_path, r) as f: train_lines f.readlines() with open(val_annotation_path, r) as f: val_lines f.readlines() # 实例化 Dataset 类还记得 dataloader.py 吗这里就是调用它 train_dataset UnetDataset(train_lines, input_shape, num_classes, True, VOCdevkit_path) val_dataset UnetDataset(val_lines, input_shape, num_classes, False, VOCdevkit_path) # 实例化 DataLoader负责开启多线程把数据打包成 Batch 喂给 GPU gen DataLoader(train_dataset, shuffleTrue, batch_sizebatch_size, num_iternum_workers, pin_memoryTrue, drop_lastTrue, collate_fnunet_dataset_collate) gen_val DataLoader(val_dataset, shuffleTrue, batch_sizebatch_size, num_iternum_workers, pin_memoryTrue, drop_lastTrue, collate_fnunet_dataset_collate)5. 核心训练循环 (The Training Loop)这是最热闹的地方。程序会在这里跑几十到上百个 Epoch世代。for epoch in range(start_epoch, end_epoch): # 如果到了解冻的时机就把主干网络的锁定解开 if epoch Freeze_Epoch and not UnFreeze_Flag: batch_size Unfreeze_batch_size # ... 重新设置优化器和学习率 ... UnFreeze_Flag True # 调用 fit_one_epoch 函数。这是实际干活的地方。 # 每一代训练都会经历读图 - 预测 - 算误差 - 找差距 - 改参数。 fit_one_epoch(model_train, model, loss_history, optimizer, epoch, gen, gen_val, end_epoch, Cuda, dice_loss, focal_loss, cls_weights, num_classes)6. 关键函数fit_one_epoch里的逻辑意思在train.py内部或调用的工具函数中这几行是灵魂optimizer.zero_grad()擦除黑板。把上一轮计算的梯度差距清空。outputs model_train(images)模型考试。模型根据输入的图给出它认为的分割结果。loss loss_calc(outputs, pngs)对答案。计算预测图和真实标签图之间的差距。loss.backward()反向传播。把差距一层层传回去告诉前面的神经元“你刚才猜错了”。optimizer.step()调整。根据刚才传回来的信息微调模型内部的参数螺丝。2.predict.py(预测/推理入口)当你训练好模型后运行这个文件。它会启动一个窗口或循环让你输入图片并查看分割结果。3.unet.py(预测逻辑包装)注意根目录下的这个unet.py不同于nets/里的。它是一个“工具类”专门用于加载模型权重并对单张图片进行预处理方便predict.py调用。4.voc_annotation.py(数据索引生成器)在训练前必须运行。它会扫描VOCdevkit并生成train.txt等文件告诉程序“这 800 张图用来学习剩下的 200 张用来考试。5.get_miou.py(性能体检脚本)用于量化评估你的模型到底有多优秀它会计算出平均交并比mIoU这是语义分割领域最权威的指标三、函数级别的“黑话”解释__init__积木准备。比如定义“我要一个卷积层”、“我要一个池化层”。forward数据流动。定义图片进入网络后先过哪个门再进哪个房。Dataset/DataLoader搬运工。负责从硬盘一张张地把图片塞给 GPU。Optimizer(优化器)修理工。根据误差不断旋转模型内部的“螺丝钉”权重让误差越来越小。四、基于模型修改自己的图片1. 准备数据集最关键的一步你需要按照VOC 格式整理你的图片。将你的数据放入VOCdevkit/VOC2007文件夹中JPEGImages存放所有的原始彩色图片建议统一为.jpg格式。SegmentationClass存放所有的标签图片Mask必须是.png格式。注意标签图里的像素值必须是数字。比如背景像素是 0目标 A 是 1目标 B 是 2。在视觉上这些图片看起来可能是全黑的这是正常的。2. 修改类别文件 (model_data/cls_classes.txt)打开这个文件把里面的内容换成你自己的分类目标。注意第一行通常保留为background背景后面紧跟着你的目标。示例如果你要识别肺部结节文件内容应为​ background nodule ​3. 生成索引文件 (voc_annotation.py)模型训练前需要知道哪些图是用来“学习”的哪些是用来“考试”的。打开根目录下的voc_annotation.py。修改其中的classes_path指向你刚刚改好的cls_classes.txt。运行该脚本。它会在根目录生成2007_train.txt和2007_val.txt里面记录了图片的路径。4. 配置训练参数 (train.py)这是你的控制中心。在运行之前检查以下几个核心参数num_classes改为你的总类别数背景 目标数量。backbone根据你的显卡配置选择vgg或resnet50。model_path如果你是第一次训练建议填从零开始或者指向作者提供的预训练权重收敛更快。input_shape根据你的图片大小修改通常设为[512, 512]。5. 开始训练在终端或 PyCharm 中直接运行train.py。观察日志你会看到Loss误差在不断下降。产出物训练完成后在logs/文件夹下会生成一系列.pth文件。这些就是你辛苦训练出来的“大脑”。6. 使用你的模型进行预测当你训练好模型后想要测试一张新图片修改根目录下的unet.py注意不是nets/里的将model_path指向你刚才生成的最好的.pth文件。将classes_path指向你的cls_classes.txt。运行predict.py。根据提示输入图片的路径你就能看到模型对你新图片的分割效果了。给初学者的 3 个避坑小贴士尺寸一致性确保你的标签图 (Mask) 和原始图片的长宽比例是一致的否则模型会学得很痛苦。显存溢出 (OOM)如果运行train.py报错提示内存不足请尝试调小batch_size比如从 4 改为 2。标签格式很多初学者直接用彩色分割图当标签这是不对的。标签图必须是单通道的索引图。你可以用PIL库通过代码检查一下标签图的像素值是否在[0, num_classes-1]范围内。