在脑部 MRI 数据集上微调 MedGemma

MedGemma 是一组 Gemma 3 变体，旨在在医学文本与图像理解方面表现出色。目前该集合包含两个强大的变体：一款 4B 多模态 版本和一款 27B 纯文本 版本。

MedGemma 4B 模型将 SigLIP 图像编码器（在多样且去标识化的医学数据集上预训练，如胸部 X 光、皮肤病图像、眼科图像和组织病理切片）与在大规模医学数据上训练的大语言模型（LLM）相结合。

在本教程中，我们将学习如何在脑部 MRI 数据集上微调 MedGemma 4B 模型，以完成图像分类任务。目标是让更小的 MedGemma 4B 模型能够高效分类脑部 MRI 扫描，并以更高的准确率预测脑癌。

搭建环境

RunPod 是运行 GPU 工作负载的优秀平台，提供预配置环境并支持 JupyterLab。您可以启动一个 pod，立即使用偏好的编辑器开始编码。以下是环境的搭建步骤：

1. 设置 RunPod

登录 RunPod，创建一个配备 A100 GPU 和最新 PyTorch 镜像的 pod。然后点击“Deploy On-Demand”按钮启动 pod。

编辑 pod，添加以下用于集成 Hugging Face 与 Kaggle 的环境变量：

• HF_TOKEN
• KAGGLE_USERNAME
• KAGGLE_KEY

将存储容量增加到 100GB，以容纳数据集和模型检查点

2. 安装 Python 包

pod 运行后，点击“Connect”启动 JupyterLab 实例。创建一个新的 Python 笔记本，运行以下命令安装所需的 Python 包：

! pip install --upgrade --quiet transformers bitsandbytes datasets evaluate peft trl scikit-learn kaggle技术支持

加载与处理数据集

我们将使用 Kaggle 上的 Brain Cancer MRI 数据集。

3. 下载数据集

使用 Kaggle CLI 下载并解压数据集：

!kaggle datasets download -d orvile/brain-cancer-mri-dataset --unzip技术支持

Dataset URL: https://www.kaggle.com/datasets/orvile/brain-cancer-mri-dataset
License(s): CC-BY-SA-4.0
Downloading brain-cancer-mri-dataset.zip to /workspace
 90%|████████████████████████████████████▉    | 130M/144M [00:00<00:00, 231MB/s]
100%|█████████████████████████████████████████| 144M/144M [00:01<00:00, 108MB/s]技术支持

4. 加载数据集

将其加载为 Hugging Face 数据集，划分训练与验证集，并展示数据结构：

from datasets import load_dataset

data_dir = "./Brain_Cancer raw MRI data/Brain_Cancer"

# Define proportions for train and validation splits
train_size = 0.8  
validation_size = 0.2  


data = load_dataset("imagefolder", data_dir=data_dir, split="train")

# Split the dataset into train and validation sets
data = data.train_test_split(
    train_size=train_size,
    test_size=validation_size,
    shuffle=True,
    seed=42,
)

# Rename the 'test' split to 'validation'
data["validation"] = data.pop("test")

# Display dataset details
print(data)技术支持

输出：

DatasetDict({
    train: Dataset({
        features: ['image', 'label'],
        num_rows: 4844
    })
    validation: Dataset({
        features: ['image', 'label'],
        num_rows: 1212
    })
})技术支持

查看训练集中的第一张图像及其对应标签：

data["train"][0]["image"]技术支持

print(data["train"][0]["label"])技术支持

1技术支持

5. 处理数据集

在处理数据集之前，务必先检查标签名称，以确保分类任务的正确处理。

BRAIN_CANCER_CLASSES = data["train"].features["label"].names
print("Detected classes:", BRAIN_CANCER_CLASSES)技术支持

检测到的类别为：['brain_glioma', 'brain_menin', 'brain_tumor']。

Detected classes: ['brain_glioma', 'brain_menin', 'brain_tumor']技术支持

为提升分类流程，我们将为这些类别标签添加前缀（A、B、C），以便更好地组织并与自定义提示格式对齐。

BRAIN_CANCER_CLASSES = ['A: brain glioma', 'B: brain menin', 'C: brain tumor']技术支持

