def flush_and_chunk_episode(example_strings, class_ids, chunk_sizes): """Removes flushed examples from an episode and chunks it. This function: 1) splits the batch of examples into a "flush" chunk and some number of additional chunks (as determined by `chunk_sizes`), 2) throws away the "flush" chunk, and 3) removes the padded dummy examples from the additional chunks. For example, in the context of few-shot learning, where episodes are composed of a support set and a query set, `chunk_size = (150, 100, 50)` would be interpreted as describing a "flush" chunk of size 150, a "support" chunk of size 100, and a "query" chunk of size 50. Args: example_strings: 1-D Tensor of dtype str, tf.train.Example protocol buffers. class_ids: 1-D Tensor of dtype int, class IDs (absolute wrt the original dataset). chunk_sizes: tuple of ints representing the sizes of the flush and additional chunks. Returns: A tuple of episode chunks of the form `((chunk_0_example_strings, chunk_0_class_ids), (chunk_1_example_strings, chunk_1_class_ids), ...)`. """ example_strings_chunks = tf.split(example_strings, num_or_size_splits=chunk_sizes)[1:] class_ids_chunks = tf.split(class_ids, num_or_size_splits=chunk_sizes)[1:] return tuple( filter_dummy_examples(strings, ids) for strings, ids in zip(example_strings_chunks, class_ids_chunks))
def _shift_and_log_scale_fn(x, output_dims, *args, **kwargs): del args del kwargs if output_dims != self.num_dims - num_masked: raise ValueError('Expected {} output_dims, got {}.'.format( self.num_dims - num_masked, output_dims)) return tf.split(shift_and_log_scale_model(x), 2, axis=-1)
def fit_gaussian(embeddings, damping=1e-7, full_covariance=False): """Fits a unimodal Gaussian distribution to `embeddings`. Args: embeddings: A [batch_size, embedding_dim] tf.Tensor of embeddings. damping: The scale of the covariance damping coefficient. full_covariance: Whether to use a full or diagonal covariance. Returns: Parameter estimates (means and log variances) for a Gaussian model. """ if full_covariance: num, dim = tf.split(tf.shape(input=embeddings), num_or_size_splits=2) num, dim = tf.squeeze(num), tf.squeeze(dim) sample_mean = tf.reduce_mean(input_tensor=embeddings, axis=0) centered_embeddings = embeddings - sample_mean sample_covariance = tf.einsum('ij,ik->kj', centered_embeddings, centered_embeddings) # Outer product. sample_covariance += damping * tf.eye(dim) # Positive definiteness. sample_covariance /= tf.cast(num, dtype=tf.float32) # Scale by N. return sample_mean, sample_covariance else: sample_mean, sample_variances = tf.nn.moments(x=embeddings) log_variances = tf.math.log(sample_variances + damping * tf.ones_like(sample_variances)) return sample_mean, log_variances
def get_z(h, latent_shape, name="latent", dist="normal"): h = tf.layers.dense(h, latent_shape * 2, name=name, activation=None, reuse=tf.AUTO_REUSE) mu, log_sigma = tf.split(h, 2, axis=-1) sigma = 0.1 + 0.9 * tf.nn.softplus(log_sigma) if dist == "multi": return tf.contrib.distributions.MultivariateNormalDiag(mu, sigma) else: return tf.distributions.Normal(mu, sigma)
def _get_class_labels_and_predictions(labels, logits, num_classes, multi_label): """Returns list of per-class-labels and list of per-class-predictions. Args: labels: A `Tensor` of size [n, k]. In the multi-label case, values are either 0 or 1 and k = num_classes. Otherwise, k = 1 and values are in [0, num_classes). logits: A `Tensor` of size [n, `num_classes`] representing the logits of each pixel and semantic class. num_classes: Number of classes. multi_label: Boolean which defines if we are in a multi_label setting, where pixels can have multiple labels, or not. Returns: class_labels: List of size num_classes, where each entry is a `Tensor' of size [batch_size, height, width] of type float with values of 0 or 1 representing the ground truth labels. class_predictions: List of size num_classes, each entry is a `Tensor' of size [batch_size, height, width] of type float with values of 0 or 1 representing the predicted labels. """ class_predictions = [None] * num_classes if multi_label: class_labels = tf.split(labels, num_or_size_splits=num_classes, axis=1) class_logits = tf.split(logits, num_or_size_splits=num_classes, axis=1) for c in range(num_classes): class_predictions[c] = tf.cast(tf.greater(class_logits[c], 0), dtype=tf.float32) else: class_predictions_flat = tf.argmax(logits, 1) class_labels = [None] * num_classes for c in range(num_classes): class_labels[c] = tf.cast(tf.equal(labels, c), dtype=tf.float32) class_predictions[c] = tf.cast(tf.equal(class_predictions_flat, c), dtype=tf.float32) return class_labels, class_predictions
def call(self, observation, step_type=(), network_state=()): del step_type # unused. output = tf.cast(tf.nest.flatten(observation)[0], tf.float32) for layer in self._mlp_layers: output = layer(output) shift, log_scale_diag = tf.split(output, 2, axis=-1) log_scale_diag = tf.clip_by_value(log_scale_diag, -20, 2) base_distribution = tfp.distributions.MultivariateNormalDiag( loc=shift, scale_diag=tf.exp(log_scale_diag)) distribution = SquashToSpecDistribution( base_distribution, self._single_action_spec) distribution = tf.nest.pack_sequence_as(self.output_spec, [distribution]) return distribution, network_state
def preprocess_spatial_observation(input_obs, spec, categorical_embedding_dims=16, non_categorical_scaling='log'): with tf.name_scope('preprocess_spatial_obs'): features = Lambda(lambda x: tf.split(x, x.get_shape()[1], axis=1))(input_obs) for f in spec.features: if f.is_categorical: features[f.index] = Lambda(lambda x: tf.squeeze(x, axis=1))(features[f.index]) features[f.index] = Embedding(f.scale, categorical_embedding_dims)(features[f.index]) features[f.index] = Permute((3, 1, 2))(features[f.index]) else: if non_categorical_scaling == 'log': features[f.index] = Lambda(lambda x: tf.log(x + 1e-10))(features[f.index]) elif non_categorical_scaling == 'normalize': features[f.index] = Lambda(lambda x: x / f.scale)(features[f.index]) return features
def preprocess_observation(input_obs, spec): if spec.is_spatial: features = Lambda( lambda x: tf.split(x, x.get_shape()[1], axis=1))(input_obs) for f in spec.features: if f.type == FeatureType.CATEGORICAL: features[f.index] = Lambda( lambda x: one_hot_encode(x, f.scale))( features[f.index]) else: features[f.index] = Lambda(lambda x: x / f.scale)( features[f.index]) return features else: return input_obs
def fit_gaussian_mixture(embeddings, responsibilities, damping=1e-7, full_covariance=False): """Fits a unimodal Gaussian distribution `embeddings`. Args: embeddings: A [batch_size, embedding_dim] tf.Tensor of embeddings. responsibilities: The per-component responsibilities. damping: The scale of the covariance damping coefficient. full_covariance: Whether to use a full or diagonal covariance. Returns: Parameter estimates for a Gaussian mixture model. """ num, dim = tf.split(tf.shape(input=embeddings), num_or_size_splits=2) num, dim = tf.squeeze(num), tf.squeeze(dim) num_classes = responsibilities.shape[1] mixing_proportion = tf.einsum('jk->k', responsibilities) mixing_proportion /= tf.cast(num, dtype=tf.float32) mixing_logits = tf.math.log(mixing_proportion) sample_mean = tf.einsum('ij,ik->jk', responsibilities, embeddings) sample_mean /= tf.reduce_sum( input_tensor=responsibilities, axis=0)[:, tf.newaxis] centered_embeddings = ( embeddings[:, tf.newaxis, :] - sample_mean[tf.newaxis, :, :]) if full_covariance: sample_covariance = tf.einsum('ijk,ijl->ijkl', centered_embeddings, centered_embeddings) # Outer product. sample_covariance += damping * tf.eye(dim) # Positive definiteness. weighted_covariance = tf.einsum('ij,ijkl->jkl', responsibilities, sample_covariance) weighted_covariance /= tf.reduce_sum( input_tensor=responsibilities, axis=0)[:, tf.newaxis, tf.newaxis] return ( _split_and_squeeze(sample_mean, num_splits=num_classes), _split_and_squeeze(weighted_covariance, num_splits=num_classes), [mixing_logits], ) else: avg_x_squared = ( tf.matmul(responsibilities, embeddings**2, transpose_a=True) / tf.reduce_sum(input_tensor=responsibilities, axis=0)[:, tf.newaxis]) avg_means_squared = sample_mean**2 avg_x_means = ( sample_mean * tf.matmul(responsibilities, embeddings, transpose_a=True) / tf.reduce_sum(input_tensor=responsibilities, axis=0)[:, tf.newaxis]) sample_variances = ( avg_x_squared - 2 * avg_x_means + avg_means_squared + damping * tf.ones(dim)) log_variances = tf.math.log(sample_variances) return ( _split_and_squeeze(sample_mean, num_splits=num_classes), _split_and_squeeze(log_variances, num_splits=num_classes), [mixing_logits], )
def _split_and_squeeze(tensor, num_splits, axis=0): return [ tf.squeeze(t) for t in tf.split(tensor, axis=axis, num_or_size_splits=num_splits) ]
def _split_mode_params(params): return [ tf.squeeze(p) for p in tf.split(params, axis=0, num_or_size_splits=self.num_modes) ]
def discriminator(self, z): probabilities = self.discriminator_net(z) return tf.split(probabilities, 2, axis=-1)[0]