【深度学习|语义分割之UNet】继承自 PyTorch 的 nn.Module的UNet——基于编码器-解码器结构的语义分割任务的卷积神经网络。附代码及解读。
【深度学习|语义分割之UNet】继承自 PyTorch 的 nn.Module的UNet——基于编码器-解码器结构的语义分割任务的卷积神经网络。附代码及解读。
文章目录
欢迎宝子们点赞、关注、收藏！欢迎宝子们批评指正！
祝所有的硕博生都能遇到好的导师！好的审稿人！好的同门！顺利毕业！
大多数高校硕博生毕业要求需要参加学术会议，发表EI或者SCI检索的学术论文会议论文：
可访问艾思科蓝官网，浏览即将召开的学术会议列表。会议入口：https://ais.cn/u/mmmiUz
前言
这段代码定义了一个 U-Net 网络的实现，U-Net 是一种广泛应用于图像分割任务的神经网络架构。它的特点是采用了 编码器-解码器结构（即上下采样结构），用于从输入图像中提取特征并进行像素级别的预测。该实现继承自 PyTorch 的 nn.Module，并且在此基础上做了一些扩展，比如支持 **检查点（checkpointing）**以降低内存使用。
我们将逐行解释这段代码的细节。
1. 导入必要的部分
from .unet_parts import *
- 这行代码从
unet_parts.py中导入所有定义的网络模块。通常这个文件包含了 U-Net 所需的基本模块，如卷积块（DoubleConv），下采样块（Down），上采样块（Up），和输出卷积层（OutConv）等。
2. UNet 类定义
class UNet(nn.Module):
def __init__(self, num_classes, bilinear=False):
super(UNet, self).__init__()
self.n_channels = 3
self.n_classes = num_classes
self.bilinear = bilinear
UNet继承自nn.Module，是 PyTorch 的所有神经网络的基类。- 在
__init__方法中，初始化了：
n_channels = 3：输入图像的通道数，通常是 RGB 图像，所以设置为 3。
n_classes = num_classes：输出的类别数，通常是图像分割中的类别数。
bilinear = bilinear：一个布尔值，决定是否使用双线性插值来做上采样。如果为True，则上采样时使用双线性插值，否则使用转置卷积。
3. 定义 U-Net 的各层
self.inc = (DoubleConv(3, 64)) # 输入的卷积块
self.down1 = (Down(64, 128)) # 第一层下采样
self.down2 = (Down(128, 256)) # 第二层下采样
self.down3 = (Down(256, 512)) # 第三层下采样
factor = 2 if bilinear else 1 # 根据是否使用双线性插值调整factor的大小
self.down4 = (Down(512, 1024 // factor)) # 第四层下采样
self.up1 = (Up(1024, 512 // factor, bilinear)) # 第1层上采样
self.up2 = (Up(512, 256 // factor, bilinear)) # 第2层上采样
self.up3 = (Up(256, 128 // factor, bilinear)) # 第3层上采样
self.up4 = (Up(128, 64, bilinear)) # 第4层上采样
self.outc = (OutConv(64, num_classes)) # 输出卷积
self.inc是输入的卷积块，通常包含两次卷积操作以及激活函数。DoubleConv(3, 64)表示输入通道数为 3，输出通道数为 64。self.down1, self.down2, self.down3, self.down4是下采样块（编码器部分），用于逐步降低特征图的空间分辨率，增加通道数以提取更深层次的特征。每一层都会包含卷积、激活和池化操作，具体的实现由Down类定义。factor = 2 if bilinear else 1：根据是否启用双线性插值，上采样时使用的倍数进行调整。self.up1, self.up2, self.up3, self.up4是上采样块（解码器部分），用于逐步恢复特征图的空间分辨率。Up 类会在上采样过程中使用转置卷积或双线性插值。self.outc是输出卷积层，将上采样后的特征图转化为最终的类别预测，输出通道数为num_classes，即类别数。
4. 前向传播方法
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) # 第1层上采样，跳跃连接
x = self.up2(x, x3) # 第2层上采样，跳跃连接
x = self.up3(x, x2) # 第3层上采样，跳跃连接
x = self.up4(x, x1) # 第4层上采样，跳跃连接
logits = self.outc(x) # 最后的输出卷积
return logits
forward 方法定义了数据流在网络中的前向传播过程，具体步骤如下：
- 输入
x首先通过初始卷积块self.inc提取低级特征。 - 然后数据通过四层下采样块（
self.down1到self.down4）逐步减少空间分辨率并增加特征通道。 - 在解码器部分，通过四层上采样块（
self.up1到self.up4）恢复空间分辨率。在每层上采样时，都使用跳跃连接（skip connections）将编码器部分的特征传递到解码器部分。 - 最后通过
self.outc将恢复后的特征图通过卷积转化为最终的类别预测输出。
5. 检查点功能
def use_checkpointing(self):
self.inc = torch.utils.checkpoint(self.inc)
self.down1 = torch.utils.checkpoint(self.down1)
self.down2 = torch.utils.checkpoint(self.down2)
self.down3 = torch.utils.checkpoint(self.down3)
self.down4 = torch.utils.checkpoint(self.down4)
self.up1 = torch.utils.checkpoint(self.up1)
self.up2 = torch.utils.checkpoint(self.up2)
self.up3 = torch.utils.checkpoint(self.up3)
self.up4 = torch.utils.checkpoint(self.up4)
self.outc = torch.utils.checkpoint(self.outc)
use_checkpointing方法启用了检查点（checkpointing）功能，以减少训练过程中内存的使用。检查点技术的基本思想是在前向传播时不保存中间激活，而是在反向传播时重新计算它们。这样做可以显著减少 GPU 内存的占用，特别是在训练大规模网络时。这里使用了torch.utils.checkpoint对每一层进行包装。
总结
- 该段代码实现了一个 U-Net 网络，适用于语义分割任务。
- U-Net 采用了 编码器-解码器结构，通过逐步下采样和上采样恢复空间分辨率，并利用 跳跃连接（skip connections）将编码器的特征传递给解码器，从而保持了细粒度的空间信息。
- 模型支持 双线性插值 上采样或 转置卷积，并且提供了 检查点功能，以降低内存开销。
U-Net 网络因其高效的结构和优异的性能，广泛应用于医学图像分割、遥感图像分割等任务中。
第二届人工智能、系统与网络安全国际学术会议 (AISNS 2024）
- 【ACM独立出版，最快3个月检索 | 往届已实现EI检索】
- 大会时间：2024年12月20-22日
- 大会地点：中国 · 湘潭
- 大会官网：www.aisns.org
第二届人工智能与自动化控制国际学术会议（AIAC 2024）
- 时间地点：2024年12月20-22日 中国-广州
- IEEE独立出版，EI快稳检索，广东工业大学主办，往届已EI&Scopus双检索
- 会议官网：http://www.icaiac.org/
第四届信号处理与通信技术国际学术会议（SPCT 2024）
- 时间地点：2024年12月27-29日，中国·深圳·香港中文大学（深圳校区）
- 香港中文大学（深圳）主办
- 设置优秀评选，ACM独立出版，EI Compendex、Scopus检索
- 会议官网：www.icspct.com
2024年智能通信、感知与电磁学术会议（ICSE 2024）
- 2024年12月27-29日 中国-广州
- 24年年度最后一个电磁、感知类IEEE学术会议！！！
- 中山大学主办，IEEE出版，学生投稿优惠400
- 会议官网：www.icicse.net
第五届神经网络、信息与通信工程国际学术会议（NNICE 2025）
- 大会官网：www.icnnice.com
- 大会时间：2025年1月10-12日
- IEEE独立出版，往届均已见刊检索
- IEEE出版|广东工业大学主办|往届均已见刊检索
第五届生物信息学与智能计算国际学术研讨会（BIC 2025）
- 2025年1月10-12日 中国·沈阳
- https://www.ic-bic.net
- EI、SCOPUSACM出版 | 往届均已见刊检索
- 辽宁大学主办 | 检索稳定