From aa4efbba57600c882bc54f6828020ba4d0b0b2da Mon Sep 17 00:00:00 2001 From: Srinivas Vasudevan Date: Tue, 16 Oct 2018 16:39:22 -0700 Subject: [PATCH] Update uses of distributions / bijectors to use TensorFlow Probability. --- docs/tex/api/model-development.tex | 8 ++++---- .../tutorials/automated-transformations.tex | 8 ++++---- edward/inferences/inference.py | 13 ++++++++++++- edward/inferences/klqp.py | 17 +++++++++++++---- edward/inferences/map.py | 12 ++++++++++-- edward/models/dirichlet_process.py | 16 +++++++++++++--- edward/models/empirical.py | 19 +++++++++++++++---- edward/models/param_mixture.py | 13 +++++++++++-- edward/models/point_mass.py | 12 ++++++++++-- edward/models/random_variable.py | 2 +- edward/util/random_variables.py | 13 ++++++++++++- examples/deep_exponential_family.py | 3 ++- notebooks/automated_transformations.ipynb | 5 +++-- tests/inferences/ar_process_test.py | 8 ++++---- .../inference_auto_transform_test.py | 13 +++++++------ tests/util/transform_test.py | 6 +++--- 16 files changed, 124 insertions(+), 44 deletions(-) diff --git a/docs/tex/api/model-development.tex b/docs/tex/api/model-development.tex index 095fc85dc..8733e3d6c 100644 --- a/docs/tex/api/model-development.tex +++ b/docs/tex/api/model-development.tex @@ -7,14 +7,14 @@ \subsubsection{Developing Custom Random Variables} Oftentimes we'd like to implement our own random variables. To do so, write a class that inherits the \texttt{RandomVariable} class in \texttt{edward.models} and -the \texttt{Distribution} class in \texttt{tf.contrib.distributions} (in that +the \texttt{Distribution} class in \texttt{tfp.distributions} (in that order). A template is provided below. \begin{lstlisting}[language=Python] from edward.models import RandomVariable -from tensorflow.contrib.distributions import Distribution +import tensorflow_probability as tfp -class CustomRandomVariable(RandomVariable, Distribution): +class CustomRandomVariable(RandomVariable, tfp.distributions.Distribution): def __init__(self, *args, **kwargs): super(CustomRandomVariable, self).__init__(*args, **kwargs) @@ -39,7 +39,7 @@ \subsubsection{Developing Custom Random Variables} in the Github repository. For more details and more methods one can implement, see the API documentation in TensorFlow's -\href{https://www.tensorflow.org/api_docs/python/tf/contrib/distributions/Distribution}{\texttt{Distribution} class}. +\href{https://www.tensorflow.org/probability/api_docs/python/tfp/distributions}{\texttt{Distribution} class}. \subsubsection{Advanced settings} diff --git a/docs/tex/tutorials/automated-transformations.tex b/docs/tex/tutorials/automated-transformations.tex index 6e0b359aa..643da9ecd 100644 --- a/docs/tex/tutorials/automated-transformations.tex +++ b/docs/tex/tutorials/automated-transformations.tex @@ -99,18 +99,18 @@ \subsubsection{Automated Transformations in Inference} and then inverting the transformation: \begin{lstlisting}[language=Python] -from tensorflow.contrib.distributions import bijectors +import tensorflow_probability as tfp x_unconstrained = inference.transformations[x] # transformed prior x_transform = x_unconstrained.bijector # transformed prior's transformation -qx_constrained = ed.transform(qx, bijectors.Invert(x_transform)) +qx_constrained = ed.transform(qx, tfp.bijectors.Invert(x_transform)) \end{lstlisting} The set of transformations is given by \texttt{inference.transformations}, which is a dictionary with keys given by any constrained latent variables and values given by their transformed distribution. We use the -\href{https://www.tensorflow.org/api_docs/python/tf/distributions/bijectors}{\texttt{bijectors}} +\href{https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors}{\texttt{bijectors}} module in \texttt{tf.distributions} in order to handle invertible transformations. @@ -158,7 +158,7 @@ \subsubsection{Automated Transformations in Inference} transformation. \begin{lstlisting}[language=Python] -from tensorflow.contrib.distributions import bijectors +import tensorflow_probability as tfp x_unconstrained = inference.transformations[x] # transformed prior x_transform = x_unconstrained.bijector # transformed prior's transformation diff --git a/edward/inferences/inference.py b/edward/inferences/inference.py index 28614223f..dd1e8e575 100644 --- a/edward/inferences/inference.py +++ b/edward/inferences/inference.py @@ -13,7 +13,18 @@ from edward.util import check_data, check_latent_vars, get_session, \ get_variables, Progbar, transform -from tensorflow.contrib.distributions import bijectors +try: + import tensorflow_probability as tfp + bijectors = tfp.bijectors +except Exception as e: + print("{0}. Can not import TensorFlow Probability, " + "defaulting to TensorFlow.".format(e)) + try: + from tensorflow.contrib.distributions import bijectors + except Exception as e2: + raise ImportError( + "{0}. Your TensorFlow version is not supported.".format(e2)) + @six.add_metaclass(abc.ABCMeta) diff --git a/edward/inferences/klqp.py b/edward/inferences/klqp.py index 3cfbc9cea..a584f9110 100644 --- a/edward/inferences/klqp.py +++ b/edward/inferences/klqp.py @@ -11,9 +11,19 @@ try: from edward.models import Normal - from tensorflow.contrib.distributions import kl_divergence + import tensorflow_probability as tfp + kl_divergence = tfp.distributions.kl_divergence + FULLY_REPARAMETERIZED = tfp.distributions.FULLY_REPARAMETERIZED except Exception as e: - raise ImportError("{0}. Your TensorFlow version is not supported.".format(e)) + print("{0}. Can not import TensorFlow Probability, " + "defaulting to TensorFlow.".format(e)) + try: + from edward.models import Normal + from tensorflow.contrib.distributions import kl_divergence + FULLY_REPARAMETERIZED = tf.contrib.distributions.FULLY_REPARAMETERIZED + except Exception as e2: + raise ImportError( + "{0}. Your TensorFlow version is not supported.".format(e2)) class KLqp(VariationalInference): @@ -136,8 +146,7 @@ def build_loss_and_gradients(self, var_list): Normal. """ is_reparameterizable = all([ - rv.reparameterization_type == - tf.contrib.distributions.FULLY_REPARAMETERIZED + rv.reparameterization_type == FULLY_REPARAMETERIZED for rv in six.itervalues(self.latent_vars)]) is_analytic_kl = all([isinstance(z, Normal) and isinstance(qz, Normal) for z, qz in six.iteritems(self.latent_vars)]) diff --git a/edward/inferences/map.py b/edward/inferences/map.py index 406d461d4..1901c9133 100644 --- a/edward/inferences/map.py +++ b/edward/inferences/map.py @@ -10,9 +10,17 @@ from edward.util import copy, transform try: - from tensorflow.contrib.distributions import bijectors + import tensorflow_probability as tfp + bijectors = tfp.bijectors except Exception as e: - raise ImportError("{0}. Your TensorFlow version is not supported.".format(e)) + print("{0}. Can not import TensorFlow Probability, " + "defaulting to TensorFlow.".format(e)) + try: + from tensorflow.contrib.distributions import bijectors + except Exception as e2: + raise ImportError( + "{0}. Your TensorFlow version is not supported.".format(e2)) + class MAP(VariationalInference): diff --git a/edward/models/dirichlet_process.py b/edward/models/dirichlet_process.py index 9b6efe767..2a8349839 100644 --- a/edward/models/dirichlet_process.py +++ b/edward/models/dirichlet_process.py @@ -5,13 +5,23 @@ import tensorflow as tf from edward.models.random_variable import RandomVariable -from tensorflow.contrib.distributions import Distribution try: from edward.models.random_variables import Bernoulli, Beta - from tensorflow.contrib.distributions import NOT_REPARAMETERIZED + import tensorflow_probability as tfp + Distribution = tfp.distributions.Distribution + NOT_REPARAMETERIZED = tfp.distributions.NOT_REPARAMETERIZED except Exception as e: - raise ImportError("{0}. Your TensorFlow version is not supported.".format(e)) + print("{0}. Can not import TensorFlow Probability, " + "defaulting to TensorFlow.".format(e)) + try: + from edward.models.random_variables import Bernoulli, Beta + from tensorflow.contrib.distributions import Distribution + from tensorflow.contrib.distributions import NOT_REPARAMETERIZED + except Exception as e2: + raise ImportError( + "{0}. Your TensorFlow version is not supported.".format(e2)) + class distributions_DirichletProcess(Distribution): diff --git a/edward/models/empirical.py b/edward/models/empirical.py index 7da9b8265..c36b424ca 100644 --- a/edward/models/empirical.py +++ b/edward/models/empirical.py @@ -5,12 +5,23 @@ import tensorflow as tf from edward.models.random_variable import RandomVariable -from tensorflow.contrib.distributions import Distribution try: - from tensorflow.contrib.distributions import FULLY_REPARAMETERIZED + import tensorflow_probability as tfp + Distribution = tfp.distributions.Distribution + Categorical = tfp.distributions.Categorical + FULLY_REPARAMETERIZED = tfp.distributions.FULLY_REPARAMETERIZED except Exception as e: - raise ImportError("{0}. Your TensorFlow version is not supported.".format(e)) + print("{0}. Can not import TensorFlow Probability, " + "defaulting to TensorFlow.".format(e)) + try: + from tensorflow.contrib.distributions import Distribution + Categorical = tf.distributions.Categorical + FULLY_REPARAMETERIZED = tf.contrib.distributions.FULLY_REPARAMETERIZED + except Exception as e2: + raise ImportError( + "{0}. Your TensorFlow version is not supported.".format(e2)) + class distributions_Empirical(Distribution): @@ -104,7 +115,7 @@ def _sample_n(self, n, seed=None): return tf.tile(input_tensor, multiples) else: probs = tf.ones([self.n]) / tf.cast(self.n, dtype=tf.float32) - cat = tf.contrib.distributions.Categorical(probs) + cat = Categorical(probs) indices = cat._sample_n(n, seed) tensor = tf.gather(input_tensor, indices) return tensor diff --git a/edward/models/param_mixture.py b/edward/models/param_mixture.py index 4f2f3f5b1..26dab74a9 100644 --- a/edward/models/param_mixture.py +++ b/edward/models/param_mixture.py @@ -6,12 +6,21 @@ import tensorflow as tf from edward.models.random_variable import RandomVariable -from tensorflow.contrib.distributions import Distribution try: from edward.models.random_variables import Categorical + import tensorflow_probability as tfp + Distribution = tfp.distributions.Distribution except Exception as e: - raise ImportError("{0}. Your TensorFlow version is not supported.".format(e)) + print("{0}. Can not import TensorFlow Probability, " + "defaulting to TensorFlow.".format(e)) + try: + from edward.models.random_variables import Categorical + from tensorflow.contrib.distributions import Distribution + except Exception as e2: + raise ImportError( + "{0}. Your TensorFlow version is not supported.".format(e2)) + class distributions_ParamMixture(Distribution): diff --git a/edward/models/point_mass.py b/edward/models/point_mass.py index b63031b6d..afda67278 100644 --- a/edward/models/point_mass.py +++ b/edward/models/point_mass.py @@ -8,9 +8,17 @@ from tensorflow.contrib.distributions import Distribution try: - from tensorflow.contrib.distributions import FULLY_REPARAMETERIZED + import tensorflow_probability as tfp + FULLY_REPARAMETERIZED = tfp.distributions.FULLY_REPARAMETERIZED except Exception as e: - raise ImportError("{0}. Your TensorFlow version is not supported.".format(e)) + print("{0}. Can not import TensorFlow Probability, " + "defaulting to TensorFlow.".format(e)) + try: + FULLY_REPARAMETERIZED = tf.contrib.distributions.FULLY_REPARAMETERIZED + except Exception as e2: + raise ImportError( + "{0}. Your TensorFlow version is not supported.".format(e2)) + class distributions_PointMass(Distribution): diff --git a/edward/models/random_variable.py b/edward/models/random_variable.py index 65f69e85e..f3e72d59b 100644 --- a/edward/models/random_variable.py +++ b/edward/models/random_variable.py @@ -41,7 +41,7 @@ class RandomVariable(object): `RandomVariable` assumes use in a multiple inheritance setting. The child class must first inherit `RandomVariable`, then second inherit a - class in `tf.contrib.distributions`. With Python's method resolution + class in `tfp.distributions`. With Python's method resolution order, this implies the following during initialization (using `distributions.Bernoulli` as an example): diff --git a/edward/util/random_variables.py b/edward/util/random_variables.py index 3a581505a..f434b71ec 100644 --- a/edward/util/random_variables.py +++ b/edward/util/random_variables.py @@ -15,7 +15,18 @@ from tensorflow.python.framework.ops import set_shapes_for_outputs from tensorflow.python.util import compat -tfb = tf.contrib.distributions.bijectors +try: + import tensorflow_probability as tfp + tfb = tfp.bijectors +except Exception as e: + print("{0}. Can not import TensorFlow Probability, " + "defaulting to TensorFlow.".format(e)) + try: + tfb = tf.contrib.distributions.bijectors + except Exception as e2: + raise ImportError( + "{0}. Your TensorFlow version is not supported.".format(e2)) + def check_data(data): diff --git a/examples/deep_exponential_family.py b/examples/deep_exponential_family.py index 3b57415b9..b2c71e925 100644 --- a/examples/deep_exponential_family.py +++ b/examples/deep_exponential_family.py @@ -71,6 +71,7 @@ class objects visual import numpy as np import os import tensorflow as tf +import tenosrflow_probability as tfp from datetime import datetime from edward.models import Gamma, Poisson, Normal, PointMass, \ @@ -126,7 +127,7 @@ def lognormal_q(shape, name=None): "scale", shape, initializer=tf.random_normal_initializer(stddev=0.1)) rv = TransformedDistribution( distribution=Normal(loc, tf.maximum(tf.nn.softplus(scale), min_scale)), - bijector=tf.contrib.distributions.bijectors.Exp()) + bijector=tfp.bijectors.Exp()) return rv diff --git a/notebooks/automated_transformations.ipynb b/notebooks/automated_transformations.ipynb index e6c04c15b..138142716 100644 --- a/notebooks/automated_transformations.ipynb +++ b/notebooks/automated_transformations.ipynb @@ -80,7 +80,8 @@ "import tensorflow as tf\n", "\n", "from edward.models import Empirical, Gamma, Normal\n", - "from tensorflow.contrib.distributions import bijectors" + "import tensorflow_probability as tfp\n", + "bijectors = tfp.bijectors", ] }, { @@ -172,7 +173,7 @@ "`inference.transformations`, which is a dictionary with keys\n", "given by any constrained latent variables and values given by their\n", "transformed distribution. We use the\n", - "[`bijectors`](https://www.tensorflow.org/versions/master/api_docs/python/tf/distributions/bijectors)\n", + "[`bijectors`](https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors)\n", "module in `tf.distributions` in order to handle invertible\n", "transformations.\n", "`qx_unconstrained` is a random variable distributed\n", diff --git a/tests/inferences/ar_process_test.py b/tests/inferences/ar_process_test.py index 88c63e4f6..6940452a8 100644 --- a/tests/inferences/ar_process_test.py +++ b/tests/inferences/ar_process_test.py @@ -5,16 +5,15 @@ import edward as ed import numpy as np import tensorflow as tf +import tensorflow_probability as tfp from edward.models import Normal, PointMass from scipy.optimize import minimize from edward.models import RandomVariable -from tensorflow.contrib.distributions import Distribution -from tensorflow.contrib.distributions import FULLY_REPARAMETERIZED -class AutoRegressive(RandomVariable, Distribution): +class AutoRegressive(RandomVariable, tfp.distributions.Distribution): # a 1-D AR(1) process # a[t + 1] = a[t] + eps with eps ~ N(0, sig**2) def __init__(self, T, a, sig, *args, **kwargs): @@ -29,7 +28,8 @@ def __init__(self, T, a, sig, *args, **kwargs): if 'allow_nan_stats' not in kwargs: kwargs['allow_nan_stats'] = False if 'reparameterization_type' not in kwargs: - kwargs['reparameterization_type'] = FULLY_REPARAMETERIZED + kwargs['reparameterization_type'] = ( + tfp.distributions.FULLY_REPARAMETERIZED) if 'validate_args' not in kwargs: kwargs['validate_args'] = False if 'name' not in kwargs: diff --git a/tests/inferences/inference_auto_transform_test.py b/tests/inferences/inference_auto_transform_test.py index 2aa1c9b6f..32a159dff 100644 --- a/tests/inferences/inference_auto_transform_test.py +++ b/tests/inferences/inference_auto_transform_test.py @@ -5,11 +5,11 @@ import edward as ed import numpy as np import tensorflow as tf +import tensorflow_probability as tfp from edward.models import (Empirical, Gamma, Normal, PointMass, TransformedDistribution, Beta, Bernoulli) from edward.util import transform -from tensorflow.contrib.distributions import bijectors class test_inference_auto_transform_class(tf.test.TestCase): @@ -20,7 +20,7 @@ def test_auto_transform_true(self): # automated transformation on latter (assuming it is softplus). x = TransformedDistribution( distribution=Normal(0.0, 0.5), - bijector=tf.contrib.distributions.bijectors.Softplus()) + bijector=tfp.bijectors.Softplus()) x.support = 'nonnegative' qx = Normal(loc=tf.Variable(tf.random_normal([])), scale=tf.nn.softplus(tf.Variable(tf.random_normal([])))) @@ -36,7 +36,8 @@ def test_auto_transform_true(self): n_samples = 10000 x_mean, x_var = tf.nn.moments(x.sample(n_samples), 0) x_unconstrained = inference.transformations[x] - qx_constrained = transform(qx, bijectors.Invert(x_unconstrained.bijector)) + qx_constrained = transform( + qx, tfp.bijectors.Invert(x_unconstrained.bijector)) qx_mean, qx_var = tf.nn.moments(qx_constrained.sample(n_samples), 0) stats = sess.run([x_mean, qx_mean, x_var, qx_var]) self.assertAllClose(info_dict['loss'], 0.0, rtol=0.2, atol=0.2) @@ -49,7 +50,7 @@ def test_auto_transform_false(self): # automated transformation; it should fail. x = TransformedDistribution( distribution=Normal(0.0, 0.5), - bijector=tf.contrib.distributions.bijectors.Softplus()) + bijector=tfp.bijectors.Softplus()) x.support = 'nonnegative' qx = Normal(loc=tf.Variable(tf.random_normal([])), scale=tf.nn.softplus(tf.Variable(tf.random_normal([])))) @@ -115,7 +116,7 @@ def test_hmc_custom(self): with self.test_session() as sess: x = TransformedDistribution( distribution=Normal(1.0, 1.0), - bijector=tf.contrib.distributions.bijectors.Softplus()) + bijector=tfp.bijectors.Softplus()) x.support = 'nonnegative' qx = Empirical(tf.Variable(tf.random_normal([1000]))) @@ -140,7 +141,7 @@ def test_hmc_default(self): with self.test_session() as sess: x = TransformedDistribution( distribution=Normal(1.0, 1.0), - bijector=tf.contrib.distributions.bijectors.Softplus()) + bijector=tfp.bijectors.Softplus()) x.support = 'nonnegative' inference = ed.HMC([x]) diff --git a/tests/util/transform_test.py b/tests/util/transform_test.py index f541f5008..4c09659b3 100644 --- a/tests/util/transform_test.py +++ b/tests/util/transform_test.py @@ -5,12 +5,12 @@ import edward as ed import numpy as np import tensorflow as tf +import tensorflow_probability as tfp from collections import namedtuple from edward.models import ( Beta, Dirichlet, DirichletProcess, Gamma, MultivariateNormalDiag, Normal, Poisson, TransformedDistribution) -from tensorflow.contrib.distributions import bijectors class test_transform_class(tf.test.TestCase): @@ -24,14 +24,14 @@ def assertSamplePosNeg(self, sample): def test_args(self): with self.test_session(): x = Normal(-100.0, 1.0) - y = ed.transform(x, bijectors.Softplus()) + y = ed.transform(x, tfp.bijectors.Softplus()) sample = y.sample(10).eval() self.assertTrue((sample >= 0.0).all()) def test_kwargs(self): with self.test_session(): x = Normal(-100.0, 1.0) - y = ed.transform(x, bijector=bijectors.Softplus()) + y = ed.transform(x, bijector=tfp.bijectors.Softplus()) sample = y.sample(10).eval() self.assertTrue((sample >= 0.0).all())