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

基于计图深度学习框架(Jittor)[1]，四川大学智能医学中心开发了医学图像处理模型库 SCU-JittorMed，并开源了三个基于Jittor的医学图像处理方法，具体包括：

Jittor-NODE: 基于 Jittor 实现的 Neural ODE Solver ，支持 Euler 与 RK4 两种固定步长的 Solver ，并实现了 nmODE 模型。
Jittor-SimplTrans:一种高效的Transformer变体，在保证性能的同时有效降低了可学习参数量。
Jittor-AugmentBN: 一种针对医学图像小批量训练的归一化方法，提升了模型在小批量样本训练时的准确率与稳定性。

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

Part 1

Jittor-NODE

1.1 nmODE 方法介绍

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

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

其中是随时间 t 演化的隐藏状态， f 是由神经网络建模的非线性变换， W 是可学习参数。 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. nmODE 与 UNet 以及 convODE 的分割结果对比

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

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

Part 2

Jittor-SimplTrans

2.1 SimplTrans 方法介绍

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

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

其中

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

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

图 3. 原始 Attention 与 Simplified 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 Normalization （ NABN ） [4] 。

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

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

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

图 6. NABN 算法流程

3.2 NABN 实验结果

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

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

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

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

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

表 6. U-Net 与 BN 和 NABN 在 MSD 肝脏分割数据集上的分割结果对比。 Dice (%) 值越高越好， HD (95%) 值越低越好

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

图 7. Grad-CAM 可解释性实验

Part 4

计图开源

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

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

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

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

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

参考资料

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.
Zhang Yi. nmODE: neural memory ordinary differential equation. Artificial Intelligence Review, 56(12), 2023: 14403-14438.
Lei Xu, Haiying Luo, and Zhang Yi. Simplified Transformer. Neurocomputing , 2025. Accepted.
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往期回顾

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