M3D：利用多模态大型语言模型推进3D医学图像分析

<a href=https://arxiv.org/abs/2404.00578>**论文** | 数据 | 模型 | 训练 | 基准测试 | 在线演示

M3D是首创且全面的3D医学分析多模态大型语言模型系列工作，包括：

M3D-Data：最大规模的开源3D医学数据集，包含120K图像-文本对和662K指令-回复对；
M3D-LaMed：具有M3D-CLIP预训练视觉编码器的多功能多模态模型，能够执行图像-文本检索、报告生成、视觉问答、定位和分割等任务；
M3D-Bench：最全面的自动评估基准，涵盖8个任务。

通知

📢 [2024.06.12]

🔥🔥🔥 我们发布了一个在线演示。欢迎大家现在就来试试！
我们发布了一个轻量但强大的模型，M3D-LaMed-Phi-3-4B。初步测试表明，它的表现优于M3D-LaMed-Llama-2-7B。我们正在进行详细实验。请先尝试使用。
我们发现之前的M3D-LaMed-Llama-2-7B模型在分割任务中存在问题。我们已修复此问题，并将在接下来几天内重新发布新模型。

新闻

[2024.06.14] 我们发布了一个轻量但强大的模型M3D-LaMed-Phi-3-4B和一个在线演示。
[2024.04.28] 我们发布了数据、代码和模型。

快速开始

在这里，我们可以基于Hugging Face轻松使用我们的模型。

import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import simple_slice_viewer as ssv
import SimpleITK as sikt

device = torch.device('cuda') # 'cpu', 'cuda'
dtype = torch.bfloat16 # 或 bfloat16, float16, float32

model_name_or_path = 'GoodBaiBai88/M3D-LaMed-Phi-3-4B'
proj_out_num = 256

# 准备你的3D医学图像：
# 1. 图像形状需要处理为1*32*256*256，考虑调整大小等方法。
# 2. 图像需要归一化到0-1，考虑最小-最大归一化。
# 3. 图像格式需要转换为.npy
# 4. 虽然我们没有在2D图像上训练，但理论上2D图像可以插值到1*32*256*256的形状作为输入。
image_path = "./Data/data/examples/example_03.npy"

