模块#

MONAI 旨在促进医疗影像分析中深度学习在不同粒度上的应用。本文档提供了模块的概述并重点介绍了关键功能。

核心代码库被设计为一个轻量级、灵活且全面的 API 库,适用于不同专业水平的用户。构建块易于理解和使用,它们被精心解耦,可以轻松集成到现有的 PyTorch 程序和更大的系统中。通过利用工作流和捆绑包 API,用户还可以快速建立高效且稳健的模型训练或评估流程,用于各种特定领域的应用。

总体架构和模块如下图所示

architecture overview

I/O、处理与增强#

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

支持数组和字典风格数据的变换#

3d transform examples

这支持基本的图像变换,以及更复杂的预处理流程,例如跨不同模态和模型监督输入的同步操作。[数组和字典示例]

各种基于图像块的采样机制#

2d transform examples

实现了高级的图像块采样方法,用于选择性预处理,例如根据用户指定的采样权重图进行加权、类别平衡采样。输出可以是序列或迭代器模式,这允许使用不同类型的洗牌策略。

集成第三方库的图像 I/O#

内置了多种后端,可以支持各种格式。它易于扩展,可用于自定义格式读取器。

monai.data.MetaTensor#

核心数据结构将 PyTorch 原生 Tensor API 与元数据处理相结合,以便深度学习模型和流程可以轻松地纳入元信息。[MetaTensor]

基于 GPU 的加速#

提供了实现以确保底层硬件资源的最佳利用。[快速训练指南]

确定性与可复现性#

通过 Randomizable API 进行细粒度的局部控制以及使用 set_determinism 进行全局控制都可以实现这一点。

解批处理与可逆变换#

invert transform 模型输出的小批量数据可以被解批处理,独立进行后处理,包括根据跟踪的元数据和已应用的操作将输出反转到预处理的早期步骤。[可逆变换演示]

增强可用性#

此外,还提供了 DataStats 变换、dev_collate 等实用工具,以及作为 PyTorch 扩展的可视化方法,以提高整体可调试性。

数据集与数据加载#

遵循 PyTorch 的设计模式,MONAI 扩展了 DatasetDataLoader API,在领域特定可用性和流程性能方面进行了重要增强。

缓存 I/O 和变换数据以加速训练#

数据驱动方法需要进行许多(可能数千个)周期的训练数据读取和预处理。MONAI 提供了基于多线程缓存的数据集来加速这一过程[数据集实验]。缓存可以是持久的、动态的(SmartCacheDataset),并且可以在不同实验中重复使用[SmartCache 示例]。下图展示了与常规 PyTorch 程序相比,训练速度的提升。

cachedataset speed

ThreadDataLoader 对比 DataLoader#

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

threaddataloader

一个 ThreadDataLoader 示例在 脾脏快速训练教程 中。

公共数据集#

为了快速开始使用流行的训练数据,MONAI 提供了几个即时可用的 Dataset 类(如 MedNISTDatasetDecathlonDatasetTciaDataset),这些类支持数据下载,并支持使用变换生成训练/评估分割。[公共数据集教程] 预定义数据集的通用工作流程

pre-defined dataset

数据集类型扩展#

Dataset API 的其他扩展包括:用于关联多个数据源的 ZipDataset,用于处理图像级和图像块级预处理的 PatchDataset,用于多模态输入的 CSVDataset,以及用于交叉验证数据准备的 partition_dataset

可微分组件、网络、损失函数与优化器#

一些深度神经网络架构已被证明在医学影像分析任务中特别有效。MONAI 实现了参考网络,旨在兼顾灵活性和代码可读性。

预定义层和块#

网络层和块通常被实现为与空间 1D、2D 和 3D 输入兼容。用户可以轻松地将这些层、块和网络集成到他们定制的流程中。提供了各种实用工具来利用现有的模型权重,例如从 MMAR 进行微调 或从 MONAI model-zoo 中的捆绑包进行微调

C++/CUDA 优化模块#

