Tensorflow：累积张量的梯度

Question

  

TL; DR：，您只需跳到下面黄色框中的问题即可。

假设我有一个 Encoder-Decoder 神经网络，其编码器和解码器的权重分别为W_1和W_2。让我们将Z表示为编码器的输出。使用批次大小n训练网络，并且将根据批次的平均损失值计算所有梯度（如下图所示，L_hat是每样本损失L）。

我想要实现的是，在向后传递中，先操纵Z的梯度，然后再将其传递给编码器的权重W_1。假设是经过某种方式修改的渐变运算符，其适用于以下条件：

如上所述，在同步传递的情况下（首先计算Z的修改后的梯度，然后向下传播到W_1）非常容易实现（Jacobian乘法是使用tf.gradients中的grad_ys完成）：

def modify_grad(grad_z):
    # do some modifications

grad_z = tf.gradients(L_hat, Z)
mod_grad_z = modify_grad(grad_z)
mod_grad_w1 = tf.gradients(Z, W_1, mod_grad_z)

  

问题是，我需要在多个批次中累积张量grad_z的梯度 Z。由于它的形状是动态的（如上图所示，其中一个维度为None），因此我无法定义tf.Variable来存储它。此外，批次大小n在训练期间可能会改变。如何存储多个批次中grad_z的平均值？

PS ：我只想结合ArXiv:1810.04650的最佳优化训练，ArXiv:1609.02132的异步网络训练和ArXiv:1711.00489的批量大小调度。