图像分割经典架构Unet

  • 图像分割必备的基础知识

图像分割(image Segmentation)是深度学习在图像领域的重要应用之一,如果说识别任务是针对一张图像进行的学习、检测任务是针对图像中不同的对象进行的学习,那分割任务就是针对单一像素进行的学习。分割任务是像素级别的有监督任务,在图像分割时,我们需要对图像中的每一个像素进行分类,因此我们可以找出图像中每个对象的“精确边界”。以下面的图像为例,我们可以找出[“猫”,“狗”]两个标签类别所对应的具体对象,并且找出这些对象的精确边界。当原始图像越复杂,分割任务中需要输出的标签类别也就越多,对下面的图像,除了也可以[“猫”,“狗”]分割之外,还可以对[“草坪”,“项圈”]等标签进行分割。

从分割的类别来看,我们可以执行将不同性质的物体分开的语义分割(semantic segmentation),也可以执行将每个对象都分割开来的实例分割(instance segmentation),还可以执行使用多边形或颜色进行分割的分割方法。同时,根据分割的“细致程度”,还可以分为粗粒度分割(Coarse Segmentation)与细粒度分割(Fine Segmentation),只要拥有对应的标签,我们就可以将图像分割到非常非常精细的程度:
image.png

  • 分割任务中的标签

不难发现,分割任务具体可以做到的分割程度是由训练图像中的标签决定的,而分割图像中的标签具体是什么样呢?在分割任务中,训练数据是原始图像,原始图像中存在的所有对象都可以被标记为某一类别,而一张图像所对应的标签一般是与原始图像相同尺寸的标签矩阵,该矩阵被上色之后被称为“遮罩矩阵”(mask)。通常来说,遮罩矩阵中的颜色数量等同于这张标签上的标签类别数量。

以下面最为简单的10x10尺寸图像为例,假设一张原始图像(一个样本)的结构为(1,10,10),那它所对应的标签的结构一般也为(1,10,10),假设总共有n_samples个样本,则数据集的标签格式为(n_samples,1,10,10)。

  • 分割架构的输出

当输入标签为图像/矩阵时,你能推断出分割任务中网络对应的输出是什么吗?对于任意用于图像识别的数据集,如果数据有num_classes个标签类别,则神经网络会对每个样本输出num_classes个对应的概率,此时作为输出层的线性层则会有num_classes个神经元,作为输出层的卷积层则会输出num_classes个1x1的特征图。

相对的,在分割任务中,如果数据包含num_classes个标签类别,神经网络则需要对每个像素值输出num_classes个对应的概率。因此,分割任务中一个样本所对应的输出是该样本上所有像素值、在所有标签类别下的概率。具体来看,以10x10的花朵图为例,该数据集的标签有6大类别,因此一个样本所对应的输出就有6张概率图,一张概率图对应着一个类别,反馈着样本上所有的像素是0类、1类、2类……num_classes类的概率。当输出概率图后,一个像素在所有概率图上最大的概率所对应的类别,就是该像素的预测类别

需要注意的是,以上图像中有不严谨之处。对一个像素而言,该像素为任意类别的概率加和之后应该为1(例如,对图像上最右下角的像素点而言,所有标位灰色的概率值加和之后应该为1)。在绘制图像时,受限于绘图公式设计,图像中并没有实现“单一像素的所有类别概率加和必须为1”这一条件,但在实际使用数据、输出结果是,这一条件是一定会被满足的。

  • 分割任务中的损失函数

同时,还有几个值得注意的问题:

  1. 一张图上的存在的标签类别可能少于整个数据集中的标签类别。标签类别数量会覆盖整个数据集上的标签类别,因此一张图像上不一定包含了所有的标签类别,但每个样本都必须输出所有类别的概率图。假设一张图像上没有类别A,那我们会期待这张图像所对应的类别A的概率图上的值会全部为0,即没有任何像素属于该类别。

  2. 由于需要输出概率图,所以分割网络的输出层一般都是卷积层。以此为基础还诞生了整个网络中只有卷积或卷积相关计算的网络全卷积网络FCN(fully convolutional network)。

  3. 虽然标签和输出都转化为了图像,但我们常用的交叉熵损失等损失函数依然可以使用,只不过此时我们公式中的预测值是多个矩阵,真实标签也是多个矩阵。具体来看:

