Обучение моделей Keras с помощью TensorFlow Cloud

Посмотреть на TensorFlow.org

Запустить в Google Colab

Посмотреть исходный код на GitHub

Скачать блокнот

Вступление

TensorFlow Cloud - это пакет Python, который предоставляет API-интерфейсы для плавного перехода от локальной отладки к распределенному обучению в Google Cloud. Он упрощает процесс обучения моделей TensorFlow в облаке в один простой вызов функции, требующий минимальной настройки и никаких изменений в вашей модели. TensorFlow Cloud автоматически решает задачи, связанные с облаком, такие как создание экземпляров виртуальных машин и стратегии распространения для ваших моделей. Это руководство продемонстрирует, как взаимодействовать с Google Cloud через TensorFlow Cloud, а также широкий спектр функций, предоставляемых в TensorFlow Cloud. Начнем с простейшего варианта использования.

Настраивать

Мы начнем с установки TensorFlow Cloud и импорта пакетов, которые нам понадобятся в этом руководстве.

pip install -q tensorflow_cloud

import tensorflow as tf
import tensorflow_cloud as tfc

from tensorflow import keras
from tensorflow.keras import layers

Обзор API: первый сквозной пример

Начнем со сценария обучения модели Keras, такого как следующий CNN:

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

model = keras.Sequential(
    [
        keras.Input(shape=(28, 28)),
        # Use a Rescaling layer to make sure input values are in the [0, 1] range.
        layers.experimental.preprocessing.Rescaling(1.0 / 255),
        # The original images have shape (28, 28), so we reshape them to (28, 28, 1)
        layers.Reshape(target_shape=(28, 28, 1)),
        # Follow-up with a classic small convnet
        layers.Conv2D(32, 3, activation="relu"),
        layers.MaxPooling2D(2),
        layers.Conv2D(32, 3, activation="relu"),
        layers.MaxPooling2D(2),
        layers.Conv2D(32, 3, activation="relu"),
        layers.Flatten(),
        layers.Dense(128, activation="relu"),
        layers.Dense(10),
    ]
)

model.compile(
    optimizer=keras.optimizers.Adam(),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=keras.metrics.SparseCategoricalAccuracy(),
)

model.fit(x_train, y_train, epochs=20, batch_size=128, validation_split=0.1)

Чтобы обучить эту модель в Google Cloud, нам просто нужно добавить вызов run() в начале скрипта перед импортом:

tfc.run()

Вам не нужно беспокоиться о задачах, связанных с облаком, таких как создание экземпляров виртуальных машин и стратегии распространения при использовании TensorFlow Cloud. API включает интеллектуальные значения по умолчанию для всех параметров - все настраивается, но многие модели могут полагаться на эти значения по умолчанию.

После вызова run() TensorFlow Cloud:

Подготовьте свой скрипт Python или блокнот для распространения.
Преобразуйте его в образ Docker с необходимыми зависимостями.
Выполните задание по обучению на виртуальной машине с графическим процессором GCP.
Потоковая передача соответствующих журналов и информации о вакансиях.

Конфигурация виртуальной машины по умолчанию - 1 главный и 0 рабочих с 8 ядрами ЦП и 1 графическим процессором Tesla T4.

Конфигурация Google Cloud

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

Создать проект GCP;
Включить службы платформы AI;
Создать учетную запись службы;
Скачать ключ авторизации;
Создайте корзину Cloud Storage.

Подробные инструкции по первоначальной настройке можно найти в README TensorFlow Cloud , а дополнительный пример настройки показан в блоге TensorFlow .

Общие рабочие процессы и облачное хранилище

В большинстве случаев вы захотите получить свою модель после обучения в Google Cloud. Для этого крайне важно перенаправить сохранение и загрузку в облачное хранилище во время удаленного обучения. Мы можем направить TensorFlow Cloud в нашу корзину Cloud Storage для различных задач. Сегмент хранилища можно использовать для сохранения и загрузки больших наборов обучающих данных, хранения журналов обратных вызовов или весов моделей, а также для сохранения файлов обученных моделей. Для начала давайте настроим fit() для сохранения модели в облачном хранилище и настроим мониторинг TensorBoard для отслеживания прогресса обучения.

def create_model():
    model = keras.Sequential(
        [
            keras.Input(shape=(28, 28)),
            layers.experimental.preprocessing.Rescaling(1.0 / 255),
            layers.Reshape(target_shape=(28, 28, 1)),
            layers.Conv2D(32, 3, activation="relu"),
            layers.MaxPooling2D(2),
            layers.Conv2D(32, 3, activation="relu"),
            layers.MaxPooling2D(2),
            layers.Conv2D(32, 3, activation="relu"),
            layers.Flatten(),
            layers.Dense(128, activation="relu"),
            layers.Dense(10),
        ]
    )

    model.compile(
        optimizer=keras.optimizers.Adam(),
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=keras.metrics.SparseCategoricalAccuracy(),
    )
    return model