为了进一步加速领域特定的例程,MONAI 引入了 C++/CUDA 实现作为 PyTorch 原生实现的扩展。MONAI 使用 PyTorch 构建 C++ 扩展的两种方式 提供模块

  • 通过 setuptools,用于包括 ResamplerConditional random field (CRF)Fast bilateral filtering using the permutohedral lattice 在内的模块。

  • 通过即时 (JIT) 编译,用于 Gaussian mixtures 模块。这种方法允许根据用户指定的参数和本地系统环境进行动态优化。下图显示了 MONAI 的高斯混合模型应用于组织和手术工具分割的结果:高斯混合模型作为后处理步骤

损失函数与优化器#

文献中常用的各种应用损失函数已被(重新)实现,例如 DiceLossGeneralizedDiceLossTverskyLossDiceFocalLoss。数值优化和相关实用工具包括 NovogradLearningRateFinder。下图展示了学习率搜索过程。

learning rate finder plot

评估#

为了运行模型推理和评估模型质量,MONAI 提供了相关广泛使用方法的参考实现。目前,包含了几种流行的评估指标和推理模式

滑动窗口推理#

对于大体积的模型推理,滑动窗口方法是一个流行的选择,可以在实现高性能的同时具有灵活的内存需求(_另外,请查看关于模型并行训练的最新研究)。它还支持 overlapblending_mode 配置来处理重叠窗口,以获得更好的性能。

sliding window scheme

医学任务指标#

已实现了各种有用的评估指标来衡量医学影像特定模型的质量。这些指标包括 Mean DiceROCAUCConfusion MatricesHausdorff DistanceSurface DistanceOcclusion Sensitivity。API 还支持多进程计算

报告生成#

提供了 MetricsSaver 来写入最终的指标汇总报告:mean(均值)、median(中位数)、max(最大值)、min(最小值)、<int>percentile(百分位数)、std(标准差)。

metrics report example

可视化#

除了简单的点和曲线绘制,还提供了直观的接口,用于在 TensorBoard 中将多维数据可视化为 GIF 动画。这可以通过可视化体积输入、分割图和中间特征图等,提供对模型的快速定性评估。一个带有可视化的可运行示例可在 UNet 训练示例 中找到。为了与 ignite 程序一起工作,MONAI 还提供了几个 ignite handler,用于使用 TensorBoardMLFlow 可视化训练曲线和指标,更多详细信息可在 TensorBoard 和 MLFlow handlers 示例 中找到。

为了方便地将 3D 图像可视化为 2D 图像帧,MONAI 基于 matplotlib 库提供了 matshow3d 实用工具。它可以绘制指定维度上的图像帧,以脾脏 3D 图像为例:matshow3d(volume=image, figsize=(100, 100), every_n=10, frame_dim=-1 show=True, cmap="gray")

matshow3d example

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

blend example

有关 TensorBoard utilitymatshow3dblend_images 的更多详细信息,请查看可视化教程

为了可视化已训练分类模型的类别激活图,MONAI 为 2D 和 3D 模型提供了 CAM、GradCAM、GradCAM++ API

CAM visualization example

以上示例是通过计算肺部 CT 病灶分类模型的 GradCAM/GradCAM++ 生成的。

工作流#

MONAI 引擎和工作流支持快速启动训练和评估实验。

这些功能将领域特定组件与通用机器学习过程解耦。它们还为更高层应用(如 AutoML、联邦学习)提供了一套统一的 API。工作流的训练器和评估器与 pytorch-ignite 的 EngineEvent-Handler 机制兼容。

通用工作流流程#

工作流和一些 MONAI 事件 handler 如下所示[工作流示例]

workflow pipeline

EnsembleEvaluator#

一个典型的集成过程被实现为一个即用型工作流[交叉验证和模型集成教程]

  1. 将所有训练数据集分割成 K 折。

  2. 使用每 K-1 折的数据训练 K 个模型。

  3. 使用所有 K 个模型在测试数据上执行推理。

  4. 计算加权平均值或投票选择最常见的值作为最终结果。

model ensemble

解批处理数据以实现灵活的后处理#