二分类交叉熵损失——


L = − ( y log ⁡ p ( x ) + ( 1 − y ) log ⁡ ( 1 − p ( x ) ) ) L = -\left( y\log p(x) + (1 - y)\log(1 - p(x)) \right) L=(ylogp(x)+(1y)log(1p(x)))

此时y就是标签矩阵,而 p ( x ) p(x) p(x)就是对应的概率图,由于都是相同尺寸的图像,所以在进行计算时并无问题。


多分类交叉熵损失,总共有K个类别——


L = − ∑ k = 1 K y k ∗ log ⁡ ( P k ( x ) ) L = -\sum_{k=1}^Ky^*_k\log(P^k(x)) L=k=1Kyklog(Pk(x))

在标签为序列的识别任务中, P k ( x ) P^k(x) Pk(x)是一个样本的类别为k的概率,并且 p k ( x ) = S o f t m a x ( 网络输出 ) p^k(x) = Softmax(网络输出) pk(x)=Softmax(网络输出) y ∗ y^* y是由真实标签做独热编码后的向量。例如,在3分类情况下,真实标签 y i y_i yi为2时, y ∗ y^* y为[ y 1 ∗ y^*_{1} y1, y 2 ∗ y^*_{2} y2, y 3 ∗ y^*_{3} y3],取值分别为:

y 1 ∗ y^*_{1} y1 y 2 ∗ y^*_{2} y2 y 3 ∗ y^*_{3} y3
0 0 0 1 1 1 0 0 0

而在标签为概率图的分割任务中, P k ( x ) P^k(x) Pk(x)是所有像素的类别为k的概率(即概率图),并且 y ∗ y^* y是将标签矩阵转化为独热形态后的独热矩阵。以之前的花朵图像为例,在num_classes个标签分类的情况下,我们会分割出num_classes个独热矩阵,如果一个像素的真实标签为该类别,则该像素在这张独热矩阵上的标签为1,否则标签为0。

  1. 标签与输出概率图的尺寸必须一致,但有时候标签和输出概率图的尺寸可以略小于原始图像

图像分割是对每个样本上的每个像素进行分类,因此最为严谨的情况下每个像素都应该有对应的标签,但对像素级别进行分割的根本目的是为了描绘出物体的轮廓。然而,当原始图像的尺寸很大、像素很高时,我们并不需要非常高清的图像才能够展示出物体的轮廓;同时,在许多时候我们需要的只是“大致轮廓”,而不需要非常精确的边缘,因此标签图像可以略小于原始图像。这样的标签可以反馈出每一类对象在原始图像中的大致轮廓,但不能精确地反馈出原始图像的每一个像素点的类别。




标签需要被放置到损失函数中使用,为了能够和标签图像进行像素值一一对应的计算,分割架构输出的概率图尺寸必须与标签尺寸一致。因此当标签小于原始图像时,概率图也会小于原始图像。并且,标签和概率图越小,其投射到原始图像上的轮廓就越不准确,依据我们的使用场景,我们可以调整网络架构的输出:

较为粗糙的分割
(概率图/标签可以大幅度小于原始图像)
更加精准地分割
(概率图/标签等于原始图像,或略小于原始图像)
drawingdrawing

现在你已经了解了关于图像分割架构的一些基本特点了。有了这些基础知识,我们在解析各种分割架构的时候就不容易进入迷雾之中,下一节我们将展开聊聊经典分割架构Unet的结构。

2 经典架构Unet

Unet是一个博采众长的架构,它于2015年被德国弗莱堡大学的研究团队提出,并在原始论文当中被用于生物医学影像的分割。但事实上,它既可以完成分类任务,也可以完成分割任务,还可以被当做无监督算法使用,这是因为它汲取了大量其他网络的精华结构、并且有机地将这些架构融合在了一起。具体地来看,Unet的架构图如下。我们可以从图例、架构上的数字取得不少关键信息。

在这里插入图片描述

