模块#
MONAI 旨在促进医学图像分析领域多粒度的深度学习应用。本文档概述了各模块及其核心功能。
核心代码库被设计为一个轻量级、灵活且全面的 API 库,适用于不同专业水平的用户。构建模块易于理解和使用,它们经过精心解耦,可以轻松集成到现有的 PyTorch 程序和大型系统中。通过利用工作流(workflow)和包(bundle)API,用户还可以针对各种特定领域的应用快速建立高效且稳健的模型训练或评估流水线。
整体架构和模块如下图所示

I/O、处理与增强#
医学图像需要专门的 I/O、预处理和增强方法。它们通常遵循特定格式,使用特定的协议进行处理,且数据数组通常是高维的。monai.transforms 和 monai.data 模块包含了一套针对各种深度学习应用的领域特定 API。
数组和字典格式的数据变换#

这不仅支持基础的图像变换,还支持更复杂的预处理流水线,例如跨不同模态和模型监督输入的同步操作。[数组和字典示例]
多种基于图像块(patch)的采样机制#

实现了用于选择性预处理的高级补丁(patch)采样方法,例如根据用户指定的采样权重图进行加权、类别平衡采样。输出可以是序列或迭代器模式,支持不同类型的洗牌(shuffling)策略。
集成第三方库的图像 I/O#
内置了多种后端以支持各种格式。对于自定义格式读取器,它具有良好的扩展性。
monai.data.MetaTensor#
核心数据结构结合了 PyTorch 原生 Tensor API 与元数据处理,使深度学习模型和流水线能够轻松整合元信息。[MetaTensor]
基于 GPU 的加速#
提供了多种实现方式,以确保底层硬件资源得到最佳利用。[快速训练指南]
确定性与可重复性#
通过 Randomizable API 进行细粒度的局部控制,以及通过 set_determinism 进行全局控制,可以实现确定性和可重复性。
解包(Decollating)与可逆变换#
模型输出的 mini-batch 数据可以被解包并独立进行后处理,包括根据跟踪的元数据和已应用的操作将输出逆变换回预处理的早期步骤。[逆变换演示]
增强的可用性#
此外,还提供了诸如 DataStats 变换、dev_collate 以及可视化方法等工具作为 PyTorch 的扩展,以提高整体的可调试性。
数据集与数据加载#
遵循 PyTorch 的设计模式,MONAI 扩展了 Dataset 和 DataLoader API,在领域特定可用性和流水线性能方面进行了重大增强。
缓存 I/O 和变换数据以加速训练#
数据驱动的方法需要对训练数据进行多次(可能数千次)Epoch 的读取和预处理。MONAI 提供了基于多线程缓存的数据集来加速此过程 [数据集实验]。缓存可以是持久化和动态的(SmartCacheDataset),并且可以在不同的实验中重用 [SmartCache 示例]。下图展示了与常规 PyTorch 程序相比的训练速度提升。

ThreadDataLoader 与 DataLoader#
如果变换非常轻量,特别是当我们将所有数据缓存到内存(RAM)中时,PyTorch DataLoader 的多进程机制可能会导致不必要的 IPC 时间并降低 GPU 利用率。MONAI 提供了 ThreadDataLoader,它在单独的线程中执行变换。

ThreadDataLoader 示例位于 Spleen 快速训练教程中。
公共数据集#
为了快速上手流行训练数据,MONAI 提供了几个即插即用的数据集类(如 MedNISTDataset, DecathlonDataset, TciaDataset),其中包括数据下载,并支持通过变换生成训练/评估集划分。[公共数据集教程] 预定义数据集的通用工作流

数据集类型扩展#
Dataset API 的其他扩展包括:用于关联多个数据源的 ZipDataset,用于处理图像级和块级预处理的 PatchDataset,用于多模态输入的 CSVDataset,以及用于交叉验证数据准备的 partition_dataset。
可微组件、网络、损失函数与优化器#
一些深度神经网络架构在医学图像分析任务中表现出显著的有效性。MONAI 实现参考网络旨在兼顾灵活性与代码可读性。
预定义层与块#
网络层和块的实现通常兼容空间 1D、2D 和 3D 输入。用户可以轻松地将层、块和网络集成到他们的自定义流水线中。提供了多种工具来利用现有的模型权重,例如来自 MONAI model-zoo 的包。
C++/CUDA 优化模块#
为了进一步加速领域特定的常规操作,引入了 MONAI C++/CUDA 实现作为 PyTorch 原生实现的扩展。MONAI 使用两种构建 C++ 扩展的方法来提供这些模块:
通过
setuptools,适用于包含Resampler,Conditional random field (CRF),Fast bilateral filtering using the permutohedral lattice的模块。通过即时(JIT)编译,用于
Gaussian mixtures模块。这种方法允许根据用户指定的参数和本地系统环境进行动态优化。下图展示了 MONAI 高斯混合模型应用于组织和手术工具分割的结果:
损失函数与优化器#
针对各种应用常用的损失函数已从文献中(重新)实现,例如 DiceLoss, GeneralizedDiceLoss, TverskyLoss, DiceFocalLoss。数值优化和相关工具包括 Novograd 和 LearningRateFinder。下图显示了一个学习率搜索过程。

