kaolin.ops.batch

Batching

Batching data in 3D can be tricky due to the heterogeneous sizes.

For instance, point clouds can have different number of points, which means we can’t always just concatenate the tensors on a batch axis.

Kaolin supports different batching strategies:

Exact

Exact batching is the logical representation for homogeneous data.

For instance, if you sample the same numbers of points from a batch of meshes, you would just have a single tensor of shape \((\text{batch_size}, \text{number_of_points}, 3)\).

Padded

Heterogeneous tensors are padded to identical dimensions with a constant value so that they can be concatenated on a batch axis. This is similar to padding for the batching of image data of different shapes.

Note

The last dimension must always be of the size of the element, e.g. 3 for 3D points (element of point clouds) or 1 for a grayscale pixel (element of grayscale textures).

For instance, for two textures \(T_0\) and \(T_1\) of shape \((32, 32, 3)\) and \((64, 16, 3)\) the batched tensor will be of shape \((2, max(32, 64), max(32, 16), 3) = (2, 64, 32, 3)\) and the padding value will be \(0\). \(T_0\) will be padded on the 1st axis by \(32\) while \(T_1\) will be padded on the 2nd axis by \(16\).

You can also enforce a specific maximum shape (if you want to have a fix memory consumption or use optimization like cudnn algorithm selection).

For instance, you can force \(T_0\) and \(T_1\) to be batched with a maximum shape of \((128, 128)\), the batched tensor will be of shape \((2, 128, 128, 3)\), \(T_0\) will be padded on the 1st axis and 2nd axis by 96 and \(T_1\) will be padded on the 1st axis by \(64\) and on the 2nd axis by \(112\).

For more information on how to do padded batching check kaolin.ops.batch.list_to_padded()

Packed

Heterogeneous tensors are reshaped to 2D \((-1, \text{last_dimension})\) and concatenated on the first axis. This is similar to packed sentences in NLP.

Note

The last dimension must always be of the size of the element, e.g. 3 for 3D points (element of point clouds) or 1 for a grayscale pixel (element of grayscale textures).

For instance, for two textures \(T_0\) and \(T_1\) of shape \((32, 32, 3)\) and \((64, 16, 3)\) The batched tensor will be of shape \((32 * 32 + 64 * 16, 3)\). \(T_0\) will be reshaped to \((32 * 32, 3)\) and \(T_1\) will be reshaped \((64 * 16, 3)\), before being concatenated on the first axis.

For more information on how to do padded batching check kaolin.ops.batch.list_to_packed()

Related attributes:

  • shape_per_tensor: 2D torch.LongTensor stores the shape of each sub-tensor except the last dimension in the padded tensor. E.g., in the example above shape_per_tensor would be torch.LongTensor([[32, 32], [64, 16]]). Refer to kaolin.ops.batch.get_shape_per_tensor() for more information.

  • first_idx: 1D torch.LongTensor stores the first index of each subtensor and the last index + 1 on the first axis in the packed tensor. E.g., in the example above first_idx would be torch.LongTensor([0, 1024, 2048]). This attribute are used for delimiting each subtensor into the packed tensor, for instance, to slice or index. Refer to kaolin.ops.batch.get_first_idx() for more information.

API

kaolin.ops.batch.fill_max_shape(shape_per_tensor, partial_max_shape=None)

Fills partial definition of shape to be at least as big as each shape in shape_per_tensor.

if the i-th dimension is -1 then the i-th output will be shape_per_tensor[:,i].max().

Parameters
Returns

the max_shape fully defined, of same size than partial_max_shape.

Return type

(torch.Tensor)

Example

>>> partial_max_shape = (6, -1, -1)
>>> shape_per_tensor = torch.LongTensor([[2, 3, 5],
...                                      [3, 4, 2]])
>>> fill_max_shape(shape_per_tensor, partial_max_shape)
tensor([6, 4, 5])
kaolin.ops.batch.get_first_idx(numel_per_tensor)

Returns the first indices of each tensor in the packed tensor.

See first_idx definition for more information.

Parameters

numel_per_tensor (torch.LongTensor) – The number of elements (vertices, faces, points…) in each unbatched tensor, as a 1D tensor.

