пятница, 21 сентября 2018 г.

TensorFlow: стремительное исполнение (eager execution), продвинутое использование

Использование объектов состояния в течение стремительного исполнения

При графовом исполнении состояние программы (например, переменные) сохраняются в глобальных коллекциях и их жизненный цикл управляется tf.Session объектом. В ходе стремительного исполнения жизненный цикл объектов состояния определяется жизненным циклом соотвествующих им Python объектов.

Переменные и объекты

В ходе стремительного исполнения переменные сохраняются до тех пор пока не удалена последняя ссылка на объект.

with tf.device("gpu:0"):
  v = tfe.Variable(tf.random_normal([1000, 1000]))
  v = None  # v больше не использует GPU память

Объектно-ориентированное сохранение

tfe.Checkpoint может сохранять в контрольные точки и восстанавливать из контрольных точек tfe.Variables:

x = tfe.Variable(10.)

checkpoint = tfe.Checkpoint(x=x)  # save as "x"

x.assign(2.)   # Назначаем новое значение переменным и сохраняем
save_path = checkpoint.save('./ckpt/')

x.assign(11.)  # Изменяем переменную после сохранения

# Восстановление значений их контрольной точки
checkpoint.restore(save_path)

print(x)  # => 2.0

Для сохранения и загрузки моделей tfe.Checkpoint сохраняет внутренние объекты состояния без необходимости скрытых переменных. Чтобы записать состояние модели оптимизатор и глобальный шаг передают их (переменные) в tfe.Checkpoint:

model = MyModel()
optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
checkpoint_dir = ‘/path/to/model_dir’
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
root = tfe.Checkpoint(optimizer=optimizer,
                  model=model,
                  optimizer_step=tf.train.get_or_create_global_step())

root.save(file_prefix=checkpoint_prefix)
# или
root.restore(tf.train.latest_checkpoint(checkpoint_dir))

Объектно-ориентированные метрики

tfe.metrics хранятся как объекты. Обновляем метрику передачей новых данных вызываемому объекту и извлекаем результат, используя tfe.metrics.result метод, например:

m = tfe.metrics.Mean("loss")
m(0)
m(5)
m.result()  # => 2.5
m([8, 9])
m.result()  # => 5.5

Сводки и TensorBoard

TensorBoard - это инструмент визуализации для понимания, отладки и оптимизирования процесса тренировки модели. Он использует события сводок, которые записываются в ходе исполнения программы.

tf.contrib.summary совместим с обоими окружениями: стремительного и графового исполнения. Операции сводки, такие как tf.contrib.summary.scalar, вставляются в ходе создания модели. Например, для записи сводок единожды каждые 100 глобальных шагов:

writer = tf.contrib.summary.create_file_writer(logdir)
# возвращает переменную глобального шага
global_step=tf.train.get_or_create_global_step()  

writer.set_as_default()

for _ in range(iterations):
  global_step.assign_add(1)
  # Необходимо включить record_summaries метод
  with tf.contrib.summary.record_summaries_every_n_global_steps(100):
    # код модели располагается здесь
    tf.contrib.summary.scalar('loss', loss)
     ...

Сложные темы автоматического дифференцирования

Динамические модели

tf.GradientTape может быть использована в динамических моделях. Этот пример для алгоритма поиска линии обратного отслеживания выглядит как обычный код NumPy, за исключением того, что присутствуют градиенты и это диффернцируемо, несмотря на сложный поток контроля:

def line_search_step(fn, init_x, rate=1.0):
  with tf.GradientTape() as tape:
    # Переменные автоматически записаны, а вручную смотрим тензор
    tape.watch(init_x)
    value = fn(init_x)
  grad = tape.gradient(value, init_x)
  grad_norm = tf.reduce_sum(grad * grad)
  init_value = value
  while value > init_value - rate * grad_norm:
    x = init_x - rate * grad
    value = fn(x)
    rate /= 2.0
  return x, value

Дополнительные функции для вычисления градиентов

