Open mkretsch327 opened 3 years ago
ColCarroll
commented
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.
mkretsch327
commented
3 years ago Thanks @ColCarroll ! I wish I had a better idea to help on root cause, will keep an eye out here.
cserpell
commented
3 years ago 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.
cserpell
commented
3 years ago 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.
cserpell
commented
3 years ago Hi, is there any hint on this? It continues happening, according to this example I did. Is #1382 related?
cserpell
commented
3 years ago 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')
aciobanusebi
commented
3 years ago 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.
MathiasGruber
commented
3 years ago Also facing this issue using tfp==0.13.0 and tf==2.5.0
jdbates
commented
2 years ago 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})
Martinc4321
commented
2 years ago 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 )
frazane
commented
2 years ago 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.
willemvdb42
commented
2 years ago 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.
aegonwolf
commented
1 year ago I am using TF2.7 and TFP0.15. If I pass the
save_traces=Falseargument 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
DistributionLambdalayer 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 symbolictf.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) Implementget_configandfrom_configin the layer/model class, and pass the object to thecustom_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
DistributionLambdalayer 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
tensorflowrepo!Installed package versions
Python version: 3.8.8
Script to recreate
Complete stack trace