Fine-tuner un modèle avec Trainer

Trainer fournit une boucle d’entraînement complète pour les modèles Transformers. Ce tutoriel illustre une classification de texte ; commencez avec un petit échantillon avant de consommer des ressources importantes.

Installer les dépendances

pip install -U transformers datasets evaluate accelerate torch

Charger et tokeniser les données

from datasets import load_dataset
from transformers import AutoTokenizer

model_id = "distilbert/distilbert-base-uncased"
dataset = load_dataset("stanfordnlp/imdb")
tokenizer = AutoTokenizer.from_pretrained(model_id)

def tokenize(batch):
    return tokenizer(batch["text"], truncation=True)

tokenized = dataset.map(tokenize, batched=True)

Conservez un jeu de validation indépendant. Vérifiez les classes, les doublons et les fuites entre entraînement et validation avant de lancer le calcul.

Configurer le modèle et Trainer

from transformers import (
    AutoModelForSequenceClassification,
    DataCollatorWithPadding,
    Trainer,
    TrainingArguments,
)

model = AutoModelForSequenceClassification.from_pretrained(model_id, num_labels=2)
args = TrainingArguments(
    output_dir="./outputs/imdb-distilbert",
    eval_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    num_train_epochs=2,
    weight_decay=0.01,
    load_best_model_at_end=True,
    report_to="none",
    seed=42,
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=tokenized["train"],
    eval_dataset=tokenized["test"],
    processing_class=tokenizer,
    data_collator=DataCollatorWithPadding(tokenizer),
)
trainer.train()
print(trainer.evaluate())

Évaluer avant de publier

Ajoutez une métrique adaptée au produit, inspectez une matrice de confusion et testez des exemples hors distribution. Versionnez le modèle de base, le dataset, les paramètres et le code. Ne publiez pas automatiquement les données, poids ou journaux s’ils contiennent des informations sensibles.

Un score global peut masquer une classe faible. Définissez le seuil d’acceptation avant l’entraînement et comparez le modèle fine-tuné au modèle de base ainsi qu’à une règle simple.