评估#
为了运行模型推理并评估模型质量,MONAI 提供了针对广泛使用方法的参考实现。目前包含了一些流行的评估指标和推理模式:
滑动窗口推理#
对于大容量数据的模型推理,滑动窗口方法是实现高性能且兼顾灵活内存需求的热门选择(_此外,请查阅关于模型并行训练的最新研究)。它还支持 overlap 和 blending_mode 配置,以处理重叠窗口并获得更好的性能。

医学任务评估指标#
实现了各种有用的评估指标来衡量特定于医学图像模型的质量。这些指标包括 Mean Dice, ROCAUC, Confusion Matrices, Hausdorff Distance, Surface Distance, Occlusion Sensitivity。这些 API 也支持多进程计算。
报告生成#
提供了 MetricsSaver 来编写最终的指标摘要报告:mean, median, max, min, <int>percentile, std。

可视化#
除了简单的点和曲线绘图,还提供了直观的接口,可将多维数据作为 GIF 动画在 TensorBoard 中可视化。这可以通过可视化体积输入、分割图和中间特征图,对模型进行快速的定性评估。一个包含可视化的可运行示例可在 UNet 训练示例中找到。为了与 Ignite 程序协作,MONAI 还提供了几个 Ignite 处理器(handler),用于使用 TensorBoard 或 MLFlow 可视化训练曲线和指标,更多详细信息请参考 TensorBoard 和 MLFlow 处理器示例。
为了轻松地将 3D 图像可视化为 2D 图像帧,MONAI 提供了基于 matplotlib 库的工具 matshow3d。它可以在指定的维度上绘制图像帧,以脾脏 3D 图像为例:matshow3d(volume=image, figsize=(100, 100), every_n=10, frame_dim=-1, show=True, cmap="gray")

MONAI 还提供了 blend_images 工具,将 image 和 label 混合为 RGB 彩色图像,以便使用指定的 cmap 模式和权重更好地可视化分割区域。以脾脏分割 image 和相应的 label 为例。

关于 TensorBoard utility、matshow3d 和 blend_images 的更多详细信息,请查看可视化教程。
为了可视化已训练分类模型的类激活映射,MONAI 为 2D 和 3D 模型提供了 CAM、GradCAM、GradCAM++ API。

上述示例是通过计算来自肺部 CT 病灶分类模型的 GradCAM/GradCAM++生成的。
工作流#
MONAI 引擎和工作流能够快速启动训练和评估实验。
这些特性将领域特定组件与通用机器学习过程解耦。它们还为更高级的应用(如 AutoML、联邦学习)提供了一套统一的 API。工作流的训练器和评估器与 pytorch-ignite 的 Engine 和 Event-Handler 机制兼容。
通用工作流流水线#
工作流和部分 MONAI 事件处理器如下所示 [工作流示例]

EnsembleEvaluator#
一个典型的集成过程已实现为即用型工作流 [交叉验证和模型集成教程]
将所有训练数据集分成 K 折。
使用每 K-1 折数据训练 K 个模型。
使用所有 K 个模型对测试数据执行推理。
计算加权平均值,或通过投票选出最常见的值作为最终结果。

解包批量数据以实现灵活的后处理#
decollate batch 自 MONAI v0.6 引入,它简化了后处理变换,并为处理不同形状的数据批次提供了灵活的后续操作。它可以将批次数据(如模型预测)解包为 Tensor 列表,带来以下好处:
为每个项目独立启用后处理变换——随机变换可以为批次中的每个预测项应用不同的方式。
简化变换 API 并减少输入验证负担,因为预处理和后处理变换现在只需要支持“通道优先”(channel-first)输入格式。
为预测和不同形状的逆变换数据启用
Invertd变换,因为数据项以列表形式存在,而不是堆叠在单个 Tensor 中。允许在灵活的指标计算中同时使用 batch-first Tensor 和通道优先的 Tensor 列表。[解包批次教程]
decollate batch 的典型过程如下所示(以 batch_size=N 的模型预测和标签为例)

