def _get_output_shape(rank, strides, padding, dilations, input_shape,
output_size, filter_shape, output_padding):
"""Compute the `output_shape` and `strides` arg used by `conv_transpose`."""
if output_padding is None:
output_padding = (None, ) * rank
else:
output_padding = nn_util_lib.prepare_tuple_argument(
output_padding, rank, 'output_padding')
for stride, out_pad in zip(strides, output_padding):
if out_pad >= stride:
raise ValueError('Stride {} must be greater than output '
'padding {}.'.format(strides, output_padding))
assert len(filter_shape) == rank
assert len(strides) == rank
assert len(output_padding) == rank
event_shape = []
for i in range(-rank, 0):
event_shape.append(
_deconv_output_length(input_shape[i - 1],
filter_shape[i],
padding=padding,
output_padding=output_padding[i],
stride=strides[i],
dilation=dilations[i]))
event_shape.append(output_size)
batch_shape = input_shape[:-rank - 1]
output_shape = prefer_static.concat([batch_shape, event_shape], axis=0)
strides = (1, ) + strides + (1, )
return output_shape, strides
def __init__(
self,
input_size,
output_size, # keras::Conv::filters
# Conv specific.
filter_shape, # keras::Conv::kernel_size
rank=2, # keras::Conv::rank
strides=1, # keras::Conv::strides
padding='VALID', # keras::Conv::padding; 'CAUSAL' not implemented.
# keras::Conv::data_format is not implemented
dilations=1, # keras::Conv::dilation_rate
# Weights
init_kernel_fn=None, # tfp.experimental.nn.initializers.glorot_uniform()
init_bias_fn=None, # tf.initializers.zeros()
make_kernel_bias_fn=nn_util_lib.make_kernel_bias,
dtype=tf.float32,
batch_shape=(),
# Misc
activation_fn=None,
name=None):
"""Constructs layer.
Note: `data_format` is not supported since all nn layers operate on
the rightmost column. If your channel dimension is not rightmost, use
`tf.transpose` before calling this layer. For example, if your channel
dimension is second from the left, the following code will move it
rightmost:
```python
inputs = tf.transpose(inputs, tf.concat([
[0], tf.range(2, tf.rank(inputs)), [1]], axis=0))
```
Args:
input_size: ...
In Keras, this argument is inferred from the rightmost input shape,
i.e., `tf.shape(inputs)[-1]`. This argument specifies the size of the
second from the rightmost dimension of both `inputs` and `kernel`.
Default value: `None`.
output_size: ...
In Keras, this argument is called `filters`. This argument specifies the
rightmost dimension size of both `kernel` and `bias`.
filter_shape: ...
In Keras, this argument is called `kernel_size`. This argument specifies
the leftmost `rank` dimensions' sizes of `kernel`.
rank: An integer, the rank of the convolution, e.g. "2" for 2D
convolution. This argument implies the number of `kernel` dimensions,
i.e., `kernel.shape.rank == rank + 2`.
In Keras, this argument has the same name and semantics.
Default value: `2`.
strides: An integer or tuple/list of n integers, specifying the stride
length of the convolution.
In Keras, this argument has the same name and semantics.
Default value: `1`.
padding: One of `"VALID"` or `"SAME"` (case-insensitive).
In Keras, this argument has the same name and semantics (except we don't
support `"CAUSAL"`).
Default value: `'VALID'`.
dilations: An integer or tuple/list of `rank` integers, specifying the
dilation rate to use for dilated convolution. Currently, specifying any
`dilations` value != 1 is incompatible with specifying any `strides`
value != 1.
In Keras, this argument is called `dilation_rate`.
Default value: `1`.
init_kernel_fn: ...
Default value: `None` (i.e.,
`tfp.experimental.nn.initializers.glorot_uniform()`).
init_bias_fn: ...
Default value: `None` (i.e., `tf.initializers.zeros()`).
make_kernel_bias_fn: ...
Default value: `tfp.experimental.nn.util.make_kernel_bias`.
dtype: ...
Default value: `tf.float32`.
batch_shape: ...
Default value: `()`.
activation_fn: ...
Default value: `None`.
name: ...
Default value: `None` (i.e., `'ConvolutionV2'`).
"""
filter_shape = nn_util_lib.prepare_tuple_argument(
filter_shape, rank, arg_name='filter_shape')
batch_shape = (np.array([], dtype=np.int32) if batch_shape is None else
prefer_static.reshape(batch_shape, shape=[-1]))
batch_ndims = prefer_static.size(batch_shape)
if tf.get_static_value(batch_ndims) == 0:
# In this branch, we statically know there are no batch dims.
kernel_shape = filter_shape + (input_size, output_size)
bias_shape = [output_size]
apply_kernel_fn = _make_convolution_fn(rank, strides, padding,
dilations)
else:
# In this branch, there are either static/dynamic batch dims or
# dynamically no batch dims.
kernel_shape = prefer_static.concat(
[batch_shape, filter_shape, [input_size, output_size]], axis=0)
bias_shape = prefer_static.concat(
[batch_shape, tf.ones(rank), [output_size]], axis=0)
apply_kernel_fn = lambda x, k: nn_util_lib.convolution_batch( # pylint: disable=g-long-lambda
x,
k,
rank=rank,
strides=strides,
padding=padding,
data_format='NBHWC',
dilations=dilations)
kernel, bias = make_kernel_bias_fn(kernel_shape, bias_shape,
init_kernel_fn, init_bias_fn,
batch_ndims, batch_ndims, dtype)
self._make_kernel_bias_fn = make_kernel_bias_fn # For tracking.
super(ConvolutionV2, self).__init__(kernel=kernel,
bias=bias,
apply_kernel_fn=apply_kernel_fn,
dtype=dtype,
activation_fn=activation_fn,
name=name)
def __init__(
self,
input_size,
output_size, # keras::Conv::filters
# Conv specific.
filter_shape, # keras::Conv::kernel_size
rank=2, # keras::Conv::rank
strides=1, # keras::Conv::strides
padding='VALID', # keras::Conv::padding; 'CAUSAL' not implemented.
# keras::Conv::data_format is not implemented
dilations=1, # keras::Conv::dilation_rate
# Weights
init_kernel_fn=None, # tfp.experimental.nn.initializers.glorot_uniform()
init_bias_fn=None, # tf.initializers.zeros()
make_posterior_fn=nn_util_lib.make_kernel_bias_posterior_mvn_diag,
make_prior_fn=nn_util_lib.make_kernel_bias_prior_spike_and_slab,
posterior_value_fn=tfd.Distribution.sample,
unpack_weights_fn=unpack_kernel_and_bias,
dtype=tf.float32,
# Penalty.
penalty_weight=None,
posterior_penalty_fn=kl_divergence_monte_carlo,
# Misc
activation_fn=None,
seed=None,
name=None):
"""Constructs layer.
Note: `data_format` is not supported since all nn layers operate on
the rightmost column. If your channel dimension is not rightmost, use
`tf.transpose` before calling this layer. For example, if your channel
dimension is second from the left, the following code will move it
rightmost:
```python
inputs = tf.transpose(inputs, tf.concat([
[0], tf.range(2, tf.rank(inputs)), [1]], axis=0))
```
Args:
input_size: ...
In Keras, this argument is inferred from the rightmost input shape,
i.e., `tf.shape(inputs)[-1]`. This argument specifies the size of the
second from the rightmost dimension of both `inputs` and `kernel`.
Default value: `None`.
output_size: ...
In Keras, this argument is called `filters`. This argument specifies the
rightmost dimension size of both `kernel` and `bias`.
filter_shape: ...
In Keras, this argument is called `kernel_size`. This argument specifies
the leftmost `rank` dimensions' sizes of `kernel`.
rank: An integer, the rank of the convolution, e.g. "2" for 2D
convolution. This argument implies the number of `kernel` dimensions,
i.e., `kernel.shape.rank == rank + 2`.
In Keras, this argument has the same name and semantics.
Default value: `2`.
strides: An integer or tuple/list of n integers, specifying the stride
length of the convolution.
In Keras, this argument has the same name and semantics.
Default value: `1`.
padding: One of `"VALID"` or `"SAME"` (case-insensitive).
In Keras, this argument has the same name and semantics (except we don't
support `"CAUSAL"`).
Default value: `'VALID'`.
dilations: An integer or tuple/list of `rank` integers, specifying the
dilation rate to use for dilated convolution. Currently, specifying any
`dilations` value != 1 is incompatible with specifying any `strides`
value != 1.
In Keras, this argument is called `dilation_rate`.
Default value: `1`.
init_kernel_fn: ...
Default value: `None` (i.e.,
`tfp.experimental.nn.initializers.glorot_uniform()`).
init_bias_fn: ...
Default value: `None` (i.e., `tf.initializers.zeros()`).
make_posterior_fn: ...
Default value:
`tfp.experimental.nn.util.make_kernel_bias_posterior_mvn_diag`.
make_prior_fn: ...
Default value:
`tfp.experimental.nn.util.make_kernel_bias_prior_spike_and_slab`.
posterior_value_fn: ...
Default valye: `tfd.Distribution.sample`
unpack_weights_fn:
Default value: `unpack_kernel_and_bias`
dtype: ...
Default value: `tf.float32`.
penalty_weight: ...
Default value: `None` (i.e., weight is `1`).
posterior_penalty_fn: ...
Default value: `kl_divergence_monte_carlo`.
activation_fn: ...
Default value: `None`.
seed: ...
Default value: `None` (i.e., no seed).
name: ...
Default value: `None` (i.e.,
`'ConvolutionVariationalFlipoutV2'`).
"""
filter_shape = nn_util_lib.prepare_tuple_argument(
filter_shape, rank, arg_name='filter_shape')
kernel_shape = filter_shape + (input_size, output_size)
self._make_posterior_fn = make_posterior_fn # For variable tracking.
self._make_prior_fn = make_prior_fn # For variable tracking.
batch_ndims = 0
super(ConvolutionVariationalFlipoutV2, self).__init__(
posterior=make_posterior_fn(kernel_shape, [output_size],
init_kernel_fn, init_bias_fn,
batch_ndims, batch_ndims, dtype),
prior=make_prior_fn(kernel_shape, [output_size], init_kernel_fn,
init_bias_fn, batch_ndims, batch_ndims, dtype),
apply_kernel_fn=_make_convolution_fn(rank, strides, padding,
dilations),
posterior_value_fn=posterior_value_fn,
unpack_weights_fn=unpack_weights_fn,
dtype=dtype,
penalty_weight=penalty_weight,
posterior_penalty_fn=posterior_penalty_fn,
activation_fn=activation_fn,
seed=seed,
name=name)
def __init__(
self,
input_size,
output_size, # keras::Conv::filters
# Conv specific.
filter_shape, # keras::Conv::kernel_size
rank=2, # keras::Conv::rank
strides=1, # keras::Conv::strides
padding='VALID', # keras::Conv::padding; 'CAUSAL' not implemented.
# keras::Conv::data_format is not implemented
dilations=1, # keras::Conv::dilation_rate
output_padding=None, # keras::ConvTranspose::output_padding
# Weights
init_kernel_fn=None, # tfp.experimental.nn.initializers.glorot_uniform()
init_bias_fn=None, # tf.initializers.zeros()
make_kernel_bias_fn=nn_util_lib.make_kernel_bias,
dtype=tf.float32,
# Misc
activation_fn=None,
name=None):
"""Constructs layer.
Note: `data_format` is not supported since all nn layers operate on
the rightmost column. If your channel dimension is not rightmost, use
`tf.transpose` before calling this layer. For example, if your channel
dimension is second from the left, the following code will move it
rightmost:
```python
inputs = tf.transpose(inputs, tf.concat([
[0], tf.range(2, tf.rank(inputs)), [1]], axis=0))
```
Args:
input_size: ...
In Keras, this argument is inferred from the rightmost input shape,
i.e., `tf.shape(inputs)[-1]`. This argument specifies the size of the
second from the rightmost dimension of both `inputs` and `kernel`.
Default value: `None`.
output_size: ...
In Keras, this argument is called `filters`. This argument specifies the
rightmost dimension size of both `kernel` and `bias`.
filter_shape: ...
In Keras, this argument is called `kernel_size`. This argument specifies
the leftmost `rank` dimensions' sizes of `kernel`.
rank: An integer, the rank of the convolution, e.g. "2" for 2D
convolution. This argument implies the number of `kernel` dimensions,
i.e., `kernel.shape.rank == rank + 2`.
In Keras, this argument has the same name and semantics.
Default value: `2`.
strides: An integer or tuple/list of n integers, specifying the stride
length of the convolution.
In Keras, this argument has the same name and semantics.
Default value: `1`.
padding: One of `"VALID"` or `"SAME"` (case-insensitive).
In Keras, this argument has the same name and semantics (except we don't
support `"CAUSAL"`).
Default value: `'VALID'`.
dilations: An integer or tuple/list of `rank` integers, specifying the
dilation rate to use for dilated convolution. Currently, specifying any
`dilations` value != 1 is incompatible with specifying any `strides`
value != 1.
In Keras, this argument is called `dilation_rate`.
Default value: `1`.
output_padding: An `int` or length-`rank` tuple/list representing the
amount of padding along the input spatial dimensions (e.g., depth,
height, width). A single `int` indicates the same value for all spatial
dimensions. The amount of output padding along a given dimension must be
lower than the stride along that same dimension. If set to `None`
(default), the output shape is inferred.
In Keras, this argument has the same name and semantics.
Default value: `None` (i.e., inferred).
init_kernel_fn: ...
Default value: `None` (i.e.,
`tfp.experimental.nn.initializers.glorot_uniform()`).
init_bias_fn: ...
Default value: `None` (i.e., `tf.initializers.zeros()`).
make_kernel_bias_fn: ...
Default value: `tfp.experimental.nn.util.make_kernel_bias`.
dtype: ...
Default value: `tf.float32`.
activation_fn: ...
Default value: `None`.
name: ...
Default value: `None` (i.e., `'ConvolutionTranspose'`).
"""
filter_shape = nn_util_lib.prepare_tuple_argument(
filter_shape, rank, 'filter_shape')
kernel_shape = filter_shape + (output_size, input_size
) # Note transpose.
batch_ndims = 0
kernel, bias = make_kernel_bias_fn(kernel_shape, [output_size],
init_kernel_fn, init_bias_fn,
batch_ndims, batch_ndims, dtype)
super(ConvolutionTranspose,
self).__init__(kernel=kernel,
bias=bias,
apply_kernel_fn=_make_convolution_transpose_fn(
rank, strides, padding, dilations,
filter_shape, output_size, output_padding),
dtype=dtype,
activation_fn=activation_fn,
name=name)