Я пытаюсь вычислить медиану тензора для всех партий и всех процессов. Однако я получил ошибку TypeError: Expected int32, got None of type 'NoneType' instead. при вызове optimizer.compute_gradients(loss). Кажется, что вычисление градиентов плохо работает с операцией horovod allgather. Простая иллюстрация того, чего я хотел бы достичь, состоит в следующем:
my_tensor = compute_my_tensor() # has dimension (None, N) for example
gathered_my_tensor = hvd.allgather(my_tensor) # this will concatenate my_tensor of all processes along first dimension
median = get_median(gathered_my_tensor)
loss = compute_loss(my_tensor, median) # say subtract each element of my_tensor by the median, then an L2 loss
optimizer = tf.compat.v1.train.AdamOptimizer()
optimizer = hvd.DistributedOptimizer(optimizer)
gradients = optimizer.compute_gradients(loss) # this step throws the error
Очевидно, что операция allgather() не дифференцируема, поэтому она встречает значение None при вычислении градиентов? У кого-нибудь есть идеи, как решить эту проблему?