Этот материал посвящён ускорению обучения моделей с использованием бессерверных заданий. В частности, речь пойдёт о том, как запускать обучение с применением Pytorch, GPU и платформы Vertex.

В предыдущей статье я упомянул о том, что локальное обучение огромных моделей — это не всегда хорошо в условиях ограниченности ресурсов. Иногда так поступают просто потому, что другого варианта нет. Но случается и так, что в распоряжении ML‑специалиста имеется некий облачный провайдер, вроде Google Cloud Platform, инструменты которого способны значительно ускорить обучение моделей. А именно:

Провайдер может предоставить доступ к современнейшим машинам, характеристики которых (память, GPU и прочее) соответствуют нуждам клиента.
Провайдер может позволить одновременный запуск нескольких обучающих заданий, после завершения которых можно выбрать из того, что у них получилось, наилучшим образом обученную модель.

Не стоит и говорить о том, что вывод обучения в облако облегчит жизнь обычным компьютерам. Однажды я запустил недельную задачу обучения модели на своём ноутбуке и уехал в отпуск. Когда я вернулся, то, от перегрева батареи, тачпад буквально выскакивал из корпуса.

Мы рассмотрим конкретный пример дообучения модели BERT на комментариях в социальных сетях. Модель выполняет анализ тональности текста. Как мы увидим — обучение подобных моделей на CPU — это неоптимальный подход к решению весьма трудоёмкой задачи. Затем мы поговорим о том, как воспользоваться Google Cloud Platform для ускорения обучения с применением GPU, заплатив за это всего 60 центов.

Что такое BERT?

BERT расшифровывается как «Bidirectional Encoder Representations from Transformers». Компания Google перевела эту модель в разряд опенсорсных в 2018 году. BERT, в основном, используется для решения NLP‑задач, так как она была обучена восприятию смысла предложений и даёт доступ к мощным словесным эмбеддингам (представлениям). Отличие BERT от других моделей — таких, как Word2Vec и Glove, заключается в том, что она, для обработки текстов, использует трансформеры. Трансформеры (за подробностями о них можете обратиться к моей предыдущей статье) — это семейство нейросетевых архитектур, которые, что немного роднит их с RNN, могут обрабатывать последовательности в обоих направлениях. Это даёт им, например, возможность захвата контекста, окружающего некое слово.

Что такое анализ тональности текста?

Анализ тональности текста (Sentiment Analysis) — это специфическая задача из сферы NLP, цель которой заключается в классификации текстов по категориям, связанным с их тональностью. Тональность обычно определяют как положительную, отрицательную и нейтральную. Этот вид анализа текста очень часто используют для анализа ответов на различные вопросы, публикаций в социальных сетях, обзоров продуктов и так далее.

Дообучение BERT на данных из социальных сетей

Загрузка и подготовка данных

Набор данных, который мы будем использовать, взят с Kaggle. Загрузить его можно здесь (он распространяется по лицензии CC BY 4.0). В своих экспериментах я ограничился наборами данных Facebook и Twitter.

Следующий фрагмент кода берёт csv‑файлы и разделяет их на фрагменты, формируя в заданном месте три файла, содержащих наборы данных для обучения, валидации и тестирования модели. Рекомендую сохранять эти файлы в Google Cloud Storage.

Запустить скрипт можно такой командой:

python make_splits --output-dir gs://your-bucket/

Вот код скрипта:

import pandas as pd
import argparse
import numpy as np
from sklearn.model_selection import train_test_split


def make_splits(output_dir):
    df=pd.concat([
        pd.read_csv("data/farisdurrani/twitter_filtered.csv"),
        pd.read_csv("data/farisdurrani/facebook_filtered.csv")
    ])
    df = df.dropna(subset=['sentiment'], axis=0)
    df['Target'] = df['sentiment'].apply(lambda x: 1 if x==0 else np.sign(x)+1).astype(int)

    df_train, df_ = train_test_split(df, stratify=df['Target'], test_size=0.2)
    df_eval, df_test = train_test_split(df_, stratify=df_['Target'], test_size=0.5)

    print(f"Files will be saved in {output_dir}")
    df_train.to_csv(output_dir + "/train.csv", index=False)
    df_eval.to_csv(output_dir + "/eval.csv", index=False)
    df_test.to_csv(output_dir + "/test.csv", index=False)

    print(f"Train : ({df_train.shape}) samples")
    print(f"Val : ({df_eval.shape}) samples")
    print(f"Test : ({df_test.shape}) samples")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--output-dir')
    args, _ = parser.parse_known_args()
    make_splits(args.output_dir)

