在Tensorflow中做Mandlebrot示例的更有效方法

Question

在tensorflow網站上有一個示例 ，其中TF用於計算Mandlebrot集。 以下是相關代碼段：

# Compute the new values of z: z^2 + x
zs_ = zs*zs + xs

# Have we diverged with this new value?
not_diverged = tf.abs(zs_) < 4

# Operation to update the zs and the iteration count.
#
# Note: We keep computing zs after they diverge! This
#       is very wasteful! There are better, if a little
#       less simple, ways to do this.
#
step = tf.group(
  zs.assign(zs_),
  ns.assign_add(tf.cast(not_diverged, tf.float32))
  )

相關的引言是“有更好的方法，即使不太簡單，也可以做到這一點”。 誰知道有什么更好的方法嗎？

我搞亂了TF中的光線跟蹤，遇到了90％的像素會聚的情況，但是我一直在重新計算它們，因為我不知道如何在不犧牲張量的情況下更新整個張量的子集加快在任何地方使用向量運算的好處。

Answer 1

我想出了一種方法。

關鍵是使用tf.where查找尚未發散的像素， tf.gather_nd將這些像素拉入較小的陣列，對這些特定像素執行更新步驟，然后使用tf.scatter_nd （或另一個“分散” ”）將稀疏更新應用於表示狀態的某些變量。

在我的用例中，這節省了大量時間。