tensorflow / tensorflow/probability

Bijectors Permute() and RealNVP() failing with tf.keras Model in TF2.0-nightly

Open
#355 12 comments 4 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

keras layers
Dominant language
Jupyter Notebook
Stars
4.4k
Forks
1.1k
No merged PRs in 30d

Description

When using bijectors tfb.Permute and tfb.RealNVP to transform an input to an output in a keras model (using either of forward() or inverse() transformations), one runs into multiple (possibly related) errors, with TensorFLow 2.0.0-dev20190408 and TensorFlow Probability 0.7.0-dev.

To demonstrate, when trying to transform tf.keras input to output using one of these bijectors or any chaining of them (testing for both forward() and inverse() transforms):

bijector = permute # no error here
bijector = realnvp
bijector = realnvp(permute)
bijector = permute(realnvp)

one gets the following 3 errors (respectively):

ValueError: Trying to share variable real_nvp_default_template/dense/kernel, but specified shape (1, 784) and found shape (392, 784).
NotImplementedError: Rightmost dimension must be known prior to graph execution.
ValueError: The last dimension of the inputs to `Dense` should be defined. Found `None`.
Click here to see the full test code 
import tensorflow as tf
import tensorflow_probability as tfp
import tensorflow_datasets as tfds

tf.config.gpu.set_per_process_memory_growth(True)

tfk = tf.keras
tfkl = tf.keras.layers
tfpl = tfp.layers
tfd = tfp.distributions
tfb = tfp.bijectors

class BijectorForward(tfkl.Layer):
    def __init__(self, bijector):
        super(BijectorForward, self).__init__()
        self.bijector = bijector

    def call(self, input):
        return self.bijector.forward(input)

class BijectorInverse(tfkl.Layer):
    def __init__(self, bijector):
        super(BijectorInverse, self).__init__()
        self.bijector = bijector

    def call(self, input):
        return self.bijector.inverse(input)

dim_z = 28**2
shape_x = (28, 28, 1)

def get_bijectors():
    permute = tfb.Permute(tf.concat([tf.range(dim_z//2, dim_z), tf.range(0, dim_z//2)], axis=0))

    additive_cf = tfb.real_nvp_default_template(
        [dim_z, dim_z//2], shift_only=True, activation=None)
    realnvp = tfb.RealNVP(
        num_masked=dim_z//2, shift_and_log_scale_fn=additive_cf, is_constant_jacobian=True)

    return permute, realnvp

def construct_keras_bijector_models(bijector):
    input_z = tfkl.Input(shape=(dim_z,))
    output_x = BijectorForward(bijector)(input_z)
    model_forward = tfk.models.Model(inputs=input_z, outputs=output_x)

    input_x = tfkl.Input(shape=shape_x)
    output_z = BijectorInverse(bijector)(input_x)
    model_inverse = tfk.models.Model(inputs=input_x, outputs=output_z)

    return model_forward, model_inverse

def test1():
    permute, realnvp = get_bijectors()
    bijector = permute
    construct_keras_bijector_models(bijector)

def test2():
    permute, realnvp = get_bijectors()
    bijector = realnvp
    construct_keras_bijector_models(bijector)

def test3():
    permute, realnvp = get_bijectors()
    bijector = realnvp(permute)
    construct_keras_bijector_models(bijector)

def test4():
    permute, realnvp = get_bijectors()
    bijector = permute(realnvp)
    construct_keras_bijector_models(bijector)

import traceback

for test in [test1, test2, test3, test4]:
    print('\n------{}------\n'.format(test.__name__))
    try:
        test()
    except Exception as e:
        print(e)
        #traceback.print_exc()
Click here to see the corresponding errors 

------test1------

2019-04-08 17:58:42.998488: I tensorflow/core/platform/cpu_feature_guard.cc:142] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AV
X2 AVX512F FMA
2019-04-08 17:58:43.012348: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library libcuda.so.1
2019-04-08 17:58:43.157452: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x50cbf50 executing computations on platform CUDA. Devices:
2019-04-08 17:58:43.157505: I tensorflow/compiler/xla/service/service.cc:175]   StreamExecutor device (0): GeForce GTX 1080 Ti, Compute Capability 6.1
2019-04-08 17:58:43.179127: I tensorflow/core/platform/profile_utils/cpu_utils.cc:94] CPU Frequency: 3500000000 Hz
2019-04-08 17:58:43.180183: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x4cdb210 executing computations on platform Host. Devices:
2019-04-08 17:58:43.180232: I tensorflow/compiler/xla/service/service.cc:175]   StreamExecutor device (0): <undefined>, <undefined>
2019-04-08 17:58:43.180913: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1551] Found device 0 with properties:
name: GeForce GTX 1080 Ti major: 6 minor: 1 memoryClockRate(GHz): 1.607
pciBusID: 0000:65:00.0
totalMemory: 10.92GiB freeMemory: 10.51GiB
2019-04-08 17:58:43.180944: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1674] Adding visible gpu devices: 0
2019-04-08 17:58:43.181018: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library libcudart.so.10.0
2019-04-08 17:58:43.186741: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1082] Device interconnect StreamExecutor with strength 1 edge matrix:
2019-04-08 17:58:43.186780: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1088]      0
2019-04-08 17:58:43.186794: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1101] 0:   N
2019-04-08 17:58:43.187269: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1222] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10222
 MB memory) -> physical GPU (device: 0, name: GeForce GTX 1080 Ti, pci bus id: 0000:65:00.0, compute capability: 6.1)

