Issue 1:
System information

Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes

OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 16.04

Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: N/A

TensorFlow installed from (source or binary): pip

TensorFlow version (use command below): 1.14.0rc0

Python version: 3.6.8

Bazel version (if compiling from source): N/A

GCC/Compiler version (if compiling from source): N/A

CUDA/cuDNN version: N/A

GPU model and memory: N/A

Describe the current behavior

I am encountering serialization issues when trying to dump the config from a tf.keras.Model object without complex things like custom layers (or even Lambdas...)

The code worked well with tf 1.13.1 however in tf 1.1.4, json/yaml serialization fails, and to_yaml and model_from_yaml fails as well

Describe the expected behavior

We should be able to serialize the config from a tf.keras.Model into json/yaml and load it back, especially without using custom layers

I checked the changelog for 1.14 and did not see any documented changes on model serialization for tf 1.14

Is serializing just the config file something that should not be done in 1.14 / 2.0 ?

Code to reproduce the issue

# %%
import tensorflow as tf
import json
from ruamel import yaml

# %%
def mini_cnn(input_shape, num_classes):
inputs = tf.keras.layers.Input(shape=input_shape)

x = tf.keras.layers.Conv2D(16, (3, 3), padding="same", use_bias=False)(inputs)
x = tf.keras.layers.BatchNormalization(axis=-1)(x)
x = tf.keras.layers.Activation("relu")(x)
x = tf.keras.layers.MaxPooling2D((2, 2), strides=2)(x)
x = tf.keras.layers.Conv2D(32, (3, 3), padding="same", use_bias=False)(x)
x = tf.keras.layers.BatchNormalization(axis=-1)(x)
x = tf.keras.layers.Activation("relu")(x)
x = tf.keras.layers.MaxPooling2D((2, 2), strides=2)(x)
x = tf.keras.layers.Conv2D(64, (3, 3), padding="same", use_bias=False)(x)
x = tf.keras.layers.BatchNormalization(axis=-1)(x)
x = tf.keras.layers.Activation("relu")(x)
x = tf.keras.layers.MaxPooling2D((2, 2), strides=2)(x)
x = tf.keras.layers.Flatten()(x)
x = tf.keras.layers.Dense(64, use_bias=False)(x)
x = tf.keras.layers.BatchNormalization(axis=-1)(x)
x = tf.keras.layers.Activation("relu")(x)
x = tf.keras.layers.Dropout(0.25)(x)

logits = tf.keras.layers.Dense(num_classes, use_bias=True, name="logits")(x)
probas = tf.keras.layers.Activation(activation="softmax", name="probas")(logits)

model = tf.keras.Model(inputs=inputs, outputs=probas, name="model")

return model

# %%
m = mini_cnn((64, 64, 3), 2)

## %%
try:
c = m.get_config()
m2 = tf.keras.models.model_from_config(c)
except:
print("c error")

try:
y = m.to_yaml()
m2 = tf.keras.models.model_from_yaml(y)
except:
print("y error")

try:
j = m.to_json()
m2 = tf.keras.models.model_from_json(j)
except:
print("j error")

Other info / logs

YAML traceback (seems to come from batchnorm)

Traceback (most recent call last):
File "run.py", line 49, in <module>
m2 = tf.keras.models.model_from_yaml(y)
File "/my-env/lib/python3.6/site-packages/tensorflow/python/keras/saving/model_config.py", line 78, in model_from_yaml
return deserialize(config, custom_objects=custom_objects)
File "/my-env/lib/python3.6/site-packages/tensorflow/python/keras/layers/serialization.py", line 89, in deserialize
printable_module_name='layer')
File "/my-env/lib/python3.6/site-packages/tensorflow/python/keras/utils/generic_utils.py", line 192, in deserialize_keras_object
list(custom_objects.items())))
File "/my-env/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py", line 1131, in from_config
process_node(layer, node_data)
File "/my-env/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py", line 1087, in process_node
layer(flat_input_tensors[0], **kwargs)
File "/my-env/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py", line 591, in __call__
self._maybe_build(inputs)
File "/my-env/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py", line 1881, in _maybe_build
self.build(input_shapes)
File "/my-env/lib/python3.6/site-packages/tensorflow/python/keras/layers/normalization.py", line 286, in build
raise ValueError('Duplicate axis: %s' % self.axis)
ValueError: Duplicate axis: ListWrapper([3, 3])
None

