计图开源:医学图像处理模型库SCU-JittorMed

2025/05/28 Tutorial

计图开源:医学图像处理模型库SCU-JittorMed

基于计图深度学习框架(Jittor)[1],四川大学智能医学中心开发了医学图像处理模型库SCU-JittorMed,并开源了三个基于Jittor的医学图像处理方法,具体包括:
  • Jittor-NODE: 基于Jittor实现的Neural ODE Solver,支持EulerRK4两种固定步长的Solver,并实现了nmODE模型。

  • Jittor-SimplTrans:一种高效的Transformer变体,在保证性能的同时有效降低了可学习参数量。

  • Jittor-AugmentBN: 一种针对医学图像小批量训练的归一化方法,提升了模型在小批量样本训练时的准确率与稳定性。

以下将分别介绍这三种方法的核心思想与实验效果。

Part 1

Jittor-NODE

1.1 nmODE 方法介绍

ResNet 是一种经典的离散深度神经网络。从数学角度来看,ResNet中的残差块可以被理解为对连续动力系统的一种离散近似。具体而言,每个残差单元形式为:

这实际上等价于对以下常微分方程进行一步欧拉法(Euler method)的数值近似:

其中是随时间演化的隐藏状态,f 是由神经网络建模的非线性变换,是可学习参数。ResNet 中的跨层连接可以看作是在时间轴上以固定步长推进。随着时间步长趋近于0,网络层数趋于连续,自然过渡到了神经常微分方程模型(Neural Ordinary Differential Equation, NODE)。NODE 不再通过离散的网络层来更新状态,而是通过求解连续时间的微分方程:

四川大学智能医学中心章毅教授提出的nmODE[2]NODE的一种特殊变体,它显式地将初始状态与外部输入分离开来,为动态系统建模提供了更大的灵活性。

nmODE通过将激活函数替换为非单调函数(如sin² ),有效提升网络的表达能力。通过构建隐式映射方程并引入常微分方程求解,使得每个外部输入x都映射至唯一吸引子y*,从而得到稳定的记忆表达。

此外,nmODE将记忆机制作为吸引子嵌入网络动力学中,建立输入与稳定状态(吸引子)之间的非线性映射。该模型将学习神经元与记忆神经元分离,结构清晰,具有全局吸引子的特性。

1. nmODE模型架构

1.2 nmODE 实验效果

自然图像上的实验: CIFAR-10 分类任务中,nmODE 相比扩展模型 ANODE 展现出更强的表达能力,准确率达到 70.73%,显著优于 ANODE  60.60% (如表1所示)

1. CIFAR-10图像分类准确率

医学图像分割实验:nmODE 18类头颈部器官分割中整体优于 UNet  convODE,验证了其在医学场景下的有效性 (如表2所示)。

2. nmODEUNet以及convODE的分割结果对比

可视化分析:图2直观展示了nmODE对于特征的增强能力。其中,每个部分的四张图像分别代表CT切片、nmODE的外部输入、nmODE的输出和标签。

2. 脑干、眼、喉和脊髓影像数据分割的nmODE特征可视化。

Part 2

Jittor-SimplTrans

2.1 SimplTrans 方法介绍

Transformer 虽性能强大,但其注意力机制带来较高的参数量和计算开销。为此,四川大学章毅教授团队提出简化结构 SimplTrans[3],通过数学分析去除冗余参数,保留核心特性,在多项任务中实现高效建模与优异性能。

原始Transformer的注意力机制本质上是一个带参数的双线性映射:

其中

我们注意到,从矩阵理论出发,发现几乎所有的方阵都是对角化且可逆的,因此可将A视为单位阵,从而得到无参数形式(Simplified Attention):

这种形式保留了双线性结构,消除了多余参数。简化后的注意力模块包括:输入先经双线性映射获得注意力得分矩阵,再乘以线性投影, 通过残差连接与LayerNorm得到输出。相比原始Attention,其结构显著简化。

3. 原始AttentionSimplified Attention对比

2.2 SimplTrans实验结果

图像分类(MNIST/CIFAR-10/CIFAR-100)任务上,Simplified Attention取得更高准确率,且参数量和FLOPs更低。

3. MNIST数据集上的实验结果

 WMT 2018 EN-CH 翻译中,Simplified Attention 的表现优于原始 Attention,参数量减少超 40%。其注意力图融合水平与垂直特征,表现出更强的多样性。

4. Attention注意力图可视化结果

Part 3

Jittor-AugmentBN

3.1 NABN 方法介绍

