如何修复tensorflow中的“变量不存在梯度”错误

Question

我正在用Tensorflow 2.1.15和Tensorflow概率0.13构建一个变分自动编码器（VAE）模型。
我的问题是我无法避免得到“变量不存在梯度”

代码如下：

import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
tfb = tfp.bijectors
tfpl = tfp.layers

import numpy as np

from tensorflow.keras import Sequential, Model
from tensorflow.keras.layers import (Dense, Flatten, Reshape, Concatenate, Conv2D, 
                                     UpSampling2D, BatchNormalization)

print( "tf.version", tf.__version__, "tfp.version", tfp.__version__ )

prior_distribution = tfd.MultivariateNormalDiag( loc = np.zeros( ( 2, ) ), scale_diag = np.ones( ( 2, ) ) )

kl_regularizer = tfpl.KLDivergenceRegularizer(
    distribution_b = prior_distribution,
    use_exact_kl = False,
    test_points_fn = lambda t : t.sample( 3 ),
    test_points_reduce_axis = None,
    weight = 1.0
    )

def reconstruction_loss(batch_of_images, decoding_dist):
    return - tf.reduce_mean( decoding_dist.log_prob( batch_of_images ) )

image_array = np.random.uniform( 0, 1, ( 100, 36, 36, 3 ) ).astype( np.float64 ) / 255.0
train_dataset = image_array[ 0 : 80 ]
test_dataset = image_array[ 80 : 100 ]

encoder = tf.keras.Sequential( [
    Conv2D( 3, 3, padding = "same", input_shape = ( 36, 36, 3 ) ),
    Flatten( input_shape = ( 36, 36, 3 ) ),
    Dense( tfpl.MultivariateNormalTriL.params_size( 2 ), name = "encoder_dense" ),
    tfpl.MultivariateNormalTriL( 2, activity_regularizer = kl_regularizer, dtype = tf.float64, name = "encoder_outdist" )
])

decoder = tf.keras.Sequential( [
    Dense( 36 * 36 * 3, input_shape = ( 2, ), name = "decoder_dense" ),
    Reshape( ( 36, 36, 3 ), name = "decoder_reshape" ),
    Conv2D( 3, 3, padding = "same", name = "decoder_conv2d" ),
    Flatten( name = "decoder_flatten" ),
    tfpl.IndependentBernoulli( event_shape = ( 36, 36, 3 ), convert_to_tensor_fn = tfd.Bernoulli.logits )
])

vae = tf.keras.Model( inputs=encoder.inputs, outputs=decoder( encoder.outputs ) )
vae.compile( optimizer = "adam", loss = reconstruction_loss )

print( "vae.summary()", vae.summary() )

vae.fit( train_dataset, epochs = 5, batch_size = 20 )

下面是在Google Colab上执行此操作的结果：

tf.version 2.5.0 tfp.version 0.13.0
Model: "model_37"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d_38_input (InputLayer) [(None, 36, 36, 3)]       0         
_________________________________________________________________
conv2d_38 (Conv2D)           (None, 36, 36, 3)         84        
_________________________________________________________________
flatten_38 (Flatten)         (None, 3888)              0         
_________________________________________________________________
encoder_dense (Dense)        (None, 5)                 19445     
_________________________________________________________________
encoder_outdist (Multivariat multiple                  0         
_________________________________________________________________
sequential_77 (Sequential)   multiple                  11748     
=================================================================
Total params: 31,277
Trainable params: 31,277
Non-trainable params: 0
_________________________________________________________________
vae.summary() None
Epoch 1/5
WARNING:tensorflow:Gradients do not exist for variables ['decoder_dense/kernel:0', 'decoder_dense/bias:0', 'decoder_conv2d/kernel:0', 'decoder_conv2d/bias:0'] when minimizing the loss.
WARNING:tensorflow:Gradients do not exist for variables ['decoder_dense/kernel:0', 'decoder_dense/bias:0', 'decoder_conv2d/kernel:0', 'decoder_conv2d/bias:0'] when minimizing the loss.
4/4 [==============================] - 1s 12ms/step - loss: 0.0196
Epoch 2/5
4/4 [==============================] - 0s 10ms/step - loss: 0.0092
Epoch 3/5
4/4 [==============================] - 0s 12ms/step - loss: 0.0023
Epoch 4/5
4/4 [==============================] - 0s 10ms/step - loss: 0.0029
Epoch 5/5
4/4 [==============================] - 0s 10ms/step - loss: 0.0029
<tensorflow.python.keras.callbacks.History at 0x7fd95cb15e50>

我相信，当这种情况发生时，它也会扰乱优化。我不知道如何解决这个问题。你能帮忙吗