MLP API

MLP API provides tooling for implementing MLP encoders / decoders used in the Lightplane renderer and splatter.

class lightplane.DecoderParams(mlp_params: Tensor, n_hidden_trunk: Tensor, n_hidden_opacity: Tensor, n_hidden_color: Tensor, color_chn: int)[source]

Class configuring the learnable parameters of the decoder from Lightplane Renderer.

The decoder comprises a learnable function that predicts color and an opacity value given a grid feature sampled at every point along a rendering ray.

Specifically, the decoder function consists of three MLPs: trunk_mlp, opacity_mlp, and color_mlp. The three MLPs predict opacity and color as follows:

  1. use_separate_color_grid==False:

    grid -> f_i -> trunk_mlp -> e_i -> e_i + ray_encoding -> color_mlp -> c_i
                                -> opacity_mlp -> o_i

If the renderer uses a single grid for both opacity and color, an MLP trunk_mlp maps the grid-sampled feature f_i to a trunk feature e_i, which is later converted to opacity and color with a pair of additional color and opacity MLP heads color_mlp and opacity_mlp. The trunk feature e_i is summed with ray_encoding before color_mlp to make the predicted color viewpoint dependent.

  1. use_separate_color_grid==True:

    grid       -> f_i  -> opacity_mlp -> o_i
color_grid -> cf_i -> cf_i + ray_encoding -> color_mlp -> c_i

If the renderer uses a separate color grid (use_separate_color_grid==True), the trunk MLP will be omitted The opacity_mlp and color_mlp predict the opacity o_i and color values c_i, respectively, given an opacity/color features (f_i and cf_i) sampled from the corresponding grid grid and color_grid.

The parameters of the three MLPs are stored in the mlp_params attribute. Here, mlp_params is a 1D tensor which concatenates the flattened weight matrices and bias vectors of the three MLPs in the following order:

mlp_params = torch.cat(
    [
        weights_trunk[0].flatten(),
        ...
        weights_trunk[-1].flatten(),
        biases_trunk[0],
        ...
        biases_trunk[-1],
        weights_opacity[0].flatten(),
        ...
        weights_opacity[-1].flatten(),
        biases_opacity[0],
        ...
        baises_opacity[-1],
        weights_color[0].flatten(),
        ...
        weights_color[n].flatten(),
        biases_color[0],
        ...
        biases_color[-1],
    ]
)

Here, weights_XXX[i] correspond to a (M, N) tensor storing the weight matrix of the i-th MLP layer. Similarly, biases_XXX[i] is a (N,) tensor storing the bias vector.

The MLP multiplies the input features from the right, i.e.:

output[i+1] = input[i] @ weights_XXX[i] + biases_XXX[i]

Hence, M / N is the input / output channel dimension.

In addition to the mlp_params, the DecoderParams class stores the number of hidden units each MLP. Specifically, n_hidden_trunk, n_hidden_opacity, and n_hidden_color are tensors of shape (n_layers+1,) that store the number of input channels followed by the output channel number of each layer in the trunk, opacity, and color MLPs, respectively.

Note

One can convert the 1D mlp_params tensor to the more-interpretable list of weight matrices and bias tensors using the flattened_decoder_params_to_list function.

Note

Since the Triton language of Lightplane’s GPU kernel constraints the number of rendering channels to at least 16, the color_chn attribute is used to store the effective number of rendered output channels. If the effective number of rendered channels is less than 16, the MLP parameters are padded with zeros to match the minimum size.

mlp_params

The parameters for the Lightplane Rendering decoder.

Type:

torch.Tensor

n_hidden_trunk

(n_layers+1,) Long tensor storing the number of input channels followed by the number of hidden units in each layer of the trunk_mlp. Note that this tensor can be empty if the trunk MLP is not used.

Type:

torch.Tensor

n_hidden_opacity

(n_layers+1,) Long tensor storing the number of input channels followed by the number of hidden units in each layer of the opacity_mlp.

Type:

torch.Tensor

n_hidden_color

(n_layers+1,) Long tensor storing the number of input channels followed by the number of hidden units in each layer of the color_mlp.

Type:

torch.Tensor

color_chn

The number of rendered channels.

Type:

int

class lightplane.SplatterParams(mlp_params: Tensor, n_hidden: Tensor)[source]

Class representing learnable parameters of the MLP from Lightplane Splatter.

