网站建设短信,石家庄seo网站优化公司,房屋平面设计图,wordpress留言发送邮件在本文中#xff0c;主要是对3D UNet 进行一个学习和梳理。对于3D UNet 网上的资料和GitHub直接获取的代码很多#xff0c;不需要自己从0开始。那么本文的目的是啥呢#xff1f;
本文就是想拆解下其中的结构#xff0c;看看对于一个3D的UNet#xff0c;和2D的UNet#x…在本文中主要是对3D UNet 进行一个学习和梳理。对于3D UNet 网上的资料和GitHub直接获取的代码很多不需要自己从0开始。那么本文的目的是啥呢
本文就是想拆解下其中的结构看看对于一个3D的UNet和2D的UNet究竟有什么不同如果是你自己构建有什么样的经验和技巧可以学习。
3D的UNet的论文地址3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation
对于2D的UNet感兴趣的小伙伴可以先跳转去这里【BraTS】Brain Tumor Segmentation 脑部肿瘤分割2UNet的复现相信阅读完你会对这个模型心中已经有了结构。
对本系列的其他篇章点击下面链接
【3D 图像分割】基于 Pytorch 的 VNet 3D 图像分割1综述篇【3D 图像分割】基于 Pytorch 的 VNet 3D 图像分割2基础数据流篇【3D 图像分割】基于 Pytorch 的 VNet 3D 图像分割6数据预处理【3D 图像分割】基于 Pytorch 的 VNet 3D 图像分割7数据预处理【3D 图像分割】基于 Pytorch 的 VNet 3D 图像分割8CT肺实质分割【3D 图像分割】基于 Pytorch 的 VNet 3D 图像分割9patch 的 crop 和 merge 操作
一、 3D UNet 结构剖析
unet无论是2D还是3D从整体结构上进行划分大体可以分位以下两个阶段
下采样的阶段也就是U的左边encoder负责对特征提取上采样的阶段也就是U的右边decoder负责对预测恢复。
如下图展示的这样
其中
蓝色框表示的是特征图绿色长箭头是concat操作橘色三角是convbnrelu的组合红色的向下箭头是max pool黄色的向上箭头是up conv最后的紫色三角是conv恢复了最终的输出特征图
对于模型构建这块可以在论文中看看作者是如何描述网络结构的
Like the standard u-net, it has an analysis and a synthesis path each with four resolution steps.In the analysis path, each layer contains two 3 × 3 × 3 convolutions each followed by a rectified linear unit (ReLu), and then a 2 × 2 × 2 max pooling with strides of two in each dimension.In the synthesis path, each layer consists of an upconvolution of 2 × 2 × 2 by strides of two in each dimension, followed by two 3 × 3 × 3 convolutions each followed by a ReLu.Shortcut connections from layers of equal resolution in the analysis path provide the essential high-resolution features to the synthesis path.In the last layer a 1×1×1 convolution reduces the number of output channels to the number of labels which is 3 in our case.
从论文中的网络结构示意图也可以发现
水平看每一个小块基本都是三个特征图最后一层除外水平看每个特征图之间都是橘色三角是convbnrelu的组合最后一层除外encoder阶段连接各个水平块的是下采样decoder阶段连接各个水平块的是反卷积upconvolution还有就是绿色长箭头的concat和最后的conv输出特征图。
二、 3D UNet 复现
复线在3D UNet前可以先参照下相对简单且很深渊源的2D UNet结构。其中被多次使用的一个水平块中也是两个convbnrelu的组合2D UNet的构建如下所示
class ConvBlock2d(nn.Module):def __init__(self, in_ch, out_ch):super(ConvBlock2d, self).__init__()# 第1个3*3的卷积层self.conv1 nn.Sequential(nn.Conv2d(in_ch, out_ch, kernel_size3, stride1, padding1),nn.BatchNorm2d(out_ch),nn.ReLU(inplaceTrue),)# 第2个3*3的卷积层self.conv2 nn.Sequential(nn.Conv2d(out_ch, out_ch, kernel_size3, stride1, padding1),nn.BatchNorm2d(out_ch),nn.ReLU(inplaceTrue),)# 定义数据前向流动形式def forward(self, x):x self.conv1(x)x self.conv2(x)return x而在3D UNet的一个水平块中同样是两个convbnrelu的组合如下所示
is_elu False
def activateELU(is_elu, nchan):if is_elu:return nn.ELU(inplaceTrue)else:return nn.PReLU(nchan)def ConvBnActivate(in_channels, middle_channels, out_channels):# This is a block with 2 convolutions# The first convolution goes from in_channels to middle_channels feature maps# The second convolution goes from middle_channels to out_channels feature mapsconv nn.Sequential(nn.Conv3d(in_channels, middle_channels, stride1, kernel_size3, padding1),nn.BatchNorm3d(middle_channels),activateELU(is_elu, middle_channels),nn.Conv3d(middle_channels, out_channels, stride1, kernel_size3, padding1),nn.BatchNorm3d(out_channels),activateELU(is_elu, out_channels),)return conv2.1、模块搭建
可以发现nn.Conv2d变成了nn.Conv3dnn.BatchNorm2d变成了nn.BatchNorm3d。遵照这个规则构建下采样MaxPool3d、上采样反卷积ConvTranspose3d以及最后紫色一层卷积输出特征层FinalConvolution如下
def DownSample():# It halves the spatial dimensions on every axes (x,y,z)return nn.MaxPool3d(kernel_size2, stride2)def UpSample(in_channels, out_channels):# It doubles the spatial dimensions on every axes (x,y,z)return nn.ConvTranspose3d(in_channels, out_channels, kernel_size2, stride2)def FinalConvolution(in_channels, out_channels):return nn.Conv3d(in_channels, out_channels, kernel_size1)除此之外绿色长箭头concat操作是在水平方向上也就是列上进行组合如下所示
def CatBlock(x1, x2):return torch.cat((x1, x2), 1)至此构建模型所需要的各个组块都准备完毕了。接下来就是构建模型将各个组块搭起来。其中有个规律
除encoder中第一convbnrelu外每一次前都需要下采样decoder中每一个convbnrelu前都需要上采样并且decoder中第一个conv操作需要进行concat操作DownSample的channel不变特征图尺寸变小UpSample的channel不变特征图尺寸变大
那就把这些规则根据图示给加上组合后的一个类就如下所示
import torch
import torch.nn as nn
import torch.nn.functional as Fclass UNet3D(nn.Module):def __init__(self, num_out_classes2, input_channels1, init_feat_channels32):super().__init__()# Encoder layers definitionsself.down_sample DownSample()self.init_conv ConvBnActivate(input_channels, init_feat_channels, init_feat_channels*2)self.down_conv1 ConvBnActivate(init_feat_channels*2, init_feat_channels*2, init_feat_channels*4)self.down_conv2 ConvBnActivate(init_feat_channels*4, init_feat_channels*4, init_feat_channels*8)self.down_conv3 ConvBnActivate(init_feat_channels*8, init_feat_channels*8, init_feat_channels*16)# Decoder layers definitionsself.up_sample1 UpSample(init_feat_channels*16, init_feat_channels*16)self.up_conv1 ConvBnActivate(init_feat_channels*(168), init_feat_channels*8, init_feat_channels*8)self.up_sample2 UpSample(init_feat_channels*8, init_feat_channels*8)self.up_conv2 ConvBnActivate(init_feat_channels*(84), init_feat_channels*4, init_feat_channels*4)self.up_sample3 UpSample(init_feat_channels*4, init_feat_channels*4)self.up_conv3 ConvBnActivate(init_feat_channels*(42), init_feat_channels*2, init_feat_channels*2)self.final_conv FinalConvolution(init_feat_channels*2, num_out_classes)# Softmaxself.softmax F.softmaxdef forward(self, image):# Encoder Part ## B x 1 x Z x Y x Xlayer_init self.init_conv(image)# B x 64 x Z x Y x Xmax_pool1 self.down_sample(layer_init)# B x 64 x Z//2 x Y//2 x X//2layer_down2 self.down_conv1(max_pool1)# B x 128 x Z//2 x Y//2 x X//2max_pool2 self.down_sample(layer_down2)# B x 128 x Z//4 x Y//4 x X//4layer_down3 self.down_conv2(max_pool2)# B x 256 x Z//4 x Y//4 x X//4max_pool_3 self.down_sample(layer_down3)# B x 256 x Z//8 x Y//8 x X//8layer_down4 self.down_conv3(max_pool_3)# B x 512 x Z//8 x Y//8 x X//8# Decoder part #layer_up1 self.up_sample1(layer_down4)# B x 512 x Z//4 x Y//4 x X//4cat_block1 CatBlock(layer_down3, layer_up1)# B x (256512) x Z//4 x Y//4 x X//4layer_conv_up1 self.up_conv1(cat_block1)# B x 256 x Z//4 x Y//4 x X//4layer_up2 self.up_sample2(layer_conv_up1)# B x 256 x Z//2 x Y//2 x X//2cat_block2 CatBlock(layer_down2, layer_up2)# B x (128256) x Z//2 x Y//2 x X//2layer_conv_up2 self.up_conv2(cat_block2)# B x 128 x Z//2 x Y//2 x X//2layer_up3 self.up_sample3(layer_conv_up2)# B x 128 x Z x Y x Xcat_block3 CatBlock(layer_init, layer_up3)# B x (64128) x Z x Y x Xlayer_conv_up3 self.up_conv3(cat_block3)# B x 64 x Z x Y x Xfinal_layer self.final_conv(layer_conv_up3)# B x 2 x Z x Y x Xreturn self.softmax(final_layer, dim1)2.2、模型初测
定义好了模型还不算完分阶段测试下构建的网络是不是和我们所预想的一样。我们给他一个输入测试下是否与我们最初的想法是一致的是否报错等等问题如下这样
DEVICE torch.device(cuda if torch.cuda.is_available() else cpu) # 没gpu就用cpu
print(DEVICE)# Tensors for 3D Image Processing in PyTorch
# Batch x Channel x Z x Y x X
# Batch size BY x Number of channels x (BY Z dim) x (BY Y dim) x (BY X dim)if __name__ __main__:from torchsummary import summarymodel UNet3D(num_out_classes3, input_channels3, init_feat_channels32)# print(model)summary(model, input_size(3, 128, 128, 64), batch_size-1, devicecpu)打印的内容如下
----------------------------------------------------------------Layer (type) Output Shape Param #
Conv3d-1 [-1, 32, 128, 128, 64] 2,624BatchNorm3d-2 [-1, 32, 128, 128, 64] 64PReLU-3 [-1, 32, 128, 128, 64] 32Conv3d-4 [-1, 64, 128, 128, 64] 55,360BatchNorm3d-5 [-1, 64, 128, 128, 64] 128PReLU-6 [-1, 64, 128, 128, 64] 64MaxPool3d-7 [-1, 64, 64, 64, 32] 0Conv3d-8 [-1, 64, 64, 64, 32] 110,656BatchNorm3d-9 [-1, 64, 64, 64, 32] 128PReLU-10 [-1, 64, 64, 64, 32] 64Conv3d-11 [-1, 128, 64, 64, 32] 221,312BatchNorm3d-12 [-1, 128, 64, 64, 32] 256PReLU-13 [-1, 128, 64, 64, 32] 128MaxPool3d-14 [-1, 128, 32, 32, 16] 0Conv3d-15 [-1, 128, 32, 32, 16] 442,496BatchNorm3d-16 [-1, 128, 32, 32, 16] 256PReLU-17 [-1, 128, 32, 32, 16] 128Conv3d-18 [-1, 256, 32, 32, 16] 884,992BatchNorm3d-19 [-1, 256, 32, 32, 16] 512PReLU-20 [-1, 256, 32, 32, 16] 256MaxPool3d-21 [-1, 256, 16, 16, 8] 0Conv3d-22 [-1, 256, 16, 16, 8] 1,769,728BatchNorm3d-23 [-1, 256, 16, 16, 8] 512PReLU-24 [-1, 256, 16, 16, 8] 256Conv3d-25 [-1, 512, 16, 16, 8] 3,539,456BatchNorm3d-26 [-1, 512, 16, 16, 8] 1,024PReLU-27 [-1, 512, 16, 16, 8] 512ConvTranspose3d-28 [-1, 512, 32, 32, 16] 2,097,664Conv3d-29 [-1, 256, 32, 32, 16] 5,308,672BatchNorm3d-30 [-1, 256, 32, 32, 16] 512PReLU-31 [-1, 256, 32, 32, 16] 256Conv3d-32 [-1, 256, 32, 32, 16] 1,769,728BatchNorm3d-33 [-1, 256, 32, 32, 16] 512PReLU-34 [-1, 256, 32, 32, 16] 256ConvTranspose3d-35 [-1, 256, 64, 64, 32] 524,544Conv3d-36 [-1, 128, 64, 64, 32] 1,327,232BatchNorm3d-37 [-1, 128, 64, 64, 32] 256PReLU-38 [-1, 128, 64, 64, 32] 128Conv3d-39 [-1, 128, 64, 64, 32] 442,496BatchNorm3d-40 [-1, 128, 64, 64, 32] 256PReLU-41 [-1, 128, 64, 64, 32] 128ConvTranspose3d-42 [-1, 128, 128, 128, 64] 131,200Conv3d-43 [-1, 64, 128, 128, 64] 331,840BatchNorm3d-44 [-1, 64, 128, 128, 64] 128PReLU-45 [-1, 64, 128, 128, 64] 64Conv3d-46 [-1, 64, 128, 128, 64] 110,656BatchNorm3d-47 [-1, 64, 128, 128, 64] 128PReLU-48 [-1, 64, 128, 128, 64] 64Conv3d-49 [-1, 3, 128, 128, 64] 195Total params: 19,077,859
Trainable params: 19,077,859
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 12.00
Forward/backward pass size (MB): 8544.00
Params size (MB): 72.78
Estimated Total Size (MB): 8628.78
----------------------------------------------------------------其中我们测试的参数量是19,077,859论文中说的参数量The architecture has 19069955 parameters in total. 有略微的差别。
后面再调用模型进行一次前向传播loss运算和反向回归。如果这里都通过了那么后面构建训练代码就更简单了很多。如下
if __name__ __main__:input_channels 3num_out_classes 2init_feat_channels 32batch_size 4model UNet3D(num_out_classesnum_out_classes, input_channelsinput_channels, init_feat_channelsinit_feat_channels)# B x C x Z x Y x X# 4 x 1 x 64 x 64 x 64input_batch_size (batch_size, input_channels, 128, 128, 64)input_example torch.rand(input_batch_size)unet model.to(DEVICE)input_example input_example.to(DEVICE)output unet(input_example)# output output.cpu().detach().numpy()# Expected output shape# B x N x Z x Y x X# 4 x 2 x 64 x 64 x 64expected_output_shape (batch_size, num_out_classes, 128, 128, 64)print(Output shape {}.format(output.shape))assert output.shape expected_output_shape, Unexpected output shape, check the architecture!expected_gt_shape (batch_size, 128, 128, 64)ground_truth torch.ones(expected_gt_shape)ground_truth ground_truth.long().to(DEVICE)# Defining loss fnce_layer torch.nn.CrossEntropyLoss()# Calculating lossce_loss ce_layer(output, ground_truth)print(CE Loss {}.format(ce_loss))# Back propagationce_loss.backward()输出内容如下
Output shape torch.Size([4, 2, 128, 128, 64])
CE Loss 0.68233871459960942.3、疑问汇总
在GitHub上一篇关于3D UNet的仓库获得了1.6k 星星。链接地址在这里pytorch-3dunet
在这个GitHub里面增加了很多的注释也带来了一些心中的疑惑。
2.3.1、什么时候使用softmax什么时候使用sigmoid
选择使用softmax或sigmoid作为输出层的依据取决于您的任务类型和具体情况。 如果您的任务是对每个像素进行多类别分类(语义分割)例如图像分割任务那么您可以使用softmax作为输出层。softmax将为每个像素分配一个概率分布表示该像素属于每个类别的概率这样可以确保每个像素的预测结果归一化并且所有通道的概率之和为1。这种方法通常用于分割器官或病变等结构。 如果您的任务是对每个像素进行二元分类例如肿瘤检测任务那么您可以使用sigmoid作为输出层。sigmoid将为每个像素分配一个0到1之间的值表示该像素属于正类的概率。这种方法通常用于检测二元结构如肿瘤。但是二元分类任务使用softmax也是可以的。
总之选择哪种输出层取决于您的任务类型和具体情况。
2.3.2、训练阶段是不需要softmax/sigmoid只在推理阶段使用呢 if True applies the final normalization layer (sigmoid or softmax), otherwise the networks returns the output from the final convolution layer; use False for regression problems, e.g. de-noising 在训练阶段输出层的特征图通常不需要经过sigmoid或softmax函数处理因为在计算损失函数时通常会使用原始的特征图和标签图进行比较。 在推理阶段输出层的特征图需要经过sigmoid或softmax函数处理以将特征图转换为像素级别的预测结果。对于分割一个类别的任务您可以使用sigmoid函数将特征图转换为像素级别的二进制掩码表示每个像素属于结节的概率。对于分割多个类别的任务您可以使用softmax函数将特征图转换为像素级别的类别标签。
因此在推理阶段您需要将输出层的特征图通过sigmoid或softmax函数进行处理以获得像素级别的预测结果。
在上面的GitHub有个训练的提示如下这样 Training loss shape of target When training with binary-based losses, i.e.: BCEWithLogitsLoss, DiceLoss, BCEDiceLoss, GeneralizedDiceLoss: The target data has to be 4D (one target binary mask per channel).When training with WeightedCrossEntropyLoss, CrossEntropyLoss, PixelWiseCrossEntropyLoss the target dataset has to be 3D, see also pytorch documentation for CE loss: https://pytorch.org/docs/master/generated/torch.nn.CrossEntropyLoss.html final_sigmoid in the model config section applies only to the inference time (validation, test): When training with BCEWithLogitsLoss, DiceLoss, BCEDiceLoss, GeneralizedDiceLoss set final_sigmoidTrue;When training with cross entropy based losses (WeightedCrossEntropyLoss, CrossEntropyLoss, PixelWiseCrossEntropyLoss) set final_sigmoidFalse so that Softmax normalization is applied to the output.
2.3.3、在训练阶段真的不可以加入sigmoid或softmax吗
万事没有一个太绝对了。在训练阶段使用了sigmoid或softmax也是可以的以获得类似于推理阶段的预测结果。这种方法称为“软标签”可以帮助模型更好地学习特征和提高分割结果的质量。因为sigmoid或softmax类似于一个规范化层可以降低提高收敛效率
使用软标签时您需要将每个像素的标签从硬标签0或1转换为概率分布。对于分割一个类别的任务您可以使用sigmoid函数将标签转换为0到1之间的值表示该像素属于结节的概率。对于分割多个类别的任务您可以使用softmax函数将标签转换为每个类别的概率分布。
请注意使用软标签会增加模型的训练难度和计算复杂度。因此
如果您的数据集足够大且质量良好您可以不使用软标签来训练模型也就是训练阶段不使用sigmoid或softmax但是如果您的数据集较小或存在噪声数据使用软标签可能会提高模型的性能和分割结果的质量。也就是训练阶段使用sigmoid或softmax。
在论文3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation论文中第3.2章节介绍了如何使用软标签。 2.3.4、out_channels 的数量要不要加背景层 out_channels (int): number of output segmentation masks; Note that the of out_channels might correspond to either different semantic classes or to different binary segmentation mask. It’s up to the user of the class to interpret the out_channels and use the proper loss criterion during training (i.e. CrossEntropyLoss (multi-class) or BCEWithLogitsLoss (two-class) respectively) 我的理解是有多少个目标类out_channels 就是多少不需要加背景类。但是我也看到就只有一个类别但是做了加1操作的。这点我再了解下。如果你有什么心得欢迎评论区交流。
三、总结
UNet网络的结构无论是二维的还是三维的都是比较容易理解的这可能也是为什么那么受欢迎的原因之一吧。如果你看过之前那篇关于2D UNet的过程再看本篇应该就简单的很多。觉得本篇更简单一些呢。
我觉得本篇最大的价值就是
逐模块的分析了结构对后续的模型构建提供了思路构建完模型需要先预测试两种方式可选对模型的优势和劣势分析。
如果你阅读的过程中发现了问题和疑问欢迎评论区交流。
