Обучение моделей на облачных TPU в Google Colab за 4 простых шага

Я обучил модель нейронного машинного перевода (NMT) на TPU и теперь чувствую себя волшебником ...

У вас есть простая старая модель TensorFlow, которая требует слишком больших вычислительных ресурсов для обучения на стандартном рабочем ноутбуке. Я понял. Я тоже был там, и, честно говоря, наблюдать, как мой ноутбук два раза подряд ломался после попытки обучить на нем модель, было больно смотреть.

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

TL; DR: В этой статье показано, как легко обучить любую модель TensorFlow на TPU с очень небольшим изменением кода.

Что такое ТПУ?

Tensor Processing Unit (TPU) - это ускоритель, созданный инженерами Google Brain по индивидуальному заказу, который специализируется на обучении глубоких и дорогостоящих в вычислительном отношении моделей машинного обучения.

Давайте посмотрим на вещи в перспективе, чтобы дать вам представление о том, насколько крутым и мощным является TPU. Стандартный процессор Intel MacBook Pro может выполнять несколько операций за такт. Стандартный серийный графический процессор может выполнять десятки тысяч операций за цикл. Современный TPU может выполнять сотни тысяч операций за цикл (иногда до 128K OPS).

Чтобы понять масштаб, представьте, что вы используете эти устройства для печати книги. ЦП может печатать посимвольно. Графический процессор может печатать несколько слов за раз. ТПУ? Ну, он может печатать целую страницу за раз. Это потрясающая скорость и мощность, которые теперь в нашем распоряжении; привет Google за предоставление всем доступа к высокопроизводительному оборудованию.

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

Предварительная обработка MNIST

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

Давайте сначала начнем с извлечения и предварительной обработки нашего набора данных. Это не должно быть большой проблемой:

import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical
# This can be any dataset
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.reshape([x_train.shape[0], 784])
x_test = x_test.reshape([x_test.shape[0], 784])
x_train = x_train / 255
x_test = x_test / 255
y_train = to_categorical(y_train, 10)
y_test = to_categorical(y_test, 10)
print (x_train.shape, y_train.shape)
print (x_test.shape, y_test.shape)
# >> (60000, 784), (60000, 10)
# >> (10000, 784), (10000, 10)

Я постарался сделать это руководство как можно более полным, чтобы вы могли тренировать свои модели с невероятной скоростью и чувствовать себя на вершине мира. Чтобы обучить нашу модель на TPU, мы сделаем 4 шага:

1. Подключиться к доступному экземпляру TPU
2. Инициализировать распределенную стратегию обучения
3. Постройте нашу модель в соответствии с указанной стратегией
4. Обучите модель и почувствуйте себя Мастером

Давайте прямо в это погрузимся!

Примечание. Не стесняйтесь предварительно обрабатывать наборы данных с помощью tf.data или конвертировать их в TFRecords , если вам это нравится. Эти шаги не являются обязательными, но могут быть полезны при работе с экземплярами с нехваткой памяти в случае, если вы работаете с очень большими наборами данных.

Подключение к TPU

Когда я возился с TPU на Colab, подключение к одному было самым утомительным. Поиск в Интернете и просмотр руководств заняли несколько часов, но, наконец, я смог это сделать.

Сначала нам нужно подключить наш ноутбук Colab, работающий локально, к среде выполнения TPU. Чтобы изменить время выполнения, просто щелкните вкладку Runtime на панели навигации. Отобразится раскрывающееся меню, в котором вы можете выбрать параметр Change runtime type. Появится всплывающее окно, в котором вы можете выбрать параметр TPU из раскрывающегося списка.

Чтобы подключиться к экземпляру TPU, нам нужно создать TPUClusterResolver, который принимает доступное устройство и подготавливает его для нас:

Это должно подключить наш текущий сеанс Colab к доступному экземпляру TPU. Перейдем к инициализации стратегии и, наконец, к вызову TPU.

Примечание. Если вы испытываете трудности с подключением к экземпляру TPU даже после нескольких попыток и повторных запусков, вы можете вместо этого изменить тип среды выполнения на GPU . Код по-прежнему компилируется и работает эффективно.