Полученные данные должны выглядеть примерно так:

https://miro.medium.com/v2/resize:fit:542/1*Iwp1_8fVAbzTVs4HmkFFIQ.png — Данные, используемые в эксперименте

Использование предварительно обученной маленькой модели BERT

В нашем эксперименте будет использоваться облегчённый вариант модели BERT — BERT‑Tiny. Эта модель уже обучена на огромном объёме данных. Но это — необязательно данные из социальных сетей, и вовсе не факт то, что обучение модели велось с целью анализа тональности текстов. Именно поэтому мы и будем дообучать модель.

BERT‑Tiny содержит всего всего 2 слоя с размерностью в 128 элементов. Полный список подобных моделей можно посмотреть здесь — на тот случай, если вы решите выбрать модель побольше.

Начнём с создания файла main.py, в котором подключены все необходимые модули:

import pandas as pd
import argparse
import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text as text
import logging
import os
os.environ["TFHUB_MODEL_LOAD_FORMAT"] = "UNCOMPRESSED"

def train_and_evaluate(**params):
    pass
    # будем это дополнять по мере продвижения по статье

Кроме того — внесём требования к пакетам в отдельный файл requirements.txt:

transformers==4.40.1
torch==2.2.2
pandas==2.0.3
scikit-learn==1.3.2
gcsfs

Теперь мы, для обучения модели, загрузим две сущности:

Токенизатор, который отвечает за разбиение текстовых входных данных на токены, на которых обучалась модель BERT.
Саму модель.

И то и другое можно найти на HuggingFace. Их тоже можно выгрузить в Cloud Storage. Именно это я и сделал. В результате команды их загрузки выглядят так:

# Загрузка предварительно обученного токенизатора и модели BERT
tokenizer = BertTokenizer.from_pretrained('models/bert_uncased_L-2_H-128_A-2/vocab.txt')
model = BertModel.from_pretrained('models/bert_uncased_L-2_H-128_A-2')

Теперь добавим в наш файл следующее:

class SentimentBERT(nn.Module):
    def __init__(self, bert_model):
        super().__init__()
        self.bert_module = bert_model
        self.dropout = nn.Dropout(0.1)
        self.final = nn.Linear(in_features=128, out_features=3, bias=True) 
        
        # Раскомментируйте, если хотите лишь переобучить определённые слои.
        # self.bert_module.requires_grad_(False)
        # for param in self.bert_module.encoder.parameters():
        #     param.requires_grad = True
        
    def forward(self, inputs):
        ids, mask, token_type_ids = inputs['ids'], inputs['mask'], inputs['token_type_ids']
        # print(ids.size(), mask.size(), token_type_ids.size())
        x = self.bert_module(ids, mask, token_type_ids)
        x = self.dropout(x['pooler_output'])
        out = self.final(x)
        return out

Теперь ненадолго притормозим. Если речь идёт об использовании существующей модели в проекте, на который она изначально не рассчитана, то тут имеется несколько вариантов:

Трансферное обучение: веса модели «замораживают» и используют её как средство для «извлечения признаков». Затем можно присоединить к модели дополнительные слои, идущие после её собственных слоёв. Этот подход часто используется в сфере компьютерного зрения, где модели, вроде VGG, Xception и прочих подобных, можно использовать для решения разных задач, обучая собственную модель на небольших наборах данных.
Дообучение: «размораживают» все веса модели или какую‑то их часть, после чего дообучают модель на собственном наборе данных. Это — подход, который обычно выбирают, обучая собственные большие языковые модели.

Подробности о трансферном обучении и о дообучении можно найти здесь.

В нашем случае решено было «разморозить» всю модель. Вы вполне можете решить «заморозить» один или несколько слоёв предварительно обученной модели BERT и оценить воздействие этого шага на эффективность работы результирующей модели.

