diff --git a/tensorflow_probability/python/layers/conv_variational.py b/tensorflow_probability/python/layers/conv_variational.py index e97704deb8..e6293fff59 100644 --- a/tensorflow_probability/python/layers/conv_variational.py +++ b/tensorflow_probability/python/layers/conv_variational.py @@ -197,7 +197,6 @@ def build(self, input_shape): self.kernel_prior = self.kernel_prior_fn( dtype, kernel_shape, 'kernel_prior', self.trainable, self.add_variable) - self._built_kernel_divergence = False if self.bias_posterior_fn is None: self.bias_posterior = None @@ -212,7 +211,6 @@ def build(self, input_shape): self.bias_prior = self.bias_prior_fn( dtype, (self.filters,), 'bias_prior', self.trainable, self.add_variable) - self._built_bias_divergence = False self.input_spec = tf.keras.layers.InputSpec( ndim=self.rank + 2, axes={channel_axis: input_dim}) @@ -234,20 +232,16 @@ def call(self, inputs): outputs = self._apply_variational_bias(outputs) if self.activation is not None: outputs = self.activation(outputs) - if not self._built_kernel_divergence: - self._apply_divergence(self.kernel_divergence_fn, - self.kernel_posterior, - self.kernel_prior, - self.kernel_posterior_tensor, - name='divergence_kernel') - self._built_kernel_divergence = True - if not self._built_bias_divergence: - self._apply_divergence(self.bias_divergence_fn, - self.bias_posterior, - self.bias_prior, - self.bias_posterior_tensor, - name='divergence_bias') - self._built_bias_divergence = True + self._apply_divergence(self.kernel_divergence_fn, + self.kernel_posterior, + self.kernel_prior, + self.kernel_posterior_tensor, + name='divergence_kernel') + self._apply_divergence(self.bias_divergence_fn, + self.bias_posterior, + self.bias_prior, + self.bias_posterior_tensor, + name='divergence_bias') return outputs def compute_output_shape(self, input_shape): diff --git a/tensorflow_probability/python/layers/conv_variational_test.py b/tensorflow_probability/python/layers/conv_variational_test.py index 71a08de4df..9d505d7488 100644 --- a/tensorflow_probability/python/layers/conv_variational_test.py +++ b/tensorflow_probability/python/layers/conv_variational_test.py @@ -21,7 +21,6 @@ # Dependency imports import numpy as np -import tensorflow.compat.v1 as tf1 import tensorflow.compat.v2 as tf import tensorflow_probability as tfp @@ -180,12 +179,12 @@ def __call__(self, *args, **kwargs): @test_util.test_all_tf_execution_regimes class ConvVariational(object): - def maybe_transpose_inputs(self, inputs): + def maybe_transpose_tensor(self, tensor): if self.data_format == 'channels_first': - order = channels_last_to_first(list(range(inputs.shape.rank))) - return tf.transpose(a=inputs, perm=order) + order = channels_last_to_first(list(range(tensor.shape.rank))) + return tf.transpose(a=tensor, perm=order) else: - return inputs + return tensor def _testKerasLayer(self, layer_class): # pylint: disable=invalid-name def kernel_posterior_fn(dtype, shape, name, trainable, add_variable_fn): @@ -241,7 +240,7 @@ def _testKLPenaltyKernel(self, layer_class): # pylint: disable=invalid-name elif layer_class in (tfp.layers.Convolution3DReparameterization, tfp.layers.Convolution3DFlipout): inputs = tf.random.uniform([2, 3, 3, 3, 1], seed=1) - inputs = self.maybe_transpose_inputs(inputs) + inputs = self.maybe_transpose_tensor(inputs) # No keys. input_dependent_losses = layer.get_losses_for(inputs=None) @@ -273,7 +272,7 @@ def _testKLPenaltyBoth(self, layer_class): # pylint: disable=invalid-name elif layer_class in (tfp.layers.Convolution3DReparameterization, tfp.layers.Convolution3DFlipout): inputs = tf.random.uniform([2, 3, 3, 3, 1], seed=1) - inputs = self.maybe_transpose_inputs(inputs) + inputs = self.maybe_transpose_tensor(inputs) # No keys. input_dependent_losses = layer.get_losses_for(inputs=None) @@ -307,7 +306,7 @@ def _testConvSetUp(self, layer_class, batch_size, depth=None, inputs = tf.random.uniform([batch_size, depth, height, width, channels], seed=seed()) kernel_size = (2, 2, 2) - inputs = self.maybe_transpose_inputs(inputs) + inputs = self.maybe_transpose_tensor(inputs) kernel_shape = kernel_size + (channels, filters) kernel_posterior = MockDistribution( @@ -547,7 +546,7 @@ def _testRandomConvFlipout(self, layer_class): # pylint: disable=invalid-name inputs = tf.random.uniform([batch_size, depth, height, width, channels], seed=seed()) kernel_size = (2, 2, 2) - inputs = self.maybe_transpose_inputs(inputs) + inputs = self.maybe_transpose_tensor(inputs) kernel_shape = kernel_size + (channels, filters) bias_size = (filters,) @@ -597,26 +596,30 @@ def _testRandomConvFlipout(self, layer_class): # pylint: disable=invalid-name np.prod(outputs_one_.shape)) def _testLayerInSequential(self, layer_class): # pylint: disable=invalid-name - with self.cached_session() as sess: - if layer_class in (tfp.layers.Convolution1DReparameterization, - tfp.layers.Convolution1DFlipout): - inputs = tf.random.uniform([2, 3, 1]) - elif layer_class in (tfp.layers.Convolution2DReparameterization, - tfp.layers.Convolution2DFlipout): - inputs = tf.random.uniform([2, 3, 3, 1]) - elif layer_class in (tfp.layers.Convolution3DReparameterization, - tfp.layers.Convolution3DFlipout): - inputs = tf.random.uniform([2, 3, 3, 3, 1]) - inputs = self.maybe_transpose_inputs(inputs) + if layer_class in (tfp.layers.Convolution1DReparameterization, + tfp.layers.Convolution1DFlipout): + inputs = tf.random.uniform([2, 3, 1]) + outputs = tf.random.uniform([2, 1, 2]) + elif layer_class in (tfp.layers.Convolution2DReparameterization, + tfp.layers.Convolution2DFlipout): + inputs = tf.random.uniform([2, 3, 3, 1]) + outputs = tf.random.uniform([2, 1, 1, 2]) + elif layer_class in (tfp.layers.Convolution3DReparameterization, + tfp.layers.Convolution3DFlipout): + inputs = tf.random.uniform([2, 3, 3, 3, 1]) + outputs = tf.random.uniform([2, 1, 1, 1, 2]) + inputs = self.maybe_transpose_tensor(inputs) + outputs = self.maybe_transpose_tensor(outputs) + + net = tf.keras.Sequential([ + layer_class(filters=2, kernel_size=3, data_format=self.data_format), + layer_class(filters=2, kernel_size=1, data_format=self.data_format)]) - net = tf.keras.Sequential([ - layer_class(filters=2, kernel_size=3, data_format=self.data_format), - layer_class(filters=2, kernel_size=1, data_format=self.data_format)]) - output = net(inputs) + net.compile(loss='mse', optimizer='adam') + net.fit(inputs, outputs, batch_size=2, epochs=3, steps_per_epoch=2) - # Verify that the network runs without errors - sess.run(tf1.global_variables_initializer()) - sess.run(output) + batch_output = self.evaluate(net(inputs)) + self.assertAllEqual(outputs.shape, batch_output.shape) def testKerasLayerConvolution1DReparameterization(self): self._testKerasLayer(tfp.layers.Convolution1DReparameterization) diff --git a/tensorflow_probability/python/version.py b/tensorflow_probability/python/version.py index 46e7510935..49de77855f 100644 --- a/tensorflow_probability/python/version.py +++ b/tensorflow_probability/python/version.py @@ -24,7 +24,7 @@ # stable release (indicated by `_VERSION_SUFFIX = ''`). Outside the context of a # release branch, the current version is by default assumed to be a # 'development' version, labeled 'dev'. -_VERSION_SUFFIX = 'rc0' +_VERSION_SUFFIX = 'rc1' # Example, '0.4.0-dev' __version__ = '.'.join([