U-Net模型

张炸炸CIUM

已于 2022-08-04 16:39:10 修改

阅读量4k

点赞数 8

文章标签：深度学习 pytorch python

于 2022-08-04 10:00:29 首次发布

原文链接：https://blog.csdn.net/weixin_38346042/article/details/125258332

版权

1.定义DoubleConv模块

在网络结构中，卷积基本上都是成对使用的，所以就定义了一个DoubleConv类

网络结构图：

在这里插入图片描述

U-Net原理图

1.1DoubleConv 模块代码实现

class DoubleConv(nn.Sequential):
	def __init__(self,in_channels,out_channels,mid_channels=None):
		if mid_channels is None:
			mid_channels=out_channels
			super(DoubleConv,self).__init__(
	              nn.Conv2d(in_channels,mid_channels,kernel_size=3,padding=1,bias=False),
	              nn.BatchNorm2d(mid_channels),
	              nn.ReLU(inplace=True),
	              nn.Conv2d(mid_channels,out_channels,kernel_size=3,padding=1,bias=False),
	              nn.BatchNorm2d(out_channels),
	              nn.ReLU(inplace = True)
			)

in_channels : 指的是输入特征层的channels
out_channels: 指的是经过DoubleConv层后输出特征层的channels
mid_channels: 指第一个卷积层输出的channels
通过父类的构造函数搭建DoubleConv,其中Conv2d的Kernel=3,padding=1,设置padding=1经过卷积后不会改变特征层的大小，这也是现在主流的实现方式
由于我们会使用到BN，因此将bias设置为Flase

1.2定义Down 模块

Down模块包括下采样(MaxPool) + 2个Conv2d,因为在网络左侧encoder部分基本上都是通过下采样(MaxPool) + 2个Conv2d搭建的。

Down 模块代码实现：

class Down(nn.Sequential):
	def __init__(self,in_channels,out_channels):
		super(Down,self).__init__(
			nn.MaxPool2d(2,stride=2),
			DoubleConv(in_channels,out_channels)
	)

1.3 定义UP 模块

UP模块包含：上采样(w,h翻倍)+concat拼接+2个Conv2d, 在UNet网络的右半部分decoder也就是解码器部分，基本上都是有各个上采样+concat拼接+2个Conv2d模块组成，因此这里定义了一个UP模块

UP 模块代码实现：

class Up(nn.Module):
	def __init__(self,in_channels,out_channels,bilinear=True):
		super(Up,self).__init__()
		if bilinear:
			self.up=nn.Upsample(scale_factor=2,mode='bilinear',align_corners=True)
			self.conv=DoubleConv(in_channels,out_channels,in_channels // 2)
		else:
			self.up = nn.ConvTranspose2d(in_channels,in_channels //2,kernel_size=2,stride=2)
			self.conv=DoubleConv(in_channels,out_channels)
	
	def forward(self,x1,x2):
		x1=self.up(x1)
		#[N,C,H,W]
		diff_y=x2.size()[2] - x1.size()[2]
		diff_x=x2.size()[3] - x1.size()[3]
		# padding_left,padding_right,padding_top,padding_bottom
		x1=F.pad(x1,[diff_x //2,diff_x - diff_x //2,
		             diff_y //2,diff_y - diff_y//2])
		x=torch.cat([x2,x1],dim=1)
		x=self.conv(x)
		return x

输入参数bilinear=True,表示默认情况下是会使用双线性差值
in_channels :指的是concat拼接后的channels,或者是up这个模块第一个卷积输入的channels
如果bilinear=False, 则使用论文中提到的转置卷积进行上采样，长宽翻倍，channels会减半，此时DoubleConv(in_channels,out_channels) 中mid_channels和out_channels是一样的。因为在原论文中这两个卷积的channels是一样的。参考：https://blog.csdn.net/weixin_38346042/article/details/125246551?spm=1001.2014.3001.5502

如果采用双线性插值进行上采样的话（经过双线性插值自身不会改变channels），上采样后面跟着的两个卷积的channels是不一样的，比如通过第一个卷积后channels会减半，通过第二个卷积后，channels又会减半。这样做的目的是为了经过双线性插值后得到的channels和我们要concat拼接的特征层的channels保持一致。

在这里插入图片描述

forward正向传播过程，x1指的是需要上采样的特征层，x2指的是要concat拼接的特征层。

在这里插入图片描述

流程是：首先我们需要对x1进行上采样x1=self.up(x1)得到采样之后的特征层，按理说可以直接通过concat拼接，将两个特征层在我们的channels维度进行拼接就可以了。然后再通过两个卷积层得到我们的特征输出。

def forward(self,x1,x2):
		x1=self.up(x1)
		#[N,C,H,W]
		diff_y=x2.size()[2] - x1.size()[2] # H
		diff_x=x2.size()[3] - x1.size()[3] # W
		# padding_left,padding_right,padding_top,padding_bottom
		x1=F.pad(x1,[diff_x //2,diff_x - diff_x //2,
		             diff_y //2,diff_y - diff_y//2])
		x=torch.cat([x2,x1],dim=1)
		x=self.conv(x)
		return x

但是这里作者有多做了一步，对我们上采样之后的x1进行了padding,目的是为了防止我们输入的图片如果不是16的整数倍的话，通过下采样得到的x1与我们要拼接的x2的高度和宽度是不一致的。

在我们unet的搭建过程下采样了4次也就是16倍，如果高和宽不是16的整数倍，那么在下采样的过程中，可能面临着向下取整的情况。比如在某个位置特征层大小是7x7的，经过下采样它的高和宽就变为3x3了，在通过上采样是6x6，很明显与7x7的特征层就无法拼接了。为了防止出现该问题，对x1进行了padding,padding之后就能保住x1和x2的高和宽是一致的了，这样就可以进行concat拼接了。

定义OutConv 模块

OutConv对应的是最后一个1x1的卷积层，通过这个1x1的卷积层之后就得到我们最终的输出了。这个1x1的卷积它是没有BN和ReLU激活函数的。

class OutConv(nn.Sequential):
	def __init__(self,in_channels,num_classes):
		super(OutConv,self).__init__(
			nn.Conv2d(in_channels,num_classes,kernel_size=1)
  )

UNet 网络整体搭建过程

class UNet(nn.Module):
    def __init__(self,
                 in_channels: int = 1,
                 num_classes: int = 2,
                 bilinear: bool = True,
                 base_c: int = 64):
        super(UNet, self).__init__()
        self.in_channels = in_channels
        self.num_classes = num_classes
        self.bilinear = bilinear

        self.in_conv = DoubleConv(in_channels, base_c)
        self.down1 = Down(base_c, base_c * 2)
        self.down2 = Down(base_c * 2, base_c * 4)
        self.down3 = Down(base_c * 4, base_c * 8)
        factor = 2 if bilinear else 1
        self.down4 = Down(base_c * 8, base_c * 16 // factor)
        self.up1 = Up(base_c * 16, base_c * 8 // factor, bilinear)
        self.up2 = Up(base_c * 8, base_c * 4 // factor, bilinear)
        self.up3 = Up(base_c * 4, base_c * 2 // factor, bilinear)
        self.up4 = Up(base_c * 2, base_c, bilinear)
        self.out_conv = OutConv(base_c, num_classes)

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        x1 = self.in_conv(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.out_conv(x)

        return {"out": logits}

in_channels,输入的图片如果是彩色的,in_channels=3,如果使用的是黑白的`in_channels=1
bilinear 默认是True,根据测试无论采用双线性插值还是采用转置卷积计算，他们的结果其实是差不多的。那么你采用双线性差值其实会更高效点。
base_c就是网络中第一个卷积层输出的channels,在unet网络中各层的channels都是翻倍的，比如64,128,256,512,所以就定义了一个base channel，默认为64，也可以根据自己的想法去调整channels的大小。我这边训练的时候将base_c设置为32，发现得到的结果也没啥区别，但设置为32网络的参数会降低、训练速度会得到提升。
然后经过down1,down2,dow3,down4,4个下采样过程，注意到down1~down3通过两层卷积后channels都翻倍了，但down4经过两层卷积后channels没有翻倍。因为我们要保证经过双线性差值上采样得到的channels要与拼接的特征的channels保持一致，所以这里就没有进行翻倍。因此对输出的channels除了一个factor,如果采用双线性差值factor=2，如果采用转置卷积的话（本身会对w,h翻倍，channels减半，而双线性插值不会影响channels），因此factor=1。

 factor = 2 if bilinear else 1
 self.down4 = Down(base_c * 8, base_c * 16 // factor)

完整的UNet搭建代码：

from typing import Dict
import torch
import torch.nn as nn
import torch.nn.functional as F


class DoubleConv(nn.Sequential):
    def __init__(self, in_channels, out_channels, mid_channels=None):
        if mid_channels is None:
            mid_channels = out_channels
        super(DoubleConv, self).__init__(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )


class Down(nn.Sequential):
    def __init__(self, in_channels, out_channels):
        super(Down, self).__init__(
            nn.MaxPool2d(2, stride=2),
            DoubleConv(in_channels, out_channels)
        )


class Up(nn.Module):
    def __init__(self, in_channels, out_channels, bilinear=True):
        super(Up, self).__init__()
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
        else:
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
            self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
        x1 = self.up(x1)
        # [N, C, H, W]
        diff_y = x2.size()[2] - x1.size()[2]
        diff_x = x2.size()[3] - x1.size()[3]

        # padding_left, padding_right, padding_top, padding_bottom
        x1 = F.pad(x1, [diff_x // 2, diff_x - diff_x // 2,
                        diff_y // 2, diff_y - diff_y // 2])

        x = torch.cat([x2, x1], dim=1)
        x = self.conv(x)
        return x


class OutConv(nn.Sequential):
    def __init__(self, in_channels, num_classes):
        super(OutConv, self).__init__(
            nn.Conv2d(in_channels, num_classes, kernel_size=1)
        )


class UNet(nn.Module):
    def __init__(self,
                 in_channels: int = 1,
                 num_classes: int = 2,
                 bilinear: bool = True,
                 base_c: int = 64):
        super(UNet, self).__init__()
        self.in_channels = in_channels
        self.num_classes = num_classes
        self.bilinear = bilinear

        self.in_conv = DoubleConv(in_channels, base_c)
        self.down1 = Down(base_c, base_c * 2)
        self.down2 = Down(base_c * 2, base_c * 4)
        self.down3 = Down(base_c * 4, base_c * 8)
        factor = 2 if bilinear else 1
        self.down4 = Down(base_c * 8, base_c * 16 // factor)
        self.up1 = Up(base_c * 16, base_c * 8 // factor, bilinear)
        self.up2 = Up(base_c * 8, base_c * 4 // factor, bilinear)
        self.up3 = Up(base_c * 4, base_c * 2 // factor, bilinear)
        self.up4 = Up(base_c * 2, base_c, bilinear)
        self.out_conv = OutConv(base_c, num_classes)

    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        x1 = self.in_conv(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.out_conv(x)

        return {"out": logits}

张炸炸CIUM

关注

8
点赞
踩
45

收藏

觉得还不错? 一键收藏
0
评论
U-Net模型

在网络结构中，卷积基本上都是成对使用的，所以就定义了一个DoubleConv类网络结构图：1.1DoubleConv 模块代码实现)in_channels : 指的是输入特征层的channelsout_channels: 指的是经过DoubleConv层后输出特征层的channelsmid_channels: 指第一个卷积层输出的channels。......
复制链接

扫一扫