UNet神经网络

0 引言

随着深度学习领域中各类算法的迅速发展&＃xff0c;卷积神经网络&＃xff08;CNN&＃xff09;被广泛应用在了分类任务上&＃xff0c;输出的结果是整个图像的类标签。在生物医学领域&＃xff0c;医生需要对病人的病灶区域进行病理分析&＃xff0c;这时需要一种更先进的网络模型&＃xff0c;即能通过少量的图片训练集&＃xff0c;就能实现对像素点类别的预测&＃xff0c;并且可以对像素点进行着色绘图&＃xff0c;形成更复杂、严谨的判断。于是U-Net网络被设计了出来。

1 U-Net概念及原理

U-Net网络结构最早由Ronneberger等人于2015年提出。该图像的核心思想是引入了跳跃连接&＃xff0c;使得图像分割的精度大大提升。
U-Net网络的主要结构包括了解码器、编码器、瓶颈层三个部分。

• 编码器&＃xff1a;包括了四个程序块。每个程序块都包括 $3\times 3$ 的卷积&＃xff08;使用Relu激活函数&＃xff09;&＃xff0c;步长为 $2$ 的 $2\times 2$ 的池化层&＃xff08;下采样&＃xff09;。每个程序块处理后&＃xff0c;特征图逐步减小。

• 解码器&＃xff1a; 与编码器部分对称&＃xff0c;也包括四个程序块&＃xff0c;每个程序块包括步长为 $2$ 的 $2\times 2$ 的上采样操作&＃xff0c;然后与编码部分进行特征映射级联&＃xff08;Concatenate&＃xff09;&＃xff0c;即拼接&＃xff0c;最后通过两个 $3\times 3$ 的卷积&＃xff08;Relu&＃xff09;。

• 瓶颈层&＃xff1a;包含两个 $3\times 3$ 的卷积层。

最后经过一个 $1\times 1$的卷积层得到最后的输出。

如图所示&＃xff0c;该网络模型形似字母“U”&＃xff0c;故称为U-Net。

整体过程&＃xff1a;
先对图片进行卷积和池化。比如说一开始输入的图片大小是 $224\times 224$&＃xff0c;进过四次池化后&＃xff0c;分别得到 $112\times 112$ , $56\times 56$ , $28\times 28$, $14\times 14$ 四个不同尺寸的特征图。然后对 $14\times 14$ 的特征图做上采样&＃xff0c;得到 $28\times 28$ 的特征图。将这个 $28\times 28$的特征图与之前池化得到的 $28\times 28$ 特征图进行通道上的拼接&＃xff08;concat)&＃xff0c;然后再对拼接之后的特征图做卷积和上采样&＃xff0c;得到 $56\times 56$ 的特征图&＃xff0c;然后再与之前的 $56\times 56$ 拼接&＃xff0c;卷积然后再上采样&＃xff0c;经过四次就就可以得到一个与原输入图像大小相同的图片了。

在本图片上的U-Net中&＃xff0c;它输入大小为 $572\times 572$, 而输出大小为 $388\times 388$, 那是因为它在卷积过程中没有加padding层所造成的。

2 代码

import torch
import torch.nn as nn
import torch.nn.functional as F# Double Convolution
class DoubleConv2d(nn.Module):def __init__(self, inputChannel, outputChannel):super(DoubleConv2d, self).__init__()self.conv &＃61; nn.Sequential(nn.Conv2d(inputChannel, outputChannel, kernel_size&＃61;3, padding&＃61;1),nn.BatchNorm2d(outputChannel),nn.ReLU(True),nn.Conv2d(outputChannel, outputChannel, kernel_size&＃61;3, padding&＃61;1),nn.BatchNorm2d(outputChannel),nn.ReLU(True))def forward(self, x):out &＃61; self.conv(x)return out# Down Sampling
class DownSampling(nn.Module):def __init__(self):super(DownSampling, self).__init__()self.down &＃61; nn.MaxPool2d(kernel_size&＃61;2)def forward(self, x):out &＃61; self.down(x)return out# Up Sampling
class UpSampling(nn.Module):# Use the deconvolutiondef __init__(self, inputChannel, outputChannel):super(UpSampling, self).__init__()self.up &＃61; nn.Sequential(nn.ConvTranspose2d(inputChannel, outputChannel, kernel_size&＃61;2, stride&＃61;2),nn.BatchNorm2d(outputChannel))def forward(self, x, y):x &＃61;self.up(x)diffY &＃61; y.size()[2] - x.size()[2]diffX &＃61; y.size()[3] - x.size()[3]x &＃61; F.pad(x, [diffX // 2, diffX - diffX // 2,diffY // 2, diffY - diffY // 2])out &＃61; torch.cat([y, x], dim&＃61;1)return outclass Unet(nn.Module):def __init__(self):super(Unet, self).__init__()self.layer1 &＃61; DoubleConv2d(1, 64)self.layer2 &＃61; DoubleConv2d(64, 128)self.layer3 &＃61; DoubleConv2d(128, 256)self.layer4 &＃61; DoubleConv2d(256, 512)self.layer5 &＃61; DoubleConv2d(512, 1024)self.layer6 &＃61; DoubleConv2d(1024, 512)self.layer7 &＃61; DoubleConv2d(512, 256)self.layer8 &＃61; DoubleConv2d(256, 128)self.layer9 &＃61; DoubleConv2d(128, 64)self.layer10 &＃61; nn.Conv2d(64, 2, kernel_size&＃61;3, padding&＃61;1) # The last output layerself.down &＃61; DownSampling()self.up1 &＃61; UpSampling(1024, 512)self.up2 &＃61; UpSampling(512, 256)self.up3 &＃61; UpSampling(256, 128)self.up4 &＃61; UpSampling(128, 64)def forward(self, x):conv1 &＃61; self.layer1(x)down1 &＃61; self.down(conv1)conv2 &＃61; self.layer2(down1)down2 &＃61; self.down(conv2)conv3 &＃61; self.layer3(down2)down3 &＃61; self.down(conv3)conv4 &＃61; self.layer4(down3)down4 &＃61; self.down(conv4)conv5 &＃61; self.layer5(down4)up1 &＃61; self.up1(conv5, conv4)conv6 &＃61; self.layer6(up1)up2 &＃61; self.up2(conv6, conv3)conv7 &＃61; self.layer7(up2)up3 &＃61; self.up3(conv7, conv2)conv8 &＃61; self.layer8(up3)up4 &＃61; self.up4(conv8, conv1)conv9 &＃61; self.layer9(up4)out &＃61; self.layer10(conv9)return out# Test partmynet &＃61; Unet()
# device &＃61; torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# mynet.to(device)
input &＃61; torch.rand(3, 1, 572, 572)
# output &＃61; mynet(input.to(device))
output &＃61; mynet(input)
print(output.shape) # (3,2,572,572)


https://www.jianshu.com/p/a73f74992b1a
https://arxiv.org/pdf/1505.04597v1.pdf
https://blog.csdn.net/qq_34107425/article/details/110184747
https://blog.csdn.net/weixin_41857483/article/details/120768804