config error

Traceback (most recent call last):
File "run.py", line 42, in <module>
m2 = tf.keras.models.model_from_config(c)
File "/my-env/lib/python3.6/site-packages/tensorflow/python/keras/saving/model_config.py", line 55, in model_from_config
return deserialize(config, custom_objects=custom_objects)
File "/my-env/lib/python3.6/site-packages/tensorflow/python/keras/layers/serialization.py", line 81, in deserialize
layer_class_name = config['class_name']
KeyError: 'class_name'
None

here is the yaml file that was generated:

backend: tensorflow
class_name: Model
config:
input_layers:
- [input_1, 0, 0]
layers:
- class_name: InputLayer
config:
batch_input_shape: !!python/tuple [null, 64, 64, 3]
dtype: float32
name: input_1
sparse: false
inbound_nodes: []
name: input_1
- class_name: Conv2D
config:
activation: linear
activity_regularizer: null
bias_constraint: null
bias_initializer:
class_name: Zeros
config: {dtype: float32}
bias_regularizer: null
data_format: channels_last
dilation_rate: &id001 !!python/tuple [1, 1]
dtype: float32
filters: 16
kernel_constraint: null
kernel_initializer:
class_name: GlorotUniform
config: {dtype: float32, seed: null}
kernel_regularizer: null
kernel_size: !!python/tuple [3, 3]
name: conv2d
padding: same
strides: &id002 !!python/tuple [1, 1]
trainable: true
use_bias: false
inbound_nodes:
- - - input_1
- 0
- 0
- {}
name: conv2d
- class_name: BatchNormalization
config:
axis: !!python/object/new:tensorflow.python.training.tracking.data_structures._ListWrapper
listitems: [3]
state:
_external_modification: false
_last_wrapped_list_snapshot: [3]
_non_append_mutation: false
_self_extra_variables: []
_self_trainable: true
_storage: [3]
beta_constraint: null
beta_initializer:
class_name: Zeros
config: {dtype: float32}
beta_regularizer: null
center: true
dtype: float32
epsilon: 0.001
gamma_constraint: null
gamma_initializer:
class_name: Ones
config: {dtype: float32}
gamma_regularizer: null
momentum: 0.99
moving_mean_initializer:
class_name: Zeros
config: {dtype: float32}
moving_variance_initializer:
class_name: Ones
config: {dtype: float32}
name: batch_normalization
scale: true
trainable: true
inbound_nodes:
- - - conv2d
- 0
- 0
- {}
name: batch_normalization
- class_name: Activation
config: {activation: relu, dtype: float32, name: activation, trainable: true}
inbound_nodes:
- - - batch_normalization
- 0
- 0
- {}
name: activation
- class_name: MaxPooling2D
config:
data_format: channels_last
dtype: float32
name: max_pooling2d
padding: valid
pool_size: !!python/tuple [2, 2]
strides: !!python/tuple [2, 2]
trainable: true
inbound_nodes:
- - - activation
- 0
- 0
- {}
name: max_pooling2d
- class_name: Conv2D
config:
activation: linear
activity_regularizer: null
bias_constraint: null
bias_initializer:
class_name: Zeros
config: {dtype: float32}
bias_regularizer: null
data_format: channels_last
dilation_rate: *id001
dtype: float32
filters: 32
kernel_constraint: null
kernel_initializer:
class_name: GlorotUniform
config: {dtype: float32, seed: null}
kernel_regularizer: null
kernel_size: !!python/tuple [3, 3]
name: conv2d_1
padding: same
strides: *id002
trainable: true
use_bias: false
inbound_nodes:
- - - max_pooling2d
- 0
- 0
- {}
name: conv2d_1
- class_name: BatchNormalization
config:
axis: !!python/object/new:tensorflow.python.training.tracking.data_structures._ListWrapper
listitems: [3]
state:
_external_modification: false
_last_wrapped_list_snapshot: [3]
_non_append_mutation: false
_self_extra_variables: []
_self_trainable: true
_storage: [3]
beta_constraint: null
beta_initializer:
class_name: Zeros
config: {dtype: float32}
beta_regularizer: null
center: true
dtype: float32
epsilon: 0.001
gamma_constraint: null
gamma_initializer:
class_name: Ones
config: {dtype: float32}
gamma_regularizer: null
momentum: 0.99
moving_mean_initializer:
class_name: Zeros
config: {dtype: float32}
moving_variance_initializer:
class_name: Ones
config: {dtype: float32}
name: batch_normalization_1
scale: true
trainable: true
inbound_nodes:
- - - conv2d_1
- 0
- 0
- {}
name: batch_normalization_1
- class_name: Activation
config: {activation: relu, dtype: float32, name: activation_1, trainable: true}
inbound_nodes:
- - - batch_normalization_1
- 0
- 0
- {}
name: activation_1
- class_name: MaxPooling2D
config:
data_format: channels_last
dtype: float32
name: max_pooling2d_1
padding: valid
pool_size: !!python/tuple [2, 2]
strides: !!python/tuple [2, 2]
trainable: true
inbound_nodes:
- - - activation_1
- 0
- 0
- {}
name: max_pooling2d_1
- class_name: Conv2D
config:
activation: linear
activity_regularizer: null
bias_constraint: null
bias_initializer:
class_name: Zeros
config: {dtype: float32}
bias_regularizer: null
data_format: channels_last
dilation_rate: *id001
dtype: float32
filters: 64
kernel_constraint: null
kernel_initializer:
class_name: GlorotUniform
config: {dtype: float32, seed: null}
kernel_regularizer: null
kernel_size: !!python/tuple [3, 3]
name: conv2d_2
padding: same
strides: *id002
trainable: true
use_bias: false
inbound_nodes:
- - - max_pooling2d_1
- 0
- 0
- {}
name: conv2d_2
- class_name: BatchNormalization
config:
axis: !!python/object/new:tensorflow.python.training.tracking.data_structures._ListWrapper
listitems: [3]
state:
_external_modification: false
_last_wrapped_list_snapshot: [3]
_non_append_mutation: false
_self_extra_variables: []
_self_trainable: true
_storage: [3]
beta_constraint: null
beta_initializer:
class_name: Zeros
config: {dtype: float32}
beta_regularizer: null
center: true
dtype: float32
epsilon: 0.001
gamma_constraint: null
gamma_initializer:
class_name: Ones
config: {dtype: float32}
gamma_regularizer: null
momentum: 0.99
moving_mean_initializer:
class_name: Zeros
config: {dtype: float32}
moving_variance_initializer:
class_name: Ones
config: {dtype: float32}
name: batch_normalization_2
scale: true
trainable: true
inbound_nodes:
- - - conv2d_2
- 0
- 0
- {}
name: batch_normalization_2
- class_name: Activation
config: {activation: relu, dtype: float32, name: activation_2, trainable: true}
inbound_nodes:
- - - batch_normalization_2
- 0
- 0
- {}
name: activation_2
- class_name: MaxPooling2D
config:
data_format: channels_last
dtype: float32
name: max_pooling2d_2
padding: valid
pool_size: !!python/tuple [2, 2]
strides: !!python/tuple [2, 2]
trainable: true
inbound_nodes:
- - - activation_2
- 0
- 0
- {}
name: max_pooling2d_2
- class_name: Flatten
config: {data_format: channels_last, dtype: float32, name: flatten, trainable: true}
inbound_nodes:
- - - max_pooling2d_2
- 0
- 0
- {}
name: flatten
- class_name: Dense
config:
activation: linear
activity_regularizer: null
bias_constraint: null
bias_initializer:
class_name: Zeros
config: {dtype: float32}
bias_regularizer: null
dtype: float32
kernel_constraint: null
kernel_initializer:
class_name: GlorotUniform
config: {dtype: float32, seed: null}
kernel_regularizer: null
name: dense
trainable: true
units: 64
use_bias: false
inbound_nodes:
- - - flatten
- 0
- 0
- {}
name: dense
- class_name: BatchNormalization
config:
axis: !!python/object/new:tensorflow.python.training.tracking.data_structures._ListWrapper
listitems: [1]
state:
_external_modification: false
_last_wrapped_list_snapshot: [1]
_non_append_mutation: false
_self_extra_variables: []
_self_trainable: true
_storage: [1]
beta_constraint: null
beta_initializer:
class_name: Zeros
config: {dtype: float32}
beta_regularizer: null
center: true
dtype: float32
epsilon: 0.001
gamma_constraint: null
gamma_initializer:
class_name: Ones
config: {dtype: float32}
gamma_regularizer: null
momentum: 0.99
moving_mean_initializer:
class_name: Zeros
config: {dtype: float32}
moving_variance_initializer:
class_name: Ones
config: {dtype: float32}
name: batch_normalization_3
scale: true
trainable: true
inbound_nodes:
- - - dense
- 0
- 0
- {}
name: batch_normalization_3
- class_name: Activation
config: {activation: relu, dtype: float32, name: activation_3, trainable: true}
inbound_nodes:
- - - batch_normalization_3
- 0
- 0
- {}
name: activation_3
- class_name: Dropout
config: {dtype: float32, name: dropout, noise_shape: null, rate: 0.25, seed: null,
trainable: true}
inbound_nodes:
- - - activation_3
- 0
- 0
- {}
name: dropout
- class_name: Dense
config:
activation: linear
activity_regularizer: null
bias_constraint: null
bias_initializer:
class_name: Zeros
config: {dtype: float32}
bias_regularizer: null
dtype: float32
kernel_constraint: null
kernel_initializer:
class_name: GlorotUniform
config: {dtype: float32, seed: null}
kernel_regularizer: null
name: logits
trainable: true
units: 2
use_bias: true
inbound_nodes:
- - - dropout
- 0
- 0
- {}
name: logits
- class_name: Activation
config: {activation: softmax, dtype: float32, name: probas, trainable: true}
inbound_nodes:
- - - logits
- 0
- 0
- {}
name: probas
name: model
output_layers:
- [probas, 0, 0]
keras_version: 2.2.4-tf

One can see the error with ListWrapper

Note:

Lambda layers serialization seem to fail, whereas they worked in tf1.13 as well

Issue 2:
Currently, the caching that Bijector objects do to avoid unnecessary calculations does not work when the bijector is used in a TransformedDistribution object. I believe the culprit is the reshaping that the distribution object does here; when we call Bijector.inverse on the output, the Bijector object cannot tell that this is merely a reshaped version of what it calculated previously.

My use case is to sample from a TransformedDistribution and then later calculate the log probability of that sample.

This issue is particularly a problem when using a bijector whose inverse is numerically delicate (in my case, I'm chaining together softplus bijectors and my own custom affine bijector).

I'm willing to work on a fix for this problem, but I'm not sure what the best way to do it is (adding caching to the TransformedDistribution code might work, but that seems like code duplication).

System information

OS Platform and Distribution: Ubuntu 16.04

TensorFlow installed from: Source

TensorFlow version: v1.1.0-rc2-773-g7fa0cf3 (commit 7fa0cf3)

Bazel version: 0.4.5