传知代码-KAN卷积：医学图像分割新前沿

代码以及视频讲解

本文所涉及所有资源均在传知代码平台可获取

概述

在本文中深入探讨KAN卷积在医学图像分割领域的创新应用，特别是通过引入Tokenized KAN Block（Tok Kan）这一突破性设计，将深度学习中的图像分割技术推向了新的高度。KAN作为一种能够替代传统MLP（多层感知机）的网络结构，以其独特的优势在多个领域展现出强大的潜力。而在医学图像分割这一复杂且关键的领域，KAN卷积更是凭借其高效处理图像特征的能力，成为了研究的热点。本文将U-Net结构中的卷积部分替换成了KAN卷积，将MLP部分用KANLinear取代，同时融入了类似Vision Transformer（VIT）的移位思想，使得模型在捕捉图像全局信息的同时，也能精准定位局部细节。该创新型能够支撑起一篇论文。

核心创新点

1. 将KAN卷积引入分割网络中

Kolmogorov–Arnold Networks（KAN）通常不是直接指代一种具体的卷积神经网络架构，但在这里我们可以理解为一种特殊的卷积或特征提取机制，可能基于Kolmogorov-Arnold表示定理（也称为超位置定理），该定理提供了多变量函数可以通过一系列一元函数和固定二元函数的组合来表示的理论基础。因此，KAN卷积可能意味着一种高度非线性的、能够捕捉复杂依赖关系的卷积操作。将其引入U-Net中，可以显著提升模型对图像特征的提取能力，特别是那些需要高级抽象和复杂交互的特征。

2. KANLinear替换MLP

将传统的MLP层替换为KANLinear层，可能意味着这一层结合了KAN的某些特性（如非线性处理能力或复杂的函数逼近能力）和线性变换的简洁性。这种替换可能使模型在保持高效计算的同时，能够更灵活地处理特征之间的复杂关系，进一步增强模型的特征表示能力,同时能够减少模型复杂度。KANLinear层可能通过其独特的机制，更好地整合和转换来自不同层级的特征，从而有助于模型在全局和局部信息之间做出更精准的权衡。

3. 融入移位思想

虽然Vision Transformer（VIT）本身并不直接包含“移位”操作，但其自注意力机制能够捕捉图像中的全局依赖关系，这一点与移位操作在促进信息流动和增强全局感受野方面的作用相似。在U-Net中融入类似VIT的思想，可能意味着引入了一种能够跨越空间位置直接交互特征的机制（如自注意力模块），或者通过某种形式的特征重排（类似于移位但更灵活）来增强模型的全局理解能力。这种设计使得模型在保持对局部细节敏感的同时，也能够有效地整合全局信息，从而在处理复杂图像任务时展现出更高的性能。本文中设计了沿着两个方向的移位。

模块介绍

KAN

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

UNext模块

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

本文主要结构

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

主要代码

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_()

数据集

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

结果展示

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

分割图：

运行过程：
从附件下载文件，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

运行结果：

本文在附件的readme中提供了除了BUSI之外的数据集，可以自行尝试增加工作量。

参考文献

1. UNeXt: MLP-based Rapid Medical Image Segmentation Network
2. KAN: Kolmogorov–Arnold Networks

源码下载