整个架构图呈现对称的“U字型”,因此该网络被称为Unet。查看架构细节,不难发现Unet其实拥有多重身份和多重标签:

  1. Unet是一个全卷积网络(FCN),它没有使用除了卷积、池化和转置卷积之外的任何层,架构中的数据始终都是四维的图像。

  2. Unet是一个Encoder-Decoder,它的结构与深度卷积自动编码器高度相似:输入大图像、使用卷积层和池化层组成的Encoder将图像压缩,形成数据隐式表示Laten Representation,接着又使用由卷积层和转置卷积层组成的Decoder将图像放大,最终输出大图像。这一结构左右对称、两边大中间小的结构是自动编码器的典型结构。

灵魂拷问:Encoder-Decoder(也就是Autoencoder,自动编码器)是无监督算法,Unet是有监督算法,难道因为结构一致就可以说Unet是一个自动编码器吗?


事实上,自动编码器算法不是有监督,却胜似有监督。以降噪自动编码器为例,首先我们具有一组原始数据A,为了加大自动编码器的训练难度、防止自动编码器直接把输入信息原封不动搬到输出层,我们需要在原始数据A基础上加上噪音,构成带噪音的数据A’。
在这里插入图片描述


在训练过程中,我们对编码器输入A’,让架构从A’生成数据B,并且我们对架构的要求是:数据B要尽量与无噪音的数据A相似,这样架构就拥有了“降噪”的能力。这一过程看似平常,但其实已经等同于一个有监督的过程了:

不难发现,自动编码器中损失函数衡量A与B之间的差异,而有监督算法当中损失函数衡量预测标签yhat与真实标签y之间的差异,无论我们面对的任务是给图像上色、还是给图像补全、还是其他任务,只要认为A是真实标签、A’是特征矩阵,那自动编码器就可以被当成一个有监督算法使用。如果想要为黑白图像上色,你应该准备彩色图像作为原始数据A,黑白图像作为输入数据A’,如果你想要将冬天的照片变成夏天,那你应该准备夏天的照片作为原始数据A,冬天的照片作为原始数据A’。至此,专用于“数据表示”的无监督算法就成为了可以“依葫芦画瓢”的有监督算法了,真是妙哉。因此,认为Unet是一个自动编码器没有任何问题。

  1. Unet是一个分割网络,因此它所使用的标签是遮罩,要输出的是概率图,使用的损失函数是二分类交叉熵损失。从最终输出的结果来看,Unet被创造的原始论文中所使用的分割图像应该是二分类的分割图像。同时,网络输入图像尺寸与输出图像尺寸差异较大,因此可以判断原始论文中所使用的生物医学影像并不需要太高的轮廓精度,检查原始论文,会进一步印证我们的观点:

  2. Unet使用了跳跃链接,直接将Encoder中尚未压缩完毕的图像传向Decoder,这样可以将更多原始信息直接传向输出方向,帮助生成与原始数据更相似的图像。经过跳跃链接穿到Decoder的图像需要与经过编码后的数据合并,再进行卷积,这样Decoder在输出最终的概率图时可以参考的信息就变得更加丰富了。甚至,我们说Unet就是使用了跳跃链接的、有监督版本的自动编码器。

接下来,让我们来复现一下Unet架构。

在这三天的课程当中,我们不断强调了一个事实:在深度学习的后期,许多时候架构中的具体层已经不是最为关键的内容了,相对的,数据如何在架构中流通才是我们真正关心的内容。观察架构图,我们可以将Unet总结为如下结构:

输入 → 【双卷积 + 池化】x4 → 【双卷积】 → 【转置卷积 + 双卷积】x4 → 【1x1卷积】→ 输出

    [------Encoder------]  [bottleneck] [-------Decoder-------]   [Output]

因此毫无疑问的,我们可以先定义一个双卷积的结构。在原始论文中规定,无论是在Encoder还是Decoder当中,每个卷积层后面都跟ReLU激活函数。同时,作为卷积架构的惯例,我们在每个卷积层后面、ReLU激活函数之前使用Batch Normalization。需要注意的是,Unet架构中使用的不是令特征图尺寸保持不变的卷积层,而是每经过一个卷积就会将特征图长宽缩小2的卷积层。