------test2------

WARNING: Logging before flag parsing goes to stderr.
W0408 17:58:43.255850 139912563205952 deprecation.py:323] From ~/.env/tf2gpu/lib/python3.6/site-packages/tensorflow_probability/python/bijectors/real_nvp.p
y:291: dense (from tensorflow.python.layers.core) is deprecated and will be removed in a future version.
Instructions for updating:
Use keras.layers.dense instead.
W0408 17:58:43.256973 139912563205952 deprecation.py:506] From ~/.env/tf2gpu/lib/python3.6/site-packages/tensorflow/python/ops/init_ops.py:1257: calling Va
rianceScaling.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.
Instructions for updating:
Call initializer instance with the dtype argument instead of passing it to the constructor
Trying to share variable real_nvp_default_template/dense/kernel, but specified shape (1, 784) and found shape (392, 784).

originally defined at:
  File "bijectors-keras-errors.py", line 59, in test2
    permute, realnvp = get_bijectors()
  File "bijectors-keras-errors.py", line 36, in get_bijectors
    [dim_z, dim_z//2], shift_only=True, activation=None)
  File "~/.env/tf2gpu/lib/python3.6/site-packages/tensorflow_probability/python/bijectors/real_nvp.py", line 303, in real_nvp_default_template
    return tf.compat.v1.make_template("real_nvp_default_template", _fn)
  File "~/.env/tf2gpu/lib/python3.6/site-packages/tensorflow/python/ops/template.py", line 154, in make_template
    **kwargs)
  File "~/.env/tf2gpu/lib/python3.6/site-packages/tensorflow/python/ops/template.py", line 217, in make_template_internal
    create_graph_function=create_graph_function_)


------test3------

Rightmost dimension must be known prior to graph execution.

------test4------

The last dimension of the inputs to `Dense` should be defined. Found `None`.

originally defined at:
  File "bijectors-keras-errors.py", line 69, in test4
    permute, realnvp = get_bijectors()
  File "bijectors-keras-errors.py", line 36, in get_bijectors
    [dim_z, dim_z//2], shift_only=True, activation=None)
  File "~/.env/tf2gpu/lib/python3.6/site-packages/tensorflow_probability/python/bijectors/real_nvp.py", line 303, in real_nvp_default_template
    return tf.compat.v1.make_template("real_nvp_default_template", _fn)
  File "~/.env/tf2gpu/lib/python3.6/site-packages/tensorflow/python/ops/template.py", line 154, in make_template
    **kwargs)
  File "~/.env/tf2gpu/lib/python3.6/site-packages/tensorflow/python/ops/template.py", line 217, in make_template_internal
    create_graph_function=create_graph_function_)

So I'd like to end up with a tf.keras model that can be .fit() to optimize parameters of real_nvp_default_templates MLP while using the whole chain of bijectors as a transformation of models' input to its output (or as a part of some larger transformation comprising of other trainable parameters).

Is this a bug or just a currently missing feature (as I'm on TF2 nightly)? Am I assuming some non-supported use-case here? Is there an obvious way to achieve what I need that I'm missing here?

Thanks for any response and for all the great work! TFP rocks!

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by running the reproduction in bijectors-keras-errors.py with the reported TensorFlow and TensorFlow Probability versions, then inspect tensorflow_probability/python/bijectors/real_nvp.py around real_nvp_default_template. Compare the Permute, RealNVP, and chained forward/inverse cases; done means the supported tf.keras usage and expected behavior are established, with the reported errors either reproduced and covered or clearly documented.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.