接下来，我们创建一个自定义提示词，在微调时用于引导模型。提示中包含更新后的类别标签。

options = "\n".join(BRAIN_CANCER_CLASSES)
PROMPT = f"What is the most likely type of brain cancer shown in the MRI image?\n{options}"技术支持

为准备数据集进行微调，我们将创建一个名为“messages”的新列。该列包含结构化数据，表示用户查询（提示）与助手响应（正确标签）。

def format_data(example: dict[str, any]) -> dict[str, any]:
    example["messages"] = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                },
                {
                    "type": "text",
                    "text": PROMPT,
                },
            ],
        },
        {
            "role": "assistant",
            "content": [
                {
                    "type": "text",
                    "text": BRAIN_CANCER_CLASSES[example["label"]],
                },
            ],
        },
    ]
    return example

# Apply the formatting to the dataset
formatted_data = data.map(format_data)技术支持

为验证格式是否正确，我们可以查看首个样本的 message 列：

formatted_data["train"][0]["messages"]技术支持

得到的格式化数据点如下所示：

[{'content': [{'text': None, 'type': 'image'},
   {'text': 'What is the most likely type of brain cancer shown in the MRI image?\nA: brain glioma\nB: brain menin\nC: brain tumor',
    'type': 'text'}],
  'role': 'user'},
 {'content': [{'text': 'B: brain menin', 'type': 'text'}],
  'role': 'assistant'}]技术支持

在脑部 MRI 数据集上微调 MedGemma

本节将把 MedGemma 4B Instruct 模型在脑部 MRI 数据集上进行微调。过程包括下载模型与处理器、设置 LoRA 适配器、配置训练器并启动训练。

6. 加载模型与分词器

由于 MedGemma 是门控模型，您需要使用 API 密钥登录 Hugging Face CLI。这也便于将微调后的模型保存到 Hugging Face Hub。若需复习，请参阅我们的课程 Working with Hugging Face: Your Guide to the Hub。

from huggingface_hub import login
import os

hf_token = os.environ.get("HF_TOKEN")
login(hf_token)技术支持

我们使用 Transformers 库加载 MedGemma 4B Instruct 模型及其处理器。模型被配置为使用 bfloat16 精度，以便在 GPU 上高效计算。

import torch
from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig

model_id = "google/medgemma-4b-it"

# Check if GPU supports bfloat16
if torch.cuda.get_device_capability()[0] < 8:
    raise ValueError("GPU does not support bfloat16, please use a GPU that supports bfloat16.")

model_kwargs = dict(
    attn_implementation="eager",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs)
processor = AutoProcessor.from_pretrained(model_id)

# Use right padding to avoid issues during training
processor.tokenizer.padding_side = "right"技术支持

7. 设置模型

为高效微调 MedGemma 4B Instruct 模型，我们将使用低秩适配（LoRA），这是一种参数高效的微调方法。

借助 LoRA，我们只需训练少量额外参数即可适配大型模型，大幅降低计算成本，同时保持性能。

from peft import LoraConfig

peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.05,
    r=16,
    bias="none",
    target_modules="all-linear",
    task_type="CAUSAL_LM",
    modules_to_save=[
        "lm_head",
        "embed_tokens",
    ],
)技术支持

为在训练中同时处理图像与文本输入，我们定义自定义的整理（collation）函数。该函数将数据集样本处理为适合模型的格式，包括对文本分词及图像预处理。

def collate_fn(examples: list[dict[str, any]]):
    texts = []
    images = []
    for example in examples:
        images.append([example["image"]])
        texts.append(
            processor.apply_chat_template(
                example["messages"], add_generation_prompt=False, tokenize=False
            ).strip()
        )

    # Tokenize the texts and process the images
    batch = processor(text=texts, images=images, return_tensors="pt", padding=True)

    # The labels are the input_ids, with the padding and image tokens masked in
    # the loss computation
    labels = batch["input_ids"].clone()

    # Mask image tokens
    image_token_id = [
        processor.tokenizer.convert_tokens_to_ids(
            processor.tokenizer.special_tokens_map["boi_token"]
        )
    ]
    # Mask tokens that are not used in the loss computation
    labels[labels == processor.tokenizer.pad_token_id] = -100
    labels[labels == image_token_id] = -100
    labels[labels == 262144] = -100

    batch["labels"] = labels
    return batch技术支持