易于集成到流行工作流中#
除了基于 pytorch-ignite 的 monai.engines 外,大多数 MONAI 模块可以独立使用或与其他软件包结合。例如,MONAI 可以轻松集成到诸如 PyTorch-Lightning 和 MLflow 等流行框架中。
包(Bundle)#
MONAI 包的目标是定义一个打包的模型,其中包含允许用户和程序理解如何使用该模型及其用途所需的关键信息。一个包包含单个网络的存储权重(以 Pickle 格式的状态字典形式)以及可选的 Torchscript 对象和/或 ONNX 对象。还包含额外的 JSON 文件以存储有关模型的元数据、用于构建训练、推理和后处理变换序列的信息、纯文本描述、法律信息以及模型创建者希望包含的其他数据。更多详细信息可在包规范中找到。
包的主要好处是定义模型包,并通过结构化配置支持构建基于 Python 的工作流:
包含所有必要信息的自包含模型包。
结构化配置可用于轻松重建或原型化深度学习工作流。
通过将参数设置与 Python 代码分离,配置文件提供了良好的可读性和可用性。
配置文件可以描述灵活的工作流和组件,允许不同的底层 Python 实现。
联邦学习和 AutoML 等更高级别的学习范式可以从组件细节中解耦。
典型的包示例可以包括:
ModelName
┣━ configs
┃ ┗━ metadata.json
┣━ models
┃ ┣━ model.pt
┃ ┣━ *model.ts
┃ ┗━ *model.onnx
┗━ docs
┣━ *README.md
┗━ *license.txt
有关包配置定义、语法和示例的详细信息位于配置语法中。一份分步入门教程笔记本可以帮助用户快速建立一个包。[包示例, model-zoo]
联邦学习#
使用 MONAI 包配置,我们可以使用 MONAI 的 MonaiAlgo 类(这是抽象 ClientAlgo 类的一种联邦学习(FL)实现)来执行来自 MONAI model zoo 的包。请注意,ClientAlgo 作为定义在任何联邦学习平台上运行的算法的抽象基类提供。MonaiAlgo 实现了运行联邦学习实验所需的主要功能,即 train(), get_weights() 和 evaluate(),它们可以使用单 GPU 或多 GPU 训练来运行。此外,它还提供了对组件生命周期管理(如 initialize(), abort() 和 finalize())的实现。MONAI FL 客户端还允许使用 MonaiAlgoStats 类计算包配置中定义的数据集上的汇总数据统计信息(如强度直方图)。这些统计信息可以在 FL 服务器上共享和可视化。NVIDIA FLARE(由 NVIDIA 开发的联邦学习平台)已经构建了与 ClientAlgo 的集成件,以允许在联邦环境中轻松地使用 MONAI 包进行实验。我们的[联邦学习教程]展示了单 GPU 和多 GPU 训练以及联邦统计工作流的示例。
Auto3dseg#

Auto3DSeg 是一种针对大规模 3D 医学图像分割的全面解决方案。它利用 MONAI 和 GPU 的最新进展,高效地开发和部署具有最先进性能的算法。它首先分析数据集的强度、维度和分辨率等全局信息,然后基于数据统计和算法模板以 MONAI 包格式生成算法。接下来,所有算法启动模型训练以获得具有最佳验证性能的检查点。最后,集成模块通过对训练后的检查点进行排序来选择算法,并创建集成预测。
该解决方案为初学者和高级研究人员提供了不同水平的用户体验。它已在不同模态的大规模 3D 医学影像数据集上进行了测试。
GPU 加速、性能分析与优化#
MONAI 提供了最先进的性能优化方法,包括:
自动混合精度 (AMP)#
在训练或评估期间,只需在 SupervisedTrainer 或 SupervisedEvaluator 中设置 amp=True/False 即可启用/禁用 AMP。示例基准测试结果如下 [AMP 训练教程]
在配备 CUDA 11 和 PyTorch 1.6 的 NVIDIA V100 GPU 上开启/关闭 AMP 进行训练

在配备 CUDA 11 和 PyTorch 1.6 的 NVIDIA A100 GPU 上开启/关闭 AMP 进行训练

包括 DLProf, Nsight, NVTX 和 NVML 在内的多种工具可以与 MONAI 一起使用来识别性能瓶颈。[分析教程]
分布式训练#
MONAI 的分布式数据并行 API 与原生 PyTorch 分布式模块、pytorch-ignite 分布式模块、Horovod、XLA 和 SLURM 平台兼容。[分布式训练教程]

快速训练教程结合了 AMP 与 CacheDataset, GPU cache, GPU transforms, ThreadDataLoader 以及网络和优化器的调整,与常规 PyTorch 实现相比,可以实现显著的速度提升。