В этом посте разберем работу с PyTorch 2.0 и Hugging Face Transformers на примере fine-tune модели BERT для классификации текста.

PyTorch 2.0 лучше по производительности, скорости работы, более удобный для Python, но при этом остается таким же динамическим, как и ранее.

Разберем следующие шаги:

1. Настройка окружения и установка PyTorch 2.0.

2. Загрузка и подготовка датасета.

3. Fine-tune и оценка модели BERT с помощью Hugging Face Trainer.

4. Запуск инференса и тестирование модели.

Краткое введение: PyTorch 2.0

PyTorch 2.0, или, точнее, 1.14, полностью обратно совместим с предыдущими версиями. Он не потребует каких‑либо изменений в существующем коде PyTorch, но может оптимизировать код, если добавить model = torch.compile(model). Команда PyTorch так объясняет появление новой версии в своем FAQ: «Мы выпустили значительные новые функции, которые, на наш взгляд, меняют то, как вы используете PyTorch, поэтому мы назвали это 2.0 вместо 1.14.»

Среди этих новых функций: полная поддержка TorchDynamo, AOTAutograd, PrimTorch и TorchInductor. Это позволяет PyTorch 2.0 достигнуть ускорения времени обучения в 1,3-2 раза на более 40 архитектурах моделей от HuggingFace Transformers. Подробнее о PyTorch 2.0 можно узнать на официальном "GET STARTED".

Примечание: Этот туториал был создан и запущен на инстансе AWS EC2 g5.xlarge, включая GPU NVIDIA A10G.

1. Настройка окружения и установка PyTorch 2.0

Первый шаг - установить PyTorch 2.0 и библиотеки от Hugging Face, transformers и datasets.

# Установка PyTorch 2.0 с cuda 11.7
!pip install "torch>=2.0" --extra-index-url https://download.pytorch.org/whl/cu117 --upgrade --quiet 

Также ставим последнюю версию transformers , которая включает нативную интеграцию PyTorch 2.0 в Trainer.

# Установка transformers и dataset
!pip install "transformers==4.27.1" "datasets==2.9.0" "accelerate==0.17.1" "evaluate==0.4.0" tensorboard scikit-learn
# Установка git-lfs для загрузки модели и логов в hugging face hub
!sudo apt-get install git-lfs --yes

В этом примере для версионирования моделей мы будем использовать Hugging Face Hub. Чтобы загрузить модель на Hub, вначале необходимо зарегистрироваться на Hugging Face. Для входа в свою учетную запись и сохранения токена (ключа доступа) на диске используем login из пакета huggingface_hub.

from huggingface_hub import login

login(
  token="", # ADD YOUR TOKEN HERE
  add_to_git_credential=True
)

2. Загрузка и подготовка датасета

Будем обучать модель классификации текста на датасете BANKING77. Датасет BANKING77 содержит текстовые обращения от клиентов из области банковского/финансового сектора. Он состоит из 13 083 обращений, размеченных на 77 интентов (классов).

Для загрузки BANKING77 мы будем использовать метод load_dataset() из библиотеки ???? Datasets.

from datasets import load_dataset

# Dataset id from huggingface.co/dataset
dataset_id = "banking77"

# Load raw dataset
raw_dataset = load_dataset(dataset_id)

print(f"Train dataset size: {len(raw_dataset['train'])}")
print(f"Test dataset size: {len(raw_dataset['test'])}")

Посмотрим на пример из набора данных.

from random import randrange

random_id = randrange(len(raw_dataset['train']))
raw_dataset['train'][random_id]
# {'text': "I can't get google pay to work right.", 'label': 2}

Для обучения модели нужно предварительно преобразовать текст в токены. Это делается токенизатором. Также он преобразует токены в соответствующие им идентификаторы в предобученном словаре. Подробнее о токенизации можно узнать в главе 6 от Hugging Face Course.

from transformers import AutoTokenizer

# Model id to load the tokenizer
model_id = "bert-base-uncased"
# Load Tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Tokenize helper function
def tokenize(batch):
    return tokenizer(batch['text'], padding='max_length', truncation=True, return_tensors="pt")

# Tokenize dataset
raw_dataset =  raw_dataset.rename_column("label", "labels") # to match Trainer
tokenized_dataset = raw_dataset.map(tokenize, batched=True,remove_columns=["text"])

print(tokenized_dataset["train"].features.keys())
# dict_keys(['input_ids', 'token_type_ids', 'attention_mask','lable'])

3. Fine-tune и оценка модели BERT с помощью Hugging Face Trainer

