TensorFlow: Error Saving Model - ValueError: Cannot pickle Tensor -- its value is not known statically:

Viewed 293

I'm looking to save a regression model using TensorFlow/Keras, but upon doing so I am getting the following error;

ValueError: Cannot pickle Tensor -- its value is not known statically: Tensor("variational_gaussian_process_15/Softplus:0", shape=(), dtype=float64).

The model is;

tf.random.set_seed(11)
# For numeric stability, set the default floating-point dtype to float64
tf.keras.backend.set_floatx('float64')
x_range = np.array([0,1])

# Build mvodel.
num_inducing_points = 40
model = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=[1078]),
    tf.keras.layers.Dense(1, activation='sigmoid'),
    tf.keras.layers.Dense(1, kernel_initializer='ones', use_bias=False),
    tfp.layers.VariationalGaussianProcess(
        num_inducing_points=num_inducing_points,
        kernel_provider=RBFKernelFn(),
        event_shape=[1],
        inducing_index_points_initializer=tf.constant_initializer(
            np.linspace(*x_range, num=num_inducing_points,
                        dtype=X_train.dtype)[..., np.newaxis]),
        unconstrained_observation_noise_variance_initializer=(
            tf.constant_initializer(np.array(0.54).astype(X_train.dtype))),
    ),
])
batch_size = 32
loss = lambda y_train, rv_y: rv_y.variational_loss(
    y_train, kl_weight=np.array(batch_size, X_train.dtype) / X_train.shape[0])
model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3), loss=loss)
print(model)
print("Start training..\n")
history = model.fit(X_train, y_train, batch_size=batch_size, epochs=50, validation_split=0.2)
print("Done.")
[print(np.squeeze(w.numpy())) for w in model.weights];

There's also a class included, which I think might be causing the issue, as atop of the error message when saving is a warning;

WARNING:tensorflow:Skipping full serialization of Keras layer <__main__.RBFKernelFn object at 0x7f872a0299a0>, because it is not built.

The RBFKernel looks as follows;

class RBFKernelFn(tf.keras.layers.Layer):
  def __init__(self, **kwargs):
    super(RBFKernelFn, self).__init__(**kwargs)
    dtype = kwargs.get('dtype', None)

    self._amplitude = self.add_variable(
            initializer=tf.constant_initializer(0),
            dtype=dtype,
            name='amplitude')
    
    self._length_scale = self.add_variable(
            initializer=tf.constant_initializer(0),
            dtype=dtype,
            name='length_scale')

  def call(self, x):
    # Never called -- this is just a layer so it can hold variables
    # in a way Keras understands.
    return x

  @property
  def kernel(self):
    return tfp.math.psd_kernels.ExponentiatedQuadratic(
      amplitude=tf.nn.softplus(0.1 * self._amplitude),
      length_scale=tf.nn.softplus(5. * self._length_scale)
    )

Lastly, I tried the following to save

0 Answers
Related