UNet

网上的资料和GitHub直接获取的代码很多#xff0c;不需要自己从0开始。

那么本文的目的是啥呢#xff1f;

本文就是想拆解下其中的结构#xff0c;看看对于一个3D的UNet#xff0c;和2D的UNet#x…在本文中主要是对3D

UNet

网上的资料和GitHub直接获取的代码很多不需要自己从0开始。

那么本文的目的是啥呢

本文就是想拆解下其中的结构看看对于一个3D的UNet和2D的UNet究竟有什么不同如果是你自己构建有什么样的经验和技巧可以学习。

U-Net:

对于2D的UNet感兴趣的小伙伴可以先跳转去这里【BraTS】Brain

Tumor

脑部肿瘤分割2UNet的复现相信阅读完你会对这个模型心中已经有了结构。

【3D

unet无论是2D还是3D从整体结构上进行划分大体可以分位以下两个阶段

下采样的阶段也就是U的左边encoder负责对特征提取上采样的阶段也就是U的右边decoder负责对预测恢复。

蓝色框表示的是特征图绿色长箭头是concat操作橘色三角是convbnrelu的组合红色的向下箭头是max

Like

水平看每一个小块基本都是三个特征图最后一层除外水平看每个特征图之间都是橘色三角是convbnrelu的组合最后一层除外encoder阶段连接各个水平块的是下采样decoder阶段连接各个水平块的是反卷积upconvolution还有就是绿色长箭头的concat和最后的conv输出特征图。

UNet

UNet结构。

其中被多次使用的一个水平块中也是两个convbnrelu的组合2D

class

padding1),nn.BatchNorm2d(out_ch),nn.ReLU(inplaceTrue),)#

nn.Sequential(nn.Conv2d(out_ch,

out_ch,

padding1),nn.BatchNorm2d(out_ch),nn.ReLU(inplaceTrue),)#

定义数据前向流动形式def

UNet的一个水平块中同样是两个convbnrelu的组合如下所示

is_elu

nn.Sequential(nn.Conv3d(in_channels,

middle_channels,

padding1),nn.BatchNorm3d(middle_channels),activateELU(is_elu,

middle_channels),nn.Conv3d(middle_channels,

out_channels,

padding1),nn.BatchNorm3d(out_channels),activateELU(is_elu,

conv2.1、模块搭建

可以发现nn.Conv2d变成了nn.Conv3dnn.BatchNorm2d变成了nn.BatchNorm3d。

遵照这个规则构建下采样MaxPool3d、上采样反卷积ConvTranspose3d以及最后紫色一层卷积输出特征层FinalConvolution如下

def

nn.ConvTranspose3d(in_channels,

out_channels,

kernel_size1)除此之外绿色长箭头concat操作是在水平方向上也就是列上进行组合如下所示

def

1)至此构建模型所需要的各个组块都准备完毕了。

接下来就是构建模型将各个组块搭起来。

其中有个规律

除encoder中第一convbnrelu外每一次前都需要下采样decoder中每一个convbnrelu前都需要上采样并且decoder中第一个conv操作需要进行concat操作DownSample的channel不变特征图尺寸变小UpSample的channel不变特征图尺寸变大

import

init_feat_channels32):super().__init__()#

Encoder

init_feat_channels*2)self.down_conv1

ConvBnActivate(init_feat_channels*2,

init_feat_channels*4)self.down_conv2

ConvBnActivate(init_feat_channels*4,

init_feat_channels*8)self.down_conv3

ConvBnActivate(init_feat_channels*8,

Decoder

UpSample(init_feat_channels*16,

init_feat_channels*16)self.up_conv1

ConvBnActivate(init_feat_channels*(168),

init_feat_channels*8)self.up_sample2

init_feat_channels*8)self.up_conv2

ConvBnActivate(init_feat_channels*(84),

init_feat_channels*4)self.up_sample3

init_feat_channels*4)self.up_conv3

ConvBnActivate(init_feat_channels*(42),

init_feat_channels*2)self.final_conv

FinalConvolution(init_feat_channels*2,

num_out_classes)#

self.up_sample2(layer_conv_up1)#

256

self.up_sample3(layer_conv_up2)#

128

self.final_conv(layer_conv_up3)#

