Skip to content
This repository was archived by the owner on Jul 7, 2023. It is now read-only.

Return ed.RandomVariables in initializers and operate on them in regu… #1182

Merged
afrozenator merged 1 commit into from
Nov 7, 2018
Merged

Return ed.RandomVariables in initializers and operate on them in regu… #1182

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 24 additions & 11 deletions tensor2tensor/layers/bayes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,11 @@
from __future__ import division
from __future__ import print_function

import functools
import tensorflow as tf

from tensorflow_probability import edward2 as ed


class Softplus(tf.keras.constraints.Constraint):
"""Softplus constraint."""
Expand Down Expand Up @@ -110,11 +113,7 @@ def __call__(self, shape=None, dtype=None, partition_info=None):
raise ValueError('A TrainableInitializer must be built by a layer before '
'usage, and is currently only compatible with Bayesian '
'layers.')
noise = tf.random_normal(self.shape, dtype=self.dtype, seed=self.seed)
output = self.mean + self.stddev * noise
# TODO(trandustin): Hack to store parameters so KL reg. can operate on them.
output._parameters = (self.mean, self.stddev) # pylint: disable=protected-access
return output
return ed.Normal(loc=self.mean, scale=self.stddev)

@dusenberrymw dusenberrymw Nov 1, 2018

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was thinking that we could create and store this in the build function to remove a bit of Python overhead in eager mode, but it looks like Edward2 RVs only sample from the underlying distribution on the first access in eager mode due to rv.value using a memoization strategy by storing the sample on the first access. It seems like it could potentially simplify some use cases if rv.value called rv.distribution.sample on each access to rv.value. That being said, you definitely thought much longer on the design of Edward[1|2], so I bet there is a reason that would make this undesirable.

@dustinvtran dustinvtran Nov 1, 2018 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We memoize for usability reasons, so that you operate on RandomVariables as if they were any other Tensor. For example, if you call x = ed.Normal(...), you would expect y = x + 1 and z = x + 1 to return the same result.

On caching: There's an issue that actually arises before Ed2, which is that during __init__, Distributions call TF ops on its parameters. This breaks gradients if Distribution.__init__ isn't recorded in the gradient tape. @vdumoulin raised a point about this in Brain Chat.


def get_config(self):
return {
Expand Down Expand Up @@ -151,12 +150,11 @@ def __init__(self, mean=0., stddev=1.):
self.stddev = stddev

def __call__(self, x):
mean, stddev = x._parameters # pylint: disable=protected-access
variance2 = tf.square(self.stddev)
variance_ratio = tf.square(stddev) / variance2
regularization = tf.square(mean - self.mean) / (2. * variance2)
regularization += (variance_ratio - 1. - tf.log(variance_ratio)) / 2.
return regularization
"""Computes regularization given an ed.Normal random variable as input."""
if not isinstance(x, ed.RandomVariable):
raise ValueError('Input must be an ed.RandomVariable.')
random_variable = ed.Normal(loc=self.mean, scale=self.stddev)

@danijar danijar Nov 3, 2018 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any reason to not create this object in __init__ already instead of saving self.mean and save.stddev separately?

@dustinvtran dustinvtran Nov 3, 2018

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Breaks Eager unfortunately if self.mean and self.stddev depend on trainable variables. See the above comment about caching. It would be nice if Distribution __init__s don't add TF ops to the graph.

return random_variable.distribution.kl_divergence(x.distribution)

def get_config(self):
return {
Expand Down Expand Up @@ -276,3 +274,18 @@ def build(self, input_shape):
else:
self.bias = None
self.built = True

# TODO(trandustin): Waiting on T2T to drop dependence on
# TF<=1.12rc2. A TF commit enables tf.colocate_with to work for
# Tensor-like inputs. This lets us use the parent method instead of
# this one.
def _handle_weight_regularization(self, name, variable, regularizer):
"""Create lambdas which compute regularization losses."""

def _loss_for_variable(v):
"""Creates a regularization loss `Tensor` for variable `v`."""
with tf.name_scope(name + '/Regularizer'):
regularization = regularizer(v)
return regularization

self.add_loss(functools.partial(_loss_for_variable, variable))
31 changes: 10 additions & 21 deletions tensor2tensor/layers/bayes_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,25 +29,20 @@

class BayesTest(parameterized.TestCase, tf.test.TestCase):

# TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
# support for TF 1.10
# @tf.contrib.eager.run_test_in_graph_and_eager_modes()
@tf.contrib.eager.run_test_in_graph_and_eager_modes()
def testDenseReparameterizationKernel(self):
inputs = tf.to_float(np.random.rand(5, 3, 12))
layer = bayes.DenseReparameterization(4, activation=tf.nn.relu)
outputs1 = layer(inputs)
outputs2 = layer(inputs)
self.evaluate(tf.global_variables_initializer())
# res1, res2 = self.evaluate([outputs1, outputs2])
res1, _ = self.evaluate([outputs1, outputs2])
res1, res2 = self.evaluate([outputs1, outputs2])
self.assertEqual(res1.shape, (5, 3, 4))
self.assertAllGreaterEqual(res1, 0.)
# self.assertNotAllClose(res1, res2)
self.assertNotAllClose(res1, res2)
layer.get_config()

# TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
# support for TF 1.10
# @tf.contrib.eager.run_test_in_graph_and_eager_modes()
@tf.contrib.eager.run_test_in_graph_and_eager_modes()
def testDenseReparameterizationBias(self):
inputs = tf.to_float(np.random.rand(5, 3, 12))
layer = bayes.DenseReparameterization(4, kernel_initializer="zero",
Expand All @@ -56,15 +51,12 @@ def testDenseReparameterizationBias(self):
outputs1 = layer(inputs)
outputs2 = layer(inputs)
self.evaluate(tf.global_variables_initializer())
# res1, res2 = self.evaluate([outputs1, outputs2])
res1, _ = self.evaluate([outputs1, outputs2])
res1, res2 = self.evaluate([outputs1, outputs2])
self.assertEqual(res1.shape, (5, 3, 4))
self.assertAllGreaterEqual(res1, 0.)
# self.assertNotAllClose(res1, res2)
self.assertNotAllClose(res1, res2)

# TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
# support for TF 1.10
# @tf.contrib.eager.run_test_in_graph_and_eager_modes()
@tf.contrib.eager.run_test_in_graph_and_eager_modes()
def testDenseReparameterizationDeterministic(self):
inputs = tf.to_float(np.random.rand(5, 3, 12))
layer = bayes.DenseReparameterization(4, kernel_initializer="zero",
Expand All @@ -73,15 +65,12 @@ def testDenseReparameterizationDeterministic(self):
outputs1 = layer(inputs)
outputs2 = layer(inputs)
self.evaluate(tf.global_variables_initializer())
# res1, res2 = self.evaluate([outputs1, outputs2])
res1, _ = self.evaluate([outputs1, outputs2])
res1, res2 = self.evaluate([outputs1, outputs2])
self.assertEqual(res1.shape, (5, 3, 4))
self.assertAllGreaterEqual(res1, 0.)
# self.assertAllClose(res1, res2)
self.assertAllClose(res1, res2)

# TODO(trandustin): Remove the hack in the code, or re-enable once T2T drops
# support for TF 1.10
# @tf.contrib.eager.run_test_in_graph_and_eager_modes()
@tf.contrib.eager.run_test_in_graph_and_eager_modes()
def testDenseReparameterizationModel(self):
inputs = tf.to_float(np.random.rand(3, 4, 4, 1))
model = tf.keras.Sequential([
Expand Down