我正在尝试理解 tensorflow 中 GPU 上的并行性,因为我需要将它应用到更丑陋的图表上。
import tensorflow as tf
from datetime import datetime
with tf.device('/device:GPU:0'):
var = tf.Variable(tf.ones([100000], dtype=tf.dtypes.float32), dtype=tf.dtypes.float32)
@tf.function
def foo():
return tf.while_loop(c, b, [i], parallel_iterations=1000) #tweak
@tf.function
def b(i):
var.assign(tf.tensor_scatter_nd_update(var, tf.reshape(i, [-1,1]), tf.constant([0], dtype=tf.dtypes.float32)))
return tf.add(i,1)
with tf.device('/device:GPU:0'):
i = tf.constant(0)
c = lambda i: tf.less(i,100000)
start = datetime.today()
with tf.device('/device:GPU:0'):
foo()
print(datetime.today()-start)
在上面的代码中,var 是一个长度为 100000 的张量,其元素如上所示更新。当我将 parallel_iterations 值从 10、100、1000、10000 更改时。即使明确提到了 parallel_iterations 变量,也几乎没有任何时间差(全部为 9.8 秒)。
我希望这些在 GPU 上并行发生。我该如何实施?