Автор: Сюэчен Ли, стажер по разработке программного обеспечения
Источник: публичный аккаунт TensorFlow.
Eager Execution упрощает построение модели в TensorFlow, а Graph Execution обеспечивает оптимизацию для более быстрого выполнения модели и эффективности хранения. В этом сообщении блога показано, как написать код TensorFlow для преобразования модели, созданной с помощью API tf.keras с использованием Eager Execution, в граф и, наконец, развернуть эту модель в Cloud TPU с поддержкой API tf.estimator.
Примечание: ссылка на tf.keraswww.tensorflow.org/guide/keras
ссылка на tf.estimatorwoohoo.tensorflow.org/expensive/ES Тим…
В качестве примера мы используем обратимые остаточные сети (RevNet, Gomez и т. д.). В следующих разделах предполагается, что читатель имеет базовое представление о сверточных нейронных сетях и TensorFlow. Вы можете найти полный код для этой статьи здесь (чтобы убедиться, что код работает во всех настройках, настоятельно рекомендуется использовать tf-nightly или tf-nightly-gpu).
RevNet
Сети RevNet аналогичны остаточным сетям (ResNet, He и т. д.), за исключением того, что они обратимы и могут реконструировать промежуточные вычисления с учетом выходных данных. Одним из преимуществ этого метода является то, что мы можем экономить память, реконструируя активации, а не сохраняя их все в памяти во время обучения (напомним, что, поскольку этого требует цепное правило, нам нужны промежуточные результаты для вычисления информации о входном градиенте). Это позволяет нам обрабатывать большие пакеты и обучать более глубокие модели, чем обычное обратное распространение в традиционных архитектурах. В частности, этот метод реализуется путем определения сети с помощью набора искусно построенных уравнений:
Верхнее и нижнее уравнения определяют прямой расчет и его инверсию соответственно. Здесь x1 и x2 — входы (отделенные от общего входа x), y1 и y2 — выходы, а F и G — ConvNets. Это позволяет нам точно реконструировать активации во время обратного распространения, поэтому нет необходимости хранить эти данные во время обучения.
Используйте tf.keras.Model для определения прямых и обратных проходов.
Предполагая, что мы используем класс ResidualInner для создания экземпляров функций F и G, мы можем определить обратимый блок кода, создав подкласс tf.keras.Model, и определить прямой проход, заменив метод вызова, показанный в уравнении выше:
1 class Residual(tf.keras.Model):
2 def __init__(self, filters):
3 super(Residual, self).__init__()
4 self.f = ResidualInner(filters=filters, strides=(1, 1))
5 self.g = ResidualInner(filters=filters, strides=(1, 1))
6
7 def call(self, x, training=True):
8 x1, x2 = tf.split(x, num_or_size_splits=2, axis=self.axis)
9 f_x2 = self.f(x2, training=training)
10 y1 = f_x2 + x1
11 g_y1 = self.g(y1, training=training)
12 y2 = g_y1 + x2
13 return tf.concat([y1, y2], axis=self.axis)
Параметр обучения здесь используется для определения состояния нормализации партии. Когда Eager Execution включен, пакетные нормализованные скользящие средние автоматически обновляются, когда training=True . При выполнении эквивалентного графа нам нужно вручную получить пакетные нормализованные обновления, используя метод get_updates_for.
Чтобы создать обратный проход с эффективным использованием памяти, нам нужно использовать tf.GradientTape в качестве диспетчера контекста для отслеживания градиентов (только при необходимости): Примечание: ссылка tf.GradientTapewoohoo.tensorflow.org/API_docs/friends…
1 def backward_grads(self, y, dy, training=True):
2 dy1, dy2 = dy
3 y1, y2 = y
4
5 with tf.GradientTape() as gtape:
6 gtape.watch(y1)
7 gy1 = self.g(y1, training=training)
8 grads_combined = gtape.gradient(
9 gy1, [y1] + self.g.trainable_variables, output_gradients=dy2)
10 dg = grads_combined[1:]
11 dx1 = dy1 + grads_combined[0]
12 x2 = y2 - gy1
13
14 with tf.GradientTape() as ftape:
15 ftape.watch(x2)
16 fx2 = self.f(x2, training=training)
17 grads_combined = ftape.gradient(
18 fx2, [x2] + self.f.trainable_variables,output_gradients=dx1)
19 df = grads_combined[1:]
20 dx2 = dy2 + grads_combined[0]
21 x1 = y1 - fx2
22
23 x = x1, x2
24 dx = dx1, dx2
25 grads = df + dg
26
27 return x, dx, grads
Вы можете найти точный набор вычислений градиента в «Алгоритме 1» статьи (мы упростили промежуточный шаг, используя переменную z1 в коде). Этот алгоритм тщательно разработан таким образом, чтобы, учитывая выходные данные и градиент потерь по отношению к выходным данным, в каждом блоке обратимого кода мы могли вычислить градиент по отношению к входным и модельным переменным и восстановить входные данные. Вызовите функцию tape.gradient(y, x), чтобы вычислить градиент y относительно x. Мы также можем использовать параметр output_gradients для явного применения правила цепочки.
Используйте Eager Execution для ускорения прототипирования
Очевидным преимуществом использования Eager Execution для прототипирования является использование императивных операций. Мы можем получить результаты немедленно, без необходимости строить график, а затем инициализировать сеанс для запуска.
Например, мы проверяем нашу модель, сравнивая обратимые градиенты обратного распространения с градиентами, рассчитанными с помощью общего обратного распространения:
1 block = Residual()
2 x = tf.random_normal(shape=(N, C, H, W))
3 dy = tf.random_normal(shape=(N, C, H, W))
4 with tf.GradientTape() as tape:
5 tape.watch(x)
6 y = block(x)
7 # Compute true grads
8 dx_true = tape.gradient(y, x, output_gradients=dy)
9
10 # Compute grads from reconstruction
11 dx, _ = block.backward_grads(x, y, dy)
12
13 # Check whether the difference is below a certain 14 threshold
thres = 1e-6
15 diff_abs = tf.reshape(abs(dx - dx_true), [-1])
16 assert all(diff_abs < thres)
В приведенном выше фрагменте dx_true — это градиент, возвращаемый общим обратным распространением, а dx — это градиент, возвращаемый после выполнения обратимого обратного распространения. Eager Execution интегрирует собственный Python, поэтому такие функции, как all и abs, можно напрямую применять к тензорам.
Сохранение и загрузка контрольных точек с помощью tf.train.Checkpoint
Чтобы контрольные точки можно было сохранять и загружать с помощью Eager Execution и Graph Execution, команда TensorFlow рекомендует использовать API tf.train.Checkpoint.
Чтобы сохранить модель, мы создаем экземпляр tf.train.Checkpoint со всеми объектами, которые мы хотим сохранить. Этот экземпляр может включать нашу модель, оптимизатор, который мы используем, график скорости обучения и глобальные шаги:
1 checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer,
2 learning_rate=learning_rate, global_step=global_step)
Мы можем сохранить и восстановить конкретный обученный экземпляр следующим образом:
1 checkpoint.save(file_prefix)
2 checkpoint.restore(save_path)
Улучшите производительность Eager Execution с помощью tf.contrib.eager.defun
Из-за накладных расходов на интерпретацию кода Python Eager Execution иногда может быть медленнее, чем выполнение эквивалентного графа. Этот разрыв в производительности можно устранить, используя tf.contrib.eager.defun для компиляции функций Python, состоящих из операций TensorFlow, в вызываемые графы TensorFlow. При обучении моделей глубокого обучения мы обычно можем применять tf.contrib.eager.defun в трех основных местах:
- Форвардные вычисления
- Расчет инверсии градиента
- применить градиент к переменной
Например, мы можем определить прямые проходы и вычисления градиента следующим образом:
1 tfe = tf.contrib.eager
2 model.call = tfe.defun(model.call)
3 model.compute_gradients = tfe.defun(model.compute_gradients)
Чтобы определить шаг применения градиента оптимизатора, нам нужно обернуть его внутри другой функции:
1 def apply_gradients(optimizer, gradients, variables, global_step=None):
2 optimizer.apply_gradients(
3 zip(gradients, variables), global_step=global_step)
4 apply_gradients = tfe.defun(apply_gradients)
tf.contrib.eager.defun находится в активной разработке, и его применение представляет собой развивающуюся технологию. Для получения дополнительной информации см. строку документации. Примечание: его ссылка на строку документацииGitHub.com/tensorflow/…
Обертывание функции Python с помощью tf.contrib.eager.defun оптимизирует всю программу, заставляя API TensorFlow выполнять вызовы внутри функции Python для построения графика, а не выполнять операции немедленно. Не все функции Python могут быть успешно преобразованы в эквивалентные графики, особенно функции с динамическим потоком управления (например, if или while в содержимом Tensor). tf.contrib.autograph — это связанный инструмент, который увеличивает площадь поверхности кода Python, которую можно преобразовать в граф TensorFlow. По состоянию на август 2018 года работа по интеграции Autograph с использованием defun все еще продолжается. Примечание: ссылка tf.contrib.autographwoohoo.tensorflow.org/expensive/auto G…
Создайте конвейер ввода с помощью TFRecords и tf.data.Dataset.
Eager Execution совместим с API tf.data.Dataset. Мы можем читать файлы TFRecords:
1 dataset = tf.data.TFRecordDataset(filename)
2 dataset = dataset.repeat(epochs).map(parser).batch(batch_size)
Чтобы повысить производительность, мы также можем использовать функцию предварительной выборки и настроить num_parallel_calls.
Поскольку набор данных состоит из пар изображений и меток, циклически перемещаться по этому набору данных в Eager Execution очень просто. В этом случае нам даже не нужно явно определять итератор:
1 for image, label in dataset:
2 logits = model(image, training=True)
3 ...
Оберните модель Keras оценщиком и выполните ее как график
Поскольку API tf.keras также поддерживает построение графиков, те же самые модели, созданные с помощью Eager Execution, также можно использовать в качестве функций построения графиков, предоставляемых оценщику, но с немного измененным кодом. Чтобы изменить пример RevNet, созданный с помощью Eager Execution, мы просто оборачиваем модель Keras в model_fn и используем эту модель в соответствии с инструкциями API tf.estimator.
1 def model_fn(features, labels, mode, params):
2 model = RevNet(params["hyperparameters"])
3 if mode == tf.estimator.ModeKeys.TRAIN:
4 optimizer = tf.train.MomentumOptimizer(learning_rate, momentum)
5 logits, saved_hidden = model(features, training=True)
6 grads, loss = model.compute_gradients(saved_hidden, labels, training=True)
7 with tf.control_dependencies(model.get_updates_for(features)):
8 train_op = optimizer.apply_gradients(zip(grads, model.trainable_variables))
9 return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
Вы можете использовать API tf.data для определения input_fn, требуемого API tf.estimator, как обычно, и считывать данные из TFRecords.
Оберните модель Keras оценщиком TPU для обучения Cloud TPU
Оберните модель и входной конвейер с помощью Estimator, чтобы модель могла работать в Cloud TPU.
Необходимые шаги следующие: Настройка определенных конфигураций для Cloud TPU Переключитесь с tf.estimator.Estimator на tf.contrib.tpu.TPUEstimator. Используйте tf.contrib.tpu.CrossShardOptimizer для переноса распространенных оптимизаторов. Примечание. Ссылка на облачный TPUGitHub.com/tensorflow/…настроить ссылкуwoohoo.tensorflow.org/API_docs/friends…Ссылка на tf.contrib.tpu.TPUEstimatorwoohoo.tensorflow.org/API_docs/friends…ссылка на tf.contrib.tpu.CrossShardOptimizerwoohoo.tensorflow.org/API_docs/friends…
Для получения конкретных инструкций ознакомьтесь со сценарием оценки TPU в папке примеров RevNet. Мы надеемся использовать tf.contrib.tpu.keras_to_tpu_model для дальнейшего упрощения процесса запуска модели Keras на TPU в будущем. Примечание. Ссылка на скрипт TPU EstimatorGitHub.com/tensorflow/…ссылка на tf.contrib.tpu.keras_to_tpu_modelGitHub.com/tensorflow/…
Необязательно: Производительность модели
По сравнению с обычным обратным распространением с tf.GradientTape в сочетании с упрощением вычисления градиента без дополнительных прямых проходов мы можем выполнить обратимое обратное распространение RevNet всего с 25% вычислительных затрат.
Синие и оранжевые кривые на рисунке представляют скорость выборки в секунду для общего обратного распространения и обратимого обратного распространения, соответственно, по мере увеличения глобальных шагов. Этот рисунок взят из RevNet-104, обученного на одном Tesla P100 с использованием смоделированных данных ImageNet с размером пакета 32.
Чтобы проверить экономию памяти, мы строим график использования памяти во время обучения. Синие и черные кривые — общее и обратимое обратное распространение соответственно. На этом рисунке записано 100 итераций обучения режима диаграммы RevNet-104 с использованием смоделированных данных ImageNet с размером пакета 128. График создается mprof во время обучения на ЦП, поэтому мы тренируемся с тем же размером пакета, используя общее обратное распространение.
в заключении
На примере RevNet мы покажем, как быстро создавать прототипы моделей машинного обучения с помощью Eager Execution и API tf.keras. Это не только упрощает создание модели, но мы также можем легко преобразовать модель в оценщик и развернуть ее в Cloud TPU для повышения производительности. Вы можете найти полный код для этой статьи здесь. Кроме того, обязательно ознакомьтесь с другими примерами использования Eager Execution. Примечание: ссылка здесьGitHub.com/tensorflow/…Другие примеры ссылокGitHub.com/tensorflow/…