Returns

first indices for each unbatched tensor in the packed tensor, and the last index + 1, as 1D tensor.

Return type

(torch.LongTensor)

Example

>>> numel_per_tensor = torch.LongTensor([2, 3, 5])
>>> get_first_idx(numel_per_tensor)
tensor([ 0,  2,  5, 10])
kaolin.ops.batch.get_shape_per_tensor(tensor_list)

Returns the shape of each tensor in the tensor list except the last dimension.

See shape_per_tensor for packed or padded for more information.

Parameters

tensor_list (sequence of torch.Tensor) – any python sequence of tensors of the identical type, number of dimensions, and last dimension size, e.g. \([(H_0, W_0, C), (H_1, W_1, C)]\).

Returns

the shape of each subtensor (except for the last dim), of shape \((len(\text{tensor_list}), \text{tensor_list[0].ndim} - 1)\).

Return type

(torch.Tensor)

Examples

>>> tensor_list = [
...         torch.zeros((1, 3, 4, 2)),
...         torch.ones((2, 5, 3, 2))
... ]
>>> get_shape_per_tensor(tensor_list)
tensor([[1, 3, 4],
        [2, 5, 3]])
kaolin.ops.batch.list_to_packed(tensor_list)

Converts a sequence of torch.Tensor into a single packed tensor.

torch.Tensor of same type, number of dimensions and last dimension size will be reshaped to \((-1, \text{last_dim})\) and concatenated on first axis. E.g.: With input of shapes \([(X_0, Y_0, Z_0, C), (X_1, Y_1, Z_1, C)]\) the output packed tensor will be of shape \(((X_0 * Y_0 * Z_0 + X_1 * Y_1 * Z_1), C)\). The output shape_per_tensor will be the tensor: \([[X_0, Y_0, Z_0], [X_1, Y_1, Z_1]]\).

Parameters

tensor_list (sequence of torch.Tensor) – any python sequence of tensors of identical type, number of dimensions, and last dimension size, e.g. \([(H_0, W_0, C), (H_1, W_1, C)]\).

Returns

the packed tensor and the associated shape_per_tensor

Return type

(torch.Tensor, torch.LongTensor)

Example

>>> a = torch.LongTensor([[0, 1, 2],
...                       [1, 2, 3]])
>>> b = torch.LongTensor([[2, 4, 5]])
>>> packed_tensor, shape_per_tensor = list_to_packed([a, b])
>>> packed_tensor
tensor([[0, 1, 2],
        [1, 2, 3],
        [2, 4, 5]])
>>> shape_per_tensor
tensor([[2],
        [1]])
kaolin.ops.batch.list_to_padded(tensor_list, padding_value, max_shape=None)

Converts a sequence of torch.Tensor into a single padded tensor.

torch.Tensor of same type, number of dimensions and last dimension size will be padded and stacked on first axis. E.g.: With input of shapes \([(X_0, Y_0, Z_0, C), (X_1, Y_1, Z_1, C)]\) the output padded tensor will be of shape \((2, max(X_0, X_1, \text{max_shape}[0]), max(Y_0, Y_1, \text{max_shape}[1]), max(Z_0, Z_1, \text{max_shape}[2]), C)\) The output shape_per_tensor with be the tensor: \([[X_0, Y_0, Z_0], [X_1, Y_1, Z_1]].\)

Parameters
  • tensor_list (sequence of torch.Tensor) – any python sequence of tensors of identical type, number of dimensions, and last dimension size, e.g. \([(H_0, W_0, C), (H_1, W_1, C)]\).

  • padding_value (float) – the value that will be used as padding.

  • max_shape (list, tuple or torch.LongTensor) – list of maximum value for each dim of the output shape (except batch and last axis), if a value is set to None then it will be the maximum value among the tensors. Default: All maximum values among the tensors.

Returns

the padded tensor and the associated shape_per_tensor.

Return type

(torch.Tensor, torch.LongTensor)

Example