The splatter comprises a learnable function that predicts a vector splatted to the output 3D feature grid. Specifically, the function is defined as follows:

MLP(feature_grid[x] + splatting_feature[u]) -> splat_vector[x]

where x corresponds to the 3D point along the the ray of pixel u, feature_grid[x] is the input shape grid sampled at point x, and splatting_feature[u] is the splatted feature at pixel u. The splatting MLP outputs splat_vector[x] which is pushed back into the output grid.

The parameters of the MLP are stored in the mlp_params attribute. Here, mlp_params is a 1D tensor which concatenates the flattened weight matrices and bias vectors of the MLP in the following order:

mlp_params = torch.cat(
    [
        weights[0].flatten(),
        ...
        weights[-1].flatten(),
        biases[0],
        ...
        biases[-1],
    ]
)

Here, weights[i] correspond to a (M, N) tensor storing the weight matrix of the i-th MLP layer. Similarly, biases[i] is a (N,) tensor storing the bias vector.

The MLP multiplies the input features from the right, i.e.:

output[i+1] = input[i] @ weights[i] + biases[i]

Hence, M / N is the input / output channel dimension.

In addition to the mlp_params, the SplatterParams class stores the number of MLP’s hidden units. Specifically, the n_hidden field is a tensor of shape (n_layers+1,) that stores the number of input channels followed by the output channel number of each layer in the MLP.

mlp_params

The parameters for the Lightplane rendering decoder.

Type:

torch.Tensor

n_hidden

(n_layers+1,) Long tensor storing the number of input channels followed by the number of hidden units in each layer of the splatting MLP.

Type:

torch.Tensor

lightplane.flatten_decoder_params(weights_trunk: Tuple[Tensor, ...], biases_trunk: Tuple[Tensor, ...], weights_opacity: Tuple[Tensor, ...], biases_opacity: Tuple[Tensor, ...], weights_color: Tuple[Tensor, ...], biases_color: Tuple[Tensor, ...], pad_color_channels_to_min_block_size: bool = True) → Tuple[Tensor, Tensor, Tensor][source]

The hepler function to flatten the decoder parameters into a single tensor, and get the number of hidden units for each layer in each MLP (n_hidden_XX).

Parameters:

  • weights_trunk – Tuple of weight matrices for trunk_mlp.

  • biases_trunk – Tuple of bias vectors for trunk_mlp.

  • weights_opacity – Tuple of weight matrices for opacity_mlp.

  • biases_opacity – Tuple of bias vectors for opacity_mlp.

  • weights_color – Tuple of weight matrices for color_mlp.

  • biases_color – Tuple of bias vectors for color_mlp.

  • pad_color_channels_to_min_block_size – If True, the MLP parameters are padded with zeros to match the minimum size of the triton minimum block size.

lightplane.flatten_splatter_params(weights: Tuple[Tensor, ...], biases: Tuple[Tensor, ...])[source]

The hepler function to flatten the splatter parameters into a single tensor, and get the number of hidden units for each layer in the MLP (n_hidden).

Parameters:

  • weights – Tuple of weight matrices for the MLP.

  • biases – Tuple of bias vectors for the MLP.

lightplane.flattened_decoder_params_to_list(mlp_params: Tensor, n_hidden_trunk: Tensor, n_hidden_opacity: Tensor, n_hidden_color: Tensor, transpose: bool = False) → Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...], Tuple[Tensor, ...], Tuple[Tensor, ...], Tuple[Tensor, ...], Tuple[Tensor, ...]][source]

This function converts the flattened MLP parameters into a list of weight matrices, and bias vectors for each MLP. It is the inverse function of flatten_decoder_params.

Parameters:

  • mlp_params – The flattened MLP parameters, i.e. 1D tensor.

  • n_hidden_trunk – (n_layers+1,) Long tensor storing the number of input channels followed by the number of hidden units in each layer of the trunk_mlp. Note that this tensor can be empty if the trunk MLP is not used.

  • n_hidden_opacity – (n_layers+1,) Long tensor storing the number of input channels followed by the number of hidden units in each layer of the opacity_mlp.

  • n_hidden_color – (n_layers+1,) Long tensor storing the number of input channels followed by the number of hidden units in each layer of the color_mlp.

  • transpose – If True, the weight matrices are transposed.

Returns:

  • weights_trunk – Weight matrices of the trunk MLP.

  • biases_trunk – Bias vectors of the trunk MLP.

  • weights_opacity – Weight matrices of the opacity MLP.

  • biases_opacity – Bias vectors of the opacity MLP.

  • weights_color – Weight matrices of the color MLP.

  • biases_color – Bias vectors of the color MLP.

lightplane.flattened_triton_decoder_to_list(mlp_params: Tensor, n_layers_trunk: int, n_layers_opacity: int, n_layers_color: int, input_chn: int, hidden_chn: int, color_chn: int)[source]

Another helper function to convert the flattened MLP parameters into a list of weight matrices, and bias vectors for each MLP. Given mlp_params, the number of layers for each MLP, input/output number of channesl, and hidden units number, this function returns the list of weight matrices and bias vectors for each MLP.

Parameters:

  • mlp_params – The flattened MLP parameters, i.e. 1D tensor.

  • n_layers_trunk – The number of layers in the trunk_mlp.

  • n_layers_opacity – The number of layers in the opacity_mlp.

  • n_layers_color – The number of layers in the color_mlp.

  • input_chn – The number of input channels.

  • hidden_chn – The number of hidden units in the MLP layers.

  • color_chn – The number of rendered channels.

lightplane.get_triton_function_input_dims(n_hidden_trunk: Tensor, n_hidden_opacity: Tensor, n_hidden_color: Tensor)[source]

Function to get the MLP layers and hidden units from n_hidden_trunk, n_hidden_opacity and n_hidden_color inside decoder_params object.

lightplane.init_decoder_params(device: device, n_layers_opacity: int, n_layers_trunk: int, n_layers_color: int, input_chn: int = 32, hidden_chn: int = 32, color_chn: int = 3, opacity_init_bias: float = 0.0, pad_color_channels_to_min_block_size: bool = True, use_separate_color_grid: bool = False) → DecoderParams[source]

The function initializes the learnable parameters of the Lightplane Renderer decoder given mlp configurations. Weights and biases of three MLPs inside decoder (trunk_mlp, opacity_mlp, and color_mlp) are initialized using Xavier initialization by function _xavier_init_mlp_params, and are flattened into a single tensor mlp_params by function flatten_decoder_params.

Since the Triton language of Lightplane’s GPU kernel constraints the number of rendering channels to at least 16, the color_chn attribute is used to store the effective number of rendered output channels. If the effective number of rendered channels is less than 16, the MLP parameters are padded with zeros to match the minimum size.

Parameters:

  • device – The device to store the parameters.

  • n_hidden_trunk – (n_layers+1,) Long tensor storing the number of input channels followed by the number of hidden units in each layer of the trunk_mlp. Note that this tensor can be empty if the trunk MLP is not used.

  • n_hidden_opacity – (n_layers+1,) Long tensor storing the number of input channels followed by the number of hidden units in each layer of the opacity_mlp.

  • n_hidden_color – (n_layers+1,) Long tensor storing the number of input channels followed by the number of hidden units in each layer of the color_mlp.

  • input_chn – The number of input channels, which is the number of channel for feature_grid.

  • hidden_chn – The number of hidden units in the MLP layers.

  • color_chn – The number of rendered channels.

  • opacity_init_bias – The initial bias value for the opacity MLP.

  • pad_color_channels_to_min_block_size – If True, the MLP parameters are padded with zeros to match the minimum size of the triton minimum block size.

  • use_separate_color_grid – If True, the renderer uses a separate color grid.

lightplane.init_splatter_params(device: device, n_layers: int, input_chn: int = 32, hidden_chn: int = 32, out_chn: int = 16) → SplatterParams[source]

The function initializes the learnable parameters of the Lightplane Splatter given mlp configurations. Weights and biases of the MLP inside LightPlane Splatter are initialized using Xavier initialization by function _xavier_init_mlp_params, and are flattened into a single tensor mlp_params by function flatten_splatter_params.

Since the outout of the mlp is a vector splatted to the output 3D feature grid, whose number of channels is the same as the output_grid, which is typically more than 16. So we do not need to pad the MLP parameters to match the minimum size of the triton minimum block size.

Parameters:

  • device – The device to store the parameters.

  • n_layers – (n_layers+1,) Long tensor storing the number of input channels followed by the number of hidden units in each layer of the mlp.

  • input_chn – The number of input channels.

  • hidden_chn – The number of hidden units in the MLP layers.

  • out_chn – The number of output channels.