python-3.x - Tensorflow 2.0 (Python) - 图形模式下的并行循环

Question

假设有不同形状的变量：

tvs = model.trainable_variables

如果想对这些变量应用一些自定义优化操作，每个变量都应该有自己的图表：

ops = []
for idx, tv in enumerate(tvs):
    @tf.function
    def op(xs):
        # Some operations
        tv.assign_add(xs[idx])
    ops.append(op)

然后在其他地方应用这些操作（例如在训练循环中）：

@tf.function
def step_ops(xs):
    for op in ops : op(xs)

其中step_ops将张量列表（例如，变量的梯度）作为输入。

问题是，op insides ops 会在 step_ops 中并行运行吗？

如果没有，如何让它并行运行？或者有没有更好的方法来实现这一点？

非常感谢。

标签： python-3.xtensorflow