После подготовки данных можно начинать обучение модели. Мы будем использовать модель bert-base-uncased. Первым шагом будет загрузка модели с помощью класса AutoModelForSequenceClassification из Hugging Face Hub. Так мы создаем модель с весами предобученной модели BERT, но с "головой" сверху специально под нашу задачу классификации. Здесь мы передаем количество классов (77) из нашего набора данных и имена меток, чтобы сделать вывод результатов более читабельным.

from transformers import AutoModelForSequenceClassification

# Model id to load the tokenizer
model_id = "bert-base-uncased"

# Prepare model labels - useful for inference
labels = tokenized_dataset["train"].features["labels"].names
num_labels = len(labels)
label2id, id2label = dict(), dict()
for i, label in enumerate(labels):
    label2id[label] = str(i)
    id2label[str(i)] = label

# Download the model from huggingface.co/models
model = AutoModelForSequenceClassification.from_pretrained(
    model_id, num_labels=num_labels, label2id=label2id, id2label=id2label
)

Будем мониторить качество модели во время обучения. Trainer поддерживает оценку во время обучения, для этого определим compute_metrics. Используем библиотеку evaluate для расчета метрики f1 на тестовом наборе данных.

import evaluate
import numpy as np

# Metric Id
metric = evaluate.load("f1")

# Metric helper method
def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return metric.compute(predictions=predictions, references=labels, average="weighted")

Последний шаг - определить гиперпараметры TrainingArguments для обучения. Здесь же добавим параметры от PyTorch 2.0 для быстрого времени обучения. Чтобы использовать последние улучшения PyTorch 2.0 передаем опцию torch_compile в TrainingArguments.

Также используем интеграцию Trainer с Hugging Face Hub для загрузки весов модели, логов и метрик в репозиторий во время обучения.

from huggingface_hub import HfFolder
from transformers import Trainer, TrainingArguments

# Id for remote repository
repository_id = "bert-base-banking77-pt2"

# Define training args
training_args = TrainingArguments(
    output_dir=repository_id,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=8,
    learning_rate=5e-5,
	num_train_epochs=3,
	# PyTorch 2.0 specifics
    bf16=True, # bfloat16 training
	torch_compile=True, # optimizations
    optim="adamw_torch_fused", # improved optimizer
    # logging & evaluation strategies
    logging_dir=f"{repository_id}/logs",
    logging_strategy="steps",
    logging_steps=200,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    save_total_limit=2,
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    # push to hub parameters
    report_to="tensorboard",
    push_to_hub=True,
    hub_strategy="every_save",
    hub_model_id=repository_id,
    hub_token=HfFolder.get_token(),

)

# Create a Trainer instance
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    compute_metrics=compute_metrics,
)

Для запуска обучения использем метод train от Trainer.

# Start training
trainer.train()

Для сравнения времени обучения мы также запустили обучение без опции torch_compile:

{'train_runtime': 696.2701, 'train_samples_per_second': 43.1, 'eval_f1': 0.928788}

Используя опцию torch_compile и оптимизацию adamw_torch_fused, время обучения уменьшается на 52.5% по сравнению с обучением без PyTorch 2.0:

{'train_runtime': 457.7964, 'train_samples_per_second': 65.55, 'eval_f1': 0.931773}

Таким образом, время обучения снизились с 696 до 457 секунд. Значение train_samples_per_second выросло с 43 до 65. Значение f1-метрики такое же или чуть лучше, чем при обучении без использования torch_compile.

PyTorch 2.0 невероятно мощен! ????

Сохраним наши результаты и токенизатор в Hugging Face Hub и создадим карточку модели.

# Save processor and create model card
tokenizer.save_pretrained(repository_id)
trainer.create_model_card()
trainer.push_to_hub()

4. Запуск инференса и тестирование модели

Посмотрим на инференс модели на текстовом примере. Финальный классификатор получаем с помощью pipeline из библиотеки transformers.

from transformers import pipeline

# load model from huggingface.co/models using our repository id
classifier = pipeline("sentiment-analysis", model=repository_id, tokenizer=repository_id, device=0)

sample = "I have been waiting longer than expected for my bank card, could you provide information on when it will arrive?"


pred = classifier(sample)
print(pred)
# [{'label': 'card_arrival', 'score': 0.9903606176376343}]

Заключение

В этом посте мы разобрались, как использовать PyTorch 2.0 для обучения модели классификации текста на наборе данных BANKING77. PyTorch 2.0 — это мощный инструмент, позволяющий ускорить время обучения. В нашем примере, запущенном на NVIDIA A10G, мы смогли достичь на 52.5% лучшей производительности.

Кроме того, мы увидели, как несложно дообучить BERT под свою задачу средствами Hugging Face и PyTorch.

Мой телеграм‑канал о DS и не только.