class DoubleConv2d(nn.Module):
    def __init__(self,in_channels, out_channels):
        super().__init__()
        self.conv = nn.Sequential(nn.Conv2d(in_channels, out_channels,3,1,0,bias=False)
                                 ,nn.BatchNorm2d(out_channels)
                                 ,nn.ReLU(inplace=True)
                                 ,nn.Conv2d(out_channels, out_channels,3,1,0,bias=False)
                                 ,nn.BatchNorm2d(out_channels)
                                 ,nn.ReLU(inplace=True)
                                 )
    def forward(self,x):
        return self.conv(x)

接下来我们来考虑如何实现该架构中的数据流。Unet中的数据并不依照从左到右进行线性流动:因为有跳跃链接的存在,Encoder中的每个双卷积的输出结果都必须被直接传输到Decoder中每个双卷积的输入层。因此每经过一次双卷积结构我们就需要保存中间结果,因此我们可以按如下方式梳理数据流:

#Encoder
#输入x

#x = 双卷积(x)
#保存x
#x = 池化层(x)

#x = 双卷积(x)
#保存x
#x = 池化层(x)

#x = 双卷积(x)
#保存x
#x = 池化层(x)

#x = 双卷积(x)
#保存x
#x = 池化层(x)
#完全等同于

l = []

for i in range(4):
    #x = 双卷积(x)
    #l.append(x)
    #x = 池化层(x)

#在这个循环中,4个卷积层的输入特征图数量和输出特征图数量不一致,因此实际上是4个不同的双卷积结构
#但相对的,池化层是完全一致的Maxpool(2),且池化层没有需要迭代的权重,因此可以被重复使用
#所以,我们可以将4个卷积层定义成一个序列,并单独定义池化层

for i in range(4):
    #x = 双卷积[i](x)
    #l.append(x)
    #x = 池化层(x)
    
#l = [x1,x2,x3,x4]

相似地,由于跳跃链接的存在,我们需要将数据与跳跃链接传过来的数据合并后,才能输入到Decoder中的每个双卷积层。并且需要注意的是,在Encoder中池化层是位于双卷积的后面,但在Decoder中转置卷积层是位于双卷积层的前面:

输入 → 【双卷积 + 池化】x4 → 【双卷积】 → 【转置卷积 + 双卷积】x4 → 【1x1卷积】→ 输出

    [------Encoder------]  [bottleneck] [-------Decoder-------]   [Output]

#Decoder
#l
#x此时是瓶颈结构的输出

#x = 转置卷积(x)
#x = 合并(x,x4)
#x = 双卷积(x)

#x = 转置卷积(x)
#x = 合并(x,x3)
#x = 双卷积(x)

#x = 转置卷积(x)
#x = 合并(x,x2)
#x = 双卷积(x)

#x = 转置卷积(x)
#x = 合并(x,x1)
#x = 双卷积(x)
#按循环写:

for i in range(4):
    #x = 转置卷积[i](x)
    #x = 合并(x,l[-i])
    #x = 双卷积[i](x)

#其中4个转置卷积的参数不同,4个双卷积的参数也不同,因此都需要构建成序列
  • 架构复现