我们使用 trl 库中的 SFTConfig 类来定义训练参数。这些参数控制微调过程，包括批大小、学习率与梯度累积步数。

from trl import SFTConfig

args = SFTConfig(
    output_dir="medgemma-brain-cancer",            
    num_train_epochs=1,                       
    per_device_train_batch_size=8,                           
    per_device_eval_batch_size=8,                            
    gradient_accumulation_steps=8,                           
    gradient_checkpointing=True,                             
    optim="adamw_torch_fused",                               
    logging_steps=0.1,                                        
    save_strategy="epoch",                                   
    eval_strategy="steps",                                   
    eval_steps=0.1,                                           
    learning_rate=2e-4,                             
    bf16=True,                                               
    max_grad_norm=0.3,                                       
    warmup_ratio=0.03,                                       
    lr_scheduler_type="linear",                              
    push_to_hub=True,                                        
    report_to="none",
    gradient_checkpointing_kwargs={"use_reentrant": False},  
    dataset_kwargs={"skip_prepare_dataset": True},           
    remove_unused_columns = False,                           
    label_names=["labels"],                                  
)技术支持

通过 SFTTrainer，我们将模型、数据集、数据整理器、训练参数与 LoRA 配置整合到一个流程中，从而简化微调过程并提升易用性。

from trl import SFTTrainer

trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=formatted_data["train"],
    eval_dataset=formatted_data["validation"].shuffle().select(range(50)), 
    peft_config=peft_config,
    processing_class=processor,
    data_collator=collate_fn,
)技术支持

8. 模型训练

当模型、数据集与训练配置准备就绪后，即可开始微调。借助 SFTTrainer，我们只需一条命令即可启动训练：

trainer.train()技术支持

训练过程大约耗时 1 小时 8 分钟。期间训练损失与验证损失随步骤稳步下降，表明模型在有效学习。

训练完成后，可使用 save_model() 方法将微调后的模型保存在本地，并推送至 Hugging Face Hub。

trainer.save_model()技术支持

模型现已发布于 kingabzpro/medgemma-brain-cancer · Hugging Face

来源 kingabzpro/medgemma-brain-cancer

模型评估

为评估 MedGemma 4B 模型的性能，我们将在验证数据集上测试基础模型与微调后的模型。流程包括清理内存、准备测试数据、生成响应，并计算准确率与 F1 分数等关键指标。

开始评估前，我们移除训练设置以释放 GPU 内存，确保测试环境干净

del model
del trainer
torch.cuda.empty_cache()技术支持

9. 测试前的设置

我们将验证集格式化为模型所需的输入结构。这涉及为每个样本创建包含用户提示的“messages”列。

def format_test_data(example: dict[str, any]) -> dict[str, any]:
    example["messages"] = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                },
                {
                    "type": "text",
                    "text": PROMPT,
                },
            ],
        },
    ]
    return example

test_data = data["validation"]
test_data = test_data.map(format_test_data)技术支持

为评估模型性能，我们使用 evaluate 库，该库为分类等任务提供了预置指标。导入库并加载所需指标后，我们从测试数据集中提取真实标签。随后定义辅助函数 compute_metrics，通过将预测与真实标签对比来计算准确率与 F1 分数。

import evaluate

accuracy_metric = evaluate.load("accuracy")
f1_metric = evaluate.load("f1")

# Ground-truth labels
REFERENCES = test_data["label"]


def compute_metrics(predictions: list[int]) -> dict[str, float]:
    metrics = {}
    metrics.update(
        accuracy_metric.compute(
            predictions=predictions,
            references=REFERENCES,
        )
    )
    metrics.update(
        f1_metric.compute(
            predictions=predictions,
            references=REFERENCES,
            average="weighted",
        )
    )
    return metrics技术支持

为确保标签处理的一致性，我们将数据集的“label”列转换为 ClassLabel 类型。这样可在标签索引与名称之间高效映射。此外，我们定义备用标签映射，以便在后处理阶段处理标签格式的差异。

from datasets import ClassLabel

test_data = test_data.cast_column(
    "label",
    ClassLabel(names=BRAIN_CANCER_CLASSES)
)

