为什么我不能对作为参数传递给 tf.function 的变量执行梯度？

问题描述

我的训练循环给了我以下警告：

警告：tensorflow：在最小化损失时，变量 ['noise:0'] 不存在梯度。

经过一些修补后，我确定这只发生在将噪声变量作为参数传递给我的损失函数（tf.function）时。下面的代码表明，当损失函数不是tf.function或者在函数中引用了全局噪声变量时，没有问题。它还表明，当噪声变量用作 tf.function 中的参数时，尝试对其执行梯度会导致错误：

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from  tensorflow_probability import distributions as tfd 
from tensorflow_probability import bijectors as tfb

constrain_positive = tfb.Shift(np.finfo(np.float64).tiny)(tfb.Exp())
noise = tfp.util.TransformedVariable(initial_value=.1,bijector=constrain_positive,dtype=np.float64,name="noise")
trainable_variables = [noise.variables[0]]
kernel = tfp.math.psd_kernels.ExponentiatedQuadratic()
optimizer = tf.keras.optimizers.Adam()
index_points = tf.constant([[0]],dtype=np.float64)
observations = tf.constant([0],dtype=np.float64)

# I can train noise when it is passed as an argument to a python function 
def loss_function_1(index_points,observations,kernel,observation_noise_variance):
    gp = tfd.GaussianProcess(kernel,index_points,observation_noise_variance=observation_noise_variance)
    return -gp.log_prob(observations)

with tf.GradientTape() as tape:
    nll_1 = loss_function_1(index_points,noise)
grad_1 = tape.gradient(nll_1,trainable_variables)
print(grad_1)
optimizer.apply_gradients(zip(grad_1,trainable_variables))

# I can train noise if it is used in a tf.function and not passed as an argument
@tf.function(autograph=False,experimental_compile=False)
def loss_function_2(index_points,kernel):
    gp = tfd.GaussianProcess(kernel,observation_noise_variance=noise)
    return -gp.log_prob(observations)

with tf.GradientTape() as tape:
    nll_2 = loss_function_2(index_points,kernel)
grad_2 = tape.gradient(nll_2,trainable_variables)
print(grad_2)
optimizer.apply_gradients(zip(grad_2,trainable_variables))

# I can train noise if it is passed as an argument to a tf.function if the tf.function
# uses the global variable
@tf.function(autograph=False,experimental_compile=False)
def loss_function_3(index_points,observation_noise_variance=noise)
    return -gp.log_prob(observations)

with tf.GradientTape() as tape:
    nll_3 = loss_function_3(index_points,noise)
grad_3 = tape.gradient(nll_3,trainable_variables)
print(grad_3)
optimizer.apply_gradients(zip(grad_3,trainable_variables))

# I cannot train noise if it is passed as an argument to a tf.function if the tf.function
# the local variable
@tf.function(autograph=False,experimental_compile=False)
def loss_function_4(index_points,observation_noise_variance=observation_noise_variance)
    return -gp.log_prob(observations)

with tf.GradientTape() as tape:
    nll_4 = loss_function_4(index_points,noise)
grad_4 = tape.gradient(nll_4,trainable_variables)
print(grad_4)
optimizer.apply_gradients(zip(grad_4,trainable_variables))

此代码打印：

[]
[]
[]
[无]

然后它返回错误信息：

ValueError：没有为任何变量提供梯度：['noise:0']。

理想情况下，我会获得 tf.function 的性能提升，所以我不想使用 loss_function_1。此外，我希望能够将不同的噪声变量传递给我的损失函数，因此我不想像在 loss_function_2 或 loss_function_3 中那样使用全局变量。

当我尝试对作为参数传递给 tf.function 的变量执行渐变时，为什么会得到 None？我该如何解决这个问题？

解决方法

你无法解决它，它的工作原理就是这样。

当您使用 tf.function 时，您将 Python 代码转换为静态图（特别是 DAG）。该图有一些输入节点和一些输出节点。 输入节点是函数的参数，输出节点是返回值。

在函数体内定义一个tf.Variable，或者等效地将一个tf.Variable作为函数参数传递，意味着每次调用静态图中都会创建一个新的变量节点并创建一个新变量每次你调用它，它就是你想要的。

实际上，当您拥有具有状态（tf.Variable 和类似的）的对象时，您不能在 tf.function 修饰的函数中定义它们，但您必须打破函数作用域并声明外部变量。

您可以使用声明全局变量的解决方案。更好的解决方案是将代码重构为更加面向对象，将变量声明为类的私有属性，以免全局公开变量对象。

我在 this article 中介绍了这种行为，您可以在其中找到有关如何重构代码以及在使用 tf.function 时如何思考的一些见解