Open mkretsch327 opened 3 years ago
Hey Matt! I can reproduce, and add that it also happens with the nightly tensorflow/tfp, and with the release candidate and tensorflow==2.5.0, but I don't yet have good leads on fixing it.
Thanks @ColCarroll ! I wish I had a better idea to help on root cause, will keep an eye out here.
Same error here, cannot go for tensorflow 2.5 because of this. I am using python 3.9, though. I cannot save any model with DistributionLamda.
If it helps, I checked that the actual call to __call__
in Layer gets completely executed and then, when it returns, the calls to __exit__
of the CallContextManagers begin to be called, but it seems that the error already happened. I tried changing the check in base_layer_utils.is_subclassed
, called from __call__
, trying to force not using autograph, based on your script (no idea what I am doing), and it didn't work.
Hi, is there any hint on this? It continues happening, according to this example I did. Is #1382 related?
As a workaround, I confirm that weights can be stored and then loaded back in another declared model:
the_model = models.Model(...)
the_model.save_weights('file_name')
...
another_model = models.Model(...)
another_model.load_weights('file_name')
I had a similar problem, but when I tried it with tfp v0.12.2 and tf version 2.4.1 it gave me another error: https://github.com/tensorflow/probability/issues/1268 I also posted there the versions I used to make it work.
Also facing this issue using tfp==0.13.0 and tf==2.5.0
Found a workaround for this issue using tf==2.5.0 and tfp==0.12.2 (have not tested other versions):
By wrapping the call to tfp.layers.DistributionLambda
in a user-defined layer class, you can pass that class to the custom_objects
argument of tf.keras.models.load_model
. The model will save and load as normal.
Here's an example snippet I wrote for a model whose output layer is a 2D Gaussian RV:
class Normal2DLayer(tf.keras.layers.layer)
def __init__(self):
super(Normal2DLayer, self).__init__()
self.loc = tf.keras.layers.Dense(2, name='loc')
self.scale = tf.keras.layers.Dense(2, name='scale')
self.rotation = tf.keras.layers.Dense(1, name='rotation') # Parameterizing a 2D Gaussian as a rotated ellipsoid
def call(self, inputs):
loc = self.loc(inputs)
scale_vec = self.scale(inputs)
theta = self.rotation(inputs)
scale_diag = tf.linalg.diag(scale_vec)
rotation_mat = tf.reshape(tf.stack([tf.math.cos(theta), -tf.math.sin(theta), tf.math.sin(theta), tf.math.cos(theta)], axis=1), (tf.shape(theta)[0], 2, 2))
scale_mat = tf.matmul(rotation_mat, scale_diag)
params = [loc, scale_mat]
density = tfp.layers.DistributionLambda(
make_distribution_fn = lambda params: tfd.MultivariateNormalLinearOperator(loc=params[0], scale=tf.linalg.LinearOperatorFullMatrix(params[1])),
)(params)
return density
def get_config(self):
config = dict()
return config
@classmethod
def from_config(cls, config):
return cls(**config)
From here, as long as I pass the class definition to the custom_objects
argument, the normal keras save
and load_model
functions work as expected:
model.save(path_to_model)
model = tf.keras.models.load_model(path_to_model, custom_objects={"Normal2DLayer" : Normal2DLayer})
I am newbie (like week of self of learning), but getting same error with this example and with other unrelated one ("Probabilistic Bayesian Neural Networks" Git repo. When I just add save() after training. ) It uses DistributionLambda too and it seems to me this DistributionLambda.class doesn't implement/override some necessary function. (python3 -c 'import tensorflow as tf; print(tf.version)' output: 2.7.0 )
I am using TF2.7 and TFP0.15. If I pass the save_traces=False
argument to model.save()
it works.
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow.keras import layers
X = np.random.randn(100,10)
y = np.random.randn(100,1)
n_features = X.shape[-1]
n_outputs = y.shape[-1]
inputs = layers.Input(shape = (n_features,), name='features')
x = layers.Dense(16, activation='relu')(inputs)
x = layers.Dense(2)(x)
outputs = tfp.layers.IndependentNormal(n_outputs, name='outputs')(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs)
model.save("mymodel", save_traces=False)
del model
model = tf.keras.models.load_model("mymodel")
model.summary()
I found this in the Tensorflow documentation
The argument save_traces has been added to model.save, which allows you to toggle SavedModel function tracing. Functions are saved to allow the Keras to re-load custom objects without the original class definitons, so when save_traces=False, all custom objects must have defined get_config/from_config methods.
So I guess as long as we use classes that have the get_config/from_config methods, we should be fine. And as we can see here, the DistributionLambda
layer that serves as base for many tfp layers, has these methods.
I have the same issue using tensorflow 2.7.0 and tensorflow-probability 0.15.0. The suggestion from @frazane to use save_traces=False
did not work for me. But the solution by @jdbates to use a custom layer did.
I am using TF2.7 and TFP0.15. If I pass the
save_traces=False
argument tomodel.save()
it works.import numpy as np import tensorflow as tf import tensorflow_probability as tfp from tensorflow.keras import layers X = np.random.randn(100,10) y = np.random.randn(100,1) n_features = X.shape[-1] n_outputs = y.shape[-1] inputs = layers.Input(shape = (n_features,), name='features') x = layers.Dense(16, activation='relu')(inputs) x = layers.Dense(2)(x) outputs = tfp.layers.IndependentNormal(n_outputs, name='outputs')(x) model = tf.keras.Model(inputs=inputs, outputs=outputs) model.save("mymodel", save_traces=False) del model model = tf.keras.models.load_model("mymodel") model.summary()
I found this in the Tensorflow documentation
The argument save_traces has been added to model.save, which allows you to toggle SavedModel function tracing. Functions are saved to allow the Keras to re-load custom objects without the original class definitons, so when save_traces=False, all custom objects must have defined get_config/from_config methods.
So I guess as long as we use classes that have the get_config/from_config methods, we should be fine. And as we can see here, the
DistributionLambda
layer that serves as base for many tfp layers, has these methods.
this worked for me, thanks a lot, the original error I had was:
OperatorNotAllowedInGraphError: Iterating over a symbolic
tf.Tensoris not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.
I am adding it here, because it seems to be different from OP's error but the problem was the same.
edit: while I was able to save it, I wasn't able to load it @frazane I actually get the error that the methods are missing:
(1) Implement
get_configand
from_configin the layer/model class, and pass the object to the
custom_objectsargument when loading the model. For more details, see: https://www.tensorflow.org/guide/keras/save_and_serialize
I presume they were added in later editions of tfp
When saving a keras model that incorporates a
DistributionLambda
layer fromtfp.layers
, I receive a stack trace ending in the following (complete stack trace at end of post). I observe this error with tfp v0.12.2 and tf version 2.5.0. However, it doesn't happen with tfp v0.12.2 and tf version 2.4.1.Please let me know if this better belongs as an issue over on the main
tensorflow
repo!Installed package versions
Python version: 3.8.8
Script to recreate
Complete stack trace