U-KAN
U-KAN精度高于Mamba、Transformer构建的Unet模型,本算法基于U-KAN进一步优化精度。
论文
U-KAN Makes Strong Backbone for Medical Image Segmentation and Generation
模型结构
KAN已成为MLP的一个有前景的替代方案,U-KAN将新兴算法KAN的优势融入到成熟的U-Net Pipeline 中提升精度,同时增加可解释性,本算法进一步引入部分QKV注意力、KAN ffn等来增强模型的特征提取能力,弥补KAN的不足。
算法原理
U-KAN与其它图像分割算法一样,将图片数据送入模型后,依次通过预处理、特征提取,最后经过卷积预测出像素类别实现分割。
环境配置
mv U-KAN-optimize_pytorch U-KAN # 去框架名后缀
# if torch>2.0, modify /usr/local/lib/python3.10/site-packages/timm/models/layers/helpers.py: from torch._six import container_abcs -> import collections.abc as container_abcs
Docker(方法一)
docker pull image.sourcefind.cn:5000/dcu/admin/base/pytorch:2.1.0-ubuntu20.04-dtk24.04.1-py3.10
# <your IMAGE ID>为以上拉取的docker的镜像ID替换,本镜像为:a4dd5be0ca23
docker run -it --shm-size=32G -v $PWD/U-KAN:/home/U-KAN -v /opt/hyhal:/opt/hyhal:ro --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video --name ukan <your IMAGE ID> bash
cd /home/U-KAN
pip install -r Seg_UKAN/requirements.txt # requirements.txt
Dockerfile(方法二)
cd U-KAN/docker
docker build --no-cache -t ukan:latest .
docker run --shm-size=32G --name ukan -v /opt/hyhal:/opt/hyhal:ro --privileged=true --device=/dev/kfd --device=/dev/dri/ --group-add video -v $PWD/../../U-KAN:/home/U-KAN -it ukan bash
# 若遇到Dockerfile启动的方式安装环境需要长时间等待,可注释掉里面的pip安装,启动容器后再安装python库:pip install -r Seg_UKAN/requirements.txt。
Anaconda(方法三)
1、关于本项目DCU显卡所需的特殊深度学习库可从光合开发者社区下载安装:
DTK驱动:dtk24.04.1
python:python3.10
torch:2.1.0
torchvision:0.16.0
Tips:以上dtk驱动、python、torch等DCU相关工具版本需要严格一一对应。
2、其它非特殊库参照requirements.txt安装
pip install -r Seg_UKAN/requirements.txt # requirements.txt
数据集
BUSI(Breast Ultrasound Image)
本项目无需下载原始数据集,采用U-KAN作者提供的预处理数据pre-processed dataset 即可。 项目中已包含busi,解压即可使用,训练数据目录结构如下:
Seg_UKAN
├── inputs
│ ├── busi
│ ├── images
│ ├── malignant (1).png
| ├── ...
| ├── masks
│ ├── 0
│ ├── malignant (1)_mask.png
| ├── ...
数据集快速下载中心:SCNet AIDatasets ,项目中预处理数据集可从快速下载通道下载:busi_cvc_glas_preprocessed 。
训练
单机单卡
# 以公开数据集busi为基础进行算法效果优化
cd Seg_UKAN
python train.py --arch UKAN --dataset busi --input_w 256 --input_h 256 --name busi_UKAN --data_dir ./inputs
更多资料可参考源项目的README_origin
result
精度
数据集:busi,max epoch为400,训练框架:pytorch。
算法 | Dice |
---|---|
U-KAN | 78.75% |
U-KAN-optimize | 79.64% |
应用场景
算法类别
图像分割
热点应用行业
医疗,电商,制造,能源