Инициализация стратегии распределенного обучения

Теперь, когда мы изменили среду выполнения и получили список доступных TPU, нам нужно вызвать TPU, создав распределенную стратегию обучения - оболочку вокруг модели, которая делает ее совместимой с многоядерным обучением.

resolver дает нам доступ к TPU, так что мы, наконец, можем построить на нем параллельно распределенный конвейер. Это необходимый шаг, потому что TPU - это распределенный обучающий процессор, а не одноядерный, как традиционный ЦП. Используя этот метод стратегии, запрыгнуть на борт TPU очень просто! Он должен предоставить вам список доступных экземпляров:

Построение нашей модели в рамках стратегии распределенного обучения

Имея под рукой стратегию обучения, мы можем приступить к построению нашей модели, используя эту стратегию как таковую:

Это похоже на много жаргона, но все, что он делает, это берет модель regulartf.keras (которая обычно запускается на ЦП), помещает ее в TPU и автоматически распределяет ее по всем доступным ядрам TPU.

Обучение нашей модели

Это захватывающая часть процесса. Наконец-то мы можем обучить нашу модель на Cloud TPU бесплатно. Имея в своих руках такую ​​мощь, давайте воспользуемся ею и потренируемся по MNIST (я знаю ... очень антиклиматично).

Убедитесь, что количество экземпляров полностью делится на параметр steps_per_epoch, чтобы все экземпляры использовались во время
обучения. Например, в нашем обучающем наборе 60000 экземпляров. 60000
делится на 50, что означает, что все наши экземпляры передаются в модель без каких-либо остатков. Если вы нажмете запустить, через некоторое время он должен начать обучение на экземпляре TPU:

Заключительный обзор

С этим скоро должно начаться обучение. Colab загрузит TPU и загрузит на него архитектуру модели. Вскоре вы должны увидеть классический макет индикатора выполнения Keras в выводе терминала. Поздравляю! Вы только что успешно использовали TPU.

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

После прочтения этой статьи я хочу, чтобы вы знали, что тренироваться на ускоренном оборудовании может каждый, независимо от опыта. Это лишь один из многих шагов, предпринятых Google AI для демократизации машинного обучения и искусственного интеллекта в массах.

Некоторые наблюдения о тренировках в ТПУ

Модель нейронного машинного перевода, которую я написал (для TensorFlow Docs), заняла у меня меньше минуты, чтобы обучиться на TPU. Я был удивлен, потому что одна и та же модель тренировалась на процессоре более 4 часов (что, вероятно, объясняет, почему мой ноутбук дважды ломался и не хватало памяти).

Моя модель показала очень высокую точность - выше, чем я достиг при обучении на CPU. По производительности и скорости Cloud TPU не имеет себе равных, когда дело доходит до качественного обучения!

Примечание: вы можете найти мой код здесь. Удачи!

В двух словах

Google всегда приносит подарки, когда запускает новые игрушки для машинного обучения, с которыми мы можем повозиться и поиграть. Процессор TPU, безусловно, является благом для сообщества машинного обучения, поскольку он играет важную роль в демократизации ИИ - он дает каждому возможность использовать ускоренное оборудование независимо от демографии. Благодаря TPU исследования и эксперименты достигают рекордных высот, и теперь люди, как никогда раньше, занимаются машинным обучением!

Я надеюсь, что эта статья помогла вам обучить ваши модели более эффективным и элегантным способом. Если у вас есть какие-либо вопросы об использовании TPU или вы хотите поговорить в целом о технологиях и машинном обучении, не стесняйтесь отвечать на них ниже в комментариях или поймайте меня в Twitter или LinkedIn. Обычно я отвечаю в течение дня.

А пока я поймаю тебя в следующем!

Статья Ришаба Ананда

Призыв к действию… своего рода

Хотите узнать о новейших и величайших технологиях и рассказать о последних достижениях в области машинного обучения, науки о данных и технологий? Ловите другие мои статьи или подписывайтесь на них! Ваши отзывы и постоянная поддержка очень много значат и побуждают меня продолжать писать качественный контент для вашего обучения!