Как я могу распараллелить автоматическое дифференцирование с помощью tf.GradientTape?

#python #multithreading #tensorflow2.0 #gradienttape

#питон #многопоточность #тензорный поток 2,0 #градиентная лента

Вопрос:

Я хотел бы автоматически дифференцировать довольно сложную функцию, которую я хочу распараллелить.

Я использую TensorFlow 2.x и использую tf.GradientTape для дифференцирования.

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

 import pdb import tensorflow as tf import threading  # This ThreadWithResult is from https://stackoverflow.com/a/65447493/1935801 and works fine on its own class ThreadWithResult(threading.Thread):  def __init__(self, group=None, target=None, name=None, args=(), kwargs={}, *, daemon=None):  def function():  self.result = target(*args, **kwargs)  super().__init__(group=group, target=function, name=name, daemon=daemon)  def my_function(x):  return x*x   x*x*x  def my_function_threaded(x):  def square(x):  result = x*x  return result   def cube(x):  result = x*x*x  return result   t1 = ThreadWithResult(target=square, args=(x,))  t2 = ThreadWithResult(target=cube, args=(x,))   t1.start()  t2.start()   t1.join()  t2.join()   y = t1.result   t2.result   return y  x = tf.constant(3.0) print("my_function(x) =", my_function(x)) print("my_function_threaded(x) =", my_function_threaded(x))  with tf.GradientTape() as tape:  tape.watch(x)  y = my_function(x)  dy_dx = tape.gradient(y, x, unconnected_gradients=tf.UnconnectedGradients.ZERO) print("Simple dy_dx", dy_dx)  with tf.GradientTape() as tape:  tape.watch(x)  y = my_function_threaded(x)  dy_dx = tape.gradient(y, x, unconnected_gradients=tf.UnconnectedGradients.ZERO) print("Threaded dy_dx", dy_dx)  

Как видно из вывода, показанного ниже, градиенты нарушаются, когда для одного и того же простого вычисления используется нарезание резьбы.

 my_function(x) = tf.Tensor(36.0, shape=(), dtype=float32) my_function_threaded(x) = tf.Tensor(36.0, shape=(), dtype=float32) Simple dy_dx tf.Tensor(33.0, shape=(), dtype=float32) Threaded dy_dx tf.Tensor(0.0, shape=(), dtype=float32)  

Любые предложения/идеи о том, как я мог бы парализовать свою функцию в GradientTape, будут высоко оценены?

Комментарии:

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

2. Спасибо @Ник ОДелл. И, да, можно было бы рассчитать градиенты для каждого потока, а затем объединить их, если бы это было так же просто, как в примере с игрушкой. К сожалению, как вы и предполагаете, пример с игрушкой-это чрезмерное упрощение реальной вещи. На самом деле это сеть взаимозависимых моделей, и, хотя, возможно, их не невозможно объединить после использования градиентной ленты в каждом потоке, она больше не будет «автоматической». Итак, мне было интересно (и я надеялся, что) существует какой-то способ параллельного выполнения такой сложной автоматической дифференциации без необходимости программировать бухгалтерию.

3. Совершенно понятно. Я расскажу вам, что мне удалось выяснить во время отладки. Я попытался запустить t1 и t2 один за другим, и это не помогло. Так что это не проблема безопасности потоков или, по крайней мере, не просто проблема безопасности потоков.

4. На каком устройстве вы запускаете свои коды? Процессор или графический процессор?

5. На процессоре tensorflow автоматически распараллеливает операции в графике tensorflow, вам не нужно выполнять потоковую обработку самостоятельно. Вместо этого вам нужно использовать tf.function для создания графика тензорного потока. На графическом процессоре ситуация намного сложнее, и большую часть времени операции будут выполняться последовательно, если они не поддерживаются.