Давайте сохраним журналы TensorBoard и контрольные точки модели, созданные во время обучения, в нашей корзине облачного хранилища.

import datetime
import os

# Note: Please change the gcp_bucket to your bucket name.
gcp_bucket = "keras-examples"

checkpoint_path = os.path.join("gs://", gcp_bucket, "mnist_example", "save_at_{epoch}")

tensorboard_path = os.path.join(  # Timestamp included to enable timeseries graphs
    "gs://", gcp_bucket, "logs", datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
)

callbacks = [
    # TensorBoard will store logs for each epoch and graph performance for us.
    keras.callbacks.TensorBoard(log_dir=tensorboard_path, histogram_freq=1),
    # ModelCheckpoint will save models after each epoch for retrieval later.
    keras.callbacks.ModelCheckpoint(checkpoint_path),
    # EarlyStopping will terminate training when val_loss ceases to improve.
    keras.callbacks.EarlyStopping(monitor="val_loss", patience=3),
]

model = create_model()

Здесь мы загрузим наши данные напрямую из Keras. Как правило, рекомендуется хранить набор данных в корзине облачного хранилища, однако TensorFlow Cloud также может вмещать наборы данных, хранящиеся локально. Это описано в разделе, посвященном нескольким файлам этого руководства.

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

API-интерфейс TensorFlow Cloud предоставляет функцию remote() для определения того, выполняется ли код локально или в облаке. Это позволяет отдельно обозначать параметры fit() для локального и удаленного выполнения и предоставляет средства для легкой отладки без перегрузки вашей локальной машины.

if tfc.remote():
    epochs = 100
    callbacks = callbacks
    batch_size = 128
else:
    epochs = 5
    batch_size = 64
    callbacks = None

model.fit(x_train, y_train, epochs=epochs, callbacks=callbacks, batch_size=batch_size)

Epoch 1/5
938/938 [==============================] - 6s 3ms/step - loss: 0.4867 - sparse_categorical_accuracy: 0.8556
Epoch 2/5
938/938 [==============================] - 3s 3ms/step - loss: 0.0638 - sparse_categorical_accuracy: 0.9796
Epoch 3/5
938/938 [==============================] - 3s 3ms/step - loss: 0.0411 - sparse_categorical_accuracy: 0.9870
Epoch 4/5
938/938 [==============================] - 3s 3ms/step - loss: 0.0333 - sparse_categorical_accuracy: 0.9895
Epoch 5/5
938/938 [==============================] - 3s 3ms/step - loss: 0.0278 - sparse_categorical_accuracy: 0.9910
<tensorflow.python.keras.callbacks.History at 0x7fa57866c4e0>

Сохраним модель в GCS после завершения обучения.

save_path = os.path.join("gs://", gcp_bucket, "mnist_example")

if tfc.remote():
    model.save(save_path)

Мы также можем использовать это ведро хранилища для создания образа Docker вместо вашего локального экземпляра Docker. Для этого просто добавьте свою docker_image_bucket_name параметр docker_image_bucket_name .

tfc.run(docker_image_bucket_name=gcp_bucket)

После обучения модели мы можем загрузить сохраненную модель и просмотреть журналы TensorBoard, чтобы отслеживать производительность.

model = keras.models.load_model(save_path)

!tensorboard dev upload --logdir "gs://keras-examples-jonah/logs/fit" --name "Guide MNIST"

Масштабные проекты

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

Точки входа: поддержка скриптов Python и записных книжек Jupyter

Ваш вызов API run() не всегда будет содержаться в том же скрипте Python, что и код обучения вашей модели. Для этого мы предоставляем параметр entry_point . Параметр entry_point можно использовать для указания скрипта Python или записной книжки, в которой находится код обучения вашей модели. При вызове run() из того же сценария, что и ваша модель, используйте entry_point умолчанию для entry_point None .

зависимости от `pip`

Если ваш проект требует дополнительных зависимостей pip , можно указать дополнительные необходимые библиотеки, включив файл requirements.txt . Просто поместите в этот файл список всех необходимых зависимостей, и TensorFlow Cloud выполнит их интеграцию в вашу облачную сборку.

Блокноты Python

TensorFlow Cloud также можно запускать из записных книжек Python. Кроме того, указанная entry_point может быть записной книжкой, если это необходимо. Следует помнить о двух ключевых отличиях между TensorFlow Cloud на ноутбуках и скриптами:

При вызове run() из записной книжки необходимо указать сегмент облачного хранилища для создания и хранения образа Docker.
Аутентификация GCloud происходит полностью через ваш ключ аутентификации, без спецификации проекта. Пример рабочего процесса с использованием TensorFlow Cloud из записной книжки приведен в разделе «Собираем все вместе» этого руководства.

