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): 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 _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