【论文复现】KAN卷积：医学图像分割新前沿

Eternity._

发布于 2024-11-07 08:15:36

1080

发布于 2024-11-07 08:15:36

文章被收录于专栏：登神长阶

前言：在医学领域，图像分割是一项至关重要的任务。它不仅能够帮助医生更准确地诊断疾病，还能在手术规划、治疗评估以及医学研究等多个方面发挥巨大作用。随着人工智能和深度学习技术的飞速发展，医学图像分割技术也迎来了前所未有的变革。在众多新兴的技术中，KAN（Kernel Attention Networks，核注意力网络）卷积技术以其独特的优势，正逐步成为医学图像分割的新前沿。

传统医学图像分割方法，如阈值分割、边缘检测和区域增长等，虽然在一定程度上能够满足临床需求，但它们在处理复杂、多变的医学图像时，往往显得力不从心。近年来，基于深度学习的分割方法，特别是卷积神经网络（CNN），凭借其强大的特征提取和模式识别能力，在医学图像分割领域取得了显著成效。然而，随着医学图像数据的不断增加和临床需求的日益提高，如何进一步提升分割精度和效率，仍然是当前研究的热点和难点。

正是在这样的背景下，KAN卷积技术应运而生。它结合了传统卷积神经网络的优势，并在此基础上引入了注意力机制，通过动态调整卷积核的权重，实现对图像特征的精细捕捉和高效利用。这种机制不仅提高了分割的准确性，还增强了模型对图像中关键信息的敏感度，使得在复杂医学图像中也能实现精确分割。

在未来的医学图像分割领域，它将发挥更加重要的作用，为人类的健康事业做出更大的贡献。让我们共同期待这一新前沿技术的精彩表现！

1. 概述

本文致力于深度剖析KAN卷积在医学图像分割领域的革新性应用，特别是Tokenized KAN Block（Tok Kan）这一颠覆性设计的引入，为深度学习图像分割技术树立了新的标杆。KAN，作为一种潜力巨大的网络结构，能够替代传统的多层感知机（MLP），并在多个领域内展现出了其非凡的优势。尤其在医学图像分割这一既复杂又至关重要的领域，KAN卷积凭借其卓越的图像特征处理能力，成为了研究的焦点。在本文中，我们将U-Net结构中的卷积层替换为KAN卷积，MLP层则被KANLinear所替代，同时融入了Vision Transformer（ViT）的移位思想，使得模型能够在捕捉图像全局信息的同时，对局部细节进行精确识别。这一创新性的研究足以构成一篇完整的论文。

2. 核心创新点

将KAN卷积引入分割网络中：

Kolmogorov–Arnold Networks（KAN）在此处虽不直接等同于传统意义上的卷积神经网络架构，但我们可以将其构想为一种独特的卷积或特征提取机制，该机制深受Kolmogorov-Arnold表示定理（亦称超位置定理）启发。此定理为多变量函数的表示提供了理论基础，指出这些函数可以通过一元函数与固定二元函数的组合来构建。据此，KAN卷积可视为一种高度非线性且能有效捕捉复杂依赖关系的卷积操作。将其融入U-Net架构中，能显著提升模型提取图像特征的能力，尤其是针对那些需要高度抽象和复杂特征交互的任务。

KANLinear替换MLP：

将U-Net中的传统MLP层替换为KANLinear层，意味着这一层将KAN的特定优势（诸如强大的非线性处理能力或高精度的函数逼近能力）与线性变换的简洁高效相结合。这种替换策略旨在保持模型计算效率的同时，赋予其处理特征间复杂关系的更高灵活性，从而进一步增强模型的特征表征能力。此外，KANLinear层可能凭借其独特的机制，更有效地整合并转换来自不同网络层级的特征信息，帮助模型在全局信息与局部细节之间实现更精确的平衡。

融入移位思想：

尽管Vision Transformer（ViT）并不直接包含“移位”这一具体操作，但其自注意力机制在捕捉图像全局依赖关系上的能力，与移位操作在促进信息流动和扩大全局感受野方面的效果异曲同工。在U-Net架构中融入ViT的思想，可能指的是引入了一种机制，该机制允许特征在空间位置上直接进行交互（例如，通过自注意力模块），或者通过某种灵活的特征重排（虽非传统移位但效果类似）来加强模型的全局理解能力。这种设计策略使得模型能够在保持对局部细节高度敏感的同时，有效地整合全局信息，从而在应对复杂图像分割任务时展现出卓越的性能。在本文中，我们特别设计了沿两个方向进行的特征重排操作，以进一步提升模型的全局与局部信息整合能力。

3. 模块介绍

KAN

