Vue d'ensemble

Transformers 4.5, publié le 16 mars 2021, améliore l'export ONNX des modèles et apporte des améliorations au Trainer pour un entraînement plus flexible.

Fonctionnalités principales

Export ONNX

L'export ONNX permet de convertir les modèles Transformers vers le format ONNX pour un déploiement optimisé avec ONNX Runtime, offrant des gains de performance en inférence.

python 
from transformers import AutoTokenizer, AutoModel
import torch

model_name = 'distilbert-base-uncased'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# Export vers ONNX
dummy = tokenizer('Bonjour le monde', return_tensors='pt')
torch.onnx.export(
    model, (dummy['input_ids'], dummy['attention_mask']),
    'model.onnx',
    input_names=['input_ids', 'attention_mask'],
    dynamic_axes={'input_ids': {0: 'batch', 1: 'seq'}},
)

Améliorations du Trainer

Le Trainer gagne de nouvelles options : callbacks personnalisés, intégration avec Weights & Biases améliorée et support du gradient checkpointing pour réduire la consommation mémoire.

python 
from transformers import Trainer, TrainingArguments

args = TrainingArguments(
    output_dir='./resultats',
    num_train_epochs=3,
    per_device_train_batch_size=16,
    gradient_checkpointing=True,  # économie mémoire
    report_to='wandb',
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
)
trainer.train()

Sources