Fine tuning o MedGemma 4B usando o Google Colab e o Hugging Face usando imagens de mamografia

O objetivo deste guia é fornecer as ferramentas e o conhecimento para dar os primeiros passos com o MedGemma, ajustá-lo para suas necessidades específicas e avaliar seu desempenho, abrindo caminho para o desenvolvimento de aplicações inovadoras na área da saúde!

Todos os códigos utilizados neste tutorial podem ser encontrados em medgemma_experiments e o modelo ajustado está disponível no Hugging Face Hub.

1. Visão Geral do MedGemma

MedGemma é uma coleção de variantes do modelo Gemma 3, especificamente treinadas para desempenho em compreensão de texto e imagem médica. De acordo com a documentação, o seu principal objetivo é “acelerar a construção de aplicações de IA baseadas em saúde.”

1.1 Variantes e Requisitos de Hardware

O MedGemma possui duas variantes, uma multimodal (MedGemma 4B) e outra somente textual (MedGemma 27B):

• MedGemma 4B: Pode ser executado gratuitamente no Google Colab usando uma GPU T4 . Essa variante foi pre treinada em um dataset com diversos dados médicos, entre eles imagens radiológicas e dermatológicas.
• MedGemma 27B: Requer uma GPU A100 e quantização para reduzir o uso de memória no Google Colab.

1.2 Acesso e Autenticação

Nesse tutorial utilizaremos o Google Colab e o Hugging Face para fazer o fine tuning do modelo. Para isso será necessário:

1. Ter uma conta gratuita no Hugging Face.
2. Acessar a página do modelo MedGemma e aceitar as condições de uso
3. Configurar um token de acesso do Hugging Face. Para ajuste fino, o token deve ter acesso de gravação (write) para enviar o modelo ajustado para o Hugging Face Hub. Para uso geral, um token de acesso de leitura (read) é suficiente. No Google Colab, recomenda-se armazenar o token no gerenciador de Segredos.

2. Primeiros Passos com o MedGemma

Antes de modificar o modelo, é importante saber como interagir com o modelo base. Em *quick_start_with_hugging_face.ipynb *é possível acompanhar o passo a passo desde a autenticação com o Hugging Face até as respostas retornadas pelo modelo quando questionado a descrever as imagens de mamografia apresentadas.

Depois que a variante do modelo é configurada, precisamos especificar qual imagem e prompt serão passados para o modelo:

import os
from PIL import Image
from IPython.display import Image as IPImage, display, Markdown

prompt = "Describe this mammogram"  # @param {type: "string"}

image_url = "https://cdn.apollohospitals.com/health-library-prod/2021/08/Breast-Calcifications-scaled.jpg"  # @param {type: "string"}
! wget -nc -q {image_url}
image_filename = os.path.basename(image_url)
image = Image.open(image_filename)

Para realizar a chamada a API pipeline que fará a inferência do modelo é necessário estruturar a mensagem que será passada, contendo a instrução ao sistema (por exemplo, “You are an expert radiologist.” ), o prompt e a imagem. A vantagem de ter uma estrutura pré determinada é que pra consultas posteriores o usuário pode apenas definir as variáveis que deseja modificar.

system_instruction = "You are an expert radiologist."

messages = [
    {
        "role": "system",
        "content": [{"type": "text", "text": system_instruction}]
    },
    {
        "role": "user",
        "content": [
            {"type": "text", "text": prompt},
            {"type": "image", "image": image}
        ]
    }
]

from transformers import pipeline

pipe = pipeline(
    "image-text-to-text",
    model=model_id,
    model_kwargs=model_kwargs,
)

pipe.model.generation_config.do_sample = False

Resposta retornada pelo MedGemma

Na imagem acima o MedGemma responde que esta é uma imagem de mamografia normal, quando a imagem apresentada possui um achado relevante. A partir disso surgiu a ideia de fazer um ajuste no MedGemma com imagens de mamografias com calcificações e massas.

3. Ajuste Fino (Fine-tuning) do MedGemma

O ajuste tradicional de LLMs consome muitos recursos pela necessidade de se adaptar bilhões de parâmetros. Para diminuir esse consumo de recursos algumas técnicas podem ser utilizadas, como quantização, LoRA e QLoRA, por exemplo.

Neste tutorial o ajuste fino do MedGemma será realizado utilizando Transformer Reinforcement Learning (TRL) com Supervised Fine-Tuning (SFT) e Quantized Low-Rank Adaptation (QLoRA).

3.1 Requisitos de Recursos para Fine Tuning no Google Colab

Para fazer o fine tuning é necessário uma GPU que suporte o tipo de dado bfloat16 e tenha pelo menos 40 GB de memória, como uma A100 no Google Colab. Além disso, é importante que a sua chave de acesso do Hugging Face tenha permissão para escrita, uma vez que o modelo será salvo no Hugging Face Hub.

3.2 Conjunto de Dados para Ajuste Fino