Самое главное здесь — добавить полносвязный слой после модуля BERT, чтобы «связать» эту модель с нашей задачей классификации, то есть — со слоем, содержащим три блока. Это позволит многократно использовать предварительно обученную модель BERT и приспособить её под решение нашей задачи.

Создание загрузчиков данных

Для того чтобы создать загрузчики данных, нам понадобится токенизатор, который был загружен ранее. Он принимает в качестве входных данных строку и выдаёт несколько групп выходных данных, среди которых можно обнаружить токены (в нашем случае это input_ids):

https://miro.medium.com/v2/resize:fit:700/1*0ZLt7XaVbHK7GqITy6cDMg.png — Выходные данные токенизатора

Токенизатор BERT — это, в определённом смысле, особенный механизм. Из того, что он возвращает, наибольший интерес для нас представляют именно токены input_ids. Это — те токены, которые использовались для кодирования входного предложения. Это могут быть слова или части слов. Например — слово «looking» может быть составлено из двух токенов — «look» и «##ing»

Создадим теперь модуль загрузчика, который будет обрабатывать наборы данных:

class BertDataset(Dataset):
    def __init__(self, df, tokenizer, max_length=100):
        super(BertDataset, self).__init__()
        self.df=df
        self.tokenizer=tokenizer
        self.target=self.df['Target']
        self.max_length=max_length
        
    def __len__(self):
        return len(self.df)
    
    def __getitem__(self, idx):
        
        X = self.df['bodyText'].values[idx]
        y = self.target.values[idx]
        
        inputs = self.tokenizer.encode_plus(
            X,
            pad_to_max_length=True,
            add_special_tokens=True,
            return_attention_mask=True,
            max_length=self.max_length,
        )
        ids = inputs["input_ids"]
        token_type_ids = inputs["token_type_ids"]
        mask = inputs["attention_mask"]

        x = {
            'ids': torch.tensor(ids, dtype=torch.long).to(DEVICE),
            'mask': torch.tensor(mask, dtype=torch.long).to(DEVICE),
            'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long).to(DEVICE)
            }
        y = torch.tensor(y, dtype=torch.long).to(DEVICE)
        
        return x, y

Написание основного скрипта для обучения модели

В первую очередь определим две функции, которые будут отвечать за обучение модели и оценку её качества:

def train(epoch, model, dataloader, loss_fn, optimizer, max_steps=None):
    model.train()
    total_acc, total_count = 0, 0
    log_interval = 50
    start_time = time.time()

    for idx, (inputs, label) in enumerate(dataloader):
        optimizer.zero_grad()
        predicted_label = model(inputs)
        
        loss = loss_fn(predicted_label, label)
        loss.backward()
        optimizer.step()
        
        total_acc += (predicted_label.argmax(1) == label).sum().item()
        total_count += label.size(0)
        
        if idx % log_interval == 0:
            elapsed = time.time() - start_time
            print(
                "Epoch {:3d} | {:5d}/{:5d} batches "
                "| accuracy {:8.3f} | loss {:8.3f} ({:.3f}s)".format(
                    epoch, idx, len(dataloader), total_acc / total_count, loss.item(), elapsed
                )
            )
            total_acc, total_count = 0, 0
            start_time = time.time()

        if max_steps is not None:
            if idx == max_steps:
                return {'loss': loss.item(), 'acc': total_acc / total_count}
    
    return {'loss': loss.item(), 'acc': total_acc / total_count}


def evaluate(model, dataloader, loss_fn):
    model.eval()
    total_acc, total_count = 0, 0

    with torch.no_grad():
        for idx, (inputs, label) in enumerate(dataloader):
            predicted_label = model(inputs)
            loss = loss_fn(predicted_label, label)
            total_acc += (predicted_label.argmax(1) == label).sum().item()
            total_count += label.size(0)

    return {'loss': loss.item(), 'acc': total_acc / total_count}

Мы приближаемся к тому, чтобы оснастить наш главный скрипт всем необходимым функционалом. Соберём воедино всё, что у нас уже есть:

Класс BertDataset, ответственный за загрузку данных.
Модель SentimentBERT, принимающая модель Tiny-BERT и добавляющая к ней дополнительный слой, ориентированный на решение нашей задачи.
Функции train() и eval(), ответственные за обучение и оценку модели.
Функция train_and_eval(), которая всё в себе объединяет.

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

