Converting Keras to Tensorflow -- Pruning + Probability

Question

I have a code that mixes Tensorflow Probability (requires TF 2.00) with Keras Pruning, pruning weights of a first Dense Layer and providing inputs for TF probability, having both codes (Keras + TF) in the same model. The code:

from tensorflow_model_optimization.sparsity import keras as sparsity
from tensorflow.python import keras
import numpy as np
tf.disable_v2_behavior()

epochs = 50
num_train_samples = x1.shape[0]
end_step = 500
print('End step: ' + str(end_step))

tfd = tfp.distributions

input_shape=x1.shape

init = tf.global_variables_initializer()

with tf.Session() as sess:
    sess.run(init)
    pruned_model = tf.keras.Sequential([
        sparsity.prune_low_magnitude(
            tf.keras.layers.Dense(1, activation='relu'),**pruning_params),
      tfp.layers.DistributionLambda(lambda t: tfd.Normal(loc=t, scale=1))
    ])

    negloglik = lambda x, rv_x: -rv_x.log_prob(x)


    pruned_model.compile(optimizer=tf.keras.optimizers.Adam(lr=0.0001), loss=negloglik)

    callbacks = [
        pruning_callbacks.UpdatePruningStep(),
        pruning_callbacks.PruningSummaries(log_dir="D:\Python\logs2", profile_batch=0)]

    # ERROR HERE IN .fit()

    pruned_model.fit(x1,y, epochs=50, verbose=True, batch_size=16,callbacks=callbacks)

    yhat2 = pruned_model(np.array(dataframe.iloc[:,1]).T.astype(np.float32).reshape(-1,1)[650:800])
    mean02 = tf.convert_to_tensor(yhat2)
    mean2 = sess.run(mean02)    
    stddev2 = yhat2.stddev()
    mean_plus_2_std2 = sess.run(mean2 - 3. * stddev2)
    mean_minus_2_std2 = sess.run(mean2 + 3. * stddev2)

Details of the error:

  File "<ipython-input-129-a0ad4118e99e>", line 1, in <module>
    pruned_model.fit(x1,y, epochs=50, verbose=True, batch_size=16,callbacks=callbacks)

  File "C:\Users\Rubens\Anaconda3\lib\site-packages\tensorflow\python\keras\engine\training.py", line 806, in fit
    shuffle=shuffle)

  File "C:\Users\Rubens\Anaconda3\lib\site-packages\tensorflow\python\keras\engine\training.py", line 2503, in _standardize_user_data
    self._set_inputs(cast_inputs)

  File "C:\Users\Rubens\Anaconda3\lib\site-packages\tensorflow\python\training\tracking\base.py", line 456, in _method_wrapper
    result = method(self, *args, **kwargs)

  File "C:\Users\Rubens\Anaconda3\lib\site-packages\tensorflow\python\keras\engine\training.py", line 2773, in _set_inputs
    outputs = self.call(inputs, training=training)

  File "C:\Users\Rubens\Anaconda3\lib\site-packages\tensorflow\python\keras\engine\sequential.py", line 256, in call
    outputs = layer(inputs, **kwargs)

  File "C:\Users\Rubens\Anaconda3\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 594, in __call__
    self._maybe_build(inputs)

  File "C:\Users\Rubens\Anaconda3\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 1713, in _maybe_build
    self.build(input_shapes)

  File "C:\Users\Rubens\Anaconda3\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\pruning_wrapper.py", line 175, in build
    self.prunable_weights = self.layer.get_prunable_weights()

  File "C:\Users\Rubens\Anaconda3\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\prune_registry.py", line 169, in get_prunable_weights
    return [getattr(layer, weight) for weight in cls._weight_names(layer)]

  File "C:\Users\Rubens\Anaconda3\lib\site-packages\tensorflow_model_optimization\python\core\sparsity\keras\prune_registry.py", line 169, in <listcomp>
    return [getattr(layer, weight) for weight in cls._weight_names(layer)]

AttributeError: 'Dense' object has no attribute 'kernel'

My question is: how to convert a Keras layer (prune_low_magnitude) to Tensorflow, or how to convert a Tensorflow Probability layer (tfp.layers.DistributionLambda) to Keras and train the model properly ?

The notebook uses Keras==2.2.4 and Tensorflow==2.0.0a0

Answer 1

I found the solution. I installed:

! pip install --upgrade tfp-nightly
! pip install tf_nightly
! pip install tf_estimator_nightly