Многофайловые проекты

Если ваша модель зависит от дополнительных файлов, вам нужно только убедиться, что эти файлы находятся в том же каталоге (или подкаталоге) указанной точки входа. Каждый файл, который хранится в том же каталоге, что и указанная entry_point будет включен в образ Docker, а также любые файлы, хранящиеся в подкаталогах, смежных с entry_point . Это также верно для зависимостей, которые могут вам понадобиться, которые не могут быть получены через pip

В качестве примера настраиваемой точки входа и многофайлового проекта с дополнительными зависимостями пакетов, взгляните на этот многофайловый пример в облачном репозитории TensorFlow . Для краткости мы просто включим в пример вызов run() :

tfc.run(
    docker_image_bucket_name=gcp_bucket,
    entry_point="train_model.py",
    requirements="requirements.txt"
)

Конфигурация машины и распределенное обучение

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

Распределение нескольких сотрудников

Здесь мы можем использовать COMMON_MACHINE_CONFIGS для обозначения 1 главного процессора и 4 рабочих графических процессоров.

tfc.run(
    docker_image_bucket_name=gcp_bucket,
    chief_config=tfc.COMMON_MACHINE_CONFIGS['CPU'],
    worker_count=2,
    worker_config=tfc.COMMON_MACHINE_CONFIGS['T4_4X']
)

По умолчанию TensorFlow Cloud выбирает лучшую стратегию распространения для конфигурации вашего компьютера с помощью простой формулы с использованием предоставленных параметров chief_config , worker_config и worker_count .

Если указанное количество графических процессоров больше нуля, будет выбран tf.distribute.MirroredStrategy .
Если количество рабочих больше нуля, в зависимости от типа ускорителя будут выбраны tf.distribute.experimental.MultiWorkerMirroredStrategy или tf.distribute.experimental.TPUStrategy .
В противном случае будет выбран tf.distribute.OneDeviceStrategy .

Распределение TPU

Обучим ту же модель на TPU, как показано:

tfc.run(
    docker_image_bucket_name=gcp_bucket,
    chief_config=tfc.COMMON_MACHINE_CONFIGS["CPU"],
    worker_count=1,
    worker_config=tfc.COMMON_MACHINE_CONFIGS["TPU"]
)

Стратегия индивидуального распространения

Чтобы указать настраиваемую стратегию распространения, отформатируйте код обычным образом, как в соответствии с распределенным учебным руководством, и установите для distribution_strategy значение None . Ниже мы укажем нашу собственную стратегию распространения для той же модели MNIST.

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

mirrored_strategy = tf.distribute.MirroredStrategy()
with mirrored_strategy.scope():
  model = create_model()

if tfc.remote():
    epochs = 100
    batch_size = 128
else:
    epochs = 10
    batch_size = 64
    callbacks = None

model.fit(
    x_train, y_train, epochs=epochs, callbacks=callbacks, batch_size=batch_size
)

tfc.run(
    docker_image_bucket_name=gcp_bucket,
    chief_config=tfc.COMMON_MACHINE_CONFIGS['CPU'],
    worker_count=2,
    worker_config=tfc.COMMON_MACHINE_CONFIGS['T4_4X'],
    distribution_strategy=None
)

Пользовательские образы Docker

По умолчанию TensorFlow Cloud использует базовый образ Docker, предоставленный Google и соответствующий вашей текущей версии TensorFlow. Однако при необходимости вы также можете указать собственный образ Docker в соответствии с вашими требованиями к сборке. В этом примере мы укажем образ Docker из более старой версии TensorFlow:

tfc.run(
    docker_image_bucket_name=gcp_bucket,
    base_docker_image="tensorflow/tensorflow:2.1.0-gpu"
)

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

Возможно, вам будет полезно пометить свои облачные задания определенными метками или транслировать журналы вашей модели во время облачного обучения. Хорошая практика - поддерживать надлежащую маркировку всех заданий в облаке для ведения учета. Для этого run() принимает словарь меток, содержащий до 64 пар ключ-значение, которые видны из журналов сборки Cloud. Доступ к таким журналам, как производительность эпох и внутреннее сохранение модели, можно получить, используя ссылку, предоставленную при выполнении tfc.run или распечатать на локальном терминале с stream_logs флага stream_logs .

job_labels = {"job": "mnist-example", "team": "keras-io", "user": "jonah"}

tfc.run(
    docker_image_bucket_name=gcp_bucket,
    job_labels=job_labels,
    stream_logs=True
)

Собираем все вместе

Для углубленного Colab, который использует многие функции, описанные в этом руководстве, следуйте этому примеру, чтобы обучить современную модель распознавать породы собак по фотографиям с помощью извлечения функций.