import pandas as pd
import time
import torch.nn as nn
import torch
import logging
import numpy as np
import argparse

from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer, BertModel

logging.basicConfig(format='%(asctime)s [%(levelname)s]: %(message)s', level=logging.DEBUG)
logging.getLogger().setLevel(logging.INFO)

# --- КОНСТАНТЫ ---
BERT_MODEL_NAME = 'small_bert/bert_en_uncased_L-2_H-128_A-2'

if torch.cuda.is_available():
    logging.info(f"GPU: {torch.cuda.get_device_name(0)} is available.")
    DEVICE = torch.device('cuda')
else:
    logging.info("No GPU available. Training will run on CPU.")
    DEVICE = torch.device('cpu')

# --- Подготовка и токенизация данных ---
class BertDataset(Dataset):
    def __init__(self, df, tokenizer, max_length=100):
        super(BertDataset, self).__init__()
        self.df=df
        self.tokenizer=tokenizer
        self.target=self.df['Target']
        self.max_length=max_length
        
    def __len__(self):
        return len(self.df)
    
    def __getitem__(self, idx):
        
        X = self.df['bodyText'].values[idx]
        y = self.target.values[idx]
        
        inputs = self.tokenizer.encode_plus(
            X,
            pad_to_max_length=True,
            add_special_tokens=True,
            return_attention_mask=True,
            max_length=self.max_length,
        )
        ids = inputs["input_ids"]
        token_type_ids = inputs["token_type_ids"]
        mask = inputs["attention_mask"]

        x = {
            'ids': torch.tensor(ids, dtype=torch.long).to(DEVICE),
            'mask': torch.tensor(mask, dtype=torch.long).to(DEVICE),
            'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long).to(DEVICE)
            }
        y = torch.tensor(y, dtype=torch.long).to(DEVICE)
        
        return x, y

# --- Определение модели ---
class SentimentBERT(nn.Module):
    def __init__(self, bert_model):
        super().__init__()
        self.bert_module = bert_model
        self.dropout = nn.Dropout(0.1)
        self.final = nn.Linear(in_features=128, out_features=3, bias=True) 
        
    def forward(self, inputs):
        ids, mask, token_type_ids = inputs['ids'], inputs['mask'], inputs['token_type_ids']
        x = self.bert_module(ids, mask, token_type_ids)
        x = self.dropout(x['pooler_output'])
        out = self.final(x)
        return out

# --- Цикл обучения ---
def train(epoch, model, dataloader, loss_fn, optimizer, max_steps=None):
    model.train()
    total_acc, total_count = 0, 0
    log_interval = 50
    start_time = time.time()

    for idx, (inputs, label) in enumerate(dataloader):
        optimizer.zero_grad()
        predicted_label = model(inputs)
        
        loss = loss_fn(predicted_label, label)
        loss.backward()
        optimizer.step()
        
        total_acc += (predicted_label.argmax(1) == label).sum().item()
        total_count += label.size(0)
        
        if idx % log_interval == 0:
            elapsed = time.time() - start_time
            print(
                "Epoch {:3d} | {:5d}/{:5d} batches "
                "| accuracy {:8.3f} | loss {:8.3f} ({:.3f}s)".format(
                    epoch, idx, len(dataloader), total_acc / total_count, loss.item(), elapsed
                )
            )
            total_acc, total_count = 0, 0
            start_time = time.time()

        if max_steps is not None:
            if idx == max_steps:
                return {'loss': loss.item(), 'acc': total_acc / total_count}
    
    return {'loss': loss.item(), 'acc': total_acc / total_count}

# --- Цикл оценки качества модели ---
def evaluate(model, dataloader, loss_fn):
    model.eval()
    total_acc, total_count = 0, 0

    with torch.no_grad():
        for idx, (inputs, label) in enumerate(dataloader):
            predicted_label = model(inputs)
            loss = loss_fn(predicted_label, label)
            total_acc += (predicted_label.argmax(1) == label).sum().item()
            total_count += label.size(0)

    return {'loss': loss.item(), 'acc': total_acc / total_count}

