UNet 代码复现

一、UNet 算法简介

1.1 什么是 UNet 算法

UNet算法是一种用于图像分割的卷积神经网络（Convolutional Neural Network，简称CNN）架构。它由Olaf Ronneberger等人在2015年提出，主要用于解决医学图像分割的问题。

UNet算法的特点是采用了U型的网络结构，因此得名UNet。

该网络结构具有编码器（Encoder）和解码器（Decoder）两个部分。

编码器负责逐步提取输入图像的特征并降低空间分辨率。
解码器则通过上采样操作将特征图恢复到原始输入图像的尺寸，并逐步生成分割结果。

UNet算法的关键创新是在解码器中引入了跳跃连接（Skip Connections），即将编码器中的特征图与解码器中对应的特征图进行连接。这种跳跃连接可以帮助解码器更好地利用不同层次的特征信息，从而提高图像分割的准确性和细节保留能力。

UNet算法在医学图像分割领域表现出色，特别适用于小样本、不平衡数据和需要保留细节信息的任务。它已被广泛应用于肿瘤分割、器官分割、细胞分割等领域，并成为图像分割领域的重要算法之一。

1.2 UNet 的优缺点

UNet算法作为一种图像分割算法，具有以下优点和缺点：

优点：

强大的分割能力：UNet算法采用了U型的网络结构和跳跃连接机制，能够有效地捕获不同层次的特征信息，从而提高图像分割的准确性和细节保留能力。
少样本学习：相比其他深度学习方法，UNet算法对于小样本情况表现出色，可以在较少的标注数据上进行训练，并取得较好的分割效果。
可扩展性：UNet算法的网络结构简单明了，容易进行扩展和修改。可以根据具体任务的需求进行网络结构的调整，添加或删除网络层次。
缺点：

计算资源需求较高：由于UNet算法通常需要处理较大的图像输入和较深的网络结构，因此对计算资源的要求较高，包括显存和计算能力。
数据不平衡问题：如果训练数据中存在类别不平衡的情况，UNet算法可能会倾向于预测出现频率较高的类别，而忽略出现频率较低的类别。这需要在数据预处理或损失函数设计上进行相应的处理。
对于大尺寸图像的处理：由于UNet算法的网络结构和内存限制，对于大尺寸的图像，需要进行分块处理或采用其他策略来解决内存不足的问题。
综上所述，UNet算法具有强大的分割能力和适应小样本学习的优点，但同时也需要较高的计算资源，并且在数据不平衡和大尺寸图像处理方面可能存在一些挑战。

二、UNet网络架构

三、代码复现

import torch
import torch.nn as nn
def double_conv(in_channels,out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels,out_channels,kernel_size=3,stride=1,padding=1),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_channels,out_channels,kernel_size=3,stride=1,padding=1),
        nn.ReLU(inplace=True)
    )
class UNet(nn.Module):
    def __init__(self,in_channels,out_channels):
        super().__init__()
        self.conv_down1=double_conv(in_channels,64)
        self.conv_down2=double_conv(64,128)
        self.conv_down3=double_conv(128,256)
        self.conv_down4=double_conv(256,512)

        self.maxpool=nn.MaxPool2d(kernel_size=2,stride=2)    #相当下采样

        self.upsample=nn.Upsample(scale_factor=2,mode='bilinear',align_corners=True)

        self.conv_up1=double_conv(512+256,256)
        self.conv_up2=double_conv(256+128,128)
        self.conv_up3=double_conv(128+64,64)

        self.last_conv=nn.Conv2d(64,out_channels,kernel_size=1)

    def forward(self,x):
        conv1=self.conv_down1(x)        #采用V行结构3-5，2-6，1-7
        x1=self.maxpool(conv1)
        conv2=self.conv_down2(x1)
        x2=self.maxpool(conv2)
        conv3=self.conv_down3(x2)
        x3=self.maxpool(conv3)
        conv4=self.conv_down4(x3)   #编码部分完成

        x4=self.upsample(conv4)

        x5=torch.cat([x4,conv3],dim=1)
        up1=self.upsample(self.conv_up1(x5))

        x6=torch.cat([up1,conv2],dim=1)
        up2=self.upsample(self.conv_up2(x6))

        x7 = torch.cat([up2, conv1], dim=1)
        up2 = self.conv_up3(x7)

        x8=self.last_conv(up2)
        out=torch.sigmoid(x8)

        return out


unet = UNet(3,1)
output = unet(torch.randn(1,3,256,256))
print(output.shape)