Esse tutorial utiliza o conjunto de dados CBIS-DDSM, que contém imagens de mamografias, para ajustar o MedGemma para classificar as imagens de acordo com o tipo de achado, calcificações ou massas. O dataset completo é composto por 3286 amostras, das quais 2628 serão utilizadas para treinamento e 658 para validação.

O conjunto de dados é pré-processado em um formato conversacional multimodal, incluindo uma pergunta de múltipla escolha: “What is the type of abnormality present in the mammogram image?”

from typing import Any

ABNORMALITY_TYPES = [
    "calcification",
    "mass",
]

options = "\n".join(ABNORMALITY_TYPES)
PROMPT = f"What is the type of abnormality present in the 
          mammogram image?\n{options}"

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": ABNORMALITY_TYPES[example["abnormality type"]],
                },
            ],
        },
    ]
    return example

3.3 Configuração e Treinamento

O treinamento é configurado usando SFT Config, onde serão definidos o número de épocas do treinamento, a taxa de aprendizado, em qual diretório o modelo ajustado será salvo, dentre outros, como pode ser visto no trecho de código a seguir:

from trl import SFTConfig

num_train_epochs = 1  # @param {type: "number"}
learning_rate = 2e-4  # @param {type: "number"}

args = SFTConfig(
    output_dir="medgemma-4b-it-sft-lora-crc100k",            # Directory and Hub repository id to save the model to
    num_train_epochs=num_train_epochs,                       # Number of training epochs
    per_device_train_batch_size=4,                           # Batch size per device during training
    per_device_eval_batch_size=4,                            # Batch size per device during evaluation
    gradient_accumulation_steps=4,                           # Number of steps before performing a backward/update pass
    gradient_checkpointing=True,                             # Enable gradient checkpointing to reduce memory usage
    optim="adamw_torch_fused",                               # Use fused AdamW optimizer for better performance
    logging_steps=50,                                        # Number of steps between logs
    save_strategy="epoch",                                   # Save checkpoint every epoch
    eval_strategy="steps",                                   # Evaluate every `eval_steps`
    eval_steps=50,                                           # Number of steps between evaluations
    learning_rate=learning_rate,                             # Learning rate based on QLoRA paper
    bf16=True,                                               # Use bfloat16 precision
    max_grad_norm=0.3,                                       # Max gradient norm based on QLoRA paper
    warmup_ratio=0.03,                                       # Warmup ratio based on QLoRA paper
    lr_scheduler_type="linear",                              # Use linear learning rate scheduler
    push_to_hub=True,                                        # Push model to Hub
    report_to="tensorboard",                                 # Report metrics to tensorboard
    gradient_checkpointing_kwargs={"use_reentrant": False},  # Set gradient checkpointing to non-reentrant to avoid issues
    dataset_kwargs={"skip_prepare_dataset": True},           # Skip default dataset preparation to preprocess manually
    remove_unused_columns = False,                           # Columns are unused for training but needed for data collator
    label_names=["labels"],                                  # Input keys that correspond to the labels
)

Utilizando a classe SFTTrainer pode-se instanciar um “treinador” usando a configuração feita anteriormente e uma função que processa os exemplos com os textos e as imagens:

from trl import SFTTrainer

trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=data["train"],
    eval_dataset=data["validation"].shuffle().select(range(200)),  # Use subset of validation set for faster run
    peft_config=peft_config,
    processing_class=processor,
    data_collator=collate_fn,
)

4. Execução de Inferência no modelo ajustado

Para verificar o comportamento do modelo ajustado nas etapas anteriores, utilizamos uma versão modificada (*testing_fine_tuned_model.ipynb) do notebook da seção 2, com algumas adaptações para carregar o novo modelo, bem como ajustar o prompt utilizado. As imagens testadas foram as mesmas do [quick_start_with_hugging_face.ipynb](https://colab.research.google.com/github/nandaw/medgemma_experiments/blob/main/notebooks/quick_start_with_hugging_face.ipynb) .*

É importante notar que o modelo não responde ao prompt com a mesma complexidade anterior, comportamento que pode ser explicado por um fenômeno chamado *catastrophic forgetting. *Durante o processo de fine tuning o modelo é apresentado a novos dados e ajusta seus parâmetros para se especializar neles, podendo “esquecer” das habilidades adquiridas durante o treinamento original. Como os dados utilizados eram compostos apenas pelos labels simples e pelas imagens, a capacidade de resposta do novo modelo ficou reduzida.

5. Próximos Passos

Este tutorial tinha como objetivo fornecer um caminho inicial para o fine tuning do MedGemma para casos de mamografia, que não eram cobertos pelo modelo inicialmente. Para diminuir o tempo de processamento o modelo foi retreinado por apenas uma época e com uma quantidade reduzida de dados. No futuro, também seria possível usar as outras informações disponíveis nos dados, como formato das lesões e malignidade, por exemplo.

*Google Colab units are provided for this project. #AISprint