LABEL_FEATURE = test_data.features["label"]

ALT_LABELS = dict([
    (label, f"({label.replace(': ', ') ')}") for label in BRAIN_CANCER_CLASSES
])技术支持

为将模型预测映射到正确类别标签，我们定义 postprocess 函数。该函数确保预测与合适的标签匹配，同时兼容规范及替代表述。

def postprocess(prediction, do_full_match: bool = False) -> int:
    if isinstance(prediction, str):
        response_text = prediction
    else:
        response_text = prediction[0]["generated_text"]

    if do_full_match:
        return LABEL_FEATURE.str2int(response_text)

    for label in BRAIN_CANCER_CLASSES:
        # accept canonical or alternative wording
        if label in response_text or ALT_LABELS[label] in response_text:
            return LABEL_FEATURE.str2int(label)

    return -1技术支持

10. 基础模型的表现

为评估基础模型表现，我们加载预训练模型与处理器，配置生成设置，并准备测试所需的提示与图像。

import torch
from transformers import AutoModelForImageTextToText, AutoProcessor

model_kwargs = dict(
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

model = AutoModelForImageTextToText.from_pretrained(
    model_id, **model_kwargs
)

from transformers import GenerationConfig
gen_cfg = GenerationConfig.from_pretrained(model_id)
gen_cfg.update(
    do_sample          = False, 
    top_k              = None, 
    top_p              = None,
    cache_implementation = "dynamic" 
)
model.generation_config = gen_cfg

processor  = AutoProcessor.from_pretrained(args.output_dir)
tok = processor.tokenizer

model.config.pad_token_id            = tok.pad_token_id
model.generation_config.pad_token_id = tok.pad_token_id

def chat_to_prompt(chat_turns):
    return processor.apply_chat_template(
        chat_turns,
        add_generation_prompt=True,   # tells the model "your turn"
        tokenize=False                # we want raw text, not ids
    )

prompts = [chat_to_prompt(c) for c in test_data["messages"]]
images  = test_data["image"]                         # already a list of PIL images
assert len(prompts) == len(images), "1 prompt must match 1 image!"技术支持

我们创建了 batch_predict 函数以批量处理测试数据。该函数为每个提示-图像对生成预测，并进行后处理以映射到正确标签。

import torch
from typing import List, Any, Callable


def batch_predict(
    prompts,
    images,
    model,
    processor,
    postprocess,
    *,
    batch_size=64,
    device="cuda",
    dtype=torch.bfloat16,
    **gen_kwargs
):
    preds = []
    for i in range(0, len(prompts), batch_size):
        texts = prompts[i : i + batch_size]
        imgs = [[img] for img in images[i : i + batch_size]]
        enc = processor(text=texts, images=imgs, padding=True, return_tensors="pt").to(
            device, dtype=dtype
        )
        lens = enc["attention_mask"].sum(dim=1)
        with torch.inference_mode():
            out = model.generate(
                **enc,
                disable_compile=True,
                **gen_kwargs
            )
        for seq, ln in zip(out, lens):
            ans = processor.decode(seq[ln:], skip_special_tokens=True)
            preds.append(postprocess(ans))
    return preds技术支持

我们将使用 batch_predict 为基础模型在测试数据集上生成预测，然后通过 compute_metrics 进行评估。

bf_preds = batch_predict(
    model         = model,
    processor     = processor,
    prompts       = prompts,
    images        = images,
    batch_size    = 64,
    max_new_tokens= 40,          # forwarded to generate
    postprocess= postprocess, # your label-mapping function
)




bf_metrics = compute_metrics(bf_preds)
print(f"Baseline metrics: {bf_metrics}")技术支持

结果并不理想。准确率为 33%，表现较差。

Baseline metrics: {'accuracy': 0.33745874587458746, 'f1': 0.1737287617650654}技术支持

为更好地理解模型行为，我们可以针对数据集中的单个样本生成预测。这需要创建一个辅助函数来处理输入并返回模型响应。

函数 predict_one 接收一个提示与一张图像作为输入，使用模型的处理器进行处理并生成响应。该函数确保模型输出被解码为可读文本。

import torch
from typing import Union, Dict, Any, List
from transformers import AutoModelForImageTextToText, AutoProcessor


def predict_one(
    prompt,
    image,
    model,
    processor,
    *,
    device="cuda",
    dtype=torch.bfloat16,
    disable_compile=True,
    **gen_kwargs
) -> str:
    inputs = processor(text=prompt, images=image, return_tensors="pt").to(
        device, dtype=dtype
    )
    plen = inputs["input_ids"].shape[-1]
    with torch.inference_mode():
        ids = model.generate(
            **inputs,
            disable_compile=disable_compile,
            **gen_kwargs
        )
    return processor.decode(ids[0, plen:], skip_special_tokens=True)技术支持

我们使用 predict_one 为数据集中第 11 个样本生成响应。这包括准备提示并运行预测函数。

idx     = 10
chat    = test_data["messages"][idx]
prompt  = processor.apply_chat_template(
            chat,
            add_generation_prompt=True,
            tokenize=False
          )

# run the one-sample helper
answer = predict_one(
    prompt   = prompt,
    image    = test_data["image"][idx],
    model    = model,
    processor= processor,
    max_new_tokens = 40 
)

print("Model answer:", answer)技术支持

结果是一段冗长的解释，说明为何选择脑胶质瘤。该回答完全错误，连分类都错了。

Model answer: Based on the MRI image, the most likely type of brain cancer is **A: brain glioma**.

Here's why:

*   **Gliomas** are a common type of brain tumor技术支持

11. 微调模型的表现

为评估微调后的模型，我们重复评估流程，从输出目录加载模型，生成预测并计算指标。

我们从输出目录加载微调模型与处理器，运行 batch_predict 并计算指标。

model = AutoModelForImageTextToText.from_pretrained(
    args.output_dir, **model_kwargs
)
model.generation_config = gen_cfg

processor  = AutoProcessor.from_pretrained(args.output_dir)
tok = processor.tokenizer

model.config.pad_token_id            = tok.pad_token_id
model.generation_config.pad_token_id = tok.pad_token_id

af_preds = batch_predict(
    model         = model,
    processor     = processor,
    prompts       = prompts,
    images        = images,
    batch_size    = 64,
    max_new_tokens= 40,          # forwarded to generate
    postprocess= postprocess, # your label-mapping function
)
af_metrics = compute_metrics(af_preds)
print(f"Fine-tuned metrics: {af_metrics}")技术支持

这些结果相较基础模型有显著提升，凸显微调的有效性。仅用 1 个 epoch，准确率便从 33% 跃升至 89%。

Fine-tuned metrics: {'accuracy': 0.8927392739273927, 'f1': 0.892641793935792}技术支持

为进一步分析微调模型的表现，我们针对测试集中的单个样本生成预测。

idx     = 10
chat    = test_data["messages"][idx]
prompt  = processor.apply_chat_template(
            chat,
            add_generation_prompt=True,
            tokenize=False
          )

# run the one-sample helper
answer = predict_one(
    prompt   = prompt,
    image    = test_data["image"][idx],
    model    = model,
    processor= processor,
    max_new_tokens = 40              # any generate-kwargs you need
)

print("Model answer:", answer)技术支持

模型以清晰准确的方式给出结果，分类既正确又结构清晰。

Model answer: C: brain tumor技术支持

如在运行上述代码时遇到问题，请参考配套笔记本：Fine_tuning_MedGemma.ipynb

结语

MedGemma 代表了将 AI 应用于医学科学的重要进步。它能够帮助医生更快速、更准确地作出判断，从而更快完成诊断并为患者制定更有效的治疗方案。

对如 MedGemma 4B Instruct 这样的视觉-语言模型进行微调，可使其适配多种医学任务，从图像分类到融合推理能力。

在本教程中，我们学习了如何在脑部 MRI 数据集上微调视觉-语言模型以进行脑癌分类。结果令人瞩目：模型准确率从 33% 提升至 89%，这一大幅跃升凸显了微调在医疗 AI 应用中的潜力。

如果您希望进一步学习，推荐参考以下资源：

• Biomedical Image Analysis in Python 课程
• Deep Learning for Images with PyTorch 课程
• Fine-Tuning Your Own Llama 3 Model 实战