如何计算 tensorflow_probability 层上的梯度？

Question

我想使用 tensorflow_probability 计算 tf.GradientTape() 层上的梯度。这是相当简单的使用正常的，例如，密集层

inp = tf.random.normal((2,5))
layer = tf.keras.layers.Dense(10)

with tf.GradientTape() as tape:
    out = layer(inp)
    loss = tf.reduce_mean(1-out)
grads = tape.gradient(loss, layer.trainable_variables)
print(grads)
[<tf.Tensor: shape=(5, 10), dtype=float32, numpy=
 array([[ 0.04086879,  0.04086879, -0.02974391,  0.04086879,  0.04086879,
          0.04086879, -0.02974391,  0.04086879, -0.02974391, -0.07061271],
        [ 0.01167339,  0.01167339, -0.02681615,  0.01167339,  0.01167339,
          0.01167339, -0.02681615,  0.01167339, -0.02681615, -0.03848954],
        [ 0.00476769,  0.00476769, -0.00492069,  0.00476769,  0.00476769,
          0.00476769, -0.00492069,  0.00476769, -0.00492069, -0.00968838],
        [-0.00462376, -0.00462376,  0.05914849, -0.00462376, -0.00462376,
         -0.00462376,  0.05914849, -0.00462376,  0.05914849,  0.06377225],
        [-0.11682947, -0.11682947, -0.06357963, -0.11682947, -0.11682947,
         -0.11682947, -0.06357963, -0.11682947, -0.06357963,  0.05324984]],
       dtype=float32)>,
 <tf.Tensor: shape=(10,), dtype=float32, numpy=
 array([-0.05, -0.05, -0.1 , -0.05, -0.05, -0.05, -0.1 , -0.05, -0.1 ,
        -0.05], dtype=float32)>]

但是如果我使用 DenseReparameterization 来执行此操作，则 grads 不会注册。

inp = tf.random.normal((2,5))
layer = tfp.layers.DenseReparameterization(10)

with tf.GradientTape() as tape:
    out = layer(inp)
    loss = tf.reduce_mean(1-out)
grads = tape.gradient(loss, layer.trainable_variables)
print(grads)
[None, None, None]

谁能告诉我如何解决这个问题，以便对梯度进行录音和注册？

Answer 1

啊哈，就是这样！我正在使用 tf v2.1.0。显然这不适用于 tensorflow_probability。我会尽快升级。谢谢 gobrewers14。