说明文档
LVM-Med:通过二阶图匹配学习医学影像大规模自监督视觉模型(NeurIPS 2023)
我们发布了 LVM-Med 的 PyTorch 预训练模型,并展示了在 2D-3D 分割、线性/全量微调图像分类和目标检测等下游任务上的应用。
LVM-Med 使用从 55 个数据集收集的约 130 万张医学图像进行训练,采用二阶图匹配方法,统一了当前的对比学习和基于实例的自监督学习方法。
<p align="center"> <img src="assets/body_lvm_med.jpg" alt="drawing" width="650"/> </p>
<p align="center"> <img src="assets/lvm_med_teaser.gif" alt="drawing" width="800"/> </p>
目录
新闻
- 2023年12月14日:LVM-Med 训练算法已准备好发布!请发送邮件给我们申请!
- 如果您需要其他架构,请通过邮件或创建 Issue 发送请求。如果请求足够多,我们将进行训练。
- 即将推出:由 LVM-Med 训练的 ConvNext 架构。
- 即将推出:用于端到端分割的 ViT 架构,论文中报告了更好的性能。
- 2023年7月31日:发布 LVM-Med ResNet50 和 LVM-Med ViT 作为骨干网络的 ONNX 支持,位于
onnx_model文件夹中。 - 2023年7月26日:我们发布了从 LVM-Med 初始化并在 LIVECell 数据集(包含 160 万个高质量细胞)上进一步训练的 ViT 架构(ViT-B 和 ViT-H)。请参阅此表格。
- 2023年6月25日:我们发布 LVM-Med 的两个预训练模型:ResNet-50 和 ViT-B。提供下游任务的脚本。
LVM-Med 预训练模型
<table> <tr> <th>架构</th> <th>参数量 (M)</th> <th> 2D 分割 (Dice) </th> <th> 3D 分割 (3D IoU) </th> <th>权重</th> </tr> <tr> <td>ResNet-50</td> <td>25.5M</td> <td>83.05</td> <td>79.02</td> <td> <a href="https://drive.google.com/file/d/11Uamq4bT_AbTf8sigIctIAnQJN4EethW/view?usp=sharing">骨干网络</a> </td> </tr> <tr> <td>ViT-B</td> <td>86.0M</td> <td>85.80</td> <td>80.90</td> <td> <a href="https://drive.google.com/file/d/17WnE34S0ylYiA3tMXobH8uUrK_mCVPT4/view?usp=sharing">骨干网络</a> </td> </tr> </table>
下载预训练模型后,请将其放置在 lvm_med_weights 文件夹中以便使用。
- 对于 Resnet-50,我们展示 端到端 分割/分类/目标检测。
- 对于 ViT-B,我们展示使用边界框的 基于提示的 分割。
重要提示: 请查看 dataset.md 以避免在使用我们的模型时可能出现的测试数据泄露。
Segment Anything Model 相关实验
- 对于所有使用 SAM 模型的实验,我们使用 SAM 的基础架构
sam_vit_b。您可以在original repo浏览此预训练权重,并将其放在./working_dir/sam_vit_b_01ec64.pth文件夹中以便正确使用 yaml 配置。
在大数据集上进一步训练 LVM-Med
我们发布了下表中提到的其他大数据集上的进一步预训练权重。 <table> <tr> <th>架构</th> <th>参数量 (M)</th> <th>数据集名称 </th> <th>权重</th> <th>描述</th> </tr> <tr> <td>ViT-B</td> <td>86.0M</td> <td> <a href="https://www.nature.com/articles/s41592-021-01249-6">LIVECell</a> </td> <td> <a href="https://drive.google.com/file/d/1SxaGXQ4FMbG8pS2zzwTIXXgxF4GdwyEU/view?usp=sharing">骨干网络</a> </td> <td> <a href="https://github.com/duyhominhnguyen/LVM-Med/blob/main/further_training_lvm_med/README.md">链接</a></td> </tr> <tr> <td>ViT-H</td> <td>632M</td> <td> <a href="https://www.nature.com/articles/s41592-021-01249-6">LIVECell</a> </td> <td> <a href="https://drive.google.com/file/d/14IhoyBXI9eP9V2xeOV2-6LlNICKjzBaJ/view?usp=sharing">骨干网络</a> </td> <td> <a href="https://github.com/duyhominhnguyen/LVM-Med/blob/main/further_training_lvm_med/README.md">链接</a></td> </tr> </table>
环境配置
代码需要 python>=3.8,以及 pytorch>=1.7 和 torchvision>=0.8。请按照此处的说明安装 PyTorch 和 TorchVision 依赖。强烈建议安装支持 CUDA 的 PyTorch 和 TorchVision。
要设置我们的项目,请运行以下命令:
git clone https://huggingface.co/duynhm/LVM-Med
cd LVM-Med
conda env create -f lvm_med.yml
conda activate lvm_med
要使用 ResNet-50 进行分割微调,我们使用 segmentation-models-pytorch 包中的 U-Net。要安装此库,您可以执行以下操作:
git clone https://github.com/qubvel/segmentation_models.pytorch.git
cd segmentation_models.pytorch
pip install -e
cd ..
mv segmentation_models_pytorch_example/encoders/__init__.py segmentation_models.pytorch/segmentation_models_pytorch/__init__.py
mv segmentation_models_pytorch_example/encoders/resnet.py segmentation_models.pytorch/segmentation_models_pytorch/resnet.py
准备数据集
脑肿瘤数据集
您可以通过 Kaggle 的 Brain Tumor Classification (MRI) 下载 Brain 数据集,并将名称更改为 BRAIN。
VinDr 数据集
您可以从VinDr链接下载数据集,并将 vinbigdata 文件夹放入 object_detection 文件夹中。要构建数据集,下载数据集后,您可以运行 object_detection 文件夹内的 convert_to_coco.py 脚本。
python convert_to_coco.py # 注意,请检查代码中第146和158行的链接以正确构建数据集
更多信息可以在 object_detection 中找到。
其他数据集
首先,您应该将需要运行的相应数据集下载到 dataset_demo 文件夹中。为了获得与您的工作尽可能接近的结果,您可以按照我们的方式准备一些特定数据集(非预分发的):
python prepare_dataset.py -ds [dataset_name]
其中:dataset_name 是您想要准备的数据集名称。之后,您应该在 dataloader/yaml_data 中将加载的数据集路径更改为我们预定义的 yaml 文件。
目前支持 Kvasir、BUID、FGADR、MMWHS_MR_Heart 和 MMWHS_CT_Heart。
注意: 您应该将数据集名称更改为正确的格式(即 Kvasir、BUID),与我们当前支持的数据集名称一致。否则,它将无法按预期工作。
下游任务
分割
1. 端到端分割
a) 训练阶段:
使用 ResNet-50 微调下游任务
python train_segmentation.py -c ./dataloader/yaml_data/buid_endtoend_R50.yml
在 ./dataloader/yaml_data/ 中更改 .yml 配置中的数据集名称以进行其他实验。
注意:要使用 ResNet-50 应用分割模型(2D 或 3D),我们建议通过设置以下内容来标准化梯度以保持训练阶段的稳定性:
clip_value = 1
torch.nn.utils.clip_grad_norm_(net.parameters(), clip_value)
请参阅文件 /segmentation_2d/train_R50_seg_adam_optimizer_2d.py 第 129-130 行中的示例。
b) 推理:
ResNet-50 版本
python train_segmentation.py -c ./dataloader/yaml_data/buid_endtoend_R50.yml -test
对于使用 SAM 的 ViT 的端到端版本,我们将很快发布比论文中报告的结果更好的版本。
2. 使用 ViT-B 的基于提示的分割
a. 使用 SAM 微调解码器的基于提示的分割(MedSAM)。
我们运行 MedSAM 基线来比较性能:
训练
python3 medsam.py -c dataloader/yaml_data/buid_sam.yml
推理
python3 medsam.py -c dataloader/yaml_data/buid_sam.yml -test
b. 与 MedSAM 类似的基于提示的分割,但使用 LVM-Med 的编码器。
训练脚本与 MedSAM 情况类似,但通过 -lvm_encoder 指定权重模型。
训练
python3 medsam.py -c dataloader/yaml_data/buid_lvm_med_sam.yml -lvm_encoder ./lvm_med_weights/lvmmed_vit.pth
测试
python3 medsam.py -c dataloader/yaml_data/buid_lvm_med_sam.yml -lvm_encoder ./lvm_med_weights/lvmmed_vit.pth -test
您还可以查看我们的示例笔记本 Prompt_Demo.ipynb,使用基于提示的 MedSAM 和带有 LVM-Med 编码器的基于提示的 SAM 进行结果可视化。演示中每个 SAM 解码器模型的预训练权重在这里。请下载 LVM-Med 和 MedSAM 的训练模型,并将它们放入 working_dir/checkpoints 文件夹中,以运行上述笔记本文件。
c. 使用 Segment Anything Model (SAM) 进行下游任务的零样本基于提示的分割
无需任何微调,使用基于边界框提示的 SAM 模型可以通过以下方式完成:
python3 zero_shot_segmentation.py -c dataloader/yaml_data/buid_sam.yml
图像分类
我们提供使用 LVM-Med 的 ResNet-50 模型进行脑肿瘤分类和 FGADR 数据集糖尿病视网膜病变分级(论文正文表 5 和附录表 12)的训练和测试脚本。ViT 模型版本将很快更新。
a. 使用 FGADR 训练
# 使用 1 个全连接层全量微调
python train_classification.py -c ./dataloader/yaml_data/fgadr_endtoend_R50_non_frozen_1_fcn.yml
# 使用多个全连接层全量微调
python train_classification.py -c ./dataloader/yaml_data/fgadr_endtoend_R50_non_frozen_fcns.yml
# 冻结全部,仅微调 1 层全连接层
python train_classification.py -c ./dataloader/yaml_data/fgadr_endtoend_R50_frozen_1_fcn.yml
# 冻结全部,仅微调多层全连接层
python train_classification.py -c ./dataloader/yaml_data/fgadr_endtoend_R50_frozen_fcns.yml
要运行 Brain dataset,请在 ./dataloader/yaml_data/ 文件夹中选择其他配置文件 brain_xyz.yml。
b. 使用 FGADR 推理
# 使用 1 个全连接层全量微调
python train_classification.py -c ./dataloader/yaml_data/fgadr_endtoend_R50_non_frozen_1_fcn.yml -test
# 使用多个全连接层全量微调
python train_classification.py -c ./dataloader/yaml_data/fgadr_endtoend_R50_non_frozen_fcns.yml -test
# 冻结全部,仅微调 1 层全连接层
python train_classification.py -c ./dataloader/yaml_data/fgadr_endtoend_R50_frozen_1_fcn.yml -test
# 冻结全部,仅微调多层全连接层
python train_classification.py -c ./dataloader/yaml_data/fgadr_endtoend_R50_frozen_fcns.yml -test
目标检测
我们展示使用 LVM-Med ResNet-50 在 Vin-Dr 数据集上进行目标检测。我们使用 Faster-RCNN 作为网络骨干。
您可以访问 object_detection 文件夹了解更多详情。
引用
如果这项工作对您的研究有帮助,请引用本文:
@article{nguyen2023lvm,
title={LVM-Med: Learning Large-Scale Self-Supervised Vision Models for Medical Imaging via Second-order Graph Matching},
author={Nguyen, Duy MH and Nguyen, Hoang and Diep, Nghiem T and Pham, Tan N and Cao, Tri and Nguyen, Binh T and Swoboda, Paul and Ho, Nhat and Albarqouni, Shadi and Xie, Pengtao and others},
journal={arXiv preprint arXiv:2306.11925},
year={2023}
}
相关工作
我们使用并修改了 SAM 和 MedSAM 的代码用于基于提示的分割设置。LVM-Med 算法的一部分采用了 Vicregl、Deepcluster-v2 的数据变换。我们还利用 vissl 框架在我们收集的数据上训练 2D 自监督方法。感谢作者们的杰出工作!
许可证
根据 CC BY-NC-ND 2.0(署名-非商业性使用-禁止演绎 2.0 通用)许可。代码仅用于学术研究。如需商业用途,请联系 Ho_Minh_Duy.Nguyen@dfki.de
duynhm/LVM-Med
作者 duynhm
创建时间: 2023-12-17 13:49:22+00:00
更新时间: 2023-12-17 15:04:10+00:00
在 Hugging Face 上查看