从 MONAI v0.6 开始引入了 decollate batch,它简化了后处理变换,并为具有各种数据形状的批次数据提供了灵活的后续操作。它可以将批处理数据(例如模型预测)解批处理为张量列表,好处包括

  1. 使后处理变换能够独立应用于每个项——随机变换可以对批次中的每个预测项以不同方式应用。

  2. 简化了变换 API 并减少了输入验证负担,因为预处理和后处理变换现在都只需要支持“通道优先”的输入格式。

  3. 为预测和形状不同的反转数据启用 Invertd 变换,因为数据项是列表中的元素,而不是堆叠在单个张量中。

  4. 在灵活的指标计算中同时允许批次优先张量和通道优先张量列表。[解批处理教程]

典型的 decollate batch 过程如下图所示(以 batch_size=N 的模型预测和标签为例)

decollate_batch

捆绑包#

MONAI 捆绑包的目标是定义一个打包模型,其中包括允许用户和程序理解模型如何使用以及用于何种目的所需的关键信息。一个捆绑包包含单个网络的存储权重,形式为 pickled 状态字典,可选地还包括一个 Torchscript 对象和/或一个 ONNX 对象。还包括额外的 JSON 文件,用于存储有关模型的元数据、用于构建训练、推理和后处理变换序列的信息、纯文本描述、法律信息以及模型创建者希望包含的其他数据。更多详细信息可在捆绑包规范中找到。

捆绑包的关键优势在于定义模型包,并通过结构化配置支持构建基于 Python 的工作流

  • 自包含的模型包包含所有必要信息。

  • 结构化配置可用于轻松重建或原型化深度学习工作流。

  • 通过将参数设置与 Python 代码分离,配置文件可以提供良好的可读性和可用性。

  • 配置文件可以描述灵活的工作流和组件,允许不同的底层 Python 实现

  • 联邦学习和 AutoML 等更高层次的学习范式可以与组件细节解耦。

典型的捆绑包示例可以包括

  ModelName
  ┣━ configs
  ┃  ┗━ metadata.json
  ┣━ models
  ┃  ┣━ model.pt
  ┃  ┣━ *model.ts
  ┃  ┗━ *model.onnx
  ┗━ docs
     ┣━ *README.md
     ┗━ *license.txt

有关捆绑包配置定义、语法和示例的详细信息,请参阅配置语法。一个分步的入门教程 notebook 可以帮助用户快速设置捆绑包。 [捆绑包示例, model-zoo]

联邦学习#

federated-learning

利用 MONAI 捆绑包配置,我们可以使用 MONAI 的 MonaiAlgo 类(它是联邦学习 (FL) 抽象类 ClientAlgo 的一个实现),来执行 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

Auto3DSeg 是一个针对大规模 3D 医学影像分割的综合解决方案。它利用 MONAI 和 GPU 的最新进展,高效地开发和部署具有最先进性能的算法。它首先分析数据集的全局信息,如强度、维度和分辨率,然后根据数据统计和算法模板生成 MONAI 捆绑包格式的算法。接下来,所有算法开始模型训练,以获得具有最佳验证性能的检查点。最后,集成模块通过对训练过的检查点进行排序来选择算法,并创建集成预测。

该解决方案为初学者和高级研究人员提供不同级别的用户体验。它已在不同模态的大规模 3D 医学影像数据集上进行了测试。

GPU 加速、性能分析与优化#

MONAI 提供了最先进的性能优化方法,包括

自动混合精度 (AMP)#

在训练或评估期间,只需在 SupervisedTrainerSupervisedEvaluator 中设置 amp=True/False 即可启用/禁用 AMP。示例基准测试结果如下[AMP 训练教程]

在搭载 CUDA 11 和 PyTorch 1.6 的 NVIDIA V100 GPU 上开启/关闭 AMP 进行训练

amp v100 results

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

amp a100 results

MONAI 可以与 DLProfNsightNVTXNVML 等多种工具一起使用,以识别性能瓶颈。[性能分析教程]

分布式训练#

MONAI 的分布式数据并行 API 与 PyTorch 原生分布式模块、pytorch-ignite 分布式模块、Horovod、XLA 和 SLURM 平台兼容。[分布式训练教程]

distributed training results

快速训练教程 结合了 AMPCacheDatasetGPU cacheGPU transformsThreadDataLoader 以及网络和优化器的调优,与 PyTorch 的常规实现相比,可以实现显著的加速。