Saving Probablistic Layers

I built a model that ends with a probabilistic layer (tfp.layers.IndependentNormal, if it makes a difference, and tensorflow version 2.11.0) and it’s necessary for both training and inference that the model outputs a distribution. When I try to save the model, however, it throws up this error:

OperatorNotAllowedInGraphError: Iterating over a symbolic tf.Tensor is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

When I tried subclassing the layer in order to override the get_config method it worked, but the loaded model could then only output tensors and not probability distributions.

If anyone can solve either of these issues (saving a tfp.layer or getting the subclasses layer to output a distribution), that would be amazing.

Thanks!

Hi @Tzvi_Blonder, Could you please provide the standalone code to reproduce the issue.Because i have tried with a sample model that contains the last layer as tfp.layers.IndependentNormal,but i did not face any error. Thank You.

Hi @Kiran_Sai_Ramineni, thank you for responding. The following code gives the same error:

inputs = keras.Input(shape=(35,))
X = Dense(24,activation=‘relu’)(inputs)
X = Dense(tfpl.IndependentNormal.params_size(1))(X)
output = tfpl.IndependentNormal(1)(X)

model = keras.Model(inputs=inputs,
outputs=output)

model.save(model_path)

When I replace the final layer (called output) with a Dense layer, there is no problem.

Hi - after seeing it work on your colab, I tried it on my own colab and it worked; I realized that colab uses tensorflow version 2.12.0, while I was using 2.11.0. I downloaded the new version and it worked.
Thank you!