network.py
初始化神经网络中的权重和偏置
def init_weights(net, init_type='normal', gain=0.02):# 定义初始化函数def init_func(m):# 获取当前模块的类名classname = m.__class__.__name__# 如果模块有权重,并且是卷积层或全连接层if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):# 根据指定的初始化类型进行权重初始化if init_type == 'normal':init.normal_(m.weight.data, 0.0, gain) # 正态分布初始化elif init_type == 'xavier':init.xavier_normal_(m.weight.data, gain=gain) # Xavier初始化elif init_type == 'kaiming':init.kaiming_normal_(m.weight.data, a=0, mode='fan_in') # Kaiming初始化elif init_type == 'orthogonal':init.orthogonal_(m.weight.data, gain=gain) # 正交初始化else:# 如果初始化类型不支持,抛出异常raise NotImplementedError('initialization method [%s] is not implemented' % init_type)# 如果模块有偏置项,则将其初始化为0if hasattr(m, 'bias') and m.bias is not None:init.constant_(m.bias.data, 0.0)# 如果模块是BatchNorm2d层elif classname.find('BatchNorm2d') != -1:init.normal_(m.weight.data, 1.0, gain) # 权重初始化为正态分布init.constant_(m.bias.data, 0.0) # 偏置初始化为0# 打印当前使用的初始化类型print('initialize network with %s' % init_type)# 将初始化函数应用于网络的所有模块net.apply(init_func)
这段代码的作用是初始化神经网络中的权重和偏置。具体来说:
函数定义:
init_weights
函数接收一个网络对象net
,一个初始化类型init_type
(如正态分布、Xavier、Kaiming 等),以及一个增益参数gain
。初始化函数:内部定义的
init_func
函数会遍历网络中的每一个层(模块),并根据层的类型初始化权重和偏置。权重初始化:
- 对于卷积层(
Conv
)和全连接层(Linear
),根据选择的初始化方法(如正态、Xavier、Kaiming、正交等)初始化权重。- 如果有偏置项,将其初始化为零。
BatchNorm层:对于批归一化层(
BatchNorm2d
),权重初始化为正态分布,偏置初始化为零。应用初始化:最后,使用
net.apply(init_func)
将初始化函数应用到整个网络,确保每一层都按照指定方式初始化。
卷积块
class conv_block(nn.Module):def __init__(self, ch_in, ch_out):# 初始化卷积块,接收输入通道数和输出通道数super(conv_block, self).__init__()# 定义一个顺序容器,包含两个卷积层和相关操作self.conv = nn.Sequential(nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=True), # 第一个卷积层nn.BatchNorm2d(ch_out), # 批归一化层,用于提高训练稳定性nn.ReLU(inplace=True), # ReLU激活函数,使用原地操作以节省内存nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1, bias=True), # 第二个卷积层nn.BatchNorm2d(ch_out), # 再次应用批归一化nn.ReLU(inplace=True) # 再次使用ReLU激活函数)
这段代码定义了一个卷积块(
conv_block
)的神经网络模块,通常用于构建卷积神经网络(CNN)。具体作用如下:
卷积层:包含两个卷积层,每个卷积层使用 3x3 的卷积核,输入和输出的通道数可以根据构造函数的参数设置。这种设计有助于提取图像特征。
批归一化:每个卷积层后面跟随一个批归一化层(
BatchNorm2d
),用于规范化输出,提高训练的稳定性,并加速收敛。激活函数:每个卷积层后面使用 ReLU 激活函数,增加网络的非线性能力。
输入和输出通道:构造函数接受输入通道数(
ch_in
)和输出通道数(ch_out
),允许在不同层之间灵活配置。
上采样卷积块
def forward(self, x):# 将输入 x 通过卷积块处理x = self.conv(x)# 返回处理后的输出return xclass up_conv(nn.Module):def __init__(self, ch_in, ch_out):# 初始化上采样卷积块,接收输入通道数和输出通道数super(up_conv, self).__init__()# 定义一个顺序容器,包含上采样和卷积层self.up = nn.Sequential(nn.Upsample(scale_factor=2), # 上采样,将特征图尺寸扩大 2 倍nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=True), # 卷积层,调整通道数nn.BatchNorm2d(ch_out), # 批归一化层,规范化输出nn.ReLU(inplace=True) # ReLU激活函数,增加非线性)def forward(self, x):# 将输入 x 通过上采样卷积块处理x = self.up(x)# 返回处理后的输出return x
这段代码定义了一个上采样卷积块(
up_conv
),主要用于图像处理和深度学习中的卷积神经网络(CNN)。具体作用如下:
上采样:通过
nn.Upsample(scale_factor=2)
,将输入特征图的尺寸扩大一倍。这在图像分割或生成任务中很常见,用于恢复图像的空间分辨率。卷积操作:
nn.Conv2d
层负责对上采样后的特征图进行卷积处理,提取特征并调整通道数。批归一化:
nn.BatchNorm2d
层用于规范化卷积层的输出,帮助提高训练的稳定性,加速收敛。激活函数:
nn.ReLU
激活函数引入非线性,使网络能够学习更复杂的特征。
循环块
class Recurrent_block(nn.Module):def __init__(self, ch_out, t=2):# 初始化循环块,接收输出通道数和循环次数super(Recurrent_block, self).__init__()self.t = t # 循环次数self.ch_out = ch_out # 输出通道数# 定义卷积序列,包括卷积、批归一化和ReLU激活self.conv = nn.Sequential(nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1, bias=True), # 卷积层nn.BatchNorm2d(ch_out), # 批归一化层nn.ReLU(inplace=True) # ReLU激活函数)def forward(self, x):# 前向传播函数for i in range(self.t): # 根据 t 的值进行循环if i == 0:x1 = self.conv(x) # 第一次循环,直接对输入 x 进行卷积操作# 在后续循环中,将当前输入 x 和上一层的输出 x1 相加,然后进行卷积x1 = self.conv(x + x1)return x1 # 返回最终的输出
这段代码实现了一个循环块(Recurrent Block),其主要功能包括:
循环处理:该块通过指定的循环次数
t
,重复对输入进行卷积操作。每次迭代都将当前输入与上一层的输出相加,从而实现特征的累积。残差连接:在
forward
方法中,第一次循环将输入x
经过卷积得到x1
,后续循环则将x1
与输入x
相加后再进行卷积。这种方式类似于残差网络(ResNet)的设计,帮助缓解深度网络的梯度消失问题。非线性特征学习:通过卷积层、批归一化和ReLU激活函数,该块能够学习复杂的非线性特征,适用于图像处理任务,如图像重建或图像分割。
RRCNN和单层卷积块
class RRCNN_block(nn.Module):def __init__(self, ch_in, ch_out, t=2):# 初始化RRCNN块,接收输入通道数、输出通道数和循环次数super(RRCNN_block, self).__init__()# 定义两个循环块(Recurrent Block)self.RCNN = nn.Sequential(Recurrent_block(ch_out, t=t), # 第一个循环块Recurrent_block(ch_out, t=t) # 第二个循环块)# 定义1x1卷积,用于调整输入通道数self.Conv_1x1 = nn.Conv2d(ch_in, ch_out, kernel_size=1, stride=1, padding=0)def forward(self, x):# 前向传播函数x = self.Conv_1x1(x) # 通过1x1卷积调整输入的通道数x1 = self.RCNN(x) # 将调整后的输出传入循环块return x + x1 # 返回输入与循环块输出的和(残差连接)class single_conv(nn.Module):def __init__(self, ch_in, ch_out):# 初始化单卷积块,接收输入通道数和输出通道数super(single_conv, self).__init__()# 定义卷积序列,包括卷积、批归一化和ReLU激活self.conv = nn.Sequential(nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=True), # 3x3卷积nn.BatchNorm2d(ch_out), # 批归一化层nn.ReLU(inplace=True) # ReLU激活函数)def forward(self, x):# 前向传播函数x = self.conv(x) # 通过定义的卷积序列进行处理return x # 返回卷积处理后的输出
RRCNN_block:
- 该类实现了一个RRCNN(Residual Recurrent Convolutional Neural Network)块。其主要特点是通过残差连接和循环块来增强特征提取能力。
- 初始化时,定义了两个循环块,每个循环块内使用
Recurrent_block
处理输入。- 使用1x1卷积将输入的通道数调整为输出通道数,以便与循环块的输出相加,形成残差连接。
- 前向传播时,首先通过1x1卷积调整通道数,然后将调整后的特征传入循环块,最后返回输入和循环块输出的和。
single_conv:
- 该类实现了一个单层卷积块,主要用于简单的特征提取。
- 初始化时,定义了一个卷积序列,包括3x3卷积、批归一化和ReLU激活函数。
- 前向传播时,通过卷积序列处理输入数据并返回输出,适合在深度学习模型中作为基本构件进行特征提取。
注意力块
import torch
import torch.nn as nn# 定义Attention Block类,继承自nn.Module
class Attention_block(nn.Module):# 初始化方法,定义了三个输入通道数:F_g(输入g的通道数),F_l(输入x的通道数),F_int(中间层的通道数)def __init__(self, F_g, F_l, F_int):super(Attention_block, self).__init__()# g通道的卷积层,用于生成与输入g相关的特征映射self.W_g = nn.Sequential(nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True), # 1x1卷积,改变通道数nn.BatchNorm2d(F_int) # 批量归一化,保证训练稳定性)# x通道的卷积层,用于生成与输入x相关的特征映射self.W_x = nn.Sequential(nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True), # 1x1卷积,改变通道数nn.BatchNorm2d(F_int) # 批量归一化)# psi卷积层:用来生成最终的注意力图(注意力权重)self.psi = nn.Sequential(nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True), # 1x1卷积,生成1通道的注意力图nn.BatchNorm2d(1), # 批量归一化nn.Sigmoid() # Sigmoid激活函数将输出映射到[0, 1]之间)# ReLU激活函数,用于激活输入信号self.relu = nn.ReLU(inplace=True)# 前向传播方法,输入g和x分别是两个特征图def forward(self, g, x):# 对g输入进行卷积操作,得到g1g1 = self.W_g(g)# 对x输入进行卷积操作,得到x1x1 = self.W_x(x)# 将g1和x1的结果相加后通过ReLU激活psi = self.relu(g1 + x1)# 通过psi卷积层得到最终的注意力图psi = self.psi(psi)# 输出x与注意力图的逐元素相乘,进行加权return x * psi
这段代码实现了一个 Attention Block,用于深度学习中的注意力机制,通常用于提升模型的表现,尤其是在图像处理任务中。通过计算注意力图(Attention Map),它能够在空间维度上对输入进行加权,突出重要的区域并抑制不重要的区域。具体来说,这段代码是基于类似于 Spatial Attention 的机制来生成一个加权系数,应用于输入的特征图。
U-Net
class U_Net(nn.Module):def __init__(self, img_ch=3, output_ch=1):super(U_Net, self).__init__()# 最大池化层:通常用于下采样,减少空间维度self.Maxpool = nn.MaxPool2d(kernel_size=2, stride=2)# 编码器部分(左侧)# 第一层卷积块:输入通道img_ch,输出通道64self.Conv1 = conv_block(ch_in=img_ch, ch_out=64)# 第二层卷积块:输入通道64,输出通道128self.Conv2 = conv_block(ch_in=64, ch_out=128)# 第三层卷积块:输入通道128,输出通道256self.Conv3 = conv_block(ch_in=128, ch_out=256)# 第四层卷积块:输入通道256,输出通道512self.Conv4 = conv_block(ch_in=256, ch_out=512)# 第五层卷积块:输入通道512,输出通道1024self.Conv5 = conv_block(ch_in=512, ch_out=1024)# 解码器部分(右侧)# 上采样层5:输入通道1024,输出通道512self.Up5 = up_conv(ch_in=1024, ch_out=512)# 卷积块5:进行进一步的卷积操作,输入通道1024,输出通道512self.Up_conv5 = conv_block(ch_in=1024, ch_out=512)# 上采样层4:输入通道512,输出通道256self.Up4 = up_conv(ch_in=512, ch_out=256)# 卷积块4:进行进一步的卷积操作,输入通道512,输出通道256self.Up_conv4 = conv_block(ch_in=512, ch_out=256)# 上采样层3:输入通道256,输出通道128self.Up3 = up_conv(ch_in=256, ch_out=128)# 卷积块3:进行进一步的卷积操作,输入通道256,输出通道128self.Up_conv3 = conv_block(ch_in=256, ch_out=128)# 上采样层2:输入通道128,输出通道64self.Up2 = up_conv(ch_in=128, ch_out=64)# 卷积块2:进行进一步的卷积操作,输入通道128,输出通道64self.Up_conv2 = conv_block(ch_in=128, ch_out=64)# 最终卷积层:将输出的64个通道映射到output_ch个通道,通常是1通道用于二分类self.Conv_1x1 = nn.Conv2d(64, output_ch, kernel_size=1, stride=1, padding=0)
U_Net
类初始化 (__init__
方法):
- 定义了 U-Net 的结构,其中
img_ch
和output_ch
是输入和输出图像的通道数,默认输入是 3 通道(RGB 图像),输出是 1 通道(单通道分割图像,通常用于二分类)。最大池化层 (
Maxpool
):
nn.MaxPool2d(kernel_size=2, stride=2)
用于下采样图像,减小空间尺寸,保留重要特征。每经过一次MaxPool
层,图像的分辨率减半。编码器部分(
Conv1
到Conv5
):
- 这些卷积层和激活函数块负责提取输入图像的高级特征。每一层卷积块都使用了
conv_block
这个自定义模块(可以理解为由卷积、激活、批量归一化等组成的模块)。每经过一层,特征图的通道数逐渐增大,网络能够学习到越来越复杂的特征。解码器部分(
Up5
到Up2
):
- 这些层负责将编码器提取到的特征图逐渐恢复到原始图像的空间分辨率。
up_conv
是一个上采样模块,通常用于将低分辨率的特征图通过反卷积(或者上采样)扩大空间维度。上采样的过程使得网络能够逐步恢复图像的空间细节。跳跃连接(skip connections):
- 这些部分通过跳跃连接的方式将编码器部分的特征图与解码器部分的相应层连接起来,帮助网络更好地恢复细节信息。这在 U-Net 中是一个关键特性,有助于避免信息丢失。
最终卷积层(
Conv_1x1
):
- 最后通过
1x1
卷积层将解码器的输出通道数映射到目标输出通道数(output_ch
)。如果是图像分割任务,通常output_ch=1
,表示单通道的分割结果。
U-Net 神经网络的 前向传播(forward) 函数
def forward(self, x):# encoding path: 编码路径,通过一系列卷积层和池化层提取特征# 第1步:输入图像经过第一个卷积块x1 = self.Conv1(x)# 第2步:对x1进行最大池化(下采样),减小空间分辨率x2 = self.Maxpool(x1)# 第3步:通过第二个卷积块提取更深层次的特征x2 = self.Conv2(x2)# 第4步:对x2进行最大池化(下采样)x3 = self.Maxpool(x2)# 第5步:通过第三个卷积块提取更深层次的特征x3 = self.Conv3(x3)# 第6步:对x3进行最大池化(下采样)x4 = self.Maxpool(x3)# 第7步:通过第四个卷积块提取更深层次的特征x4 = self.Conv4(x4)# 第8步:对x4进行最大池化(下采样)x5 = self.Maxpool(x4)# 第9步:通过第五个卷积块提取最深层次的特征x5 = self.Conv5(x5)# decoding + concat path: 解码路径 + 拼接路径,通过反卷积和跳跃连接恢复空间分辨率# 第10步:通过反卷积层对x5进行上采样(扩大空间尺寸)d5 = self.Up5(x5)# 第11步:跳跃连接,将编码器的特征x4与解码器的特征d5拼接在一起d5 = torch.cat((x4, d5), dim=1) # 在通道维度上拼接# 第12步:通过解码器的卷积块进一步处理拼接后的特征d5 = self.Up_conv5(d5)# 第13步:对d5进行上采样d4 = self.Up4(d5)# 第14步:跳跃连接,将编码器的特征x3与解码器的特征d4拼接d4 = torch.cat((x3, d4), dim=1)# 第15步:通过解码器的卷积块进一步处理拼接后的特征d4 = self.Up_conv4(d4)# 第16步:对d4进行上采样d3 = self.Up3(d4)# 第17步:跳跃连接,将编码器的特征x2与解码器的特征d3拼接d3 = torch.cat((x2, d3), dim=1)# 第18步:通过解码器的卷积块进一步处理拼接后的特征d3 = self.Up_conv3(d3)# 第19步:对d3进行上采样d2 = self.Up2(d3)# 第20步:跳跃连接,将编码器的特征x1与解码器的特征d2拼接d2 = torch.cat((x1, d2), dim=1)# 第21步:通过解码器的卷积块进一步处理拼接后的特征d2 = self.Up_conv2(d2)# 第22步:通过1x1卷积层将输出通道数转换为目标通道数(通常是1,表示分割结果)d1 = self.Conv_1x1(d2)# 最终返回分割图像return d1
编码路径:
Conv1
到Conv5
是 U-Net 的编码部分,逐层提取输入图像的特征。每一层通过卷积操作对输入图像进行特征学习。- 在每一层之后,都使用了最大池化(
Maxpool
),该操作减小空间分辨率,增大特征图的感受野,从而帮助网络提取更具抽象性的特征。解码路径:
- 解码部分通过上采样(
Up5
到Up2
)逐步恢复图像的空间分辨率。- 每一层的解码部分都使用了跳跃连接(
torch.cat((x, d), dim=1)
),这些连接将编码部分的特征(例如x5
,x4
,x3
,x2
,x1
)与对应的解码部分特征进行拼接。
- 拼接的目的是保留编码器中的高分辨率细节,这些信息对恢复精细的分割边界非常有帮助。
卷积操作:
- 在解码路径的每个步骤中,拼接后的特征图会经过一层卷积操作(例如
Up_conv5
,Up_conv4
等),以进一步处理和融合特征,避免因为上采样操作导致的特征损失或不精细。最终卷积层:
Conv_1x1
是一个1x1
卷积操作,用于将解码器最后一层的输出映射到目标通道数(通常是1,用于二分类任务的分割结果)。这是 U-Net 的输出层,产生最终的预测结果。这段代码的作用:
- 输入: 传入一个图像
x
,它通过网络的编码路径(卷积 + 池化)逐步提取特征。- 跳跃连接: 通过跳跃连接,网络将编码路径中高分辨率的特征与解码路径中的低分辨率特征进行拼接,帮助网络恢复图像的细节。
- 输出: 网络最终通过
1x1
卷积层将解码路径的输出映射到目标通道数(通常是一个单通道的分割图像)。
R2U-Net
class R2U_Net(nn.Module):def __init__(self,img_ch=3,output_ch=1,t=2):super(R2U_Net,self).__init__()# 初始化MaxPooling层(用于下采样)self.Maxpool = nn.MaxPool2d(kernel_size=2,stride=2)# 初始化Upsample层(用于上采样)self.Upsample = nn.Upsample(scale_factor=2)# 定义RRCNN块,用于编码器的各层(这些块实现卷积操作)self.RRCNN1 = RRCNN_block(ch_in=img_ch, ch_out=64, t=t)self.RRCNN2 = RRCNN_block(ch_in=64, ch_out=128, t=t)self.RRCNN3 = RRCNN_block(ch_in=128, ch_out=256, t=t)self.RRCNN4 = RRCNN_block(ch_in=256, ch_out=512, t=t)self.RRCNN5 = RRCNN_block(ch_in=512, ch_out=1024, t=t)# 解码器部分,通过上采样将特征图的空间尺寸恢复self.Up5 = up_conv(ch_in=1024, ch_out=512)self.Up_RRCNN5 = RRCNN_block(ch_in=1024, ch_out=512, t=t)self.Up4 = up_conv(ch_in=512, ch_out=256)self.Up_RRCNN4 = RRCNN_block(ch_in=512, ch_out=256, t=t)self.Up3 = up_conv(ch_in=256, ch_out=128)self.Up_RRCNN3 = RRCNN_block(ch_in=256, ch_out=128, t=t)self.Up2 = up_conv(ch_in=128, ch_out=64)self.Up_RRCNN2 = RRCNN_block(ch_in=128, ch_out=64, t=t)# 最终的1x1卷积层,将输出通道数调整为目标输出通道数(一般为1,用于二分类)self.Conv_1x1 = nn.Conv2d(64, output_ch, kernel_size=1, stride=1, padding=0)def forward(self, x):# 编码路径# 1. 第一步,通过RRCNN块提取输入特征x1 = self.RRCNN1(x)# 2. 对x1进行下采样x2 = self.Maxpool(x1)# 3. 第二步,通过RRCNN块提取更深层次的特征x2 = self.RRCNN2(x2)# 4. 对x2进行下采样x3 = self.Maxpool(x2)# 5. 第三步,通过RRCNN块提取更深层次的特征x3 = self.RRCNN3(x3)# 6. 对x3进行下采样x4 = self.Maxpool(x3)# 7. 第四步,通过RRCNN块提取更深层次的特征x4 = self.RRCNN4(x4)# 8. 对x4进行下采样x5 = self.Maxpool(x4)# 9. 第五步,通过RRCNN块提取最深层次的特征x5 = self.RRCNN5(x5)# 解码 + 拼接路径# 10. 对x5进行上采样d5 = self.Up5(x5)# 11. 拼接来自编码器部分的x4与解码器部分的d5,增强特征d5 = torch.cat((x4, d5), dim=1)# 12. 通过RRCNN块进一步处理拼接后的特征d5 = self.Up_RRCNN5(d5)# 13. 对d5进行上采样d4 = self.Up4(d5)# 14. 拼接来自编码器部分的x3与解码器部分的d4d4 = torch.cat((x3, d4), dim=1)# 15. 通过RRCNN块进一步处理拼接后的特征d4 = self.Up_RRCNN4(d4)# 16. 对d4进行上采样d3 = self.Up3(d4)# 17. 拼接来自编码器部分的x2与解码器部分的d3d3 = torch.cat((x2, d3), dim=1)# 18. 通过RRCNN块进一步处理拼接后的特征d3 = self.Up_RRCNN3(d3)# 19. 对d3进行上采样d2 = self.Up2(d3)# 20. 拼接来自编码器部分的x1与解码器部分的d2d2 = torch.cat((x1, d2), dim=1)# 21. 通过RRCNN块进一步处理拼接后的特征d2 = self.Up_RRCNN2(d2)# 22. 使用1x1卷积将输出的通道数调整为目标通道数d1 = self.Conv_1x1(d2)# 返回最终的输出return d1
网络结构概述:
R2U_Net
是基于 U-Net 的变体,采用了 RRCNN_block(残差卷积神经网络块),使网络能够在每一层内进行残差学习,从而提高网络的训练效率和准确性。- 网络的输入为一张图像,输出为经过分割处理的图像(例如,在医学图像分割中,输出为分割后的组织区域)。
编码路径:
- 编码路径用于逐步提取输入图像的特征。每一层都通过 RRCNN_block(带有残差结构的卷积块)提取特征,同时通过 Maxpool(最大池化)减少空间分辨率,使网络能够捕捉更高层次的抽象特征。
- 编码部分有五层,逐步将特征图的通道数加深,从 64 到 1024。
解码路径:
- 解码路径通过 up_conv 层进行上采样,逐步恢复空间分辨率。
- 每一层的解码都会与编码部分的对应层特征进行拼接(即跳跃连接),以保持原始图像的空间细节,这对分割任务至关重要。
- 在每一层拼接之后,通过 RRCNN_block 进一步学习和调整特征。
跳跃连接:
- 解码层的每一层都会和编码层的特征图拼接(例如,
torch.cat((x4, d5), dim=1)
),这帮助网络在解码时恢复更细粒度的空间信息,尤其是在边界细节和小区域的分割上非常重要。最终输出:
- 网络的最后通过 Conv_1x1 层将输出图像的通道数调整为
output_ch
,通常为 1,表示分割结果。
AttU_Net
class AttU_Net(nn.Module):def __init__(self,img_ch=3,output_ch=1):super(AttU_Net,self).__init__()# Maxpool层,尺寸为2x2,用于下采样self.Maxpool = nn.MaxPool2d(kernel_size=2,stride=2)# 编码路径中的卷积块,每个块包含多个卷积层及激活函数self.Conv1 = conv_block(ch_in=img_ch, ch_out=64) # 输入图像通道为img_ch,输出64个通道self.Conv2 = conv_block(ch_in=64, ch_out=128) # 输入64个通道,输出128个通道self.Conv3 = conv_block(ch_in=128, ch_out=256) # 输入128个通道,输出256个通道self.Conv4 = conv_block(ch_in=256, ch_out=512) # 输入256个通道,输出512个通道self.Conv5 = conv_block(ch_in=512, ch_out=1024) # 输入512个通道,输出1024个通道# 解码路径中的上采样卷积块self.Up5 = up_conv(ch_in=1024, ch_out=512) # 输入1024个通道,输出512个通道self.Att5 = Attention_block(F_g=512, F_l=512, F_int=256) # Attention块,用于对特征进行加权self.Up_conv5 = conv_block(ch_in=1024, ch_out=512) # 1024个通道合并后,上采样得到512个通道self.Up4 = up_conv(ch_in=512, ch_out=256) # 输入512个通道,输出256个通道self.Att4 = Attention_block(F_g=256, F_l=256, F_int=128) # Attention块self.Up_conv4 = conv_block(ch_in=512, ch_out=256) # 512个通道合并后,上采样得到256个通道self.Up3 = up_conv(ch_in=256, ch_out=128) # 输入256个通道,输出128个通道self.Att3 = Attention_block(F_g=128, F_l=128, F_int=64) # Attention块self.Up_conv3 = conv_block(ch_in=256, ch_out=128) # 256个通道合并后,上采样得到128个通道self.Up2 = up_conv(ch_in=128, ch_out=64) # 输入128个通道,输出64个通道self.Att2 = Attention_block(F_g=64, F_l=64, F_int=32) # Attention块self.Up_conv2 = conv_block(ch_in=128, ch_out=64) # 128个通道合并后,上采样得到64个通道# 输出卷积,1x1卷积将输出的通道数转换为目标输出通道self.Conv_1x1 = nn.Conv2d(64, output_ch, kernel_size=1, stride=1, padding=0)def forward(self, x):# 编码路径x1 = self.Conv1(x) # 输入图像x经过卷积层Conv1处理,输出特征图x1x2 = self.Maxpool(x1) # 对x1进行Maxpool操作,尺寸减半x2 = self.Conv2(x2) # 输入x2经过卷积层Conv2处理,输出特征图x2x3 = self.Maxpool(x2) # 对x2进行Maxpool操作,尺寸减半x3 = self.Conv3(x3) # 输入x3经过卷积层Conv3处理,输出特征图x3x4 = self.Maxpool(x3) # 对x3进行Maxpool操作,尺寸减半x4 = self.Conv4(x4) # 输入x4经过卷积层Conv4处理,输出特征图x4x5 = self.Maxpool(x4) # 对x4进行Maxpool操作,尺寸减半x5 = self.Conv5(x5) # 输入x5经过卷积层Conv5处理,输出特征图x5# 解码路径 + 拼接路径d5 = self.Up5(x5) # 对x5进行上采样,得到d5x4 = self.Att5(g=d5, x=x4) # 使用Attention块对d5和x4进行加权,得到加权后的特征x4d5 = torch.cat((x4, d5), dim=1) # 将加权后的x4和d5拼接,合并到一起d5 = self.Up_conv5(d5) # 对拼接后的d5进行上采样,得到新的d5d4 = self.Up4(d5) # 对d5进行上采样,得到d4x3 = self.Att4(g=d4, x=x3) # 使用Attention块对d4和x3进行加权,得到加权后的特征x3d4 = torch.cat((x3, d4), dim=1) # 将加权后的x3和d4拼接,合并到一起d4 = self.Up_conv4(d4) # 对拼接后的d4进行上采样,得到新的d4d3 = self.Up3(d4) # 对d4进行上采样,得到d3x2 = self.Att3(g=d3, x=x2) # 使用Attention块对d3和x2进行加权,得到加权后的特征x2d3 = torch.cat((x2, d3), dim=1) # 将加权后的x2和d3拼接,合并到一起d3 = self.Up_conv3(d3) # 对拼接后的d3进行上采样,得到新的d3d2 = self.Up2(d3) # 对d3进行上采样,得到d2x1 = self.Att2(g=d2, x=x1) # 使用Attention块对d2和x1进行加权,得到加权后的特征x1d2 = torch.cat((x1, d2), dim=1) # 将加权后的x1和d2拼接,合并到一起d2 = self.Up_conv2(d2) # 对拼接后的d2进行上采样,得到新的d2d1 = self.Conv_1x1(d2) # 最后通过1x1卷积将输出通道数变为目标通道数return d1 # 返回最终的输出d1,即分割结果
- 编码路径(Encoder): 包括了多个卷积层和池化层,用于提取输入图像的特征并进行下采样,逐渐减小图像的空间维度,同时增加特征图的通道数。
- 解码路径(Decoder): 使用上采样卷积(
up_conv
)恢复图像的空间维度,同时通过注意力机制(Attention_block
)对每个解码层的特征进行加权,以突出更为重要的特征。- 跳跃连接(Skip Connections): 在解码过程中,来自编码路径的特征图与解码后的特征图进行拼接(
torch.cat
),这帮助恢复图像的空间分辨率并结合低级特征与高级特征。- 输出: 最终使用一个
1x1
的卷积层将通道数压缩到期望的输出通道数(例如1个通道用于二分类任务)。
R2AttU_Net
class R2AttU_Net(nn.Module):# 构造函数,初始化网络结构def __init__(self, img_ch=3, output_ch=1, t=2):super(R2AttU_Net, self).__init__()# 最大池化层,作用是降采样,减小特征图尺寸self.Maxpool = nn.MaxPool2d(kernel_size=2, stride=2)# 上采样层,作用是扩大特征图尺寸self.Upsample = nn.Upsample(scale_factor=2)# 定义五个编码模块,每个编码模块使用 RRCNN_blockself.RRCNN1 = RRCNN_block(ch_in=img_ch, ch_out=64, t=t)self.RRCNN2 = RRCNN_block(ch_in=64, ch_out=128, t=t)self.RRCNN3 = RRCNN_block(ch_in=128, ch_out=256, t=t)self.RRCNN4 = RRCNN_block(ch_in=256, ch_out=512, t=t)self.RRCNN5 = RRCNN_block(ch_in=512, ch_out=1024, t=t)# 上采样和注意力模块self.Up5 = up_conv(ch_in=1024, ch_out=512) # 上采样卷积层self.Att5 = Attention_block(F_g=512, F_l=512, F_int=256) # 注意力机制模块self.Up_RRCNN5 = RRCNN_block(ch_in=1024, ch_out=512, t=t) # 上采样后的卷积模块self.Up4 = up_conv(ch_in=512, ch_out=256)self.Att4 = Attention_block(F_g=256, F_l=256, F_int=128)self.Up_RRCNN4 = RRCNN_block(ch_in=512, ch_out=256, t=t)self.Up3 = up_conv(ch_in=256, ch_out=128)self.Att3 = Attention_block(F_g=128, F_l=128, F_int=64)self.Up_RRCNN3 = RRCNN_block(ch_in=256, ch_out=128, t=t)self.Up2 = up_conv(ch_in=128, ch_out=64)self.Att2 = Attention_block(F_g=64, F_l=64, F_int=32)self.Up_RRCNN2 = RRCNN_block(ch_in=128, ch_out=64, t=t)# 输出卷积层,输出通道数为目标图像的类别数(如分割任务中的类别数)self.Conv_1x1 = nn.Conv2d(64, output_ch, kernel_size=1, stride=1, padding=0)def forward(self, x):# 编码路径x1 = self.RRCNN1(x) # 第一层R2U-Net模块x2 = self.Maxpool(x1) # 最大池化,降采样x2 = self.RRCNN2(x2) # 第二层R2U-Net模块x3 = self.Maxpool(x2) # 最大池化x3 = self.RRCNN3(x3) # 第三层R2U-Net模块x4 = self.Maxpool(x3) # 最大池化x4 = self.RRCNN4(x4) # 第四层R2U-Net模块x5 = self.Maxpool(x4) # 最大池化x5 = self.RRCNN5(x5) # 第五层R2U-Net模块(瓶颈层)# 解码和跳跃连接路径d5 = self.Up5(x5) # 上采样x4 = self.Att5(g=d5, x=x4) # 注意力机制d5 = torch.cat((x4, d5), dim=1) # 融合编码和解码信息d5 = self.Up_RRCNN5(d5) # R2U-Net模块d4 = self.Up4(d5) # 上采样x3 = self.Att4(g=d4, x=x3) # 注意力机制d4 = torch.cat((x3, d4), dim=1) # 融合d4 = self.Up_RRCNN4(d4) # R2U-Net模块d3 = self.Up3(d4) # 上采样x2 = self.Att3(g=d3, x=x2) # 注意力机制d3 = torch.cat((x2, d3), dim=1) # 融合d3 = self.Up_RRCNN3(d3) # R2U-Net模块d2 = self.Up2(d3) # 上采样x1 = self.Att2(g=d2, x=x1) # 注意力机制d2 = torch.cat((x1, d2), dim=1) # 融合d2 = self.Up_RRCNN2(d2) # R2U-Net模块# 最终输出d1 = self.Conv_1x1(d2) # 输出层,通常是类别数为 1 的卷积层return d1 # 返回最终的输出结果
这个
R2AttU_Net
模型是一个结合了 R2U-Net 和注意力机制(Attention Mechanism)的网络,主要用于图像分割任务。其结构包括:
- 编码部分:多个
RRCNN_block
负责提取图像的特征。- 解码部分:通过上采样和融合编码部分的特征,同时使用注意力机制模块对关键区域进行加权,从而提升模型的分割精度。
- 跳跃连接:每个解码层都会与对应的编码层通过跳跃连接(skip connection)融合,以便恢复更多的空间信息。
- 最终输出:通过一个 1x1 卷积层输出最终的分割结果。
注意:
RRCNN_block
、up_conv
和Attention_block
是在代码之外定义的模块,假设它们是实现了特定功能的模块。RRCNN_block
是一种基于卷积神经网络(CNN)的改进模块,up_conv
是一个上采样卷积层,Attention_block
是一个利用注意力机制增强模型聚焦于重要特征区域的模块。
main.py
def main(config):# 设置 cudnn.benchmark 为 True,启用 cudnn 的优化选项,适用于输入图像大小固定的情况。cudnn.benchmark = True# 修改模型类型为 'R2U_Net'config.model_type = 'R2U_Net'# 检查配置中的模型类型是否在允许的选项列表内if config.model_type not in ['U_Net', 'R2U_Net', 'AttU_Net', 'R2AttU_Net']:# 如果模型类型无效,输出错误信息并退出程序print('ERROR!! model_type should be selected in U_Net/R2U_Net/AttU_Net/R2AttU_Net')print('Your input for model_type was %s' % config.model_type)return# 如果模型保存路径不存在,则创建该路径if not os.path.exists(config.model_path):os.makedirs(config.model_path)# 如果结果保存路径不存在,则创建该路径if not os.path.exists(config.result_path):os.makedirs(config.result_path)# 将结果路径设定为 model_type 子文件夹路径,用于存储不同模型的结果config.result_path = os.path.join(config.result_path, config.model_type)# 如果该路径不存在,则创建它if not os.path.exists(config.result_path):os.makedirs(config.result_path)# 修改学习率(lr),使用随机值,范围在0.0000005到0.0005之间lr = random.random() * 0.0005 + 0.0000005# 修改数据增强的概率,随机值,范围在0到0.7之间augmentation_prob = random.random() * 0.7# 随机选择训练的 epoch 数量,可能是 100, 150, 200 或 250epoch = random.choice([100, 150, 200, 250])# 随机选择衰减的比例(decay_ratio),该比例决定了衰减epoch的数量decay_ratio = random.random() * 0.8# 计算衰减epoch的数量,根据衰减比例计算decay_epoch = int(epoch * decay_ratio)# 将这些配置项更新到 config 对象中config.augmentation_prob = augmentation_probconfig.num_epochs = epochconfig.lr = lrconfig.num_epochs_decay = decay_epoch# 输出当前的配置参数,方便调试print(config)# 获取训练数据加载器(train_loader)train_loader = get_loader(image_path=config.train_path,image_size=config.image_size,batch_size=config.batch_size,num_workers=config.num_workers,mode='train',augmentation_prob=config.augmentation_prob)# 获取验证数据加载器(valid_loader),此处不使用数据增强valid_loader = get_loader(image_path=config.valid_path,image_size=config.image_size,batch_size=config.batch_size,num_workers=config.num_workers,mode='valid',augmentation_prob=0.)# 获取测试数据加载器(test_loader),此处不使用数据增强test_loader = get_loader(image_path=config.test_path,image_size=config.image_size,batch_size=config.batch_size,num_workers=config.num_workers,mode='test',augmentation_prob=0.)# 创建 Solver 实例,Solver 负责训练和测试过程的控制solver = Solver(config, train_loader, valid_loader, test_loader)# 如果模式为训练(train),则调用训练函数进行训练if config.mode == 'train':solver.train()# 如果模式为测试(test),则调用测试函数进行测试elif config.mode == 'test':solver.test()
main()
函数主要完成模型配置、路径设置、数据加载器创建以及训练/测试模式的选择。如果是训练模式,则开始训练过程;如果是测试模式,则开始模型测试。随机生成的一些参数(如学习率、epoch 数量和数据增强的概率)使得每次运行时可以尝试不同的配置,从而有助于模型的泛化能力和优化。
if __name__ == '__main__':# 创建 ArgumentParser 对象,用于处理命令行参数parser = argparse.ArgumentParser()# model hyper-parameters(模型超参数)# 解析输入的图像大小,默认值为 224parser.add_argument('--image_size', type=int, default=224)# 解析 R2U_Net 或 R2AttU_Net 模型的递归步骤数,默认为 3parser.add_argument('--t', type=int, default=3, help='t for Recurrent step of R2U_Net or R2AttU_Net')# training hyper-parameters(训练超参数)# 解析输入图像的通道数,默认为 3(表示 RGB 图像)parser.add_argument('--img_ch', type=int, default=3)# 解析输出图像的通道数,默认为 1(常见的二值分割任务)parser.add_argument('--output_ch', type=int, default=1)# 解析训练的 epoch 数量,默认为 100parser.add_argument('--num_epochs', type=int, default=100)# 解析学习率衰减的 epoch 数量,默认为 70parser.add_argument('--num_epochs_decay', type=int, default=70)# 解析每个训练批次的大小,默认为 1parser.add_argument('--batch_size', type=int, default=1)# 解析加载数据时使用的并行工作线程数量,默认为 8parser.add_argument('--num_workers', type=int, default=8)# 解析学习率,默认为 0.0002parser.add_argument('--lr', type=float, default=0.0002)# 解析 Adam 优化器的 momentum1 参数,默认为 0.5(控制梯度的一阶矩估计)parser.add_argument('--beta1', type=float, default=0.5)# 解析 Adam 优化器的 momentum2 参数,默认为 0.999(控制梯度的二阶矩估计)parser.add_argument('--beta2', type=float, default=0.999)# 解析数据增强的概率,默认为 0.4,表示图像有 40% 的概率会进行数据增强parser.add_argument('--augmentation_prob', type=float, default=0.4)# 解析训练过程中每隔多少步打印一次日志,默认为 2parser.add_argument('--log_step', type=int, default=2)# 解析每隔多少步进行一次验证,默认为 2parser.add_argument('--val_step', type=int, default=2)# misc(其他配置项)# 解析当前模式,默认为 'train',可以是 'train' 或 'test'parser.add_argument('--mode', type=str, default='train')# 解析模型类型,默认为 'U_Net',可选值为 'U_Net','R2U_Net','AttU_Net','R2AttU_Net'parser.add_argument('--model_type', type=str, default='U_Net', help='U_Net/R2U_Net/AttU_Net/R2AttU_Net')# 解析模型保存路径,默认为 './models'parser.add_argument('--model_path', type=str, default='./models')# 解析训练数据集路径,默认为 './dataset/train/'parser.add_argument('--train_path', type=str, default='./dataset/train/')# 解析验证数据集路径,默认为 './dataset/valid/'parser.add_argument('--valid_path', type=str, default='./dataset/valid/')# 解析测试数据集路径,默认为 './dataset/test/'parser.add_argument('--test_path', type=str, default='./dataset/test/')# 解析结果保存路径,默认为 './result/'parser.add_argument('--result_path', type=str, default='./result/')# 解析使用的 CUDA 设备编号,默认为 1parser.add_argument('--cuda_idx', type=int, default=1)# 解析命令行参数,并将其存储在 config 变量中config = parser.parse_args()# 调用 main 函数,传入 config 参数main(config)
这段代码的作用是设置和解析命令行参数,允许用户通过命令行灵活地配置模型训练和测试过程中的各项超参数。通过这种方式,可以在不同的实验中快速调整参数并执行训练或测试任务。