# --- Главная функция ---
def train_and_evaluate(**params):

    logging.info("running with the following params :")
    logging.info(params)

    # Загрузка предварительно обученного токенизатора и модели BERT
    # поменяйте пути на те, которые используете
    tokenizer = BertTokenizer.from_pretrained('models/bert_uncased_L-2_H-128_A-2/vocab.txt')
    model = BertModel.from_pretrained('models/bert_uncased_L-2_H-128_A-2')
    
    # Параметры обучения
    epochs = int(params.get('epochs'))
    batch_size = int(params.get('batch_size'))
    learning_rate = float(params.get('learning_rate'))
    
    #  Загрузка данных
    df_train = pd.read_csv(params.get('training_file'))
    df_eval = pd.read_csv(params.get('validation_file'))
    df_test = pd.read_csv(params.get('testing_file'))

    # Создание загрузчиков данных
    train_ds = BertDataset(df_train, tokenizer, max_length=100)
    train_loader = DataLoader(dataset=train_ds,batch_size=batch_size, shuffle=True)
    eval_ds = BertDataset(df_eval, tokenizer, max_length=100)
    eval_loader = DataLoader(dataset=eval_ds,batch_size=batch_size)
    test_ds = BertDataset(df_test, tokenizer, max_length=100)
    test_loader = DataLoader(dataset=test_ds,batch_size=batch_size)
    
    # Создание модели
    classifier = SentimentBERT(bert_model=model).to(DEVICE)
    total_parameters = sum([np.prod(p.size()) for p in classifier.parameters()])
    model_parameters = filter(lambda p: p.requires_grad, classifier.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])
    logging.info(f"Total params : {total_parameters} - Trainable : {params} ({params/total_parameters*100}% of total)")
    
    # Оптимизатор и функция потерь
    optimizer = torch.optim.Adam([p for p in classifier.parameters() if p.requires_grad], learning_rate)
    loss_fn = nn.CrossEntropyLoss()

    # При пробном запуске выполнить лишь следующее
    logging.info(f'Training model with {BERT_MODEL_NAME}')
    if args.dry_run:
        logging.info("Dry run mode")
        epochs = 1
        steps_per_epoch = 1
    else:
        steps_per_epoch = None
        
    # Вперёд!
    for epoch in range(1, epochs + 1):
        epoch_start_time = time.time()
        train_metrics = train(epoch, classifier, train_loader, loss_fn=loss_fn, optimizer=optimizer, max_steps=steps_per_epoch)
        eval_metrics = evaluate(classifier, eval_loader, loss_fn=loss_fn)
        
        print("-" * 59)
        print(
            "End of epoch {:3d} - time: {:5.2f}s - loss: {:.4f} - accuracy: {:.4f} - valid_loss: {:.4f} - valid accuracy {:.4f} ".format(
                epoch, time.time() - epoch_start_time, train_metrics['loss'], train_metrics['acc'], eval_metrics['loss'], eval_metrics['acc']
            )
        )
        print("-" * 59)
    
    if args.dry_run:
        # При пробном запуске не выполнять оценку качества модели
        return None
    
    test_metrics = evaluate(classifier, test_loader, loss_fn=loss_fn)
    
    metrics = {
        'train': train_metrics,
        'val': eval_metrics,
        'test': test_metrics,
    }
    logging.info(metrics)
    
    # Сохранение модели и архитектуры в одном файле
    if params.get('job_dir') is None:
        logging.warning("No job dir provided, model will not be saved")
    else:
        logging.info("Saving model to {} ".format(params.get('job_dir')))
        torch.save(classifier.state_dict(), params.get('job_dir'))
    logging.info("Bye bye")
    
    
if __name__ == '__main__':
    # Создание аргументов
    parser = argparse.ArgumentParser()
    parser.add_argument('--training-file', required=True, type=str)
    parser.add_argument('--validation-file', required=True, type=str)
    parser.add_argument('--testing-file', type=str)
    parser.add_argument('--job-dir', type=str)
    parser.add_argument('--epochs', type=float, default=2)
    parser.add_argument('--batch-size', type=float, default=1024)
    parser.add_argument('--learning-rate', type=float, default=0.01)
    parser.add_argument('--dry-run', action="store_true")

    # Парсинг аргументов
    args, _ = parser.parse_known_args()

    # Запуск обучения
    train_and_evaluate(**vars(args))