tf.GradientTape - это мощный интерфейс для вычисления градиентов, но существует другое Autograd-style API доступное для автоматического дифференцирования. Эти функции полезны, если записывать математический код только с тензорами и градиентными функциями и без tfe.Variables:

  • tfe.gradients_function - Возвращает функцию, которая вычисляет производные ее вводного параметра функции с учетом ее аргументов. Вводный параметр функция должна возвращать скалярное значение. Когда возвращенная функция вызвана она возвращает лист tf.Tensor объектов: один элемент для каждого аргумента вводной функции. Ввиду того, что все должно быть передано как параметр функции, это становится громоздким если присутствует зависимость на многие тренируемые параметры.
  • tfe.value_and_gradients_function - Схожа с tfe.gradients_function, но когда возвращенная функция вызвана, она возвращает значение от вводной функции вместе с листом производных вводной функции с учетом ее аргументов.

В следующем примере tfe.gradients_function принимает квадратичную функцию как аргумент и возвращает функцию, которая вычисляет частичные производные квадрата с учетом ее вводов. Чтобы вычислить производную от квадрата 3, grad (3.0) возвращает 6.

def square(x):
  return tf.multiply(x, x)

grad = tfe.gradients_function(square)

square(3.)  # => 9.0
grad(3.)    # => [6.0]

# Производная второго порядка квадрата:
gradgrad = tfe.gradients_function(lambda x: grad(x)[0])
gradgrad(3.)  # => [2.0]

# Производная третьего порядка равна None:
gradgradgrad = tfe.gradients_function(lambda x: gradgrad(x)[0])
gradgradgrad(3.)  # => [None]


# С потоком контроля:
def abs(x):
  return x if x > 0. else -x

grad = tfe.gradients_function(abs)

grad(3.)   # => [1.0]
grad(-3.)  # => [-1.0]

Кастомные градиенты

Кастомные градиенты - это простой путь переопределить градиенты при стремительном и графовом исполнении. В следующей функции определяем градиент с учетом вводов, выводов, или промежуточных результатов. Например, вот простой способ обрезать норму градиента при обратном проходе:

@tf.custom_gradient
def clip_gradient_by_norm(x, norm):
  y = tf.identity(x)
  def grad_fn(dresult):
    return [tf.clip_by_norm(dresult, norm), None]
  return y, grad_fn

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

def log1pexp(x):
  return tf.log(1 + tf.exp(x))
grad_log1pexp = tfe.gradients_function(log1pexp)

# Вычисление градиента работает хорошо при x = 0.
grad_log1pexp(0.)  # => [0.5]

# Однако, при x = 100 проваливается из-за числовой нестабильности.
grad_log1pexp(100.)  # => [nan]

Здесь log1pexp функция может быть аналитически упрощена с кастомным градиентом. Реализация ниже переиспользует значение для tf.exp(x), которое вычисляется в ходе прохода вперед, делая это более эффективным за счет удаления излишних вычислений:

@tf.custom_gradient
def log1pexp(x):
  e = tf.exp(x)
  def grad(dy):
    return dy * (1 - 1 / (1 + e))
  return tf.log(1 + e), grad

grad_log1pexp = tfe.gradients_function(log1pexp)

# Как и раньше вычисление градиента работает хорошо при x = 0.
grad_log1pexp(0.)  # => [0.5]

# И вычисление градиента также работает при x = 100.
grad_log1pexp(100.)  # => [1.0]

Производительность

Вычисление автоматически передается на GPU в ходе стремительного исполнения. Если необходим контроль над тем где происходит вычисление, тогда можно заключить это в tf.device('/gpu:0') блок (или CPU эквивалент):

import time

def measure(x, steps):
  # TensorFlow инициализирует GPU при первом использовании, 
  # исключая из учета времени
  tf.matmul(x, x)
  start = time.time()
  for i in range(steps):
    x = tf.matmul(x, x)
    # Удостоверяемся, что исполнили оп (op), 
    # а не просто включили в очередь
    _ = x.numpy()  
  end = time.time()
  return end - start

shape = (1000, 1000)
steps = 200
print("Time to multiply a {} matrix by itself {} times:".format(shape, 
                                                                steps))

# Исполняем на CPU:
with tf.device("/cpu:0"):
  print("CPU: {} secs".format(measure(tf.random_normal(shape), steps)))

# Исполняем на GPU, если доступно:
if tfe.num_gpus() > 0:
  with tf.device("/gpu:0"):
    print("GPU: {} secs".format(measure(tf.random_normal(shape), steps)))