KAN模型由数学定理Kolmogorov–Arnold启发得出，该定理由前苏联的两位数学家Vladimir Arnold和Andrey Kolmogorov提出。定理表明，任何多元连续函数都可以表示为单变量连续函数的两层嵌套叠加（一个单一变量的连续函数和一系列连续的双变量函数的组合）。这为多维函数的分解提供了理论基础，也是KAN模型设计的核心思想。

KAN模型具体解读参考：

https://www.aspiringcode.com/content?id=17248603157044&uid=218e0c4d82f346c880d3be2c19307d0f

UNext模块

UNext模块一种基于卷积和多层感知器（MLP）的医学图像分割网络，旨在解决现有模型如UNet和Transformer版本在计算复杂度、参数量以及推理速度上的不足。本文是在此基础上将卷积用KAN卷积取代，将MLP用KAN模块代替。

有关UNext模块的介绍参考：

https://www.aspiringcode.com/content?id=17233927441942&uid=fb5339efc9754267a67a10a2f4b86967

4. 本文主要结构

本文受启发与上述的两种模块，将UNext模块中的上下采样中的卷积提取特征阶段用KAN卷积取代，将Tok MLP阶段用Tok Kan取代，从而增加模型提取特征的能力和渐少模型的参数量，同时提高模型非线性表达的能力。

整体架构如下:

本文所涉及的所有资源的获取方式：

https://www.aspiringcode.com/content?id=17256347583346&uid=21d463b8d98a4ef38920589d0efd153a

5. 主要代码

KAN卷积：

class KANConvs(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, enable_standalone_scale_spline=True):
        """
        定义KAN卷积层，类似于nn.Conv2d，但包括KAN样条插值权重
        :param in_channels: 输入通道数
        :param out_channels: 输出通道数
        :param kernel_size: 卷积核大小
        :param stride: 步长
        :param padding: 填充
        :param dilation: 扩张
        :param groups: 组卷积
        :param bias: 偏置项
        :param enable_standalone_scale_spline: 是否启用独立缩放样条插值
        """
        super(KANConvs, self).__init__()

        # 基本的卷积层初始化
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.groups = groups
        self.bias = bias

        # 标准卷积层参数
        self.weight = Parameter(torch.Tensor(out_channels, in_channels // groups, *self.kernel_size))
        if bias:
            self.bias = Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias', None)

        # 样条插值权重
        self.spline_weight = Parameter(torch.Tensor(out_channels, in_channels, *self.kernel_size))

        # 是否启用独立缩放样条插值
        self.enable_standalone_scale_spline = enable_standalone_scale_spline
        if enable_standalone_scale_spline:
            self.spline_scaler = Parameter(torch.ones(out_channels, 1))
            

        # 初始化权重
        self.reset_parameters()
class ConvLayer(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(ConvLayer, self).__init__()
        self.conv = nn.Sequential(
            KANConvs(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            KANConvs(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )

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

KANLinear层：

class KANBlock(nn.Module):
    def __init__(self, dim, drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, no_kan=False):
        super().__init__()

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim)

        self.layer = KANLayer(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop, no_kan=no_kan)

        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)
        elif isinstance(m, nn.Conv2d):
            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
            fan_out //= m.groups
            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
            if m.bias is not None:
                m.bias.data.zero_()

6. 数据集

本文使用的是BUSI（Breast Ultrasound Image）是一个包含乳腺超声图像的分类和分割数据集。该数据集包括了 2018 年收集的乳腺超声波图像，涵盖了 25 至 75 岁的 600 名女性患者。数据集由 780 张图像组成，每张图像的平均大小为 500*500 像素。这些图像被划分为三类：正常、良性和恶性。而在良性和恶性乳腺超声图像中，还包含了对应胸部肿瘤的详细分割标注，为深入研究和精准诊断提供了关键信息。这份数据集不仅为乳腺癌研究提供了丰富的图像资源和宝贵支持。

7. 结果展示

在BUSI数据集中的结果展示如下：

Methods	IoU	Dice
U-Net	57.22	71.91
Att-Unet	55.18	70.22
U-Net++	57.41	72.11
U-NeXt	59.06	73.08
Rolling-UNet	61.00	74.67
U-Mamba	61.81	75.55
OURS	63.45	77.05

分割图：

运行过程： 从附件下载文件，readme中有详细步骤。数据集在readme提供的链接中。

环境配置：

  - pip:
    - addict==2.4.0
    - dataclasses==0.8
    - mmcv-full==1.2.7
    - numpy==1.19.5
    - opencv-python==4.5.1.48
    - perceptual==0.1
    - pillow==8.4.0
    - scikit-image==0.17.2
    - scipy==1.5.4
    - tifffile==2020.9.3
    - timm==0.3.2
    - torch==2.4.0
    - torchvision==0.8.2
    - typing-extensions==4.0.0
    - yapf==0.31.0

运行结果：