Всё это замечательно, но, к сожалению, обучаться модель эта будет очень долго. На самом деле — речь идёт об обучении примерно 4,7 миллионов параметров. На Macbook Pro с процессором от Intel и с 16 Гб памяти один шаг обучения займёт порядка 3 секунд.

https://miro.medium.com/v2/resize:fit:595/1*G9aqdhK_OWr1qZS8AAg6JA.png — Один шаг обучения

Шаг длительностью в 3 секунды может оказаться слишком длительным в том случае, если нужно пройти 1283 шага и выполнить обработку данных в 10 эпохах…

Короче говоря — никакого праздника без GPU.

Как задействовать Vertex AI и начать праздновать?

Короткий ответ на этот вопрос состоит из двух слов: Docker и Gcloud.

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

В работе с Google есть один приятный момент: тот, кто создаёт проект, используя аккаунт Gmail, получает бонус в $300.

И, как обычно, когда речь идёт о переносе кода на некую платформу, популярным решением является использование Docker.

Докеризация скрипта

Создадим образ Docker, рассчитанный на использование GPU. В официальном репозитории Docker имеется множество подходящих образов. Я выбрал pytorch/pytorch:2.2.2-cuda11.8-cudnn8-runtime, так как я пользуюсь Pytorch 2.2.2. Выбирая образ, обратите внимание на то, чтобы он уже поддерживал бы CUDA. Иначе вам придётся самостоятельно организовать установку всего необходимого посредством Dockerfile. И, уж поверьте мне на слово, вам это не надо, за исключением, конечно, тех случаев, когда у вас в этом есть реальная необходимость.

Нижеприведённый Dockerfile организует предустановку всех необходимых CUDA‑зависимостей и драйверов. Он обеспечит нам возможность использовать их в собственном задании обучения и запустит скрипт main.py с аргументами, переданными при запуске контейнера.

FROM pytorch/pytorch:2.2.2-cuda11.8-cudnn8-runtime

WORKDIR /src
COPY . .
RUN pip install --upgrade pip && pip install -r requirements.txt

ENTRYPOINT ["python", "main.py"]

Сборка и отправка образа в Google Cloud

После того, как образ будет готов к сборке — надо его собрать и отправить в реестр образов. Это может быть любой реестр, который вам кажется подходящим, но Google Cloud предлагает службу, называемую Artefact Registry и предназначенную для хранения образов. Воспользовавшись этим реестром, вы облегчите себе задачу хранения образов в Google Cloud.

Поместите файл build.sh, код которого приведён ниже, в корень проекта, и проверьте, чтобы Dockerfile был бы на том же уровне:

# build.sh

export PROJECT_ID=<your-project-id>
export IMAGE_REPO_NAME=pt_bert_sentiment
export IMAGE_TAG=dev
export IMAGE_URI=eu.gcr.io/$PROJECT_ID/$IMAGE_REPO_NAME:$IMAGE_TAG

gcloud builds submit --tag $IMAGE_URI .

Запустите build.sh. После нескольких минут ожидания, необходимых для сборки образа, вы должны увидеть примерно следующее:

eu.gcr.io/<your-project-id>/pt_bert_sentiment:dev SUCCESS

Создание задания в Vertex AI

После того, как образ собран и отправлен в Artefact Registry, можно предложить Vertex AI запустить этот образ на любой машине, которая нам нужна. В том числе и на такой, которая оснащена мощными GPU! Как уже было сказано, Google даёт бонус в $300 при создании GCP‑проектов. Этого вполне достаточно для запуска нашей модели.

Сведения о ценах на разные конфигурации можно найти здесь. В нашем случае всё будет выглядеть так: мы возьмём машину n1-standard-4 за $0,24 в час и присоединим к ней GPU NVIDIA T4 за $0,40 в час.

https://miro.medium.com/v2/resize:fit:700/1*p1Sc0D78xW_HgFe04Owj_Q.png — Типы машин

https://miro.medium.com/v2/resize:fit:700/1*I7ADcxcsSUxZoIcwINoRTQ.png — Ускорители

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

Вам, кроме того, нужно будет передать обучающему скрипту аргументы. Синтаксические конструкции, используемые для этого в gcloud ai custom-jobs create, строятся из двух частей:

Аргументы, имеющие отношение к самому заданию: --region, --display-name, --worker-pool-spec, --service-account и --args.
Аргументы, относящиеся к обучению: --training-file, --epochs и прочее подобное.