else:
  print("GPU: not found")

Вывод (точные числа зависят от аппаратных средств):

Time to multiply a (1000, 1000) matrix by itself 200 times:
CPU: 4.614904403686523 secs
GPU: 0.5581181049346924 secs

tf.Tensor объект может быть скопирован на другое устройство для исполнения его операций:

x = tf.random_normal([10, 10])

x_gpu0 = x.gpu()
x_cpu = x.cpu()

_ = tf.matmul(x_cpu, x_cpu)    # Runs on CPU
_ = tf.matmul(x_gpu0, x_gpu0)  # Runs on GPU:0

if tfe.num_gpus() > 1:
  x_gpu1 = x.gpu(1)
  _ = tf.matmul(x_gpu1, x_gpu1)  # Runs on GPU:1

Опорные отметки (benchmarks)

Для вычислительно-емких моделей, таких как ResNet50 тренируемый на GPU, производительность стремительного исполнения сопоставима с графовым исполнением. Но эта разница становится больше для моделей с меньшими вычислениями.

Работа с графами

Хотя стремительное исполнение делает разработку и отладку более интерактивными, TensorFlow графовое исполнение имеет преимущества распреленной тренировки, оптимизаций производительности, и производственного использования. Однако, написание графового кода может отличаться от написания обычного Python кода и быть более трудным в отладке.

Для создания и тренировки граф-сконструированных моделей Python программа сначала создает граф, представляющий вычисление, затем вызывает Session.run для отправки графа на исполнение в C++ окружение. Это предоставляет:

  • Автоматическую дифференцировку, используя статический autodiff
  • Простое использование на сервере, независимо от платформы
  • Основанные на графах оптимизации (общее устранение вложенных выражений, раскладывание констант)
  • Компиляцию и ядерный поток
  • Автоматическое распределение и репликацию (размещение на узлы распределенной системы)

Производственное применение кода, написанного для стремительного исполнения, более трудоемко: либо генерирование графа из модели, либо исполнение Python интерпретатора и кода непосредственно на сервере.

Написание совместимого кода

Код, написанный для стремительного исполнения, также построит граф в ходе графового исполнения. Делаем это просто выполняя код в новой Python сессии, где стремительное исполнение не включено.

Большинство TensorFlow операций работают при стремительном исполнении, но необходимо помнить некоторые особенности:

  • Исползуйте tf.data для обработки ввода, вместо очередей. Это быстрее и проще.
  • Используйте API объектно-ориентированных слоев - такие как tf.keras.layers и tf.keras.Model - ввиду того что они имеют явное хранилище переменных.
  • Большая часть кода моделей работает одинаково при стремительном и графовом исполнении, но есть исключения. (Например, динамические модели используют Python поток контроля для изменения вычислений, основанных на вводных данных.)
  • Как только стремительное исполнение включено с помощью tf.enable_eager_execution, оно не может быть выключено. Необходимо начать новую Python сессию для возвращения к графовому исполнению.

Лучше всего писать код подходящий и для стремительного, и для графового исполнения. Это дает возможность экспериментирования и отладки при стремительном исполнении с преимуществами распределенной производительности при графовом исполнении.

Пишите, отлаживайте, итерируйте при стремительном исполнении, затем импортируйте граф модели для производственного использования. Используйте tfe.Checkpoint для сохранения и восстановления переменных моделей, это позволяет перемещаться между окружениями стремительного и графового исполнения.

Использование стремительного исполнения в графовом окружении

Выборочно включаем стремительное исполнение в TensorFlow графовом окружении, используя tfe.py_func. Это используется, когда tf.enable_eager_execution() не вызван.

def my_py_func(x):
  x = tf.matmul(x, x)  # Можно использовать tf ops
  print(x)  # but it's eager!
  return x

with tf.Session() as sess:
  x = tf.placeholder(dtype=tf.float32)
  # Вызываем функция стремительного исполнения в графе!
  pf = tfe.py_func(my_py_func, [x], tf.float32)
  sess.run(pf, feed_dict={x: [[2.0]]})  # [[4.0]]

О базовом использовании стремительного исполнения читайте предыдущий пост.