train.py

import torch
import evaluate
import numpy as np
from datasets import load_dataset
from transformers import (
    AutoImageProcessor,
    AutoModelForImageClassification,
    TrainingArguments,
    Trainer,
)

# --- 1. CARREGAR O DATASET E PREPARÁ-LO ---

# Carrega o dataset HAM10000 do Hugging Face
ds = load_dataset("Nagabu/HAM10000")

# O dataset já vem com uma divisão 'train'. Vamos criar uma divisão de teste/validação.
# Vamos dividir o 'train' em 90% para treino e 10% para teste.
split = ds["train"].train_test_split(test_size=0.1)
train_ds = split["train"]
test_ds = split["test"]

# Carrega o processador de imagem do modelo que usaremos como base.
# Este processador ajusta o tamanho da imagem e normaliza as cores.
model_name = "google/vit-base-patch16-224-in21k"
processor = AutoImageProcessor.from_pretrained(model_name)

# Pega os nomes das classes (labels) do próprio dataset
labels = train_ds.features["label"].names
id2label = {i: label for i, label in enumerate(labels)}
label2id = {label: i for i, label in enumerate(labels)}

# Função para aplicar a transformação nas imagens
def transform(example_batch):
    # Pega uma lista de imagens e as processa
    inputs = processor([x for x in example_batch["image"]], return_tensors="pt")
    # Inclui os rótulos (labels) no resultado
    inputs["label"] = example_batch["label"]
    return inputs

# Aplica a transformação em todo o dataset
train_ds.set_transform(transform)
test_ds.set_transform(transform)

# --- 2. CONFIGURAR O MODELO E O TREINAMENTO ---

# Carrega o modelo pré-treinado.
# Ele será treinado para classificar em 7 classes (o número de labels do nosso dataset)
model = AutoModelForImageClassification.from_pretrained(
    model_name,
    num_labels=len(labels),
    id2label=id2label,
    label2id=label2id,
    ignore_mismatched_sizes=True, # Importante para carregar um modelo pré-treinado com uma nova "cabeça" de classificação
)

# Carrega a métrica de acurácia
accuracy = evaluate.load("accuracy")

# Função para calcular as métricas durante a avaliação
def compute_metrics(eval_pred):
    predictions, true_labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=true_labels)

# Define os argumentos do treinamento
training_args = TrainingArguments(
    output_dir="./meu-modelo-ham10000", # Pasta onde o modelo será salvo
    per_device_train_batch_size=8,   # Quantas imagens processar por vez. Se tiver erro de memória, diminua para 8.
    num_train_epochs=3,               # Quantas vezes o modelo verá o dataset inteiro. 3 é um bom começo.
    save_strategy="epoch",            # Salva o modelo ao final de cada época
    eval_strategy="epoch",      # Avalia o modelo ao final de cada época
    logging_strategy="epoch",         # Mostra os logs de treinamento a cada época
    load_best_model_at_end=True,      # Carrega o melhor modelo encontrado durante o treinamento no final
    metric_for_best_model="accuracy", # Usa a acurácia para decidir qual é o melhor modelo
    remove_unused_columns=False,      # Necessário para manter a coluna 'label'
)

# Cria o objeto Trainer, que gerencia todo o processo de treinamento
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=test_ds,
    tokenizer=processor, # O processador de imagem também é passado como tokenizer
    compute_metrics=compute_metrics,
)

# --- 3. INICIAR O TREINAMENTO ---

print("Iniciando o treinamento...")
trainer.train()

# --- 4. SALVAR O MODELO FINAL ---

print("Treinamento concluído. Salvando o modelo final...")
trainer.save_model("./meu-modelo-ham10000-final")
print("Modelo salvo com sucesso na pasta 'meu-modelo-ham10000-final'")