U-Net++原理与实现(含Pytorch和TensorFlow源码)

机器学习与优化算法

已于 2024-08-04 14:09:15 修改

阅读量1.6k

点赞数 24

文章标签： pytorch tensorflow 人工智能深度学习图像处理

于 2024-08-02 06:53:48 首次发布

本文链接：https://blog.csdn.net/m0_59257547/article/details/140862066

版权

U-Net++原理与实现

引言

在图像处理和计算机视觉领域，图像分割是一个至关重要的任务。分割技术被广泛应用于医学图像分析、自动驾驶、卫星图像处理等诸多领域。U-Net 及其改进版本 U-Net++ 是当前流行的图像分割神经网络结构，因其高效性和精确性而备受关注。本文旨在介绍 U-Net 和 U-Net++ 的基本原理，详细对比这两种网络结构，并探讨 U-Net++ 在实际应用中的优势。

1. U-Net简介

U-Net 是一种用于生物医学图像分割的卷积神经网络，由 Olaf Ronneberger 等人在 2015 年提出。其结构主要由编码器、解码器和跳跃连接组成。
在这里插入图片描述

1.1 编码器（Encoder）

编码器通过一系列卷积层和池化层逐步提取图像的高层次特征，同时减小特征图的空间尺寸。每个卷积层包含两个3x3卷积操作，接着是一个2x2最大池化操作。

$\text{MaxPool}(\sigma(W * X + b))$

其中， $X$ 是输入特征图， $W$ 和 $b$ 分别是卷积核权重和偏置， $\sigma$ 是激活函数，通常为 ReLU。

1.2 解码器（Decoder）

解码器通过上采样操作逐步恢复特征图的空间尺寸，并与对应编码器层的特征图进行融合。每个上采样层包含一个2x2反卷积操作，随后接两个3x3卷积操作。

$\sigma(W * \text{UpSample}(X) + b)$

1.3 跳跃连接（Skip Connections）

跳跃连接将编码器每一层的特征图直接传递给解码器对应层，帮助网络更好地捕捉细节信息和上下文特征。

$Y_{\text{decoder}} = \text{Concat}(Y_{\text{encoder}}, Y_{\text{decoder}})$

2. U-Net++详解

U-Net++ 由 Zhou 等人在 2018 年提出，是对经典 U-Net 的改进，主要在增强特征传递和多尺度特征融合方面进行了优化。
在这里插入图片描述
图：(a) U-Net++ 由一个编码器和解码器组成，它们通过一系列嵌套的密集卷积块相连。U-Net++ 的核心思想是在融合之前缩小编码器和解码器之间的特征图的语义差距。例如，通过使用具有三个卷积层的密集卷积块来缩小 (X0,0, X1,3) 之间的语义差距。在图形摘要中，黑色表示原始的 U-Net，绿色和蓝色显示跳过路径上的密集卷积块，红色表示深度监督。红色、绿色和蓝色部分区分了 U-Net++ 与 U-Net。(b) U-Net++ 中第一个跳过路径的详细分析。© 如果采用深度监督训练，则可以在推理时对 U-Net++ 进行剪枝。

2.1 密集跳跃连接

U-Net++ 引入了密集的跳跃连接，在每一级的编码器和解码器之间，以及每个子 U-Net 结构内部进行连接，增强了特征的传递和利用效率。

$Y_{i,j} = \sigma(W_{i,j} * [Y_{i-1,j}, Y_{i,j-1}] + b_{i,j})$

其中， $Y_{i,j}$ 表示第 $i$ 层第 $j$ 个子网络的输出。

2.2 嵌套和多尺度特征融合

通过嵌套的 U 形结构，U-Net++ 实现了多尺度特征融合，有效提升了网络对不同尺度细节的捕捉能力。

$Y_{i,j} = \sigma(W_{i,j} * [Y_{i-1,j}, Y_{i,j-1}, ..., Y_{i,j-n}] + b_{i,j})$

2.3 参数效率和性能

尽管增加了连接和结构，U-Net++ 通过合理设计控制参数量，保持了高效率和良好的性能，适用于医学图像等复杂场景。