批量归一化(BN)虽能提升训练效率和泛化能力,但在小批量训练中易产生统计偏差,导致不稳定(如下图所示),在医学图像分析中尤为明显。因此,优化 BN 以适应小批量场景成为亟需解决的问题。针对这一现状,四川大学智能医学中心章毅教授团队提出了在不改变BN核心结构的前提下增强BN稳定性的方法:Noise-Assisted Batch NormalizationNABN[4]

5. 传统BN与论文提出的NABN在测试集效果对比

NABN 的核心思想在于通过在批量归一化中引入随机噪声,提升小批量训练时均值与方差估计的稳定性,从而增强模型鲁棒性。

其核心包括两步:对均值加入有界噪声(Noisy Centering)和基于噪声均值计算方差(Noisy Scaling),有效模拟大批量统计效果,降低对 batch size的敏感性。算法流程如下方算法所示,算法引入的噪声由一个Noisy Mean ModuleNMM)生成,其噪声强度通过超参数θ控制。

6. NABN算法流程

3.2 NABN 实验结果

图像分类任务ResNet-20/32/44 结合 NABN 后,在所有网络深度下均优于BN及其变体,结果如表4所示。

4. NABN结合不同深度的ResNetCIFAR-10数据集上测试集误差对比

OCT  Chest X-ray 医学图像分类数据集NABN 显著提升了 ResNet-50在医学图像分类任务中的性能。

5.OCTChest X-ray医学数据集上的准确率、召回率、精准率和 F1 (%)

医学图像分割任务 MSD Liver 分割数据集中,使用 U-Net 结合 NABN 后,测试集 Dice 系数提高了约 3%95% Hausdorff DistanceHD 95显著降低。

6. U-NetBNNABNMSD肝脏分割数据集上的分割结果对比。Dice (%) 值越高越好,HD (95%) 值越低越好

可视化实验:图7展示了NABN 更准确地关注到了病灶区域,有效提高了模型解释性。图中,第一行为原始图像。第二行和第三行分别为 BN  NABN  Grad-CAM 图像。

7. Grad-CAM可解释性实验

Part 4

计图开源

上述所提到的三个方法,均在四川大学智能医学中心发布的计图医学库SCU-JittorMed 中开源,开源地址为:

https://github.com/SCU-JittorMed/

使用Jittor实现以上模型,在保证效果的同时训练和推理效率得到了有效提升。

对于SimplTrans,在训练阶段,固定批次大小为256,评估数据集为MNISTJittor版本训练用时约14.3sPyTorch版本用时约15.9s,训练速度提升10%。在推理阶段,固定批次大小为256Jittor版本推理用时约0.75sPyTorch版本用时约1.04s。推理速度提升27.8%

对于NABN方法,在OCT 医学图像分类任务中,训练阶段训练样本数为 1000,固定批次大小为64,训练轮数为100Jittor版本总训练时间约为 160s,PyTorch版本约为190s,训练速度提升15.8%。在推理阶段,评估样本数为 250Jittor版本推理用时约0.75s,PyTorch 版本约0.90s秒,推理速度提升16.7%

参考资料

  1. Shi-Min Hu, Dun Liang, Guo-Ye Yang, Guo-Wei Yang, and Wen-Yang Zhou. Jittor: a novel deep learning framework with meta-operators and unified graph execution, Science China Information Science, 63(12), 2020, article no. 222103.

  2. Zhang Yi. nmODE: neural memory ordinary differential equation. Artificial Intelligence Review,  56(12), 2023: 14403-14438.

  3. Lei Xu, Haiying Luo, and Zhang Yi. Simplified Transformer. Neurocomputing, 2025. Accepted.

  4. Shengqian Zhu, Chengrong Yu, Junjie Hu. Regularizing deep neural networks for medical image analysis with augmented batch normalization. Applied Soft Computing, 154, 2024: 111337.

GGC往期回顾
  1. 计图开源:面向高斯泼溅的解耦与编辑方法DeferredGS

  2. ICML 2025: 清华与腾讯、斯坦福等单位合作联合发布大模型复杂推理评测基准RBench

  3. 第五届“计图”人工智能算法挑战赛启动

  4. 计图团队完成华为昇腾上 DeepSeek 模型的高效适配

  5. 三维透明物体重建的网格-神经混合表示 | CVMJ Spotlight

可通过下方二维码,关注清华大学图形学实验室,了解图形学、深度学习、Jittor框架、CVMJ期刊和CVM会议的相关资讯。

图片


Search

    Table of Contents