Перед последними нужно поставить --args, чтобы указать на то, что всё, что идёт дальше, предназначено для обучающего Python-скрипта.

Например, если наш скрипт принимает 2 аргумента — x и y, то получится следующее: --args=x=1,y=2.

# job.sh

export PROJECT_ID=<your-project-id>
export BUCKET=<your-bucket-id>
export REGION="europe-west4"
export SERVICE_ACCOUNT=<your-service-account>
export JOB_NAME="pytorch_bert_training"
export MACHINE_TYPE="n1-standard-4"  # Тут можно описать необходимые GPU
export ACCELERATOR_TYPE="NVIDIA_TESLA_T4"
export IMAGE_URI="eu.gcr.io/$PROJECT_ID/pt_bert_sentiment:dev"


gcloud ai custom-jobs create \
--region=$REGION \
--display-name=$JOB_NAME \
--worker-pool-spec=machine-type=$MACHINE_TYPE,accelerator-type=$ACCELERATOR_TYPE,accelerator-count=1,replica-count=1,container-image-uri=$IMAGE_URI \
--service-account=$SERVICE_ACCOUNT \
--args=\
--training-file=gs://$BUCKET/data/train.csv,\
--validation-file=gs://$BUCKET/data/eval.csv,\
--testing-file=gs://$BUCKET/data/test.csv,\
--job-dir=gs://$BUCKET/model/model.pt,\
--epochs=10,\
--batch-size=128,\
--learning-rate=0.0001

Запуск задания на платформе Vertex AI

Запустим скрипт и перейдём в наш GCP-проект, в раздел Training, который можно найти в меню Vertex.

https://miro.medium.com/v2/resize:fit:700/1*VSmqgPskAaeQb2FPuGv4sA.png — Проверка проекта

Запустим скрипт и перейдём в консоль. Там должны отобразиться сведения о состоянии задания. Сначала это будет Pending, а потом — Training.

Для того чтобы убедиться в том, что в ходе обучения модели действительно используется GPU — можно проверить задание и его ресурсы.

https://miro.medium.com/v2/resize:fit:700/1*Os5MCjlrzZ4RuR0MFmnpfQ.png — Проверка ресурсов, используемых заданием

Это указывает на то, что мы обучаем модель с помощью GPU. А значит — можно ожидать значительного роста скорости обучения! Заглянем в логи.

https://miro.medium.com/v2/resize:fit:682/1*qsAamE-BdY2lZ3QcTGvQbQ.png — Сведения об обучении модели

Сейчас на одну эпоху ушло 10 минут, а на CPU то же самое занимает целый час! Мы передали нагрузку по обучению модели платформе Vertex и ускорили процесс обучения. Это даёт нам возможность, например, запускать и другие задания, используя различные конфигурации, и при этом не перегружать ноутбуки.

А как насчёт итоговой точности модели? Получается, что после прохождения 10 эпох она находится в районе 94-95%. Можно дать модели ещё поучиться, и посмотреть — улучшит ли это её точность (или можно добавить коллбэк, останавливающий обучение пораньше, чтобы избежать переобучения модели).

https://miro.medium.com/v2/resize:fit:671/1*QHXQq2knoduoRdtYAaf51g.png — Сведения об обучении модели

Каковы результаты обучения модели?

https://miro.medium.com/v2/resize:fit:700/1*13kcuq9rEO5Pdrqp8MtmvQ.png — Результаты обучения модели

Пришло время праздновать! 🍻

О, а приходите к нам работать? 🤗 💰

Мы в wunderfund.io занимаемся высокочастотной алготорговлей с 2014 года. Высокочастотная торговля — это непрерывное соревнование лучших программистов и математиков всего мира. Присоединившись к нам, вы станете частью этой увлекательной схватки.

Мы предлагаем интересные и сложные задачи по анализу данных и low latency разработке для увлеченных исследователей и программистов. Гибкий график и никакой бюрократии, решения быстро принимаются и воплощаются в жизнь.

Сейчас мы ищем плюсовиков, питонистов, дата-инженеров и мл-рисерчеров.

Присоединяйтесь к нашей команде

Никакого праздника без GPU: дообучение BERT на Vertex AI