2.4 Pytorch代码

import torch
import torch.nn as nn

class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ConvBlock, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, 3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.conv(x)

class UNetPlusPlus(nn.Module):
    def __init__(self, in_channels=3, out_channels=1, filters=[32, 64, 128, 256, 512]):
        super(UNetPlusPlus, self).__init__()
        self.pool = nn.MaxPool2d(2, 2)
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        self.conv0_0 = ConvBlock(in_channels, filters[0])
        self.conv1_0 = ConvBlock(filters[0], filters[1])
        self.conv2_0 = ConvBlock(filters[1], filters[2])
        self.conv3_0 = ConvBlock(filters[2], filters[3])
        self.conv4_0 = ConvBlock(filters[3], filters[4])

        self.conv0_1 = ConvBlock(filters[0] + filters[1], filters[0])
        self.conv1_1 = ConvBlock(filters[1] + filters[2], filters[1])
        self.conv2_1 = ConvBlock(filters[2] + filters[3], filters[2])
        self.conv3_1 = ConvBlock(filters[3] + filters[4], filters[3])

        self.conv0_2 = ConvBlock(filters[0]*2 + filters[1], filters[0])
        self.conv1_2 = ConvBlock(filters[1]*2 + filters[2], filters[1])
        self.conv2_2 = ConvBlock(filters[2]*2 + filters[3], filters[2])

        self.conv0_3 = ConvBlock(filters[0]*3 + filters[1], filters[0])
        self.conv1_3 = ConvBlock(filters[1]*3 + filters[2], filters[1])

        self.conv0_4 = ConvBlock(filters[0]*4 + filters[1], filters[0])

        self.final = nn.Conv2d(filters[0], out_channels, kernel_size=1)

    def forward(self, x):
        x0_0 = self.conv0_0(x)
        x1_0 = self.conv1_0(self.pool(x0_0))
        x0_1 = self.conv0_1(torch.cat([x0_0, self.up(x1_0)], 1))

        x2_0 = self.conv2_0(self.pool(x1_0))
        x1_1 = self.conv1_1(torch.cat([x1_0, self.up(x2_0)], 1))
        x0_2 = self.conv0_2(torch.cat([x0_0, x0_1, self.up(x1_1)], 1))

        x3_0 = self.conv3_0(self.pool(x2_0))
        x2_1 = self.conv2_1(torch.cat([x2_0, self.up(x3_0)], 1))
        x1_2 = self.conv1_2(torch.cat([x1_0, x1_1, self.up(x2_1)], 1))
        x0_3 = self.conv0_3(torch.cat([x0_0, x0_1, x0_2, self.up(x1_2)], 1))

        x4_0 = self.conv4_0(self.pool(x3_0))
        x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
        x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, self.up(x3_1)], 1))
        x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, self.up(x2_2)], 1))
        x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, self.up(x1_3)], 1))

        output = self.final(x0_4)
        return output

# 创建模型实例
model = UNetPlusPlus(in_channels=3, out_channels=1)

2.5 TensorFlow代码

import tensorflow as tf
from tensorflow.keras import layers, Model