model = AutoModelForCausalLM.from_pretrained(
    model_name_or_path,
    torch_dtype=dtype,
    device_map='auto',
    trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(
    model_name_or_path,
    model_max_length=512,
    padding_side="right",
    use_fast=False,
    trust_remote_code=True
)

model = model.to(device=device)

# question = "你能为这个医学图像提供一个包含发现的描述吗？"
question = "这张图像中的肝脏在哪里？请输出分割掩码。"
# question = "这张图像中的肝脏在哪里？请输出边界框。"

image_tokens = "<im_patch>" * proj_out_num
input_txt = image_tokens + question
input_id = tokenizer(input_txt, return_tensors="pt")['input_ids'].to(device=device)

image_np = np.load(image_path)
image_pt = torch.from_numpy(image_np).unsqueeze(0).to(dtype=dtype, device=device)

# generation = model.generate(image_pt, input_id, max_new_tokens=256, do_sample=True, top_p=0.9, temperature=1.0)
generation, seg_logit = model.generate(image_pt, input_id, seg_enable=True, max_new_tokens=256, do_sample=True, top_p=0.9, temperature=1.0)

generated_texts = tokenizer.batch_decode(generation, skip_special_tokens=True)
seg_mask = (torch.sigmoid(seg_logit) > 0.5) * 1.0

print('问题', question)
print('生成的文本', generated_texts[0])

image = sikt.GetImageFromArray(image_np)
ssv.display(image)
seg = sikt.GetImageFromArray(seg_mask.cpu().numpy()[0])
ssv.display(seg)

模型

模型	下载链接
M3D-CLIP	HuggingFace, ModelScope
M3D-LaMed-Phi-3-4B	HuggingFace, ModelScope
M3D-LaMed-Llama-2-7B	HuggingFace, ModelScope

安装

git clone https://github.com/BAAI-DCAI/M3D.git
pip install -r requirements.txt

数据

M3D-Data支持训练和基准测试，包含4种类型的数据：

数据集	类型	图像	文本	下载链接
M3D-Cap	3D图像-文本对	120,092	42,496	HuggingFace, ModelScope
M3D-VQA	3D图像、问题和答案	96,170	509,755	HuggingFace, ModelScope
M3D-Seg	3D图像、类别文本和分割掩码	5,772	149,196	HuggingFace, ModelScope
M3D-RefSeg	3D图像、问题、答案和分割掩码	210	2,778	HuggingFace, ModelScope

请按照每个数据集的说明进行下载和预处理。您可以在数据集描述或Data/process/中找到名为m3d_xx_data_prepare.py的预处理文件，用于M3D-Cap、M3D-Seg和M3D-RefSeg。我们建议将下载和处理后的数据集保存到Data/data/。

训练

预训练权重

要训练M3D-LaMed，您需要准备一些预训练权重以获得更好的性能和更快的收敛速度。

视觉编码器

我们建议直接从M3D-CLIP下载医学3D ViT权重pretrained_ViT.bin。或者您也可以通过以下方式自行预训练3D ViT：

sh LaMed/script/train_clip.sh

LLM

Phi-3-4B：下载并按照这里的说明操作。 Llama-2-7B：下载并按照这里的说明操作。

分割模块

SegVol：下载并按照这里的说明操作。

训练

我们的训练包括两个步骤。

步骤1：预训练
步骤2：视觉指令微调

配置

我们建议使用accelerate进行训练。它由Hugging Face开发，方便支持常见的训练策略，如分布式训练、混合精度、DeepSpeed等。首次使用时需要进行配置：

accelerate config

请按照配置指南选择适合的训练策略。我们建议使用bf16和Deepspeed来加速，ZeRO类型则取决于您的具体情况。

如果您不知道如何配置，我们提供了一个简单的配置default_config.yaml供您参考。

default_config.yaml

compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
  gradient_accumulation_steps: 1
  zero3_init_flag: false
  zero_stage: 0
distributed_type: DEEPSPEED
downcast_bf16: 'no'
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

步骤1：预训练

我们应该使用图像-文本数据来对齐视觉和语言，即只训练mm_projector并冻结视觉编码器和LLM。

请分别更新LLM路径--model_name_or_path和视觉编码器路径--pretrain_vision_model。更新--output_dir以指定模型的输出路径。然后运行以下脚本：

sh LaMed/script/pretrain.sh

步骤2：视觉指令微调

通过图像-文本对、VQA、定位和分割的多任务数据进行视觉指令微调，即只对LLM执行LoRA训练，并解冻所有其他模型。

请分别更新LLM路径--model_name_or_path、视觉编码器路径--pretrain_vision_model、步骤1保存的模型路径--pretrain_mm_mlp_adapter和分割模块路径--pretrain_seg_module。更新--output_dir以指定模型的输出路径。然后运行以下脚本：

sh LaMed/script/finetune_lora.sh

合并LoRA权重

合并model_with_lora.bin的LoRA权重，将最终模型以Hugging Face格式保存到您指定的路径：

CUDA_VISIBLE_DEVICES="" python merge_lora_weights_and_save_hf_model.py \
  --version="" \
  --model_type="" \
  --model_with_lora="PATH_TO_model_with_lora.bin" \
  --output_dir="PATH_TO_SAVED_MODEL"

基准测试

我们提出了最全面的自动评估基准，涵盖3D医疗领域的8个任务，包括图像-文本检索、报告生成、封闭式VQA、开放式VQA、指代表达理解、指代表达生成、语义分割、指代表达分割。

评估

我们可以通过运行以下命令直接评估每个任务：

CUDA_VISIBLE_DEVICES="" python Bench/eval/eval_TASK.py

我们还提供了使用LLM进行更准确的报告生成任务自动评估，修改file_path后，请运行：

CUDA_VISIBLE_DEVICES="" python Bench/eval/eval_with_llm.py

引用

如果我们的数据集或项目对您有帮助，请考虑引用：

@misc{bai2024m3d,
      title={M3D: Advancing 3D Medical Image Analysis with Multi-Modal Large Language Models}, 
      author={Fan Bai and Yuxin Du and Tiejun Huang and Max Q. -H. Meng and Bo Zhao},
      year={2024},
      eprint={2404.00578},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

@article{du2023segvol,
  title={SegVol: Universal and Interactive Volumetric Medical Image Segmentation},
  author={Du, Yuxin and Bai, Fan and Huang, Tiejun and Zhao, Bo},
  journal={arXiv preprint arXiv:2311.13385},
  year={2023}
}

致谢

我们感谢以下开源项目： LLaVA、 LISA、 SegVol。