>>> a = torch.LongTensor([[0, 1, 2],
...                       [1, 2, 3]])
>>> b = torch.LongTensor([[2, 4, 5]])
>>> padded_tensor, shape_per_tensor = list_to_padded([a, b], -1, [3])
>>> padded_tensor
tensor([[[ 0,  1,  2],
         [ 1,  2,  3],
         [-1, -1, -1]],

        [[ 2,  4,  5],
         [-1, -1, -1],
         [-1, -1, -1]]])
>>> shape_per_tensor
tensor([[2],
        [1]])
kaolin.ops.batch.packed_to_list(packed_tensor, shape_per_tensor, first_idx)

Converts a single packed tensor into a sequence of torch.Tensor.

Parameters
  • packed_tensor (torch.Tensor) – input packed tensor.

  • shape_per_tensor (torch.LongTensor) – shape_per_tensor associated to the packed tensor.

  • first_idx (torch.LongTensor) – first_idx associated to the packed tensor.

Returns

list of tensor unbatched from packed_tensor

Return type

list of torch.Tensor

Example

>>> packed_tensor = torch.arange(16).reshape(8, 2)
>>> packed_tensor
tensor([[ 0,  1],
        [ 2,  3],
        [ 4,  5],
        [ 6,  7],
        [ 8,  9],
        [10, 11],
        [12, 13],
        [14, 15]])
>>> shape_per_tensor = torch.LongTensor([[3], [4], [1]])
>>> first_idx = torch.LongTensor([0, 3, 7, 8])
>>> packed_to_list(packed_tensor, shape_per_tensor, first_idx)
[tensor([[0, 1],
        [2, 3],
        [4, 5]]), tensor([[ 6,  7],
        [ 8,  9],
        [10, 11],
        [12, 13]]), tensor([[14, 15]])]
kaolin.ops.batch.packed_to_padded(packed_tensor, shape_per_tensor, first_idx, padding_value, max_shape=None)

Converts a single packed tensor into a padded tensor.

Parameters
  • packed_tensor (torch.Tensor) – a packed tensor.

  • shape_per_tensor (torch.LongTensor) – the shape_per_tensor tensor associated to the padded tensor.

  • first_idx (torch.LongTensor) – first_idx associated to the packed tensor.

  • padding_value (float) – the value that will be used as padding.

  • max_shape (list, tuple or torch.LongTensor) – list of maximum value for each dim of the output shape (except batch and last axis), if a value is set to None then it will be the maximum value among the tensors. Default: All maximum values among the tensors.

Returns

the padded tensor.

Return type

(torch.Tensor)

kaolin.ops.batch.padded_to_list(padded_tensor, shape_per_tensor)

Converts a single padded tensor into a sequence of torch.Tensor.

Parameters
Returns

list of tensor unbatched from padded_tensor

Return type

list of torch.Tensor

Example

>>> padded_tensor = torch.LongTensor([[[0, 1, 2],
...                                    [1, 2, 3],
...                                    [-1, -1, -1]],
...                                   [[2, 4, 5],
...                                    [-1, -1, -1],
...                                    [-1, -1, -1]]])
>>> shape_per_tensor = torch.LongTensor([[2], [1]])
>>> padded_to_list(padded_tensor, shape_per_tensor)
[tensor([[0, 1, 2],
        [1, 2, 3]]), tensor([[2, 4, 5]])]
kaolin.ops.batch.padded_to_packed(padded_tensor, shape_per_tensor)

Converts a single padded tensor into a packed tensor.

Parameters
Returns

the packed tensor.

Return type

(torch.Tensor)

kaolin.ops.batch.tile_to_packed(values, numel_per_tensor)

Tiles values to a packed representation of numel_per_tensor,

Parameters
  • values (torch.Tensor) – tensor of shape \((\text{batch_size},)\) of values to be tiled.

  • numel_per_tensor (torch.LongTensor) – number of elements per tensor of the output packed tensor.

Returns

The packed tensor of tiled values of shape \((sum(\text{numel_per_tensor}), 1)\).

Return type

torch.Tensor

Example

>>> values = torch.tensor([0., 6., 7.])
>>> numel_per_tensor = torch.LongTensor([2, 2, 3])
>>> tile_to_packed(values, numel_per_tensor)
tensor([[0.],
        [0.],
        [6.],
        [6.],
        [7.],
        [7.],
        [7.]])