Keras для TPU в Google Colaboratory (бесплатно!)

Играем с официальным примером Fashion MNIST

(Обновление: обратитесь к официальной документации Tensorflow 2.1+ о том, как использовать TPU с TF на Colab)

Google начал предоставлять пользователям доступ к TPU в Google Colaboratory (Colab) БЕСПЛАТНО! Google Colab уже предоставляет бесплатный доступ к графическому процессору (1 ядро ​​K80) для всех, а TPU в 10 раз дороже. (Google Cloud в настоящее время взимает 4,50 доллара США за TPU в час и 0,45 доллара США за ядро ​​K80 в час.) Какая захватывающая новость.

Хорошие ребята из сообщества fast.ai уже поделились некоторыми тестами набора данных MNIST. Однако кажется, что MNIST слишком прост для решения, что мы не видим преимуществ TPU.

Я нашел официальный пример записной книжки «Fashion MNIST с Keras и TPU» в репозитории Github tensorflow / tpu. Я внес некоторые изменения в ноутбук и трижды запускал его в разных средах (TPU, GPU, CPU) в качестве альтернативного теста. Вот что я изменил:

1. Создал набор проверки из обучающего набора
2. Измените batch_size с 1024 на 512 и epochs с 10 на 20. (Нет особой причины. Просто хочу попробовать разные гиперпараметры.)
3. Подсчитайте результаты тестирования и проверки после обучения.

Исходный код / ​​Блокноты

А вот записные книжки (они сохраняются как Gists. Используйте раскрывающееся меню «Файл / Просмотр на Github», чтобы открыть их на Github.)

Технические подробности

Все, что требуется, - это просто преобразовать модель Keras в модель TPU с помощью tf.contrib.tpu.keras_to_tpu_model:

tpu_model = tf.contrib.tpu.keras_to_tpu_model(
    model,
    strategy=tf.contrib.tpu.TPUDistributionStrategy(
        tf.contrib.cluster_resolver.TPUClusterResolver(
            tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
    )
)

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

Полученные результаты

С 3 сверточными слоями и 2 полносвязными слоями мы видим, что TPU уже обеспечивает почти двукратную производительность с точки зрения скорости по сравнению с GPU:

Точности обучения / проверки / тестирования должны быть очень близкими в разных средах, поскольку они имеют одни и те же гиперпараметры (хотя мы не устанавливали одно и то же начальное число):

Интересно то, что точность проверки TPU после обучения отличается от того, что Керас сообщил во время обучения. Я не уверен, почему, но, вероятно, это как-то связано с тем фактом, что TPU использует вычисления со смешанной точностью, и мы перемещаем график в пост-обучение ЦП, которое, как я полагаю, использует вычисления с одинарной точностью.

Заключение

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

Дополнительные материалы

Мы кратко представили некоторые обзоры и тесты TPU в апреле: