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
: 2Dtorch.LongTensor
stores the shape of each sub-tensor except the last dimension in the padded tensor. E.g., in the example aboveshape_per_tensor
would betorch.LongTensor([[32, 32], [64, 16]])
. Refer tokaolin.ops.batch.get_shape_per_tensor()
for more information.
first_idx
: 1Dtorch.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 abovefirst_idx
would betorch.LongTensor([0, 1024, 2048])
. This attribute are used for delimiting each subtensor into the packed tensor, for instance, to slice or index. Refer tokaolin.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
shape_per_tensor (torch.Tensor) – Input shape_per_tensor, of shape \((\text{N}, \text{ndim})\).
partial_max_shape (tuple, list or torch.Tensor) – partially defined maximum shape, of size
ndim
.
- Returns
the max_shape fully defined, of same size than
partial_max_shape
.- Return type
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
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
- kaolin.ops.batch.padded_to_list(padded_tensor, shape_per_tensor)¶
Converts a single padded tensor into a sequence of torch.Tensor.
- Parameters
padded_tensor (torch.Tensor) – a padded tensor.
shape_per_tensor (torch.LongTensor) – the shape_per_tensor tensor associated to the padded tensor.
- 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
padded_tensor (torch.Tensor) – a padded tensor.
shape_per_tensor (torch.LongTensor) – the shape_per_tensor tensor associated to the padded tensor.
- Returns
the packed tensor.
- Return type
- 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
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.]])