def test_unflatten_batch_to_2d_preset(self): """Test unflattening with a preset input.""" data = 1. + np.reshape(np.arange(12, dtype=np.float32), (6, 2)) sizes = (2, 3, 1) output_true = np.array((((1., 2.), (3., 4.), (0., 0.)), ((5., 6.), (7., 8.), (9., 10.)), ((11., 12.), (0., 0.), (0., 0.))), dtype=np.float32) output_true_padded = np.pad(output_true, ((0, 0), (0, 2), (0, 0)), mode="constant") output = utils.unflatten_2d_to_batch(data, sizes, max_rows=None) output_padded = utils.unflatten_2d_to_batch(data, sizes, max_rows=5) self.assertAllEqual(output, output_true) self.assertAllEqual(output_padded, output_true_padded)
def get_unflatten(self, max_num_points=None, name=None): """ Returns the method to unflatten the segmented points. Use this instead of accessing 'self._unflatten', if the class was constructed using segmented input the '_unflatten' method is created in this method. Note: In the following, A1 to An are optional batch dimensions. Args: max_num_points: An `int`, specifies the 'V' dimension the method returns, by default uses maximum of 'sizes'. `max_rows >= max(sizes)` Returns: A method to unflatten the segmented points, which returns a `Tensor` of shape `[A1,...,An,V,D]`, zero padded. Raises: ValueError: When trying to unflatten unsorted points. """ if self._unflatten is None: self._unflatten = lambda data: unflatten_2d_to_batch( data=tf.gather(data, self._sorted_indices_batch), sizes=self.get_sizes(), max_rows=max_num_points) return self._unflatten
def test_unflatten_batch_to_2d_types(self, dtype): """Test unflattening with int and float types.""" data = np.ones(shape=(6, 2), dtype=dtype) sizes = (2, 2, 2) unflattened_true = np.ones(shape=(3, 2, 2), dtype=dtype) unflattened = utils.unflatten_2d_to_batch(data, sizes) self.assertEqual(data.dtype, unflattened.dtype.as_numpy_dtype) self.assertAllEqual(unflattened, unflattened_true)
def test_unflatten_batch_to_2d_jacobian_random( self, sizes, max_rows, num_features): """Test that the jacobian is correct.""" max_rows = np.max(sizes) if max_rows is None else max_rows total_rows = np.sum(sizes) data_init = 0.1 + np.random.uniform(size=(total_rows, num_features)) data = tf.convert_to_tensor(value=data_init) unflattened = utils.unflatten_2d_to_batch(data, sizes, max_rows) self.assert_jacobian_is_correct(data, data_init, unflattened)
def test_unflatten_batch_to_2d_random(self, sizes, max_rows, num_features): """Test unflattening with random inputs.""" max_rows = np.max(sizes) if max_rows is None else max_rows output_shape = np.concatenate( (np.shape(sizes), (max_rows,), (num_features,))) total_rows = np.sum(sizes) data = 0.1 + np.random.uniform(size=(total_rows, num_features)) unflattened = utils.unflatten_2d_to_batch(data, sizes, max_rows) flattened = tf.reshape(unflattened, (-1, num_features)) nonzero_rows = tf.where(tf.norm(tensor=flattened, axis=-1)) flattened_unpadded = tf.gather( params=flattened, indices=tf.squeeze(input=nonzero_rows, axis=-1)) self.assertAllEqual(tf.shape(input=unflattened), output_shape) self.assertAllEqual(flattened_unpadded, data)
def unflatten_2d_to_batch(data): return utils.unflatten_2d_to_batch(data, sizes, max_rows)
def test_input_type_exception_raised(self): """Check that invalid input types trigger the right exception.""" with self.assertRaisesRegexp(TypeError, "'sizes' must have an integer type."): utils.unflatten_2d_to_batch(np.ones((3, 4)), np.ones((3, )))
def pool(data, pool_map, sizes, algorithm='max', name=None): # pyformat: disable """Implements graph pooling. The features at each output vertex are computed by pooling over a subset of vertices in the input graph. This pooling window is specified by the input `pool_map`. The shorthands used below are `V1`: The number of vertices in the input data. `V2`: The number of vertices in the pooled output data. `C`: The number of channels in the data. Note: In the following, A1 to An are optional batch dimensions. Args: data: A `float` tensor with shape `[A1, ..., An, V1, C]`. pool_map: A `SparseTensor` with the same type as `data` and with shape `[A1, ..., An, V2, V1]`. The features for an output vertex `v2` will be computed by pooling over the corresponding input vertices specified by the entries in `pool_map[A1, ..., An, v2, :]`. sizes: An `int` tensor of shape `[A1, ..., An, 2]` indicating the true input sizes in case of padding (`sizes=None` indicates no padding). `sizes[A1, ..., An, 0] <= V2` specifies the padding in the (pooled) output, and `sizes[A1, ..., An, 1] <= V1` specifies the padding in the input. algorithm: The pooling function, must be either 'max' or 'weighted'. Default is 'max'. For 'max' pooling, the output features are the maximum over the input vertices (in this case only the indices of the `SparseTensor` `pool_map` are used, the values are ignored). For 'weighted', the output features are a weighted sum of the input vertices, the weights specified by the values of `pool_map`. name: A name for this op. Defaults to 'graph_pooling_pool'. Returns: Tensor with shape `[A1, ..., An, V2, C]`. Raises: TypeError: if the input types are invalid. ValueError: if the input dimensions are invalid. ValueError: if `algorithm` is invalid. """ # pyformat: enable with tf.compat.v1.name_scope( name, 'graph_pooling_pool', [data, pool_map, sizes]): data = tf.convert_to_tensor(value=data) pool_map = tf.compat.v1.convert_to_tensor_or_sparse_tensor(value=pool_map) if sizes is not None: sizes = tf.convert_to_tensor(value=sizes) utils.check_valid_graph_pooling_input(data, pool_map, sizes) if sizes is not None: sizes_output, sizes_input = tf.split(sizes, 2, axis=-1) sizes_output = tf.squeeze(sizes_output, axis=-1) sizes_input = tf.squeeze(sizes_input, axis=-1) else: sizes_output = None sizes_input = None batched = data.shape.ndims > 2 if batched: x_flat, _ = utils.flatten_batch_to_2d(data, sizes_input) pool_map_block_diagonal = utils.convert_to_block_diag_2d(pool_map, sizes) else: x_flat = data pool_map_block_diagonal = pool_map if algorithm == 'weighted': pooled = tf.sparse.sparse_dense_matmul(pool_map_block_diagonal, x_flat) elif algorithm == 'max': pool_groups = tf.gather(x_flat, pool_map_block_diagonal.indices[:, 1]) pooled = tf.math.segment_max( data=pool_groups, segment_ids=pool_map_block_diagonal.indices[:, 0]) else: raise ValueError('The pooling method must be "weighted" or "max"') if batched: if sizes_output is not None: pooled = utils.unflatten_2d_to_batch(pooled, sizes_output) else: output_shape = tf.concat((tf.shape(input=pool_map)[:-1], (-1,)), axis=0) pooled = tf.reshape(pooled, output_shape) return pooled
def upsample_transposed_convolution(data, pool_map, sizes, kernel_size, transposed_convolution_op, name=None): # pyformat: disable r"""Graph upsampling by transposed convolution. Upsamples a graph using a transposed convolution op. The map from input vertices to the upsampled graph is specified by the reverse of pool_map. The inputs `pool_map` and `sizes` are the same as used for pooling: >>> pooled = pool(data, pool_map, sizes) >>> upsampled = upsample_transposed_convolution(pooled, pool_map, sizes, ...) The shorthands used below are `V1`: The number of vertices in the inputs. `V2`: The number of vertices in the upsampled output. `C`: The number of channels in the inputs. Note: In the following, A1 to A3 are optional batch dimensions. Only up to three batch dimensions are supported due to limitations with TensorFlow's dense-sparse multiplication. Please see the documentation for `graph_pooling.pool` for a detailed interpretation of the inputs `pool_map` and `sizes`. Args: data: A `float` tensor with shape `[A1, ..., A3, V1, C]`. pool_map: A `SparseTensor` with the same type as `data` and with shape `[A1, ..., A3, V1, V2]`. `pool_map` will be interpreted in the same way as the `pool_map` argument of `graph_pooling.pool`, namely `v_i_map = [..., v_i, :]` are the upsampled vertices corresponding to vertex `v_i`. Additionally, for transposed convolution a fixed number of entries in each `v_i_map` (equal to `kernel_size`) are expected: `|v_i_map| = kernel_size`. When this is not the case, the map is either truncated or the last element repeated. Furthermore, upsampled vertex indices should not be repeated across maps otherwise the output is nondeterministic. Specifically, to avoid nondeterminism we must have `intersect([a1, ..., an, v_i, :],[a1, ..., a3, v_j, :]) = {}, i != j`. sizes: An `int` tensor of shape `[A1, ..., A3, 2]` indicating the true input sizes in case of padding (`sizes=None` indicates no padding): `sizes[A1, ..., A3, 0] <= V1` and `sizes[A1, ..., A3, 1] <= V2`. kernel_size: The kernel size for transposed convolution. transposed_convolution_op: A callable transposed convolution op with the form `y = transposed_convolution_op(x)`, where `x` has shape `[1, 1, D1, C]` and `y` must have shape `[1, 1, kernel_size * D1, C]`. `transposed_convolution_op` maps each row of `x` to `kernel_size` rows in `y`. An example: `transposed_convolution_op = tf.keras.layers.Conv2DTranspose( filters=C, kernel_size=(1, kernel_size), strides=(1, kernel_size), padding='valid', ...) name: A name for this op. Defaults to 'graph_pooling_upsample_transposed_convolution'. Returns: Tensor with shape `[A1, ..., A3, V2, C]`. Raises: TypeError: if the input types are invalid. TypeError: if `transposed_convolution_op` is not a callable. ValueError: if the input dimensions are invalid. """ # pyformat: enable with tf.compat.v1.name_scope( name, 'graph_pooling_upsample_transposed_convolution', [data, pool_map, sizes]): data = tf.convert_to_tensor(value=data) pool_map = tf.compat.v1.convert_to_tensor_or_sparse_tensor(value=pool_map) if sizes is not None: sizes = tf.convert_to_tensor(value=sizes) utils.check_valid_graph_unpooling_input(data, pool_map, sizes) if not callable(transposed_convolution_op): raise TypeError("'transposed_convolution_op' must be callable.") if sizes is not None: sizes_input, sizes_output = tf.split(sizes, 2, axis=-1) sizes_input = tf.squeeze(sizes_input, axis=-1) sizes_output = tf.squeeze(sizes_output, axis=-1) else: sizes_input = None sizes_output = None num_features = tf.compat.v1.dimension_value(data.shape[-1]) batched = data.shape.ndims > 2 if batched: x_flat, _ = utils.flatten_batch_to_2d(data, sizes_input) pool_map_block_diagonal = utils.convert_to_block_diag_2d(pool_map, sizes) else: x_flat = data pool_map_block_diagonal = pool_map x_flat = tf.expand_dims(tf.expand_dims(x_flat, 0), 0) x_upsample = transposed_convolution_op(x_flat) # Map each upsampled vertex into its correct position based on pool_map. # Select 'kernel_size' neighbors for each input vertex. Truncate or repeat # as necessary. ragged = tf.RaggedTensor.from_value_rowids( pool_map_block_diagonal.indices[:, 1], pool_map_block_diagonal.indices[:, 0]) # Take up to the first 'kernel_size' entries. ragged_k = ragged[:, :kernel_size] # Fill rows with less than 'kernel_size' entries by repeating the last # entry. last = ragged_k[:, -1:].flat_values num_repeat = kernel_size - ragged_k.row_lengths() sum_num_repeat = tf.reduce_sum(input_tensor=num_repeat) ones_ragged = tf.RaggedTensor.from_row_lengths( tf.ones((sum_num_repeat,), dtype=last.dtype), num_repeat) repeat = ones_ragged * tf.expand_dims(last, -1) padded = tf.concat([ragged_k, repeat], axis=1) pool_map_dense = tf.reshape(padded.flat_values, (-1, kernel_size)) # Map rows of 'x_upsample' to positions indicated by the # indices 'pool_map_dense'. up_scatter_indices = tf.expand_dims(tf.reshape(pool_map_dense, (-1,)), -1) up_row = tf.reshape(tf.cast(up_scatter_indices, tf.int64), (-1,)) up_column = tf.range(tf.shape(input=up_row, out_type=tf.dtypes.int64)[0]) scatter_indices = tf.concat( (tf.expand_dims(up_row, -1), tf.expand_dims(up_column, -1)), axis=1) scatter_values = tf.ones_like(up_row, dtype=x_upsample.dtype) scatter_shape = tf.reduce_max(input_tensor=scatter_indices, axis=0) + 1 scatter = tf.SparseTensor(scatter_indices, scatter_values, scatter_shape) scatter = tf.sparse.reorder(scatter) row_sum = tf.sparse.reduce_sum(tf.abs(scatter), keepdims=True, axis=-1) row_sum = tf.compat.v1.where(tf.equal(row_sum, 0.), row_sum, 1.0 / row_sum) scatter = row_sum * scatter x_upsample = tf.sparse.sparse_dense_matmul(scatter, x_upsample[0, 0, :, :]) if batched: if sizes_output is not None: x_upsample = utils.unflatten_2d_to_batch(x_upsample, sizes_output) else: output_shape = tf.concat((tf.shape(input=pool_map)[:-2], tf.shape(input=pool_map)[-1:], (num_features,)), axis=0) x_upsample = tf.reshape(x_upsample, output_shape) return x_upsample