class Unet(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.encoder_conv = nn.Sequential(DoubleConv2d(1,64)
                                          ,DoubleConv2d(64,128)
                                          ,DoubleConv2d(128,256)
                                          ,DoubleConv2d(256,512)
                                         )
        
        self.encoder_down = nn.MaxPool2d(2)
        
        self.decoder_up = nn.Sequential(nn.ConvTranspose2d(1024,512,4,2,1)
                                       ,nn.ConvTranspose2d(512,256,4,2,1)
                                       ,nn.ConvTranspose2d(256,128,4,2,1)
                                       ,nn.ConvTranspose2d(128,64,4,2,1)
                                       )
        
        self.decoder_conv = nn.Sequential(DoubleConv2d(1024,512)
                                          ,DoubleConv2d(512,256)
                                          ,DoubleConv2d(256,128)
                                          ,DoubleConv2d(128,64)
                                         )
        
        self.bottleneck = DoubleConv2d(512,1024)
        
        self.output = nn.Conv2d(64,2,3,1,1)
    
    def forward(self,x):
        
        #encoder:保存每一个DoubleConv的结果为跳跃链接做准备,同时输出codes
        skip_connection = []
        
        for idx in range(4):
            x = self.encoder_conv[idx](x)
            skip_connection.append(x)
            x = self.encoder_down(x)
        
        x = self.bottleneck(x)
        
        #调换顺序
        skip_connection = skip_connection[::-1]
        
        #decoder:codes每经过一个转置卷积,就需要与跳跃链接中的值合并
        #合并后的值进入DoubleConv
        
        for idx in range(4):
            x = self.decoder_up[idx](x)
            #转换尺寸
            skip_connection[idx] = transforms.functional.resize(skip_connection[idx],size=x.shape[-2:])
            x = torch.cat((skip_connection[idx],x),dim=1)
            x = self.decoder_conv[idx](x)
        
        x = self.output(x)
        return x
  • 架构验证
net = Unet()
summary(net,input_size=(10,1,572,572),device="cpu")

==========================================================================================
Layer (type:depth-idx) Output Shape Param #

Unet – –
├─Sequential: 1 – –
│ └─DoubleConv2d: 2-1 [10, 64, 568, 568] –
│ │ └─Sequential: 3-1 [10, 64, 568, 568] 37,696
├─MaxPool2d: 1-1 [10, 64, 284, 284] –
├─Sequential: 1 – –
│ └─DoubleConv2d: 2-2 [10, 128, 280, 280] –
│ │ └─Sequential: 3-2 [10, 128, 280, 280] 221,696
├─MaxPool2d: 1-2 [10, 128, 140, 140] –
├─Sequential: 1 – –
│ └─DoubleConv2d: 2-3 [10, 256, 136, 136] –
│ │ └─Sequential: 3-3 [10, 256, 136, 136] 885,760
├─MaxPool2d: 1-3 [10, 256, 68, 68] –
├─Sequential: 1 – –
│ └─DoubleConv2d: 2-4 [10, 512, 64, 64] –
│ │ └─Sequential: 3-4 [10, 512, 64, 64] 3,540,992
├─MaxPool2d: 1-4 [10, 512, 32, 32] –
├─DoubleConv2d: 1-5 [10, 1024, 28, 28] –
│ └─Sequential: 2-5 [10, 1024, 28, 28] –
│ │ └─Conv2d: 3-5 [10, 1024, 30, 30] 4,718,592
│ │ └─BatchNorm2d: 3-6 [10, 1024, 30, 30] 2,048
│ │ └─ReLU: 3-7 [10, 1024, 30, 30] –
│ │ └─Conv2d: 3-8 [10, 1024, 28, 28] 9,437,184
│ │ └─BatchNorm2d: 3-9 [10, 1024, 28, 28] 2,048
│ │ └─ReLU: 3-10 [10, 1024, 28, 28] –
├─Sequential: 1 – –
│ └─ConvTranspose2d: 2-6 [10, 512, 56, 56] 8,389,120
├─Sequential: 1 – –
│ └─DoubleConv2d: 2-7 [10, 512, 52, 52] –
│ │ └─Sequential: 3-11 [10, 512, 52, 52] 7,079,936
├─Sequential: 1 – –
│ └─ConvTranspose2d: 2-8 [10, 256, 104, 104] 2,097,408
├─Sequential: 1 – –
│ └─DoubleConv2d: 2-9 [10, 256, 100, 100] –
│ │ └─Sequential: 3-12 [10, 256, 100, 100] 1,770,496
├─Sequential: 1 – –
│ └─ConvTranspose2d: 2-10 [10, 128, 200, 200] 524,416
├─Sequential: 1 – –
│ └─DoubleConv2d: 2-11 [10, 128, 196, 196] –
│ │ └─Sequential: 3-13 [10, 128, 196, 196] 442,880
├─Sequential: 1 – –
│ └─ConvTranspose2d: 2-12 [10, 64, 392, 392] 131,136
├─Sequential: 1 – –
│ └─DoubleConv2d: 2-13 [10, 64, 388, 388] –
│ │ └─Sequential: 3-14 [10, 64, 388, 388] 110,848
├─Conv2d: 1-6 [10, 2, 388, 388] 1,154

Total params: 39,393,410
Trainable params: 39,393,410
Non-trainable params: 0
Total mult-adds (T): 2.35

Input size (MB): 13.09
Forward/backward pass size (MB): 19926.14
Params size (MB): 157.57
Estimated Total Size (MB): 20096.80

  • 37
    点赞
  • 34
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值