Overview

Transformers 4.5, released on March 16, 2021, improves ONNX model export and brings enhancements to the Trainer for more flexible training.

Main Features

ONNX export

ONNX export allows converting Transformers models to the ONNX format for optimized deployment with ONNX Runtime, providing inference performance gains.

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 to ONNX
dummy = tokenizer('Hello world', 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'}},
)

Trainer improvements

The Trainer gains new options: custom callbacks, improved Weights & Biases integration, and gradient checkpointing support to reduce memory consumption.

python 
from transformers import Trainer, TrainingArguments

args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=16,
    gradient_checkpointing=True,  # memory saving
    report_to='wandb',
)

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

Sources