Создайте TrainerCallback, чтобы легко загружать модели HuggingFace на Kaggle

Если вы когда-либо работали с большими моделями HuggingFace, такими как Whisper-Large на Kaggle, вы знаете, что сочетание доступности графического процессора (30 часов в неделю) и ограниченности ресурсов (16 ГБ видеопамяти) может усложнить обучение. В разгар конкурса Kaggle, посвященного транскрипции аудио, я экспериментировал с моделью HuggingFace Whisper-Large на своем локальном компьютере. Как только я обучил модель до удовлетворительного уровня, я загрузил ее как набор данных Kaggle, чтобы повторно использовать ее для своих представлений. Однако постоянная повторная загрузка модели в Kaggle оказалась обременительной задачей. Именно тогда я решил создать умный обратный вызов, который автоматически загружал бы модель в Kaggle каждый раз, когда она обновлялась на HuggingFace после оценки.

Прежде чем мы начнем, важно отметить, что в блокноте для подачи заявок на соревнования Kaggle должен быть отключен Интернет, что делает невозможным прямое использование функции HuggingFace from_pretrained для загрузки вашей последней модели. Следовательно, загрузка обученной модели в виде набора данных Kaggle становится необходимостью для бесшовной интеграции.

В этой главе мы подробно рассмотрим пошаговый процесс создания обратного вызова KaggleUploader, который позволит вам ускорить рабочий процесс и без труда загружать модели HuggingFace на Kaggle во время обучения.

Давайте начнем!

Получение токена API Kaggle

Прежде чем мы продолжим, убедитесь, что у вас есть токен Kaggle API. Если у вас его нет, создайте его, посетив настройки своей учетной записи Kaggle (https://www.kaggle.com/settings/account) и нажав кнопку Создать новый токен API. . Будет загружен файл kaggle.json, содержащий ваше имя пользователя и токен API. Сохраните этот файл в папке ~/.kaggle/kaggle.json в Linux или C:\Users‹Windows-username›.kaggle\kaggle.json в Windows. В целях безопасности установите права доступа к файлам, чтобы предоставить только вам доступ для чтения, чтобы предотвратить возможные ошибки, связанные с безопасностью файлов.

chmod 600 ~/.kaggle/kaggle.json

Установка модуля Kaggle

Убедитесь, что у вас установлен модуль Kaggle. Вы можете сделать это, выполнив следующую команду:

pip install kaggle

Тонкая настройка модели Whisper

Давайте начнем с фрагмента кода, который настраивает модель Whisper на 1000 шагов, оценивая и отправляя модель в HuggingFace Hub каждые 250 шагов. Мы сохраним аргументы обучения простыми, чтобы сосредоточиться на самом важном. В зависимости от варианта использования вам может потребоваться включить токенизатор, средство подбора данных и другие компоненты.

from transformers import WhisperForConditionalGeneration
from transformers import Seq2SeqTrainingArguments
from transformers import Seq2SeqTrainer

model_name = "openai/whisper-large-v2" 
model = WhisperForConditionalGeneration.from_pretrained(model_name)

training_args = Seq2SeqTrainingArguments(
    output_dir=f"./whisper-large-v2",  
    per_device_train_batch_size=8,
    learning_rate=1e-5,
    max_steps=1000,
    save_steps=250,
    eval_steps=250,
    logging_steps=25,
    greater_is_better=False,
    push_to_hub=True,
)

trainer = Seq2SeqTrainer(
    args=training_args,
    model=model,
    train_dataset=my_dataset["train"],
    eval_dataset=my_dataset["eval"],
    callbacks=[kaggle_uploader]
)

trainer.train()

Заметили обратный вызов kaggle_uploader? Это то, что мы создадим прямо сейчас.

Создание обратного вызова HuggingFace Trainer

Теперь давайте перейдем к интересной части — созданию TrainerCallback, автоматизирующего процесс загрузки. Следующий код подробно объясняет каждый шаг:

import tempfile
import shutil
import json
import kaggle
from pathlib import Path
from transformers import TrainerCallback
from transformers.trainer_callback import TrainerControl, TrainerState
from transformers.training_args import TrainingArguments

class KaggleUploader(TrainerCallback):
    def __init__(self, dataset_path: str, id: str, title: str, isPrivate: bool):
        self.api = kaggle.KaggleApi()
        self.api.authenticate()
        self.dataset_path = dataset_path
        self.meta_data = dict(
            id=id,
            title=title,
            isPrivate=isPrivate,
            licenses=[dict(name="other")]
        )

    def on_save(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        best_model_checkpoint = str(Path(state.best_model_checkpoint).name)
        self.upload_dataset_to_kaggle(self.dataset_path, best_model_checkpoint)
        
        return super().on_evaluate(args, state, control, **kwargs)
    
    def upload_dataset_to_kaggle(self, dataset_path, checkpoint_to_save: str):
        # latest_checkpoint = find_latest_checkpoint(dataset_path)
        checkpoint = os.path.join(dataset_path, checkpoint_to_save)

        version_notes = checkpoint_to_save
        # The checkpoint has multiple files that we don't need.
        # We only need the pytorch_model.bin file, config.json and generation_config.json
        # Copy these files to a temporary folder
        with tempfile.TemporaryDirectory() as temp_dir:
            # create a directory inside named "model"
            temp_model_dir = os.path.join(temp_dir, "model")
            os.mkdir(temp_model_dir)
            # copy the files
            for file in ["pytorch_model.bin", "config.json", "generation_config.json"]:
                shutil.copy(os.path.join(checkpoint, file), temp_model_dir)

            # create dataset-metadata.json inside the temporary directory
            with open(os.path.join(temp_model_dir, "dataset-metadata.json"), "w") as f:
                json.dump(self.meta_data, f)

            self.api.dataset_create_version(temp_model_dir, version_notes=version_notes, dir_mode="zip")

Рассмотрим подробнее каждый раздел:

Создание KaggleUploader

Класс KaggleUploader отвечает за управление загрузкой модели в Kaggle. Мы можем создать его экземпляр следующим образом:

kaggle_uploader = KaggleUploader(
  dataset_path="<my_dataset_path>",
  id="<kaggle_username>/<kaggle_dataset_name>",
  title="<Kaggle Clean Dataset Name>", 
  isPrivate=True
)

В приведенном выше коде мы указываем папку, содержащую модель, которая будет загружена в Kaggle, а также имя и заголовок набора данных Kaggle. Параметр isPrivate определяет, будет ли загруженный набор данных частным или нет. Аутентификация обрабатывается внутри класса, гарантируя, что она выполняется только один раз. Стоит отметить, что HuggingFace создает подпапку на каждой контрольной точке (checkpoint-XXXX). Таким образом, путь_набора_данных должен быть родительской папкой, содержащей эти папки контрольных точек.

class KaggleUploader(TrainerCallback):
  def __init__(self, dataset_path: str, id: str, title: str, isPrivate: bool):
    self.api = kaggle.KaggleApi()
    self.api.authenticate()
    self.dataset_path = dataset_path
    self.meta_data = dict(
        id=id,
        title=title,
        isPrivate=isPrivate,
        licenses=[dict(name="other")]
    )

Загрузка модели при сохранении

Поскольку мы хотим загружать модель в Kaggle каждый раз, когда она обновляется на HuggingFace, мы используем для этой цели метод on_save:

def on_save(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
    best_model_checkpoint = str(Path(state.best_model_checkpoint).name)
    self.upload_dataset_to_kaggle(self.dataset_path, best_model_checkpoint)
    
    return super().on_evaluate(args, state, control, **kwargs)

В этой реализации будет сохранена и загружена лучшая модель (то есть модель с наибольшей потерей при проверке). Если вы хотите вместо этого сохранить последнюю модель, вы можете использовать такую ​​функцию:

def find_latest_checkpoint(self, dataset_path):
    # checkpoints are folders of the following format: checkpoint-XXXXX
    checkpoints = [
        folder
        for folder in os.listdir(dataset_path)
        if folder.startswith("checkpoint-")
    ]
    # sort the checkpoints by their step number
    checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[-1]))
    # return the path to the latest checkpoint
    return checkpoints[-1]

def on_save(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
    last_model_checkpoint = self.find_latest_checkpoint(self.dataset_path)
    self.upload_dataset_to_kaggle(self.dataset_path, best_model_checkpoint)
    
    return super().on_evaluate(args, state, control, **kwargs)

Загрузите соответствующие файлы

Папка checkpoint-XXXX содержит несколько файлов, многие из которых не нужны для отправки на Kaggle. Мы будем загружать только файлы pytorch_model.bin, config.json и generation_config.json. Это позволит избежать загрузки слишком большого количества файлов, некоторые из которых довольно тяжелые (например, optimizer.pt), что сэкономит время на загрузку и скачивание.

def upload_dataset_to_kaggle(self, dataset_path, checkpoint_to_save: str):
    checkpoint = os.path.join(dataset_path, checkpoint_to_save)

    version_notes = checkpoint_to_save
    # The checkpoint has multiple files that we don't need.
    # We only need the pytorch_model.bin file, config.json and generation_config.json
    # Copy these files to a temporary folder
    with tempfile.TemporaryDirectory() as temp_dir:
        # create a directory inside named "model"
        temp_model_dir = os.path.join(temp_dir, "model")
        os.mkdir(temp_model_dir)
        # copy the files
        for file in ["pytorch_model.bin", "config.json", "generation_config.json"]:
            shutil.copy(os.path.join(checkpoint, file), temp_model_dir)

        # create dataset-metadata.json inside the temporary directory
        with open(os.path.join(temp_model_dir, "dataset-metadata.json"), "w") as f:
            json.dump(self.meta_data, f)

        self.api.dataset_create_version(temp_model_dir, version_notes=version_notes, dir_mode="zip")

Поскольку API Kaggle позволяет загружать только папку, мы копируем необходимые файлы во временную папку. Мы также создаем файл dataset-metadata.json, используемый Kaggle для загрузки. Важно отметить, что метод self.api.dataset_create_version работает, когда набор данных уже существует и вы хотите его обновить. Если вы хотите создать новый набор данных, вам нужно использовать метод self.api.dataset_create_new, который принимает те же аргументы, за исключением version_notes.

Автоматизируя загрузку моделей из HuggingFace в Kaggle с помощью TrainerCallback, вы можете легко интегрировать свои обученные модели в отправки Kaggle и экономить драгоценное время во время соревнований или других проектов. Удачных тренировок модели и успехов в соревнованиях!

На простом английском

Спасибо, что являетесь частью нашего сообщества! Прежде чем уйти: