-# Copyright Contributors to the OpenVDB Project
-# SPDX-License-Identifier: Apache-2.0
-#
-importmath
-fromtypingimportList,Optional,Sequence,Union
-
-importtorch
-importtorch.nnasnn
-fromtorch.profilerimportrecord_function
-
-importfvdb
-fromfvdbimportGridBatch,JaggedTensor
-
-from.vdbtensorimportVDBTensor
-
-
-deffvnn_module(module):
- # Register class as a module in fvdb.nn
- old_forward=module.forward
-
- def_forward(self,*args,**kwargs):
- withrecord_function(repr(self)):
- returnold_forward(self,*args,**kwargs)
-
- module.forward=_forward
- returnmodule
-
-
-GridOrVDBTensor=Union[fvdb.GridBatch,VDBTensor]
-ListOrInt=Union[int,List[int]]
-
-
-
-[docs]
-@fvnn_module
-classMaxPool(nn.Module):
-r"""Applies a 3D max pooling over an input signal.
-
- Args:
- kernel_size: the size of the window to take a max over
- stride: the stride of the window. Default value is :attr:`kernel_size`
-
- Note:
- For target voxels that are not covered by any source voxels, the
- output feature will be set to zero.
-
- """
-
- def__init__(self,kernel_size:ListOrInt,stride:Optional[ListOrInt]=None):
- super().__init__()
- self.kernel_size=kernel_size
- self.stride=strideorself.kernel_size
-
- defforward(self,input:VDBTensor,ref_coarse_data:Optional[GridOrVDBTensor]=None)->VDBTensor:
- ifisinstance(ref_coarse_data,VDBTensor):
- coarse_grid,coarse_kmap=ref_coarse_data.grid,ref_coarse_data.kmap
- elifisinstance(ref_coarse_data,fvdb.GridBatch):
- coarse_grid,coarse_kmap=ref_coarse_data,None
- else:
- coarse_grid,coarse_kmap=None,None
-
- new_feature,new_grid=input.grid.max_pool(
- self.kernel_size,input.data,stride=self.stride,coarse_grid=coarse_grid
- )
- new_feature.jdata[torch.isinf(new_feature.jdata)]=0.0
- returnVDBTensor(new_grid,new_feature,kmap=coarse_kmap)
-
- defextra_repr(self)->str:
- return"kernel_size={kernel_size}, stride={stride}".format(kernel_size=self.kernel_size,stride=self.stride)
-
-
-
-
-[docs]
-@fvnn_module
-classAvgPool(nn.Module):
-r"""Applies a 3D average pooling over an input signal.
-
- Args:
- kernel_size: the size of the window to take average over
- stride: the stride of the window. Default value is :attr:`kernel_size`
-
- """
-
- def__init__(self,kernel_size:ListOrInt,stride:Optional[ListOrInt]=None):
- super().__init__()
- self.kernel_size=kernel_size
- self.stride=strideorself.kernel_size
-
- defforward(self,input:VDBTensor,ref_coarse_data:Optional[GridOrVDBTensor]=None)->VDBTensor:
- ifisinstance(ref_coarse_data,VDBTensor):
- coarse_grid,coarse_kmap=ref_coarse_data.grid,ref_coarse_data.kmap
- elifisinstance(ref_coarse_data,fvdb.GridBatch):
- coarse_grid,coarse_kmap=ref_coarse_data,None
- else:
- coarse_grid,coarse_kmap=None,None
-
- new_feature,new_grid=input.grid.avg_pool(
- self.kernel_size,input.data,stride=self.stride,coarse_grid=coarse_grid
- )
- returnVDBTensor(new_grid,new_feature,kmap=coarse_kmap)
-
- defextra_repr(self)->str:
- return"kernel_size={kernel_size}, stride={stride}".format(kernel_size=self.kernel_size,stride=self.stride)
-
-
-
-
-[docs]
-@fvnn_module
-classUpsamplingNearest(nn.Module):
-r"""Upsamples the input by a given scale factor using nearest upsampling.
-
- Args:
- scale_factor: the upsampling factor
- """
-
- def__init__(self,scale_factor:ListOrInt):
- super().__init__()
- self.scale_factor=scale_factor
-
- defforward(
- self,input:VDBTensor,mask:Optional[JaggedTensor]=None,ref_fine_data:Optional[GridOrVDBTensor]=None
- )->VDBTensor:
- ifisinstance(ref_fine_data,VDBTensor):
- fine_grid,fine_kmap=ref_fine_data.grid,ref_fine_data.kmap
- elifisinstance(ref_fine_data,fvdb.GridBatch):
- fine_grid,fine_kmap=ref_fine_data,None
- else:
- fine_grid,fine_kmap=None,None
-
- new_feature,new_grid=input.grid.subdivide(self.scale_factor,input.data,mask,fine_grid=fine_grid)
- returnVDBTensor(new_grid,new_feature,kmap=fine_kmap)
-
- defextra_repr(self)->str:
- return"scale_factor={scale_factor}".format(scale_factor=self.scale_factor)
-
-
-
-
-[docs]
-@fvnn_module
-classFillFromGrid(nn.Module):
-r"""
- Fill the content of input vdb-tensor to another grid.
-
- Args:
- default_value: the default value to fill in the new grid.
- """
-
- def__init__(self,default_value:float=0.0)->None:
- super().__init__()
- self.default_value=default_value
-
- defforward(self,input:VDBTensor,other_data:Optional[GridOrVDBTensor]=None)->VDBTensor:
- ifisinstance(other_data,VDBTensor):
- other_grid,other_kmap=other_data.grid,other_data.kmap
- elifisinstance(other_data,fvdb.GridBatch):
- other_grid,other_kmap=other_data,None
- else:
- returninput
-
- new_feature=other_grid.fill_from_grid(input.data,input.grid,self.default_value)
- returnVDBTensor(other_grid,new_feature,kmap=other_kmap)
-
-
-
-
-[docs]
-@fvnn_module
-classSparseConv3d(nn.Module):
-r"""Applies a 3D convolution over an input signal composed of several input
- planes, by performing a sparse convolution on the underlying VDB grid.
-
- Args:
- in_channels: number of channels in the input tensor
- out_channels: number of channels produced by the convolution
- kernel_size: size of the convolving kernel
- stride: stride of the convolution. Default value is 1
- bias: if ``True``, adds a learnable bias to the output. Default: ``True``
- transposed: if ``True``, uses a transposed convolution operator
- """
-
- CUTLASS_SUPPORTED_CHANNELS=[
- (32,64),
- (64,128),
- (128,256),
- (32,32),
- (64,64),
- (128,128),
- (256,256),
- (128,64),
- (64,32),
- (256,128),
- (384,256),
- (192,128),
- (256,512),
- (512,256),
- (512,512),
- ]
-
-"""
- Backend for performing convolutions:
- - "default": for now it is 'igemm_mode1'
- - "legacy": the old slow implementation
- - "me": MinkowskiEngine implementation
- - "halo": 10x10x10 halo buffer implementation, stride 1, kernel 3
- - "cutlass": 4x4x6 cutlass implementation, stride 1, kernel 3, forward only, limited channels support
- - "lggs": kernel optimized for sparse structures
- - "igemm_mode0": unsorted
- - "igemm_mode1": sorted + split=1
- - "igemm_mode2": sorted + split=3
- - "dense": dense convolution
- """
- backend:str="default"
- allow_tf32:bool=True
-
- def__init__(
- self,
- in_channels:int,
- out_channels:int,
- kernel_size:Union[int,Sequence]=3,
- stride:Union[int,Sequence]=1,
- bias:bool=True,
- transposed:bool=False,
- )->None:
-
- super().__init__()
- self.in_channels=in_channels
- self.out_channels=out_channels
-
- ifisinstance(kernel_size,int):
- kernel_size=(kernel_size,)*3
- assertlen(kernel_size)==3
-
- ifisinstance(stride,int):
- stride=(stride,)*3
- assertlen(stride)==3
-
- self.kernel_size=kernel_size
- self.stride=stride
- self.transposed=transposed
-
- ifself.transposed:
- # Only change kernel size instead of module dict
- out_channels,in_channels=in_channels,out_channels
-
- self.kernel_volume=math.prod(self.kernel_size)
- ifself.kernel_volume>1:
- # Weight tensor is of shape (Do, Di, K0, K1, K2), but the underlying data is (K2, K1, K0, Di, Do)
- # so we don't need to make a copy of the permuted tensor within the conv kernel.
- weight_shape=[out_channels,in_channels]+list(self.kernel_size)
- weight=torch.zeros(*weight_shape[::-1]).permute(4,3,2,1,0)
- self.weight=nn.Parameter(weight)
- else:
- self.weight=nn.Parameter(torch.zeros(out_channels,in_channels))
-
- ifbias:
- self.bias=nn.Parameter(torch.Tensor(self.out_channels))
- else:
- self.register_parameter("bias",None)
-
- self.reset_parameters()
-
- defextra_repr(self)->str:
- s="{in_channels}, {out_channels}, kernel_size={kernel_size}"
- ifself.stride!=(1,1,1):
- s+=", stride={stride}"
- ifself.biasisNone:
- s+=", bias=False"
- ifself.transposed:
- s+=", transposed=True"
- returns.format(**self.__dict__)
-
- defreset_parameters(self)->None:
- std=1/math.sqrt((self.out_channelsifself.transposedelseself.in_channels)*self.kernel_volume)
- self.weight.data.uniform_(-std,std)
- ifself.biasisnotNone:
- self.bias.data.uniform_(-std,std)
-
- def_dispatch_conv(self,in_feature,in_grid,in_kmap,out_grid):
-
- backend=self.backend
-
- sm_arch=torch.cuda.get_device_capability()[0]+torch.cuda.get_device_capability()[1]/10
- # tf32 requires compute capability >= 8.0 (Ampere)
- ifself.allow_tf32andself.weight.is_cuda:
- assert(
- sm_arch>=8
- ),"TF32 requires GPU with compute capability >= 8.0. Please set fvdb.nn.SparseConv3d.allow_tf32 = False."
-
- # bf16 requires compute capability >= 8.0 (Ampere)
- ifself.weight.is_cudaandself.weight.dtype==torch.bfloat16:
- assertsm_arch>=8,"BF16 requires GPU with compute capability >= 8.0."
-
- # float16 requires compute capability >= 7.5 (Turing)
- ifself.weight.is_cudaandself.weight.dtype==torch.float16:
- assertsm_arch>=7.5,"FP16 requires GPU with compute capability >= 7.5."
-
- # cutlass, lggs, halo backends require compute capability >= 8.0 (Ampere)
- ifbackendin["cutlass","lggs","halo"]:
- assert(
- torch.cuda.get_device_capability()[0]>=8
- ),"cutlass, LGGS and Halo backends require GPU with compute capability >= 8.0."
-
- ifbackend=="cutlass"and(
- (notself.weight.is_cuda)or(self.in_channels,self.out_channels)notinself.CUTLASS_SUPPORTED_CHANNELS
- ):
- print(
- f"Cutlass backend does not support {self.in_channels} -> {self.out_channels} convolutions, falling back to default"
- )
- backend="default"
-
- ifbackend=="lggs"and((self.in_channels,self.out_channels)notin[(128,128)]):
- print("LGGS backend only supports 128 to 128 convolution, falling back to default")
- backend="default"
-
- ifbackend=="default":
- if(notself.weight.is_cuda)orin_feature.dtype==torch.float64:
- backend="legacy"
- else:
- backend="igemm_mode1"
-
- ifbackend=="halo"andself.stride==(1,1,1)andself.kernel_size==(3,3,3):
- assertout_gridisNoneorin_grid.is_same(out_grid)
- returnin_grid,in_grid.sparse_conv_halo(in_feature,self.weight,8),None
-
- elifbackend=="dense"andself.stride==(1,1,1):
- assertout_gridisNoneorin_grid.is_same(out_grid)
- min_coord=in_grid.ijk.jdata.min(axis=0).values
- # BWHDC -> BCDHW
- dense_feature=in_grid.write_to_dense(in_feature,min_coord=min_coord).permute(0,4,3,2,1)
- dense_feature=torch.nn.functional.conv3d(dense_feature,self.weight,padding=1,stride=1)
- # BCDHW -> BWHDC
- dense_feature=dense_feature.permute(0,4,3,2,1).contiguous()
- dense_feature=in_grid.read_from_dense(dense_feature,dense_origins=min_coord)
-
- returnin_grid,dense_feature,None
-
- else:
- # Fallback to the default implementation
- can_cache=self.stride==(1,1,1)and(out_gridisNoneorout_grid.is_same(in_grid))
-
- ifin_kmapisnotNoneandin_kmap.kernel_size==self.kernel_sizeandcan_cache:
- kmap,out_grid=in_kmap,in_grid
- else:
- ifself.transposed:
- assertout_gridisnotNone
- kmap,_=out_grid.sparse_conv_kernel_map(self.kernel_size,self.stride,in_grid)
- else:
- kmap,out_grid=in_grid.sparse_conv_kernel_map(self.kernel_size,self.stride,out_grid)
-
- out_kmap=kmapifcan_cacheelseNone
-
- backend=self._build_kmap_and_convert_backend(kmap,backend)
-
- ifnotself.transposed:
- out_feature=kmap.sparse_conv_3d(in_feature,self.weight,backend)
- else:
- out_feature=kmap.sparse_transpose_conv_3d(in_feature,self.weight,backend)
-
- returnout_grid,out_feature,out_kmap
-
- def_build_kmap_and_convert_backend(self,kmap:fvdb.SparseConvPackInfo,backend:str)->fvdb.ConvPackBackend:
- ifbackendin["legacy","me"]:
- kmap.build_gather_scatter(backend=="me")
- returnfvdb.ConvPackBackend.GATHER_SCATTER
-
- elifbackend=="cutlass":
- kmap.build_cutlass(benchmark=False)
- returnfvdb.ConvPackBackend.CUTLASS
-
- elifbackend=="igemm_mode0":
- kmap.build_implicit_gemm(
- sorted=False,split_mask_num=1,training=self.training,split_mask_num_bwd=3,use_tf32=self.allow_tf32
- )
- returnfvdb.ConvPackBackend.IGEMM
-
- elifbackend=="igemm_mode1":
- kmap.build_implicit_gemm(
- sorted=True,split_mask_num=1,training=self.training,split_mask_num_bwd=3,use_tf32=self.allow_tf32
- )
- returnfvdb.ConvPackBackend.IGEMM
-
- elifbackend=="igemm_mode2":
- kmap.build_implicit_gemm(
- sorted=True,split_mask_num=3,training=self.training,split_mask_num_bwd=3,use_tf32=self.allow_tf32
- )
- returnfvdb.ConvPackBackend.IGEMM
-
- elifbackend=="lggs":
- kmap.build_lggs()
- returnfvdb.ConvPackBackend.LGGS
-
- else:
- raiseNotImplementedError(f"Backend {backend} is not supported")
-
- defforward(
- self,
- input:VDBTensor,
- out_grid:Optional[GridBatch]=None,
- )->VDBTensor:
- in_feature,in_grid,in_kmap=input.data,input.grid,input.kmap
-
- ifself.kernel_size==(1,1,1)andself.stride==(1,1,1):
- out_feature=in_feature.jdata.matmul(self.weight.transpose(0,1))
- out_feature=in_feature.jagged_like(out_feature)
- out_grid,out_kmap=in_grid,in_kmap
-
- else:
- out_grid,out_feature,out_kmap=self._dispatch_conv(in_feature,in_grid,in_kmap,out_grid)
-
- ifself.biasisnotNone:
- out_feature.jdata=out_feature.jdata+self.bias
-
- ifout_gridisNone:
- raiseRuntimeError("Failed to compute output grid. This is a bug in the implementation.")
- returnVDBTensor(out_grid,out_feature,out_kmap)
-
-
-
-
-[docs]
-@fvnn_module
-classGroupNorm(nn.GroupNorm):
-r"""Applies Group Normalization over a VDBTensor.
- See :class:`~torch.nn.GroupNorm` for detailed information.
- """
-
- defforward(self,input:VDBTensor)->VDBTensor:
- num_channels=input.data.jdata.size(1)
- assertnum_channels==self.num_channels,"Input feature should have the same number of channels as GroupNorm"
- num_batches=input.grid.grid_count
-
- flat_data,flat_offsets=input.data.jdata,input.data.joffsets
-
- result_data=torch.empty_like(flat_data)
-
- forbinrange(num_batches):
- feat=flat_data[flat_offsets[b]:flat_offsets[b+1]]
- iffeat.size(0)!=0:
- feat=feat.transpose(0,1).reshape(1,num_channels,-1)
- feat=super().forward(feat)
- feat=feat.reshape(num_channels,-1).transpose(0,1)
-
- result_data[flat_offsets[b]:flat_offsets[b+1]]=feat
-
- returnVDBTensor(input.grid,input.grid.jagged_like(result_data),input.kmap)
-
-
-
-@fvnn_module
-classBatchNorm(nn.BatchNorm1d):
-r"""Applies Batch Normalization over a VDBTensor.
- See :class:`~torch.nn.BatchNorm1d` for detailed information.
- """
-
- defforward(self,input:VDBTensor)->VDBTensor:
- num_channels=input.data.jdata.size(1)
- assertnum_channels==self.num_features,"Input feature should have the same number of channels as BatchNorm"
- result_data=super().forward(input.data.jdata)
- returnVDBTensor(input.grid,input.grid.jagged_like(result_data),input.kmap)
-
-
-# Non-linear Activations
-
-
-@fvnn_module
-classElementwiseMixin:
- defforward(self,input:VDBTensor)->VDBTensor:
- assertisinstance(input,VDBTensor),"Input should have type VDBTensor"
- res=super().forward(input.data.jdata)# type: ignore
- returnVDBTensor(input.grid,input.data.jagged_like(res),input.kmap)
-
-
-classELU(ElementwiseMixin,nn.ELU):
-r"""
- Applies the Exponential Linear Unit function element-wise:
- .. math::
- \text{ELU}(x) = \begin{cases}
- x, & \text{ if } x > 0\\
- \alpha * (\exp(x) - 1), & \text{ if } x \leq 0
- \end{cases}
- """
-
-
-classCELU(ElementwiseMixin,nn.CELU):
-r"""
- Applies the CELU function element-wise.
-
- .. math::
- \text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x/\alpha) - 1))
- """
-
-
-classGELU(ElementwiseMixin,nn.GELU):
-r"""
- Applies the Gaussian Error Linear Units function.
-
- .. math:: \text{GELU}(x) = x * \Phi(x)
-
- where :math:`\Phi(x)` is the Cumulative Distribution Function for Gaussian Distribution.
- """
-
-
-
-[docs]
-classLinear(ElementwiseMixin,nn.Linear):
-r"""
- Applies a linear transformation to the incoming data: :math:`y = xA^T + b`.
- """
-
-
-
-
-[docs]
-classReLU(ElementwiseMixin,nn.ReLU):
-r"""
- Applies the rectified linear unit function element-wise: :math:`\text{ReLU}(x) = (x)^+ = \max(0, x)`
- """
-[docs]
-classSiLU(ElementwiseMixin,nn.SiLU):
-r"""
- Applies element-wise, :math:`\text{SiLU}(x) = x * \sigma(x)`, where :math:`\sigma(x)` is the sigmoid function.
- """
-[docs]
-classDropout(ElementwiseMixin,nn.Dropout):
-r"""
- During training, randomly zeroes some of the elements of the input tensor with probability :attr:`p`
- using samples from a Bernoulli distribution. The elements to zero are randomized on every forward call.
- """
-# Copyright Contributors to the OpenVDB Project
-# SPDX-License-Identifier: Apache-2.0
-#
-importos
-
-fromtorch.utilsimportcpp_extension
-
-importfvdb
-
-
-
-[docs]
-defFVDBExtension(name,sources,*args,**kwargs):
-"""
- Utility function for creating pytorch extensions that depend on fvdb. You then have access to all fVDB's internal
- headers to program with. Example usage:
-
- .. code-block:: python
-
- from fvdb.utils import FVDBExtension
-
- ext = FVDBExtension(
- name='my_extension',
- sources=['my_extension.cpp'],
- extra_compile_args={'cxx': ['-std=c++17']},
- libraries=['mylib'],
- )
-
- :param name: The name of the extension.
- :param sources: The list of source files.
- :param args: Other arguments to pass to :func:`torch.utils.cpp_extension.CppExtension`.
- :param kwargs: Other keyword arguments to pass to :func:`torch.utils.cpp_extension.CppExtension`.
- :return: A :class:`torch.utils.cpp_extension.CppExtension` object.
- """
-
- libraries=kwargs.get("libraries",[])
- libraries.append("fvdb")
- kwargs["libraries"]=libraries
-
- library_dirs=kwargs.get("library_dirs",[])
- library_dirs.append(os.path.dirname(fvdb.__file__))
- kwargs["library_dirs"]=library_dirs
-
- include_dirs=kwargs.get("include_dirs",[])
- include_dirs.append(os.path.join(os.path.dirname(fvdb.__file__),"include"))
-
- # We also need to add this because fvdb internally will refer to their headers without the fvdb/ prefix.
- include_dirs.append(os.path.join(os.path.dirname(fvdb.__file__),"include/fvdb"))
- kwargs["include_dirs"]=include_dirs
-
- extra_link_args=kwargs.get("extra_link_args",[])
- extra_link_args.append(f"-Wl,-rpath={os.path.dirname(fvdb.__file__)}")
- kwargs["extra_link_args"]=extra_link_args
-
- extra_compile_args=kwargs.get("extra_compile_args",{})
- extra_compile_args["nvcc"]=extra_compile_args.get("nvcc",[])
- if"--extended-lambda"notinextra_compile_args["nvcc"]:
- extra_compile_args["nvcc"].append("--extended-lambda")
- kwargs["extra_compile_args"]=extra_compile_args
-
- returncpp_extension.CUDAExtension(name,sources,*args,**kwargs)
Downsample this batch of grids using average pooling.
-
-
Parameters:
-
-
pool_factor (int or 3-tuple of ints) – How much to pool by (i,e, (2,2,2) means take average over 2x2x2 from start of window).
-
data (JaggedTensor) – Data at each voxel in this grid to be downsampled (JaggedTensor of shape [B, -1, *]).
-
stride (int) – The stride to use when pooling
-
coarse_grid (GridBatch, optional) – An optional coarse grid used to specify the output. This is mainly used
-for memory efficiency so you can chache grids. If you don’t pass it in, we’ll just create it for you.
-
-
-
Returns:
-
-
coarse_data (JaggedTensor) – a JaggedTensor of shape [B, -1, *] of downsampled data.
-
coarse_grid (GridBatch) – the downsampled grid batch.
A [num_grids, 2, 3] tensor of the bounding box of each grid in this batch where bbox[i, 0] is the minimimum ijk coordinate of the i^th grid, and bbox[i, 1] is the maximum ijk coordinate.
Get the bounding box (in voxel coordinates) of the bi^th grid in the batch.
-
-
Parameters:
-
bi (int) – The index of the grid to get the bounding box of.
-
-
Returns:
-
bbox (torch.Tensor) – A tensor, bbox, of shape [2, 3] where bbox = [[bmin_i, bmin_j, bmin_z=k],
-[bmax_i, bmax_j, bmax_k]] is the bi^th bounding box such that bmin <= ijk < bmax for all voxels
-ijk in the bi^th grid.
Return a batch of grids representing the clipped version of this batch.
-Each voxel [i, j, k] in the input batch is included in the output if it lies within ijk_min and ijk_max.
-
-
Parameters:
-
-
ijk_min (list of int triplets) – Index space minimum bound of the clip region.
-
ijk_max (list of int triplets) – Index space maximum bound of the clip region.
-
-
-
Returns:
-
clipped_grid (GridBatch) – A GridBatch representing the clipped version of this grid batch.
Return a batch of grids representing the coarsened version of this batch.
-Each voxel [i, j, k] in this grid batch maps to voxel [i / branchFactor, j / branchFactor, k / branchFactor] in the coarse batch.
-
-
Parameters:
-
coarsening_factor (int or 3-tuple of ints) – How much to coarsen by (i,e, (2,2,2) means take every other voxel from start of window).
-
-
Returns:
-
coarsened_grid (GridBatch) – A GridBatch representing the coarsened version of this grid batch.
Return a batch of grids representing the convolution of this batch with a given kernel.
-Each voxel [i, j, k] in the output batch is the sum of the voxels in the input batch [i * stride, j * stride, k * stride] to [i * stride + kernel_size, j * stride + kernel_size, k * stride + kernel_size].
-
-
Parameters:
-
-
kernel_size (int or 3-tuple of ints) – The size of the kernel to convolve with.
-
stride (int or 3-tuple of ints) – The stride to use when convolving.
-
-
-
Returns:
-
conv_grid (GridBatch) – A GridBatch representing the convolution of this grid batch.
Get the cumulative number of enabled voxels in the bi^th grid in the batch. i.e. nvox_0+nvox_1+…+nvox_i. If this grid isn’t mutable, this returns the same value as cum_voxels_at.
An integer tensor containing the cumulative number of voxels indexed by the grids in this batch.
-i.e. [nvox_0, nvox_0+nvox_1, nvox_0+nvox_1+nvox_2, …]
If this is grid is mutable, disable voxels at the specified coordinates, otherwise throw an exception.
-If the ijk values are already disabled or are not represented in this GridBatch, then this function is no-op.
-
-
Parameters:
-
ijk (JaggedTensor) – A JaggedTensor of shape [num_grids, -1, 3] of ijk coordinates to disable.
A [num_grids, 2, 3] tensor of the bounding box of the dual of each grid in this batch where bbox[i, 0] is the minimimum ijk coordinate of the i^th dual grid, and bbox[i, 1] is the maximum ijk coordinate.
Return a batch of grids representing the dual of this batch.
-i.e. The centers of the dual grid correspond to the corners of this grid batch. The [i, j, k] coordinate of the dual grid corresponds to the bottom/left/back
-corner of the [i, j, k] voxel in this grid batch.
-
-
Parameters:
-
exclude_border (bool) – Whether to exclude the border of the grid batch when computing the dual grid
-
-
Returns:
-
dual_grid (GridBatch) – A GridBatch representing the dual of this grid batch.
If this is grid is mutable, enable voxels at the specified coordinates, otherwise throw an exception.
-If the ijk values are already enabled or are not represented in this GridBatch, then this function is no-op.
-
-
Parameters:
-
ijk (JaggedTensor) – A JaggedTensor of shape [num_grids, -1, 3] of ijk coordinates to disable.
Given a GridBatch and features associated with it, return a JaggedTensor representing features for this batch of grid.
-Fill any voxels not in the GridBatch with the default value.
-
-
Parameters:
-
-
other_features (JaggedTensor) – A JaggedTensor of shape [B, -1, *] containing features associated with other_grid.
-
other_grid (GridBatch) – A GridBatch containing the grid to fill from.
-
default_value (float) – The value to fill in for voxels not in the GridBatch (default 0.0).
-
-
-
Returns:
-
filled_features (JaggedTensor) – A JaggedTensor of shape [B, -1, *] of features associated with this batch of grids.
pool_factor (int or 3-tuple of ints) – How much to pool by (i,e, (2,2,2) means take max over 2x2x2 from start of window).
-
data (JaggedTensor) – Data at each voxel in this grid to be downsampled (JaggedTensor of shape [B, -1, *]).
-
stride (int) – The stride to use when pooling
-
coarse_grid (GridBatch, optional) – An optional coarse grid used to specify the output. This is mainly used
-for memory efficiency so you can chache grids. If you don’t pass it in, we’ll just create it for you.
-
-
-
Returns:
-
-
coarse_data (JaggedTensor) – a JaggedTensor of shape [B, -1, *] of downsampled data.
-
coarse_grid (GridBatch) – the downsampled grid batch.
An integer tensor containing the number of enabled voxels per grid indexed by this batch. If this grid is not mutable, this will be the same as num_voxels.
Set the voxels in this grid batch to a dense grid with shape [num_grids, width, height, depth], otpionally masking out certain voxels
-
-
Parameters:
-
-
num_grids (int) – The number of grids in the batch
-
dense_dims (triple of ints) – The dimensions of the dense grid [width, height, depth]
-
ijk_min (triple of ints) – Index space minimum bound of the dense grid.
-
voxel_sizes (float, list, tensor) – Either a float or triple specifyng the voxel size of all the grids in the batch or a tensor of shape [num_grids, 3] specifying the voxel size for each grid.
-
origins (float, list, tensor) – Either a float or triple specifyng the world space origin of all the grids in the batch or a tensor of shape [num_grids, 3] specifying the world space origin for each grid.
-
mask (torch.Tensor) – A tensor of shape [num_grids, width, height, depth] of booleans indicating which voxels to include/exclude.
Set the voxels in this grid batch to those specified by a given set of ijk coordinates (with optional padding)
-
-
Parameters:
-
-
ijk (JaggedTensor) – A JaggedTensor of shape [num_grids, -1, 3] of ijk coordinates.
-
pad_min (triple of ints) – Index space minimum bound of the padding region.
-
pad_max (triple of ints) – Index space maximum bound of the padding region.
-
voxel_sizes (float, list, tensor) – Either a float or triple specifyng the voxel size of all the grids in the batch or a tensor of shape [num_grids, 3] specifying the voxel size for each grid.
-
origins (float, list, tensor) – Either a float or triple specifyng the world space origin of all the grids in the batch or a tensor of shape [num_grids, 3] specifying the world space origin for each grid.
Set the voxels in this grid batch to those which intersect a given triangle mesh
-
-
Parameters:
-
-
mesh_vertices (JaggedTensor) – A JaggedTensor of shape [num_grids, -1, 3] of mesh vertex positions.
-
mesh_faces (JaggedTensor) – A JaggedTensor of shape [num_grids, -1, 3] of integer indexes into mesh_vertices specifying the faces of each mesh.
-
voxel_sizes (float, list, tensor) – Either a float or triple specifyng the voxel size of all the grids in the batch or a tensor of shape [num_grids, 3] specifying the voxel size for each grid.
-
origins (float, list, tensor) – Either a float or triple specifyng the world space origin of all the grids in the batch or a tensor of shape [num_grids, 3] specifying the world space origin for each grid.
Set the voxels in this grid batch to the nearest voxel to each point in a given point cloud
-
-
Parameters:
-
-
points (JaggedTensor) – A JaggedTensor of shape [num_grids, -1, 3] of point positions.
-
voxel_sizes (float, list, tensor) – Either a float or triple specifyng the voxel size of all the grids in the batch or a tensor of shape [num_grids, 3] specifying the voxel size for each grid.
-
origins (float, list, tensor) – Either a float or triple specifyng the world space origin of all the grids in the batch or a tensor of shape [num_grids, 3] specifying the world space origin for each grid.
Set the voxels in this grid batch to those which contain a point in a given point cloud (with optional padding)
-
-
Parameters:
-
-
points (JaggedTensor) – A JaggedTensor of shape [num_grids, -1, 3] of point positions.
-
pad_min (triple of ints) – Index space minimum bound of the padding region.
-
pad_max (triple of ints) – Index space maximum bound of the padding region.
-
mesh_faces (JaggedTensor) – A JaggedTensor of shape [num_grids, -1, 3] of integer indexes into mesh_vertices specifying the faces of each mesh.
-
voxel_sizes (float, list, tensor) – Either a float or triple specifyng the voxel size of all the grids in the batch or a tensor of shape [num_grids, 3] specifying the voxel size for each grid.
-
origins (float, list, tensor) – Either a float or triple specifyng the world space origin of all the grids in the batch or a tensor of shape [num_grids, 3] specifying the world space origin for each grid.
Subdivide the grid batch and associated data tensor into a finer GridBatch and data tensor using nearest neighbor sampling.
-Each voxel [i, j, k] in this grid batch maps to voxels [i * subdivFactor, j * subdivFactor, k * subdivFactor] in the fine batch.
-Each data value in the subdividided data tensor inherits its parent value
-
-
Parameters:
-
-
subdiv_factor (int or 3-tuple of ints) – How much to subdivide by (i,e, (2,2,2) means subdivide each voxel into 2^3 voxels).
-
data (JaggedTensor) – A JaggedTensor of shape [B, -1, *] containing data associated with this batch of grids.
-
mask (JaggedTensor) – A JaggedTensor of shape [num_grids, -1, 3] of booleans indicating which voxels to subdivide.
-
fine_grid (GridBatch) – An optional fine grid used to specify the output. This is mainly used
-for memory efficiency so you can chache grids. If you don’t pass it in, we’ll just create it for you.
-
-
-
Returns:
-
-
fine_data (JaggedTensor) – A JaggedTensor of shape [B, -1, *] of data associated with the fine grid batch.
-
fine_grid (GridBatch) – A GridBatch representing the subdivided version of this grid batch.
Subdivide the grid batch into a finer grid batch.
-Each voxel [i, j, k] in this grid batch maps to voxels [i * subdivFactor, j * subdivFactor, k * subdivFactor] in the fine batch.
-
-
Parameters:
-
-
subdiv_factor (int or 3-tuple of ints) – How much to subdivide by (i,e, (2,2,2) means subdivide each voxel into 2^3 voxels).
-
mask (JaggedTensor) – A JaggedTensor of shape [num_grids, -1, 3] of booleans indicating which voxels to subdivide.
-
-
-
Returns:
-
subdivided_grid (GridBatch) – A GridBatch representing the subdivided version of this grid batch.
A tensor, total_bbox, of shape [2, 3] where total_bbox = [[bmin_i, bmin_j, bmin_z=k],
-[bmax_i, bmax_j, bmax_k]] is the bounding box such that bmin <= ijk < bmax for all voxels
-ijk in the batch.
A pair of JaggedTensors (gv, ge) of shape [num_grids, -1, 3] and [num_grids, -1, 2] where gv are the corner positions of each voxel and ge are edge indices indexing into gv. This property is useful for visualizing the grid.
A VDBTensor is a thin wrapper around a GridBatch and its corresponding feature JaggedTensor, conceptually denoting a batch of
-sparse tensors along with its topology.
-It works as the input and output arguments of fvdb’s neural network layers.
-One can simply construct a VDBTensor from a GridBatch and a JaggedTensor, or from a dense tensor using from_dense().
During training, randomly zeroes some of the elements of the input tensor with probability p
-using samples from a Bernoulli distribution. The elements to zero are randomized on every forward call.
Utility function for creating pytorch extensions that depend on fvdb. You then have access to all fVDB’s internal
-headers to program with. Example usage:
fVDB, inspired by function notation to resemble \(f(VDB)\), is a data structure for encoding and operating on sparse voxel hierarchies of features in PyTorch.
-A sparse voxel hierarchy is a coarse-to-fine hierarchy of sparse voxel grids such that every fine voxel is contained within some coarse voxel.
-fvdb supports storing PyTorch tensors at the corners and centers of voxels in a hierarchy and enables a number of differentiable operations on these tensors (e.g. trilinear interpolation, splatting, ray tracing).
We give a high-level overview of the main concepts in fVDB. Namely, how fVDB encodes sparse voxel grids with attributes, as well as how fVDB efficiently manages batches with non-uniform numbers of elements.
At its core, fVDB provides means to efficiently encode mini-batches of sparse voxel grids where each voxel contains arbitrary vector or scalar attributes encoded as torch tensors. In practice, such an encoding is implemented via the GridBatch class.
-
-
A GridBatch is an indexing structure which maps 3D ijk coordinates to integer offsets which can be used to look up attributes in a tensor. The figure below illustrates this process for a GridBatch containing a single grid.
-
-
In the figure, the GridBatch acts as an acceleration structure which encodes the sparsity pattern of the grid (also known as the topology) and can translate grid coordinates into offsets into a data tensor.
-
By separating the grid topology and data tensors, the same grid can be used to operate on many different attributes without rebuilding.
-
Every operation in fVDB is built upon this kind of query (e.g. Sparse Convolution uses this query to look up features in the neighborhood of each voxel).
Each grid in a GridBatch can have a different number of voxels (e.g. in the mini batch of four cars above, each car has a different number of voxels). This means that unlike the dense case, fVDB needs to handle parallel operations over jagged batches. I.e. batches containing different numbers of elements.
-
To handle jagged batches, fVDB provides a JaggedTensor class. Conceptually, a JaggedTensor is a list of tensors with shapes \([N_0, *], [N_1, *], \ldots, [N_{B-1}, *]\) where \(B\) is the number of elements in the batch, \(N_i\) is the number of elements in the \(i^\text{th}\) batch item and \(*\) is an arbitrary numer of additional dimensions that all match between the tensors. The figure below illustrates such a list of tensors pictorially.
-
-
In practice, JaggedTensors are represented in memory by concatenating each tensor in the list into a single jdata (for Jagged Data) tensor of shape \([N_0 + N_1 + \ldots + N_{B-1}, *]\). Additionally, each JaggedTensor stores an additional jidx tensor (for Jagged Indexes) of shape \([N_0 + N_1 + \ldots + N_{B-1}]\) containing one int per element in the jagged tensor. jidx[i] is the batch index of the \(i^\text{th}\) element of jdata. Finally, a JaggedTensor contains a joffsets tensor (for Jagged Offsets) of shape \([B, 2]\) which indicates the start and end positions of the \(i^\text{th}\) tensor in the batch.
-
-
Similarly, each GridBatch also has jidx and joffsets corresponding to the batch index of each voxel in the grid, and the start and end offsets of each voxel index in a batch.
To illustrate the use of GridBatchand JaggedTensor, consider a simple example where we build a grid from a point cloud, splat some values onto the voxels of that grid, and then sample them again using a different set of points.
-
First, we construct a minibatch of grids using the input points. These input points have corresponding color attributes.
-
importfvdb
-fromfvdb.utils.examplesimportload_car_1_mesh,load_car_2_mesh
-importtorch
-importpoint_cloud_utilsaspcu
-
-# We're going to create a minibatch of two point clouds each of which
-# has a different number of points
-pts1,clrs1=load_car_1_mesh(mode="vn")
-pts2,clrs2=load_car_2_mesh(mode="vn")
-
-# Creating JaggedTensors: one for points and one for colors
-points=fvdb.JaggedTensor([pts1,pts2])
-colors=fvdb.JaggedTensor([clrs1,clrs2])
-
-# Create a grid where the voxels each have unit sidelength
-grid=fvdb.gridbatch_from_points(points,voxel_sizes=1.0)
-
-# Indexing into a JaggedTensor returns a JaggedTensor
-print(points[0].jdata.shape)
-print(points[1].jdata.shape)
-
-
-
-
Next, we splat the colors at the points to the constructed grid, yielding per-voxel colors.
-
# Splat the colors into the grid with trilinear interpolation
-# vox_colors is a JaggedTensor of per-voxel normas
-vox_colors=grid.splat_trilinear(points,colors)
-
-
-
-
Finally, we generate a new set of noisy points and sample the grid to recover colors at those new samples.
-
# Now let's generate some random points and sample the grid at those points
-sample_points=fvdb.JaggedTensor([torch.rand(10_000,3),torch.rand(11_000,3)]).cuda()
-
-# sampled_colors is a JaggedTensor with the same shape as sample_points with
-# one color sampled from the grid at each point
-sampled_colors=grid.sample_trilinear(sample_points,vox_colors)
-
Here we describe basic operations you can perform on a GridBatch. Generally, these operations involve a GridBatch and an associated JaggedTensor of data representing the attributes/features at each voxel within the grid.
You can differentiably sample data stored at the voxels of a grid using trilinear or Bézier interpolation as follows:
-
importfvdb
-fromfvdb.utils.examplesimportload_car_1_mesh,load_car_2_mesh
-importtorch
-importnumpyasnp
-
-# We're going to create a minibatch of two meshes
-v1,f1=load_car_1_mesh(mode="vf")
-v2,f2=load_car_2_mesh(mode="vf")
-
-# Build a GridBatch from two meshes
-mesh_v_jagged=fvdb.JaggedTensor([v1,v2])
-mesh_f_jagged=fvdb.JaggedTensor([f1,f2]).int()
-grid=fvdb.gridbatch_from_mesh(mesh_v_jagged,mesh_f_jagged,voxel_sizes=0.1)
-
-# Generate some sample points by adding random gaussian noise to the center of each voxel
-world_space_centers=grid.grid_to_world(grid.ijk.float())
-# 5 samples per voxel for the first grid and 7 samples per voxel for the second
-sample_pts=fvdb.JaggedTensor([
- torch.cat([world_space_centers[0].jdata]*5),
- torch.cat([world_space_centers[1].jdata]*7)])
-sample_pts+=fvdb.JaggedTensor([
- torch.randn(grid.num_voxels_at(0)*5,3).to(grid.device)*grid.voxel_sizes[0]*0.3,
- torch.randn(grid.num_voxels_at(1)*7,3).to(grid.device)*grid.voxel_sizes[1]*0.3])
-
-# Generate RGB values per voxel as the normalized absolute grid coordinate of the voxel
-per_voxel_colors=world_space_centers.clone()
-per_voxel_colors.jdata=torch.abs(world_space_centers.jdata)/torch.norm(world_space_centers.jdata,dim=-1,keepdim=True)
-
-# Sample these RGB colors at each sample point with trilinear interpolation
-# NOTE: You can use grid.sample_bezier to sample using bezier interpolation
-sampled_colors=grid.sample_trilinear(sample_pts,per_voxel_colors)
-
You can differentiably splat data at a set of points into voxels in a grid using trilinear or Bézier interpolation as follows:
-
importfvdb
-fromfvdb.utils.examplesimportload_car_1_mesh,load_car_2_mesh
-importtorch
-importpoint_cloud_utilsaspcu
-
-# We're going to create a minibatch of two point clouds each of which
-# has a different number of points
-pts1,clrs1=load_car_1_mesh(mode="vn")
-pts2,clrs2=load_car_2_mesh(mode="vn")
-
-# JaggedTensors of points and normals
-points=fvdb.JaggedTensor([pts1,pts2])
-colors=fvdb.JaggedTensor([clrs1,clrs2])
-
-# Create a grid where the voxels each have unit sidelength
-grid=fvdb.gridbatch_from_points(points,voxel_sizes=1.0)
-
-# Splat the normals into the grid with trilinear interpolation
-# vox_normals is a JaggedTensor of per-voxel normas
-# NOTE: You can use grid.splat_bezier to splat using bezier interpolation
-vox_colors=grid.splat_trilinear(points,colors)
-
You can query whether points lie in a grid as follows:
-
importfvdb
-fromfvdb.utils.examplesimportload_car_1_mesh,load_car_2_mesh
-importtorch
-importnumpyasnp
-importpoint_cloud_utilsaspcu
-
-# We're going to create a minibatch of two meshes
-v1,f1=load_car_1_mesh(mode="vf")
-v2,f2=load_car_2_mesh(mode="vf")
-f1,f2=f1.to(torch.int32),f2.to(torch.int32)
-
-# Build a GridBatch from two meshes
-mesh_v_jagged=fvdb.JaggedTensor([v1,v2]).cuda()
-mesh_f_jagged=fvdb.JaggedTensor([f1,f2]).cuda()
-grid=fvdb.gridbatch_from_mesh(mesh_v_jagged,mesh_f_jagged,voxel_sizes=0.1)
-
-# Generate some points and check if they lie within the grid
-bbox_sizes=grid.bbox[:,1]-grid.bbox[:,0]
-bbox_origins=grid.bbox[:,0]
-pts=fvdb.JaggedTensor([
- (torch.randn(10_000,3,device='cuda')-bbox_origins[0])*bbox_sizes[0],
- (torch.randn(11_000,3,device='cuda')-bbox_origins[0])*bbox_sizes[0],
-])
-
-# Get a mask indicating which points lie in the grid
-mask=grid.points_in_active_voxel(pts)
-
-
-
We visualize the points which intersect the grid (yellow points intersect and purple points do not).
-
There are methods of a grid to help ascertain whether a provided set of axis-aligned cubes (all of the same size) each are contained within voxels of the grid or intersect with any voxel of the grid. These methods are called cubes_in_grid and cubes_intersect_grid, respectively, and return a JaggedTensor of boolean values which indicate the result.
-
In this example we create some random points to represent the centers of cubes of size 0.03 units along-a-side. Then we use these two methods to discover which of those cubes would intersect a voxel of the grid and which would be entirely enclosed by a voxel in the grid.
-
importfvdb
-fromfvdb.utils.examplesimportload_car_1_mesh,load_car_2_mesh
-importtorch
-importnumpyasnp
-importpoint_cloud_utilsaspcu
-importpolyscopeasps
-importos
-
-# We're going to create a minibatch of two meshes
-v1,f1=load_car_1_mesh(mode="vf")
-v2,f2=load_car_2_mesh(mode="vf")
-f1,f2=f1.to(torch.int32),f2.to(torch.int32)
-
-# Build a GridBatch from two meshes
-mesh_v_jagged=fvdb.JaggedTensor([v1,v2]).cuda()
-mesh_f_jagged=fvdb.JaggedTensor([f1,f2]).cuda()
-grid=fvdb.gridbatch_from_mesh(mesh_v_jagged,mesh_f_jagged,voxel_sizes=0.025)
-
-# Generate some points and check if they lie within the grid
-bbox_sizes=(grid.bbox[:,1]-grid.bbox[:,0])*grid.voxel_sizes
-bbox_origins=(grid.bbox[:,0]+grid.origins)*grid.voxel_sizes
-pts=fvdb.JaggedTensor(
- [
- (torch.rand(3000,3,device="cuda"))*bbox_sizes[0]+bbox_origins[0],
- (torch.rand(2000,3,device="cuda"))*bbox_sizes[1]+bbox_origins[1],
- ]
-)
-
-cube_size=0.03
-
-# We can check if the axis-aligned cubes intersect any voxels of the grid...
-cubes_intersect_grid=grid.cubes_intersect_grid(pts,-cube_size/2,cube_size/2)
-# ... or if they are fully contained within the grid
-cubes_in_grid=grid.cubes_in_grid(pts,-cube_size/2,cube_size/2)
-
-
-
We visualize the cubes which intersect a voxel in the grid (yellow cubes intersect and purple cubes do not).
-
-
And which cubes lie entirely within a voxel of the grid
-
There is a convenient method to convert ijk values, which are integer coordinates in the index space of a grid, to indexes, which are the linearized indices of the voxels in the GridBatch. This method, ijk_to_index, returns a JaggedTensor of indexes. If the provided ijk values do not correspond to any voxels in the grid, the returned index will be -1.
-
-importfvdb
-fromfvdb.utils.examplesimportload_car_1_mesh,load_car_2_mesh
-importtorch
-importnumpyasnp
-importpoint_cloud_utilsaspcu
-importpolyscopeasps
-importos
-
-torch.random.manual_seed(0)
-
-defgenerate_random_points(bounding_box,num_points):
- min_i,min_j,min_k=bounding_box[0]
- max_i,max_j,max_k=bounding_box[1]
-
- # Generate random integer points within the bounding box
- random_i=torch.randint(min_i,max_i,size=(num_points,))
- random_j=torch.randint(min_j,max_j,size=(num_points,))
- random_k=torch.randint(min_k,max_k,size=(num_points,))
-
- random_points=torch.stack([random_i,random_j,random_k],dim=1)
-
- returnrandom_points
-
-# We're going to create a minibatch of two meshes
-v1,f1=load_car_1_mesh(mode="vf")
-v2,f2=load_car_1_mesh(mode="vf")
-f1,f2=f1.to(torch.int32),f2.to(torch.int32)
-
-# Build a GridBatch from two meshes
-mesh_v_jagged=fvdb.JaggedTensor([v1,v2]).cuda()
-mesh_f_jagged=fvdb.JaggedTensor([f1,f2]).cuda()
-grid=fvdb.gridbatch_from_mesh(mesh_v_jagged,mesh_f_jagged,voxel_sizes=0.025)
-
-rand_pts=fvdb.JaggedTensor([generate_random_points(bbox,1000)forbboxingrid.bbox]).cuda()
-
-rand_pts_indices=grid.ijk_to_index(rand_pts)
-
-print(rand_pts_indices.jdata)
-
If we have the value of an index into the feature data and want to obtain its corresponding ijk value, this is as simple as indexing into the ijk attribute of the grid which is itself just a JaggedTensor representing the ijk values of each index. Here we get the ijk values for 1000 random indexes of a GridBatch.
-
importfvdb
-fromfvdb.utils.examplesimportload_car_1_mesh,load_car_2_mesh
-importtorch
-importnumpyasnp
-importpoint_cloud_utilsaspcu
-importpolyscopeasps
-importos
-
-# We're going to create a minibatch of two meshes
-v1,f1=load_car_1_mesh(mode="vf")
-v2,f2=load_car_2_mesh(mode="vf")
-f1,f2=f1.to(torch.int32),f2.to(torch.int32)
-
-# Build a GridBatch from two meshes
-mesh_v_jagged=fvdb.JaggedTensor([v1,v2]).cuda()
-mesh_f_jagged=fvdb.JaggedTensor([f1,f2]).cuda()
-grid=fvdb.gridbatch_from_mesh(mesh_v_jagged,mesh_f_jagged,voxel_sizes=0.025)
-
-rand_indexes=torch.randint(0,grid.total_voxels,size=(1000,)).cuda()
-print(grid.ijk.jdata[rand_indexes])
-print(grid.ijk.jidx[rand_indexes])
-
However, you may also want to obtain the batch index of the grid that the ijk coordinate belongs to. This is just as simple by using the jidx attribute of the JaggedTensor which represents the batch index of each element in the JaggedTensor. Using the same random indexes as above, we can get the batch index of the grid for each random index.
While this case of having a random set of ijk values is very contrived, it is more likely that we would have a JaggedTensor of features that correspond to each ijk value but are out-of-order in relation to another grid.ijk’s and we need to reorder these features to for this other grid. For this, there is the ijk_to_inv_index function.
-
Given a JaggedTensor of ijk values, grid_batch.ijk_to_inv_index will return a scalar integer JaggedTensor of size [B,-1], where B is the number of grids in the batch and -1 represents the number of voxels in each grid. This JaggedTensor of indexes can be used to permute the input to ijk_to_inv_index (or the feature JaggedTensor that correspond to those ijk’s) to match the ordering of the grid_batch.
-
To concisely illustrate the properties of this function,
-
-
if idx=grid.ijk_to_inv_index(misordered_ijk), then grid.ijk==misordered_ijk[idx]
-
if idx=grid.ijk_to_index(misordered_ijk), then grid.ijk[idx]==misordered_ijk
-
-
Note: If any ijk values in the grid are not present in the input to ijk_to_inv_index, the returned index will be -1 at that position.
-
In this example, let’s illustrate the more useful case where we have corresponding features and ijk values from a grid that are ordered differently from another reference grid and we want to re-order the features to match the order they should be in the reference grid.
-
importfvdb
-fromfvdb.utils.examplesimportload_car_1_mesh,load_car_2_mesh
-importtorch
-importnumpyasnp
-importpoint_cloud_utilsaspcu
-
-# We're going to create a minibatch of two meshes
-v1,f1=load_car_1_mesh(mode="vf")
-v2,f2=load_car_2_mesh(mode="vf")
-f1,f2=f1.to(torch.int32),f2.to(torch.int32)
-
-# Build a GridBatch from two meshes
-mesh_v_jagged=fvdb.JaggedTensor([v1,v2]).cuda()
-mesh_f_jagged=fvdb.JaggedTensor([f1,f2]).cuda()
-reference_grid=fvdb.gridbatch_from_mesh(mesh_v_jagged,mesh_f_jagged,voxel_sizes=0.025)
-
-# 7 random feature values for each voxel in the grid
-features=reference_grid.jagged_like(torch.rand(grid.total_voxels,7).to(grid.device))
-
-# Create a set of randomly shuffled corresponding ijk/features from our original grid/features
-shuffled_ijks=[]
-shuffled_features=[]
-
-foriinrange(reference_grid.grid_count):
- perm=torch.randperm(reference_grid.num_voxels_at(i))
- shuffled_ijks.append(reference_grid.ijk.jdata[grid.ijk.jidx==i][perm])
- shuffled_features.append(features.jdata[features.jidx==i][perm])
-
-shuffled_ijks=fvdb.JaggedTensor(shuffled_ijks)
-shuffled_features=fvdb.JaggedTensor(shuffled_features)
-
-# Get the indexes to reorder the shuffled features to match the grid's original ijk ordering
-idx=reference_grid.ijk_to_inv_index(shuffled_ijks)
-
-# Permute the shuffled features based on the ordering from `ijk_to_inv_index`
-unshuffled_features=shuffled_features.jdata[idx.jdata]
-print("Do the shuffled features that have been permuted based on `ijk_to_inv_index` match the original features? ","Yes!"iftorch.all(unshuffled_features==features.jdata)else"No!")
-
If we want to get the indexes of all the spatial neighbors of a set of ijk values, we can use the neighbor_indexes method of a GridBatch. This method receives a set of ijk values and an extent (the number of voxels away from the ijk value to consider neighbors) and returns a JaggedTensor of indexes of neighbors for each ijk value. The returned JaggedTensor has jdata of size [N,extent*2+1,extent*2+1,extent*2+1] where N is the number of requested ijk values.
-
In this example we create a set of 24 random ijk values per grid and get the indexes of all neighbors 2 voxels away from each ijk.
-
importfvdb
-fromfvdb.utils.examplesimportload_car_1_mesh,load_car_2_mesh
-importtorch
-importnumpyasnp
-importpoint_cloud_utilsaspcu
-importpolyscopeasps
-importos
-
-# We're going to create a minibatch of two meshes
-v1,f1=load_car_1_mesh(mode="vf")
-v2,f2=load_car_2_mesh(mode="vf")
-f1,f2=f1.to(torch.int32),f2.to(torch.int32)
-
-# Build a GridBatch from two meshes
-mesh_v_jagged=fvdb.JaggedTensor([v1,v2]).cuda()
-mesh_f_jagged=fvdb.JaggedTensor([f1,f2]).cuda()
-grid=fvdb.gridbatch_from_mesh(mesh_v_jagged,mesh_f_jagged,voxel_sizes=0.025)
-
-# Build a set of 24 randomly selected ijk values per grid
-rand_ijks=fvdb.JaggedTensor(
- [
- grid.ijk.jdata[
- torch.randint(int(grid.ijk.joffsets[b]),int(grid.ijk.joffsets[b+1]),(24,))
- ]
- forbinrange(grid.grid_count-1)
- ],
-)
-
-# Get the indexes of all neighbors 2 voxels away from each ijk
-# Returns a JaggedTensor with jdata of size [48, 5, 5, 5] ([ N x (extent*2+1)^3 ])
-# where each [5, 5, 5] describes the layout of neighbouring indexes of each ijk
-neighbor_idxs=grid.neighbor_indexes(rand_ijks,2)
-
-
-
We visualize the voxels we selected in red and their neighbors in blue.
-
If we want to ‘clip’ a grid, meaning remove all the voxels that are outside of a bounding box, we can use the clipped_grid method of a GridBatch. The minimum and maximum ijk extents of the bounding box used for this clipping are provided as arguments and can be specified per-grid of the batch.
-
If we want to both clip the batch of grids and the accompanying JaggedTensor of batched features, we can use the clip method of a GridBatch which can be provided the features to be clipped along with the grid.
-
In this example we clip a batch of grids to independent bounding boxes and visualize the result.
-
importos
-importfvdb
-fromfvdb.utils.examplesimportload_car_1_mesh,load_car_2_mesh,load_car_3_mesh,load_car_4_mesh
-importtorch
-importnumpyasnp
-importpoint_cloud_utilsaspcu
-
-device='cuda'
-mesh_vs=[]
-mesh_fs=[]
-mesh_load_funcs=[load_car_1_mesh,load_car_2_mesh,load_car_3_mesh,load_car_4_mesh]
-
-forfuncinmesh_load_funcs:
- v,f=func(mode="vf")
- mesh_vs.append(v)
- mesh_fs.append(f.to(torch.int32))
-
-mesh_vs=fvdb.JaggedTensor(mesh_vs)
-mesh_fs=fvdb.JaggedTensor(mesh_fs)
-
-vox_size=0.01
-
-# Build GridBatch from meshes
-grid=fvdb.gridbatch_from_mesh(mesh_vs,mesh_fs,vox_size)
-
-# Use `clipped_grid` to clip the grids outside of the specified minimum and maximum bounding boxes for each grid…
-clipped_grid=grid.clipped_grid(
- ijk_min=[[-200,-200,-200],[0,-200,-200],[-200,-200,-200],[-200,0,-200]],
- ijk_max=[[0,200,200],[200,200,200],[200,0,200],[200,200,200]])
-
-# Use `clip` to both clip the grids and the accompanying JaggedTensor of features
-features=grid.jagged_like(torch.rand(grid.total_voxels,7).to(device))
-clipped_features,clipped_grid=grid.clip(
- features=features,
- ijk_min=[[-200,-200,-200],[0,-200,-200],[-200,-200,-200],[-200,0,-200]],
- ijk_max=[[0,200,200],[200,200,200],[200,0,200],[200,200,200]])
-
-
-
-
We visualize these grids, each clipped to a different bounding box.
‘Pooling’ a grid has the effect of coarsening the resolution of the grid by a constant factor (and can be performed anisotropically for different factors in each xyz spatial dimension). When pooling a grid, the values of the new, coarser grid can be determined by the maximum or mean of the values of the voxels in the original grid that are covered by each voxel in the new grid. Maximum and mean pooling can be accomplished by the max_pool and avg_pool operators.
-
In this example, we create a grid from a mesh and then perform max and mean pooling on it to illustrate the difference between the two.
-
importos
-importfvdb
-fromfvdb.utils.examplesimportload_car_1_mesh
-importtorch
-importnumpyasnp
-importpoint_cloud_utilsaspcu
-
-vox_size=0.02
-num_pts=10_000
-
-mesh_load_funcs=[load_car_1_mesh]
-
-points=[]
-normals=[]
-
-forfuncinmesh_load_funcs:
- pts,nms=func(mode="vn")
- pmt=torch.randperm(pts.shape[0])[:num_pts]
- pts,nms=pts[pmt],nms[pmt]
- points.append(pts)
- normals.append(nms)
-
-# JaggedTensors of points and normals
-points=fvdb.JaggedTensor(points)
-normals=fvdb.JaggedTensor(normals)
-
-# Create a grid
-grid=fvdb.gridbatch_from_points(points,voxel_sizes=vox_size)
-
-# Splat the normals into the grid with trilinear interpolation
-vox_normals=grid.splat_trilinear(points,normals)
-
-# Mean Pooling of normals features
-avg_normals,avg_grid=grid.avg_pool(4,vox_normals)
-# Max Pooling of normals features
-max_normals,max_grid=grid.max_pool(4,vox_normals)
-
-
-
We visualize the original grid with values of the mesh normals as features, the features/grid after mean pooling, and the features/grid after max pooling to illustrate how these pooling modes affect the resulting features of the grid.
Subdividing a grid has the effect of increasing the resolution of the grid by a constant factor (and can be performed anisotropically for different factors in each xyz spatial dimension).
-
The most straightforward way to use the subdivide method is to provide an integer value for the subdivision factor. This will result in a grid that is subdiv_factor times the resolution of the original grid in each spatial dimension.
-
importos
-importfvdb
-fromfvdb.utils.examplesimportload_car_1_mesh
-importtorch
-importnumpyasnp
-importpoint_cloud_utilsaspcu
-
-vox_size=0.02
-num_pts=10_000
-
-mesh_load_funcs=[load_car_1_mesh]
-
-points=[]
-normals=[]
-
-forfuncinmesh_load_funcs:
- pts,nms=func(mode="vn")
- pmt=torch.randperm(pts.shape[0])[:num_pts]
- pts,nms=pts[pmt],nms[pmt]
- points.append(pts)
- normals.append(nms)
-
-# JaggedTensors of points and normals
-points=fvdb.JaggedTensor(points)
-normals=fvdb.JaggedTensor(normals)
-
-# Create a grid
-grid=fvdb.gridbatch_from_points(points,voxel_sizes=vox_size)
-
-# Splat the normals into the grid with trilinear interpolation
-vox_normals=grid.splat_trilinear(points,normals)
-
-# Subdivide by a constant factor of 2
-subdiv_normals,subdiv_grid=grid.subdivide(2,vox_normals)
-
-
-
Here we visualize the original grid on the right and the grid after subdivision by a factor of 2 on the left.
-
-
In practice, in a deep neural network like a U-Net architecture, the resolution of the grid can be decreased early in the network and then the grid’s features need to be concatenated with features of the grid after its resolution is increased again. When working with traditional, dense 2D or 3D data, cropping any mismatched outputs is straightforward to be able to concatenate these features. However, in a sparse 3D grid, this is not straightforward and it’s entirely unclear how to align the features.
-
Take this simple example to illustrate the difficulties created by changing the grid topology in this way. We take the grid created from our mesh, perform Pooling by a factor of 2 and then try to Subdivide by an equal factor of 2 to invert the Pooling operation.
Let’s visualize these results from left to right of the original grid, the grid after max pooling, and the grid after max pooling and then subdividing again by the same factor.
-
-
Notice how we have not obtained the topology of the original grid and have obtained a grid with many more voxels than the original.
-
To correctly invert the Pooling operation, we can provide the subdivide function with a fine_grid optional argument which describes the topology we want the grid to have after the subdivision. The original grid before the Pooling operation can be used as this fine_grid.
-
max_normals,max_grid=grid.max_pool(2,vox_normals)
-
-# Providing the original grid as our fine_grid target
-subdiv_normals,subdiv_grid=max_grid.subdivide(2,max_normals,fine_grid=grid)
-
-
-
-
Note the matching topology of the grid after the Subdivision operation to the original grid (though the features are different due to the Pooling operation).
-
One other useful feature of the subdivide operator is that this operation can be masked so that subdivision is only performed on a subset of the voxels in the grid.
-
The optional mask argument to subdivide is a JaggedTensor of boolean values that indicates which voxels should be subdivided. Given this mask is simply a JaggedTensor, this operation can be made differentiable in a neural network and the subdivide operator can be learned.
-
Let’s illustrate a very simple example of how to use the mask argument to only subdivide the voxels which have a feature value greater than a certain threshold.
-
# Mask the grid with the normals where only the normals with a value greater than 0.5 on the x-axis are subdivided
-mask=vox_normals.jdata[:,0]>0.5
-
-subdiv_normals,subdiv_grid=grid.subdivide(2,vox_normals,mask=mask)
-
-
-
-
-
-
Getting the number of enabled voxels per grid in a batch
-
Getting the number of voxels in the grids of a GridBatch can be easily accomplished with num_voxels:
-
importfvdb
-importtorch
-
-# Create a GridBatch of random points
-batch_size=4
-pts=fvdb.JaggedTensor([torch.rand(10_000*(i+1),3)foriinrange(batch_size)])
-grid=fvdb.GridBatch(mutable=False)
-grid.set_from_points(pts,voxel_sizes=0.02)
-
-# Get the number of voxels per grid in the batch
-forbatch,num_voxelsinenumerate(grid.num_voxels):
- print(f"Grid {batch} has {num_voxels} voxels")
-
If the grid is mutable and the number of enabled voxels has changed, you can use num_enabled_voxels to get the number of enabled voxels. If a grid is not mutable, num_enabled_voxels will return the same value as num_voxels.
-
# Create a mutable GridBatch of random points
-grid=fvdb.GridBatch(mutable=True)
-grid.set_from_points(pts,voxel_sizes=0.02)
-# Disable some voxels randomly
-grid.disable_ijk(fvdb.JaggedTensor([torch.randint(0,50,(100_000,3))for_inrange(batch_size)]))
-
-# Get the number of enabled voxels per grid in the batch
-forbatch,num_voxelsinenumerate(grid.num_enabled_voxels):
- print(f"Grid {batch} has {grid.num_enabled_voxels_at(batch)} enabled voxels and {grid.num_voxels_at(batch)} total voxels")
-
Converting between ijk (grid) coordinates and world coordinates
-
A GridBatch can contain multiple grids, each with its own coordinate system that relate the grids’ ijk integer index space coordinates to xyz floating-point world-space coordinates. These axis-aligned coordinate systems are defined by specifying a list of three-dimensional world-space origin and scale values when we define the topology of the grids in the GridBatch like so:
-
importfvdb
-importtorch
-
-# Create a GridBatch of random points
-batch_size=4
-pts=fvdb.JaggedTensor([torch.rand(10_000*(i+1),3)foriinrange(batch_size)])
-grid=fvdb.GridBatch()
-grid.set_from_points(pts,
- voxel_sizes=[[0.02,0.02,0.02],[0.03,0.03,0.03],[0.04,0.04,0.04],[0.05,0.05,0.05]],
- origins=[[-.1,-.1,-.1],[0,0,0],[.1,.1,.1],[.2,-.2,.2]])
-
-
-
When voxel_sizes and origins are not defined, it is assumed all grids have a unit scale voxel_size and an origin at [0.0,0.0,0.0].
-
We can use GridBatch’s grid_to_world function to convert between ijk index coordinates and their corresponding world-space xyz coordinates. In this example, let’s obtain the world-space position that would lie at the index-space [1,1,1] point of each grid:
-
# Convert ijk coordinates to world coordinates
-ijk=fvdb.JaggedTensor([torch.ones(1,3,dtype=torch.float)for_inrange(batch_size)])
-world_coords=grid.grid_to_world(ijk)
-foriinrange(grid.grid_count):
- print(f"World-space point that lies at index [1,1,1] for Grid {i} is positioned at {world_coords.jdata[world_coords.jidx==i].tolist()}")
-
We can also do the inverse operation and convert world-space xyz coordinates to their corresponding ijk index-space coordinates using world_to_grid. In this example, let’s find the ijk index coordinates of the voxel which would contain the world-space point located at [1.0,1.0,1.0] for each grid in our GridBatch:
-
# Convert world coordinates to ijk coordinates
-xyz=fvdb.JaggedTensor([torch.ones(1,3,dtype=torch.float)for_inrange(batch_size)])
-ijk_coords=grid.world_to_grid(xyz)
-foriinrange(grid.grid_count):
- print(f"Index-space voxel that contains the point [1.0, 1.0, 1.0] for Grid {i}{ijk_coords.jdata[ijk_coords.jidx==i].int().tolist()}")
-
While grid_to_world and world_to_grid are convenient functions for these purposes, the row-major transformation matrices used for these calculations can be obtained directly from a GridBatch for use in your own logic:
-
print(f"Grid to world matrices:\n{grid.grid_to_world_matrices}")
-print(f"World to grid matrices:\n{grid.grid_to_world_matrices}")
-
Convolving the features of a GridBatch can be accomplished with either a high-level torch.nn.Module derived class provided by fvdb.nn or with more low-level methods available with GridBatch, we will illustrate both techniques.
fvdb.nn.SparseConv3d provides a high-level torch.nn.Module class for convolution on fvdb classes that is an analogue to the use of torch.nn.Conv3d. Using this module is the recommended functionality for performing convolution with fvdb because it not only manages functionality such as initializing the weights of the convolution and calling appropriate backend implementation functions but it also provides certain backend optimizations which will be illustrated in the Low-level usage section.
-
One thing to note is fvdb.nn.SparseConv3d operates on a class that wraps a GridBatch and JaggedTensor together into a convenience object, fvdb.VDBTensor, which is used by all the fvdb.nn modules.
-
A simple example of using fvdb.nn.SparseConv3d is as follows:
-
importfvdb
-importfvdb.nnasfvdbnn
-fromfvdb.utils.examplesimportload_car_1_mesh
-importtorch
-importnumpyasnp
-importpoint_cloud_utilsaspcu
-
-num_pts=10_000
-vox_size=0.02
-
-mesh_load_funcs=[load_car_1_mesh]
-
-points=[]
-normals=[]
-
-forfuncinmesh_load_funcs:
- pts,nms=func(mode="vn")
- pmt=torch.randperm(pts.shape[0])[:num_pts]
- pts,nms=pts[pmt],nms[pmt]
- points.append(pts)
- normals.append(nms)
-
-# JaggedTensors of points and normals
-points=fvdb.JaggedTensor(points)
-normals=fvdb.JaggedTensor(normals)
-
-# Create a grid
-grid=fvdb.gridbatch_from_points(points,voxel_sizes=vox_size)
-
-# Splat the normals into the grid with trilinear interpolation
-vox_normals=grid.splat_trilinear(points,normals)
-
-# VDBTensor is a simple wrapper of a grid and a feature tensor
-vdbtensor=fvdbnn.VDBTensor(grid,vox_normals)
-
-# fvdb.nn.SparseConv3d is a convenient torch.nn.Module implementing the fVDB convolution
-conv=fvdbnn.SparseConv3d(in_channels=3,out_channels=3,kernel_size=3,stride=1,bias=False).to(vdbtensor.device)
-
-output=conv(vdbtensor)
-
-
-
Let’s visualize the original grid with normals visualized as colours alongside the result of these features after a convolution initialized with random weights:
-
-
For stride values greater than 1, the output of the convolution will be a grid with a smaller resolution than the input grid (similar in topological effect to the output of a Pooling operator). Let’s illustrate this:
-
# We would expect for stride=2 that the output grid would have half the resolution (or twice the world-space size) of the input grid
-conv=fvdbnn.SparseConv3d(in_channels=3,out_channels=3,kernel_size=3,stride=2,bias=False).to(vdbtensor.device)
-
-output=conv(vdbtensor)
-
-
-
-
Transposed convolution can be performed with fvdb.nn.SparseConv3d which can increase the resolution of the grid. It only really makes sense to perform transposed sparse convolution with a target grid topology we wish to produce with this operation (see the Pooling Operators for an explanation). Therefore, an out_grid argument must be provided in this case to specify the target grid topology:
-
# Tranposed convolution operator, stride=2
-transposed_conv=fvdbnn.SparseConv3d(in_channels=3,out_channels=3,kernel_size=3,stride=2,bias=False,transposed=True).to(vdbtensor.device)
-
-# Note the use of the `out_grid` argument to specify the target grid topology
-transposed_output=transposed_conv(output,out_grid=vdbtensor.grid)
-
-
-
Here we visuzlie the original grid, the grid after strided convolution and the grid after transposed convolution inverts the topological operation of the strided convolution to produce the same topology as the original grid with the features convolved by our two layers:
The high-level fvdb.nn.SparseConv3d class wraps several pieces of GridBatch functionality to provide a convenient torch.nn.Module for convolution. However, for a more low-level approach that accomplishes the same outcome, the GridBatch class itself can be the starting point for performing convolution on the grid and its features. We will illustrate this approach for completeness, though we do recommend the use of the fvdb.nn.SparseConv3d Module for most use-cases.
-
Using the GridBatch convolution functions directly requires a little more knowledge about the implementation under-the-hood. Due to the nature of a sparse grid, in order to make convolution performant, it is useful to pre-compute a mapping of which features in the input grid will contribute to the values of the output grid when convolved by a kernel of a particular dimension and stride. This mapping structure is called a ‘kernel map’.
-
The kernel map, as well as the functionality for using it to compute the convolution, is contained within a fvdb.SparseConvPackInfo class which can be constructed by GridBatch.sparse_conv_kernel_map. A fvdb.SparseConvPackInfo must perform a pre-computation of the kernel map based on the style expected by the backend implementation of the convolution utilized by fvdb.SparseConvPackInfo.sparse_conv_3d. Here is an example of how to construct a fvdb.SparseConvPackInfo and use it to perform a convolution:
-
importfvdb
-importfvdb.nnasfvdbnn
-fromfvdb.utils.examplesimportload_car_1_mesh
-importtorch
-importnumpyasnp
-importpoint_cloud_utilsaspcu
-
-num_pts=10_000
-vox_size=0.02
-
-mesh_load_funcs=[load_car_1_mesh]
-
-points=[]
-normals=[]
-
-forfuncinmesh_load_funcs:
- pts,nms=func(mode="vn")
- pmt=torch.randperm(pts.shape[0])[:num_pts]
- pts,nms=pts[pmt],nms[pmt]
- points.append(pts)
- normals.append(nms)
-
-# JaggedTensors of points and normals
-points=fvdb.JaggedTensor(points)
-normals=fvdb.JaggedTensor(normals)
-
-# Create a grid
-grid=fvdb.gridbatch_from_points(points,voxel_sizes=vox_size)
-
-# Splat the normals into the grid with trilinear interpolation
-vox_normals=grid.splat_trilinear(points,normals)\
-
-# Create the kernel map (housed in a SparseConvPackInfo) and the output grid's topology based on the kernel parameters
-sparse_conv_packinfo,out_grid=grid.sparse_conv_kernel_map(kernel_size=3,stride=1)
-
-# The kernel map must be pre-computed based on the backend implementation we plan on using, here we use gather/scatter, the default implementation
-sparse_conv_packinfo.build_gather_scatter()
-
-# Create random weights for our convolution kernel of size 3x3x3 that takes 3 input channels and produces 3 output channels
-kernel_weights=torch.randn(3,3,3,3,3,device=grid.device)
-
-# Perform convolution on the normals colours. Gather/scatter is used as the backend, it is the default
-conv_vox_normals=sparse_conv_packinfo.sparse_conv_3d(vox_normals,weights=kernel_weights,backend=fvdb.ConvPackBackend.GATHER_SCATTER)
-
-
-
Here we visualize the output of our convolution alongside the original grid with normals visualized as colours:
-
-
The kernel map can potentially be expensive to compute, so it is often useful to re-use the SparseConvPackInfo in the same network to perform a convolution on other features or with different weights. This optimization is something fvdb.nn.SparseConv3d attempts to do where appropriate and is one reason we recommend using fvdb.nn.SparseConv3d over this low-level approach.
We introduce several ways to construct sparse voxel grids from various data sources including point clouds, coordinate lists, triangle meshes, and deriving from other grids.
-All the examples below could be found in full version (including visualization) at examples/grid_building.py and examples/grid_subdivide_coarsen.py.
If you already have an integer list of the ijk coordinates for each voxel in the grid, you could directly build grids from the list.
-Additionally, you will have to specify the voxel sizes and voxel origins for the grid.
-An example is as follows:
The above code assumes that you want to build a grid with two batch elements, one with voxel size [0.1,0.1,0.1], and the other with voxel size [0.15,0.15,0.15] (although usually you just want all the elements in your batch to have the same size, in which case you could just pass in voxel_sizes=[0.1,0.1,0.1]).
-The same logic applies for origins that specifies the world coordinates of voxel (0,0,0), which we set to the origin here.
-The grid will be constructed on the same device as coords_jagged, which is a JaggedTensor. The JaggedTensor is built from two numpy int64 arrays with size [*,3] called coords_1 and coords_2.
You could either choose to quantize the coordinates of your point cloud into ijk coordinates yourself (e.g. using np.unique((xyz/voxel_size).floor(),axis=0)), or let fvdb handle this logic. Specifically, you could do:
Above we show two methods of building grids from points. Similar functions exist for other grid building approaches. The built grids are shown as following:
-
-
In some applications, you may want to build a dilated version of the grid by ensuring that all \(2\times 2 \times 2\) voxels around each point are included in the built grid. That said, you could build the grid by:
-
importfvdb
-fromfvdb.utils.examplesimportload_car_1_mesh,load_car_2_mesh
-
-coords_1,_=load_car_1_mesh()
-coords_2,_=load_car_2_mesh()
-
-# Assemble point clouds into JaggedTensor
-pcd_jagged=fvdb.JaggedTensor([
- coords_1.cuda(),
- coords_2.cuda()
-])
-voxel_sizes=[[0.1,0.1,0.1],[0.15,0.15,0.15]]
-
-# Build grid from containing nearest voxels to the points
-grid_b=fvdb.gridbatch_from_nearest_voxels_to_points(pcd_jagged,voxel_sizes=voxel_sizes,origins=[0.0]*3)
-
We allow building grids enclosing a triangle mesh easily. The given triangle mesh does not have to be manifold nor watertight and it will be treated as a triangle soup internally.
-An example to build grids from meshes is shown as follows:
Here mesh_1_v and mesh_1_f are the vertex array and triangle array of the mesh to build grid from, with the shape of \((V, 3)\) and \((F, 3)\). The triangle array is an integer array that indexes into the vertex array (starting from 0 for each element in the batch). Same for another mesh_2_v and mesh_2_f.
If you are comparing the performance of dense pytorch 3D tensors vs sparse grids, it is usually very helpful to build the exact same input (including grid and features). In fvdb.nn, we provide a thin wrapper class VDBTensor that works like a torch.Tensor, yet enclosing the grid topology. To convert data back and forth from dense PyTorch Tensors, we could do:
-
importtorch
-importfvdb
-fromfvdb.nnimportVDBTensor
-
-# Easy way to initialize a VDBTensor from a torch 3D tensor [B, D, H, W, C]
-dense_data=torch.ones(2,32,32,32,16).cuda()
-sparse_data=fvdb.nn.vdbtensor_from_dense(dense_data,voxel_sizes=[0.1]*3)
-dense_data_back=sparse_data.to_dense()
-asserttorch.all(dense_data==dense_data_back)
-
-
-
Here sparse_data will be a fvdb.nn.VDBTensor class, containing both feature and grid attribute.
-Such a class could be fed into all the neural network components available in fvdb.nn.
A dual grid (of a primal grid) is also a grid, with its voxel centers covering the corners of the primal grid, and corners taking the positions of the centers of the primal grid.
-A dual grid shares the same voxel sizes as its primal grid, and is just a simple shift of half the voxel size in translation.
-In the picture below, green grid is the primal grid while purple grid is its dual.
-
-
To create a dual grid from a given primal grid, use GridBatch.dual_grid():
We give an overview of ways to save and load sparse grids including how fVDB’s serialized format relates to other libraries such as OpenVDB and NanoVDB. All of the examples in this tutorial are available in the examples/io.py file of the fVDB repository.
-
In these examples we will be using tools which are part of the NanoVDB project such as nanovdb_convert. It is assumed that the NanoVDB tools are available to call (i.e. findable on the system’s $PATH). If you have not already installed these tools, you can find instructions on building NanoVDB with OpenVDB from its documentation:
Batches of sparse grids can be serialized to a NanoVDB file using the fvdb.save method. Here, we create two grids of different sizes and numbers of points and save them to a compressed NanoVDB file with specified names. The names are optional.
-
importtorch
-importfvdb
-importtempfile
-importos
-importsubprocess
-
-p=fvdb.JaggedTensor(
- [
- torch.randn(10,3),
- torch.randn(100,3),
- ]
-)
-grid=fvdb.gridbatch_from_points(
- p,voxel_sizes=[[0.1,0.1,0.1],[0.15,0.15,0.15]],origins=[0.0]*3
-)
-
-# save the grid and features to a compressed nvdb file
-path=os.path.join(tempfile.gettempdir(),"two_random_grids.nvdb")
-fvdb.save(path,grid,names=["taco1","taco2"],compressed=True)
-
-
-
We can use the nanovdb_print command line tool to show information about our saved file. Note how our grids have the INDEX class since they have no features and only store the voxel indices.
-
Thefile"/tmp/tmpwnu7qc_7/two_random_grids.nvdb"containsthefollowing2grids:
-# Name Type Class Version Codec Size File Scale # Voxels Resolution
-1taco1OnIndexINDEX32.6.0BLOSC1.453MB7.181KB(0.1,0.1,0.1)1022x43x37
-2taco2OnIndexINDEX32.6.0BLOSC2.326MB12.7KB(0.15,0.15,0.15)10033x39x38
-
-
-
We can include N-dimensional features by passing a JaggedTensor as the second argument (or data kwarg) to fvdb.save. Here, we create a grid with a single, float feature channel for our grids and save it to a nvdb file.
-
# a single, scalar float feature per grid
-feats=fvdb.JaggedTensor([torch.randn(x,1)forxingrid.num_voxels])
-
-# save the grid and features to a compressed nvdb file
-path=os.path.join(tempfile.gettempdir(),"two_random_grids.nvdb")
-fvdb.save(path,grid,feats,names=["taco1","taco2"],compressed=True)
-
-
-
Again, we can use the nanovdb_print command line tool to show information about our saved file.
-
Thefile"/tmp/tmpg9ol5841/two_random_grids.nvdb"containsthefollowing2grids:
-# Name Type Class Version Codec Size File Scale # Voxels Resolution
-1taco1float?32.6.0BLOSC1.47MB7.959KB(1,1,1)1026x30x36
-2taco2float?32.6.0BLOSC2.411MB16.48KB(1,1,1)10032x33x34
-
-
-
Note how our serialized NanoVDB grids are now of type float. fVDB will automatically map N-dimensional features to appropriate NanoVDB types. For feature sizes that don’t naturally map to any NanoVDB data types, fVDB will save the feature data as NanoVDB blind-data which will be appropriately read back as N-dimensional features by fVDB.
-
Let’s try to save the same two grids with a Vec3d type by creating a JaggedTensor of 3-dimensional double-precision features.
-
# a 3-vector double feature per grid
-feats=fvdb.JaggedTensor([torch.randn(x,3,dtype=torch.float64)forxingrid.num_voxels])
-
-# save the grid and features to a compressed nvdb file
-saved_nvdb=os.path.join(tempfile.gettempdir(),"two_random_vec3d_grids.nvdb")
-fvdb.save(saved_nvdb,grid,feats,names=["taco1","taco2"],compressed=True)
-
-
-
Thefile"/tmp/tmpwnu7qc_7/two_random_grids.nvdb"containsthefollowing2grids:
-# Name Type Class Version Codec Size File Scale # Voxels Resolution
-1taco1Vec3d?32.6.0BLOSC6.077MB28.18KB(1,1,1)1023x36x35
-2taco2Vec3d?32.6.0BLOSC7.346MB41.37KB(1,1,1)10037x40x34
-
Loading NanoVDB files is as simple as calling fvdb.load. You can optionally supply a PyTorch device you’d like the grids and features loaded onto. Here, we load the two grids we saved in the previous section onto our GPU.
-
# Load the grid and features from the compressed nvdb file
-grid_batch,features,names=fvdb.load(saved_nvdb,device=torch.device("cuda:0"))
-print("Loaded grid batch total number of voxels: ",grid_batch.total_voxels)
-print("Loaded grid batch data type: %s, device: %s"%(features.dtype,features.device))
-
While saving and loading from OpenVDB files is not directly supported by fVDB, it is possible to easily convert between NanoVDB and OpenVDB files using the nanovdb_convert command line tool. Here, we convert our previously saved NanoVDB file to an OpenVDB file.
From here, our grid can be loaded by OpenVDB tools. Roundtripping our converted cache back to NanoVDB is possible with the nanovdb_convert tool as well.
-
Loading the converted OpenVDB file into fVDB shows our familiar grids and features as we expect:
-
convert_cmd="nanovdb_convert -v -f %s%s"%(# -f flag forces overwriting existing file
- vdb_path,
- saved_nvdb,
-)
-print("nanovdb_convert roundtrip the vdb to nvdb: ",convert_cmd)
-print(subprocess.check_output(convert_cmd.split()).decode("utf-8"))
-
-# Load the nvdb file of the converted vdb
-grid_batch,features,names=fvdb.load(saved_nvdb,device=torch.device("cuda:0"))
-print("Loaded grid batch total number of voxels: ",grid_batch.total_voxels)
-print("Loaded grid batch data type: %s, device: %s"%(features.dtype,features.device))
-print("\n")
-
Mutable grids refer to GridBatch whose voxels can be turned ‘off’.
-Each voxel hence not only stores an integer offset that indexes into the external feature array, but also includes a bit switch indicating whether the voxel exist or not.
-
-
The ability to turn on (enable) voxels and turn off (disable) voxels make it easier for downstream tasks such as structural optimization and neural rendering.
-Note that even if you disable some voxels, the corresponding entry in the feature array still exists.
-Such a design keeps the feature unchanged while changing the grid topology in a flexible way.
Voxels can be disabled in batches via disable_ijk:
-
# Get the IJK coordinates to be disabled
-disable_ijk:fvdb.JaggedTensor=grid.ijk.rmask(feature.jdata[:,0]>0.5)
-
-# Disable them!
-grid.disable_ijk(disable_ijk)
-
-
-
Once disabled, those voxels will virtually disappear, meaning that all subsequent grid operations such as sampling, splatting, or ray marching, will treat those voxels as they do not exist.
-One can visualize the enable mask via:
-
enabled_mask=grid.enabled_mask
-
-
-
-
Note that in the above figure, white voxels are those still enabled, while black voxels are disabled voxels.
-To verify, we try to sample features from the grid to a set of sampled points:
Because the disabled voxels will be treated as non-existing, no voxels will contribute to the features on the points at the front. Hence those points are marked as black (i.e. feature = 0).
-
The disabled voxels could be revived at any time using enable_ijk:
-
grid.enable_ijk(disable_ijk)
-
-
-
Conducting the same feature sampling with sample_trilinear, one can get the full features being correctly sampled.
In this example, we will cover a more advanced topic to perform structure optimization from images.
-Suppose we have the following observations of a single object’s mask:
-
-
The task is to recover the underlying 3D shape from the images. Here we use the underlying representation of our VDB grid.
-Such a problem could be solved in many different ways, with one obvious one being space culling.
-However, to demonstrate the wide applicability and flexibility of fVDB, we will take the differentiable rendering approach here.
-Each voxel is hereby given an opacity value, and the entire sparse grid is volume rendered into a predicted mask.
-A simple L1 loss is compared between the given ground-truth mask and the predicted mask to force them align.
-
To begin, we create a mutable grid and the corresponding opacity (alpha) by:
-
importfvdb
-importtorch
-importmath
-
-init_resolution=96
-
-# Suppose our shape lies within the unit bounding box from [-0.5, 0.5, 0.5] to [0.5, 0.5, 0.5]
-grid=fvdb.gridbatch_from_dense(
- num_grids=1,dense_dims=[init_resolution]*3,
- voxel_sizes=[1.0/(init_resolution-1)]*3,origins=[-0.5,-0.5,-0.5],
- device="cuda",
- mutable=True
-)
-
-definv_sigmoid(x)->float:
- return-math.log(1/x-1)
-alpha=torch.full((grid.total_voxels,),inv_sigmoid(0.1),device=grid.device,requires_grad=True)
-
-
-
The structure optimization is done via torch’s Adam optimizer, with the loop being:
-
optimizer=torch.optim.Adam([alpha],lr=1.0)
-
-# Optimization loop
-foritinrange(100):
- # Subsample rays from the given camera poses.
- sub_inds=torch.randint(0,ray_orig.shape[0],(10000,),device=grid.device)
- pd_opacity=render_opacity(
- grid,torch.sigmoid(alpha),
- ray_orig=ray_orig[sub_inds],ray_dir=ray_dir[sub_inds]
- )
- gt_opacity=ray_opacity[sub_inds]
-
- # Compute L1 loss
- loss=torch.mean(torch.abs(pd_opacity-gt_opacity))
-
- optimizer.zero_grad()
- loss.backward()
- optimizer.step()
-
-
-
Here render_opacity is an approximate differentiable rendering algorithm like:
During the optimization, the voxels of the grid could be disabled or enabled freely.
-In this example, we demonstrate the following strategy similar to Instant-NGP.
-
ifit>0andit%5==0:
- # Disable voxels that are transparent
- bad_mask=torch.sigmoid(alpha)<0.1
- grid.disable_ijk(grid.ijk.rmask(bad_mask))
-
- # Randomly revive voxels at the beginning.
- ifit<20:
- enable_mask=torch.rand(grid.total_voxels,device=grid.device)<0.01
- grid.enable_ijk(grid.ijk.rmask(enable_mask))
-
-
-
Note that a way simpler strategy that only turns voxels off at a much sparse internal also works in this very simplified scenario.
-The snippet above is solely for demonstration purpose of our API.
-
The optimization procedure looks as follows. One can see that we are able to recover the ground-truth voxel structure of the provided car.
-
-
A full runnable example could be found at examples/structure_optimization.py.
In this tutorial, you will be guided on how to build a simple sparse convolutional neural network using fVDB.
-If you were using MinkowskiEngine to tackle sparse 3D data previously, we will also guide you step-by-step to help you smoothly transfer from it and enjoy speed-ups and memory-savings.
-
In our simplistic U-Net case, we want to build a Res-UNet with four layers, and each layer contains several blocks.
-First, we import basic fvdb libraries:
Here fvdb.nn is a namespace similar to torch.nn, containing a broad definition of different neural layers.
-VDBTensor is a very thin wrapper around a grid (with type GridBatch) and its corresponding feature (with type JaggedTensor), and internally makes sure that the two members align.
-It also overloads a bunch of operators such as arithmetic computations. Please refer to our API docs to learn more.
All the network layers are fully compatible with torch.nn. The only difference is that they take VDBTensor as input and return a VDBTensor.
-A full network definition could then be built as:
Please note that here, when we apply transposed convolution layers, we additionally introduce the out_grid keyword arguments.
-This is needed to guide the output domain of the network, because for perception networks, the output grid topology should align with the input topology.
-Note that fVDB will NOT cache the grids to maintain maximum flexibility.
-
To perform inference with the network, you could simply create a VDBTensor and feed it into the model:
The output soutput will carry gradients during training, and you could train the sparse network accordingly.
-Please find a fully working example at examples/perception_example.py. The same network is implemented using MinkowskiEngine for reference.
-