Xreturn

定义好了模型还不算完分阶段测试下构建的网络是不是和我们所预想的一样。

我们给他一个输入测试下是否与我们最初的想法是一致的是否报错等等问题如下这样

DEVICE

----------------------------------------------------------------Layer

(type)

----------------------------------------------------------------

Input

----------------------------------------------------------------其中我们测试的参数量是19,077,859论文中说的参数量The

architecture

后面再调用模型进行一次前向传播loss运算和反向回归。

如果这里都通过了那么后面构建训练代码就更简单了很多。

如下

__name__

UNet3D(num_out_classesnum_out_classes,

init_feat_channelsinit_feat_channels)#

64input_batch_size

torch.rand(input_batch_size)unet

unet(input_example)#

torch.ones(expected_gt_shape)ground_truth

ground_truth.long().to(DEVICE)#

Defining

propagationce_loss.backward()输出内容如下

Output

在这个GitHub里面增加了很多的注释也带来了一些心中的疑惑。

2.3.1、什么时候使用softmax什么时候使用sigmoid

选择使用softmax或sigmoid作为输出层的依据取决于您的任务类型和具体情况。

如果您的任务是对每个像素进行多类别分类(语义分割)例如图像分割任务那么您可以使用softmax作为输出层。

softmax将为每个像素分配一个概率分布表示该像素属于每个类别的概率这样可以确保每个像素的预测结果归一化并且所有通道的概率之和为1。

这种方法通常用于分割器官或病变等结构。

如果您的任务是对每个像素进行二元分类例如肿瘤检测任务那么您可以使用sigmoid作为输出层。

sigmoid将为每个像素分配一个0到1之间的值表示该像素属于正类的概率。

这种方法通常用于检测二元结构如肿瘤。

但是二元分类任务使用softmax也是可以的。

2.3.2、训练阶段是不需要softmax/sigmoid只在推理阶段使用呢

True

在训练阶段输出层的特征图通常不需要经过sigmoid或softmax函数处理因为在计算损失函数时通常会使用原始的特征图和标签图进行比较。

在推理阶段输出层的特征图需要经过sigmoid或softmax函数处理以将特征图转换为像素级别的预测结果。

对于分割一个类别的任务您可以使用sigmoid函数将特征图转换为像素级别的二进制掩码表示每个像素属于结节的概率。

对于分割多个类别的任务您可以使用softmax函数将特征图转换为像素级别的类别标签。

因此在推理阶段您需要将输出层的特征图通过sigmoid或softmax函数进行处理以获得像素级别的预测结果。

Training

https://pytorch.org/docs/master/generated/torch.nn.CrossEntropyLoss.html

final_sigmoid

2.3.3、在训练阶段真的不可以加入sigmoid或softmax吗

万事没有一个太绝对了。

在训练阶段使用了sigmoid或softmax也是可以的以获得类似于推理阶段的预测结果。

这种方法称为“软标签”可以帮助模型更好地学习特征和提高分割结果的质量。

因为sigmoid或softmax类似于一个规范化层可以降低提高收敛效率

使用软标签时您需要将每个像素的标签从硬标签0或1转换为概率分布。

对于分割一个类别的任务您可以使用sigmoid函数将标签转换为0到1之间的值表示该像素属于结节的概率。

对于分割多个类别的任务您可以使用softmax函数将标签转换为每个类别的概率分布。

如果您的数据集足够大且质量良好您可以不使用软标签来训练模型也就是训练阶段不使用sigmoid或softmax但是如果您的数据集较小或存在噪声数据使用软标签可能会提高模型的性能和分割结果的质量。

也就是训练阶段使用sigmoid或softmax。

在论文3D

就是多少不需要加背景类。

但是我也看到就只有一个类别但是做了加1操作的。

这点我再了解下。

如果你有什么心得欢迎评论区交流。

三、总结

UNet网络的结构无论是二维的还是三维的都是比较容易理解的这可能也是为什么那么受欢迎的原因之一吧。

如果你看过之前那篇关于2D

UNet的过程再看本篇应该就简单的很多。

觉得本篇更简单一些呢。

逐模块的分析了结构对后续的模型构建提供了思路构建完模型需要先预测试两种方式可选对模型的优势和劣势分析。