def conv_block(inputs, filters):
    x = layers.Conv2D(filters, 3, padding='same')(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    x = layers.Conv2D(filters, 3, padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.ReLU()(x)
    return x

def UNetPlusPlus(input_shape=(256, 256, 3), num_classes=1):
    inputs = layers.Input(shape=input_shape)

    # Encoder (Downsampling)
    conv0_0 = conv_block(inputs, 32)
    pool1 = layers.MaxPooling2D(pool_size=(2, 2))(conv0_0)
    conv1_0 = conv_block(pool1, 64)
    pool2 = layers.MaxPooling2D(pool_size=(2, 2))(conv1_0)
    conv2_0 = conv_block(pool2, 128)
    pool3 = layers.MaxPooling2D(pool_size=(2, 2))(conv2_0)
    conv3_0 = conv_block(pool3, 256)
    pool4 = layers.MaxPooling2D(pool_size=(2, 2))(conv3_0)
    conv4_0 = conv_block(pool4, 512)

    # Decoder (Upsampling)
    up1_0 = layers.UpSampling2D(size=(2, 2))(conv4_0)
    up1_0 = layers.concatenate([up1_0, conv3_0])
    conv3_1 = conv_block(up1_0, 256)

    up2_0 = layers.UpSampling2D(size=(2, 2))(conv3_0)
    up2_0 = layers.concatenate([up2_0, conv2_0])
    conv2_1 = conv_block(up2_0, 128)

    up2_1 = layers.UpSampling2D(size=(2, 2))(conv3_1)
    up2_1 = layers.concatenate([up2_1, conv2_0, conv2_1])
    conv2_2 = conv_block(up2_1, 128)

    up3_0 = layers.UpSampling2D(size=(2, 2))(conv2_0)
    up3_0 = layers.concatenate([up3_0, conv1_0])
    conv1_1 = conv_block(up3_0, 64)

    up3_1 = layers.UpSampling2D(size=(2, 2))(conv2_1)
    up3_1 = layers.concatenate([up3_1, conv1_0, conv1_1])
    conv1_2 = conv_block(up3_1, 64)

    up3_2 = layers.UpSampling2D(size=(2, 2))(conv2_2)
    up3_2 = layers.concatenate([up3_2, conv1_0, conv1_1, conv1_2])
    conv1_3 = conv_block(up3_2, 64)

    up4_0 = layers.UpSampling2D(size=(2, 2))(conv1_0)
    up4_0 = layers.concatenate([up4_0, conv0_0])
    conv0_1 = conv_block(up4_0, 32)

    up4_1 = layers.UpSampling2D(size=(2, 2))(conv1_1)
    up4_1 = layers.concatenate([up4_1, conv0_0, conv0_1])
    conv0_2 = conv_block(up4_1, 32)

    up4_2 = layers.UpSampling2D(size=(2, 2))(conv1_2)
    up4_2 = layers.concatenate([up4_2, conv0_0, conv0_1, conv0_2])
    conv0_3 = conv_block(up4_2, 32)

    up4_3 = layers.UpSampling2D(size=(2, 2))(conv1_3)
    up4_3 = layers.concatenate([up4_3, conv0_0, conv0_1, conv0_2, conv0_3])
    conv0_4 = conv_block(up4_3, 32)

    outputs = layers.Conv2D(num_classes, 1, activation='sigmoid')(conv0_4)

    model = Model(inputs=inputs, outputs=outputs)
    return model

# 创建模型实例
model = UNetPlusPlus(input_shape=(256, 256, 3), num_classes=1)

3. 对比分析

在这里插入图片描述

3.1 分割性能比较

下表对比了 U-Net 和 U-Net++ 在不同数据集上的分割性能。

数据集	U-Net 精度	U-Net++ 精度
医学图像数据集	85%	90%
卫星图像数据集	80%	88%
自动驾驶数据集	82%	89%

3.2 参数量和计算开销

下表比较了 U-Net 和 U-Net++ 在网络结构复杂度、参数数量和计算资源消耗上的差异。

指标	U-Net	U-Net++
参数数量	31M	37M
计算复杂度	62 GFLOPs	75 GFLOPs
推理时间	20 ms/张	25 ms/张

结论

U-Net++ 作为 U-Net 结构的进化版，通过密集跳跃连接和多尺度特征融合显著提高了图像分割性能，尤其在细节捕捉和特征传递方面表现优异。尽管其参数量和计算开销有所增加，但在实际应用中，U-Net++ 的优势明显，值得在高精度图像分割任务中推广使用。

参考文献

[1] U-Net: Convolutional Networks for Biomedical Image Segmentation:U-Net
[2] UNet++: A Nested U-Net Architecture for Medical Image Segmentation:U-Net++

本人诚接各种数据处理、机器学习、深度学习、图像处理、时间序列预测分析等方向的算法/项目私人订制，技术在线，价格优惠。如有需要欢迎私信博主！！！