def build(self):
D = np.prod(self._event_dims)
flow = []
for i in range(self._num_coupling_layers):
if self._use_batch_normalization:
batch_normalization_bijector = bijectors.BatchNormalization()
flow.append(batch_normalization_bijector)
real_nvp_bijector = bijectors.RealNVP(
num_masked=D // 2,
shift_and_log_scale_fn=conditioned_real_nvp_template(
hidden_layers=self._hidden_layer_sizes,
# TODO: test tf.nn.relu
activation=tf.nn.tanh),
name='real_nvp_{}'.format(i))
flow.append(real_nvp_bijector)
if i < self._num_coupling_layers - 1:
permute_bijector = bijectors.Permute(
permutation=list(reversed(range(D))),
name='permute_{}'.format(i))
# TODO(hartikainen): We need to force _is_constant_jacobian due
# to the event_dim caching. See the issue filed at github:
# https://github.com/tensorflow/probability/issues/122
permute_bijector._is_constant_jacobian = False
flow.append(permute_bijector)
# Note: bijectors.Chain applies the list of bijectors in the
# _reverse_ order of what they are inputted.
self.flow = flow
def init_bijectors(self, n_layers, hidden_layers):
with tf.variable_scope(self.name):
bijectors = []
for i in range(n_layers):
if self.flow_type == "MAF":
bijectors.append(
tfb.MaskedAutoregressiveFlow(
shift_and_log_scale_fn=tfb.
masked_autoregressive_default_template(
hidden_layers=hidden_layers,
activation=tf.nn.relu,
log_scale_min_clip=self.log_scale_min_clip,
log_scale_max_clip=self.log_scale_max_clip,
shift_only=self.shift_only,
log_scale_clip_gradient=self.
log_scale_clip_gradient,
name="MAF_template_{}".format(i)),
name="MAF_{}".format(i)))
elif self.flow_type == "IAF":
bijectors.append(
tfb.Invert(tfb.MaskedAutoregressiveFlow(
shift_and_log_scale_fn=tfb.
masked_autoregressive_default_template(
hidden_layers=hidden_layers,
activation=tf.nn.relu,
log_scale_min_clip=self.log_scale_min_clip,
log_scale_max_clip=self.log_scale_max_clip,
shift_only=self.shift_only,
log_scale_clip_gradient=self.
log_scale_clip_gradient,
name="MAF_template_{}".format(i))),
name="IAF_{}".format(i)))
elif self.flow_type == "RealNVP":
bijectors.append(
tfb.RealNVP(num_masked=self.event_size - 1,
shift_and_log_scale_fn=tfb.
real_nvp_default_template(
hidden_layers=hidden_layers,
activation=tf.nn.relu,
shift_only=self.shift_only,
name="RealNVP_template_{}".format(i)),
name="RealNVP_{}".format(i)))
else:
raise ValueError("Unknown flow type {}".format(
self.flow_type))
bijectors.append(
tfb.Permute(permutation=list(range(1, self.event_size)) +
[0]))
# bijectors.append(
# tfb.Permute(
# self.init_once(np.random.permutation(self.event_size).astype("int32"),
# name="permutation_{}".format(i))
# )
# )
flow_bijector = tfb.Chain(list(reversed(bijectors[:-1])),
validate_args=True,
name="NF_chain")
return flow_bijector
def _build(self, input_shape):
input_depth = tf.compat.dimension_value(
tensorshape_util.with_rank_at_least(input_shape, 1)[-1])
self._input_depth = input_depth
flow_parts = []
for i in range(self._num_coupling_layers):
if self._use_batch_normalization:
batch_normalization_bijector = bijectors.BatchNormalization()
flow_parts += [batch_normalization_bijector]
real_nvp_bijector = bijectors.RealNVP(
num_masked=input_depth // 2,
shift_and_log_scale_fn=feedforward_scale_and_log_diag_fn(
hidden_layer_sizes=self._hidden_layer_sizes,
activation=tf.nn.relu),
name='real_nvp_{}'.format(i))
flow_parts += [real_nvp_bijector]
if i < self._num_coupling_layers - 1:
permute_bijector = bijectors.Permute(
permutation=list(reversed(range(input_depth))),
name='permute_{}'.format(i))
flow_parts += [permute_bijector]
# bijectors.Chain applies the list of bijectors in the
# _reverse_ order of what they are inputted, thus [::-1].
self.flow = bijectors.Chain(flow_parts[::-1])
self._built = True
def _build(self, input_shape):
input_depth = tf.compat.dimension_value(
tensorshape_util.with_rank_at_least(input_shape, 1)[-1])
self._input_depth = input_depth
flow_parts = []
for i in range(self._num_coupling_layers):
if self._use_batch_normalization:
# TODO(hartikainen): Allow other normalizations, e.g.
# weight normalization?
batch_normalization_bijector = bijectors.BatchNormalization()
flow_parts += [batch_normalization_bijector]
real_nvp_bijector = bijectors.RealNVP(
fraction_masked={
True: 1.0,
False: -1.0
}[i % 2 == 0] * 0.5,
bijector_fn=FeedforwardBijectorFunction(
hidden_layer_sizes=self._hidden_layer_sizes,
activation=self._activation),
name=f'real_nvp_{i}')
flow_parts += [real_nvp_bijector]
# bijectors.Chain applies the list of bijectors in the
# _reverse_ order of what they are inputted, thus [::-1].
self.flow = bijectors.Chain(flow_parts[::-1])
self._built = True