DOCSEARCHbyAlgolia
Submit Docs

gsplat

repository·main·Indexed 26 days ago

https://github.com/nerfstudio-project/gsplat

An open-source library providing CUDA-accelerated rasterization for Gaussians with Python bindings. Designed to be faster and more memory-efficient than the original 3D Gaussian Splatting implementation, it includes support for batch rendering, experimental HiGS inference, and NVIDIA 3DGUT for nonlinear camera projections and rolling shutter effects.

Tokens
17.7K
Snippets
28
Records
103
Agent score
90%

What's inside gsplat

  1. Overview of gsplat features

    main

    Core Capabilities

    gsplat is a CUDA-accelerated differentiable rasterization library for 3D Gaussians. Key advantages include:

    • Efficiency: Up to 4x less training memory footprint and 15% less training time compared to the official implementation on Mip-NeRF 360 captures.
    • Large Scale Support: Designed for extremely large scene rendering, significantly faster than the diff-gaussian-rasterization backend.
    • Advanced Rasterization: Supports batch rasterization, N-D feature rendering, depth rendering, sparse gradient, and multi-GPU distributed rasterization.
    • Modern Techniques: Includes support for absgrad, anti-aliasing (Mip-Splatting), and 3DGS-MCMC.

  2. Understand the gsplat Sensors Kernel Layer Architecture

    main

    The gsplat/sensors/kernels/ directory is the backend implementation layer for gsplat/sensors. It manages native CUDA sources for camera and spinning-LiDAR projection, C++ torch::class_<> registrations, Torch op bindings, Python autograd wrappers, and JIT build configurations for the gsplat_sensors_cuda extension.

    Key Architectural Boundaries:

    • Not the Public API: The stateless API resides in gsplat/sensors/functional/.
    • Not the Model Surface: The nn.Module wrappers reside in gsplat/sensors/models/.
    • Kernel Boundary: The kernel boundary only sees unpacked tensors. Named return types are handled in the functional layer, while kernel Python wrappers return raw tuples.

    Data Flow for a Kernel Call:

    1. gsplat.sensors.functional.<op>(...)
    2. kernels.projective_sensor_ops.<op>(...) (cross-family dispatch)
    3. kernels.cameras.ops.<op>(...) (public Python wrapper, pose unpacking)
    4. _<Op>[Bivariate]?.apply(...) (torch.autograd.Function)
    5. torch.ops.gsplat_sensors.<op>_opencv_pinhole_<distortion>(...) (backward pass)
    6. gsplat_sensors::<op>_...(...) (C++ entry in camera_torch.cpp)
    7. <op>_<...>_launch(...) (bridge prototype in camera_params.h)
    8. <op>_<...>_kernel<<<...>>>(...) (CUDA __global__ function)
    9. Device math in camera_kernel.cuh + math.cuh

  3. Use the gsplat functional API for sensors

    main

    The gsplat.sensors.functional module (or from gsplat.sensors import functional as F) provides a stateless public API for camera and spinning-LiDAR projection. This layer is the preferred import path for downstream code. It handles the conversion of raw kernel outputs into structured dataclasses and manages argument contracts.

    Note that the functional layer does not perform its own device checking; it forwards the allow_device_transfer flag to the kernel layer. Differentiability is inherited from the kernel layer, but fields like valid_flag, valid_indices, and timestamps_us are non-differentiable.

  4. Understand the gsplat Shared Module Architecture

    main

    The gsplat library is organized into several specialized modules (e.g., gsplat/geometry, gsplat/sensors, gsplat/scene, gsplat/stage) that follow a standardized four-layer architectural pattern. This pattern separates public stateless APIs from backend CUDA implementations and stateful wrappers.

    Core Layers

    • functional/: The canonical public stateless API surface. This is where you should find and import user-facing functions. It defines argument/shape/dtype contracts and delegates execution to the kernels.
    • kernels/: The backend implementation layer. It handles native extension loading, torch.autograd.Function wrappers, and backend dispatch. It contains the C++/CUDA source code.
    • models/ (Optional): Provides torch.nn.Module wrappers for trainable parameters and higher-level domain structures (like frames or sensor states).
    • components/ (Optional): Provides stateful wrappers that do not inherit from torch.nn.Module (e.g., for managing caches, indexing, or runtime state without parameter registration).

  5. Understand the gsplat geometry module architecture

    main

    The gsplat/geometry module is split into two primary roles to separate the user-facing API from CUDA implementation details:

    1. functional/: The public Python API. It defines the contract for arguments (shapes, dtypes, types), exports operations like quaternion and pose functions, and delegates execution to the backend.
    2. kernels/: The backend implementation layer. It handles loading the gsplat_geometry_cuda extension, defines torch.autograd.Function wrappers, manages native CUDA/C++ sources, and provides autograd glue.

  6. Understand the gsplat Scene Design and Scope

    main

    The gsplat scene architecture is designed around a minimal abstract Scene and a concrete GaussianScene.

    Included in the current scope:

    • A minimal abstract Scene class.
    • A concrete GaussianScene implementation.
    • Raw Gaussian tensor storage within the splats attribute.
    • Optional signal sidecars.
    • Exclusive components managed via component_index.
    • Sidecar update hooks for handling topology changes.
    • Support for AV checkpoint evaluation from saved splats.

    Not currently standardized (limitations):

    • Hierarchical scene access.
    • Scene-owned rendering helpers or optimizer logic.
    • Generalized component mutation APIs (e.g., set_component()).
    • Trainer checkpoint persistence of the full GaussianScene.state_dict().

  8. Understand the geometry module design and structure

    main

    The gsplat.geometry module is organized to separate public functional interfaces from backend implementation details. When extending or using the module, follow these structural constraints:

    • Public API: All public geometry functions are located in the functional/ directory.
    • Backend Implementations: Implementation details and CUDA kernels are located in the kernels/ directory.
    • Shared Headers: Use the include/ directory for shared public or host-side headers.
    • CUDA Helpers: Keep CUDA device helpers within their owning sources under kernels/cuda/csrc/.

    The module is designed to prioritize geometry concepts over implementation details, ensuring a clear distinction between public interfaces and backend-specific code.

  9. Understand CameraModel and LidarModel design constraints

    main

    When working with sensor models in gsplat, adhere to the following design patterns:

    • Pose Handling: CameraModel methods do not accept a Frame object. Instead, they require an explicit pose: Pose or dynamic_pose: DynamicPose along with explicit timestamp keyword arguments. The pose object is passed through the model layer to the functional layer without modification.
    • Immutability: CameraModel.transform returns a new instance of the model; it does not perform in-place mutations.
    • Quaternion Convention: The public Python API uses the wxyz quaternion convention. For compatibility with gsplat/geometry SE(3) primitives, use wxyz_to_xyzw or xyzw_to_wxyz from kernels/common/utils.py.
    • Device Management: Observation tensors on Frame subclasses (e.g., ImageFrame.image, LidarFrame.distance_m, LidarFrame.intensity) are registered as nn.Module buffers. This ensures that calling .to(device) on a frame correctly propagates the tensors to the specified device.
    • Data Types: Pose dataclasses (Pose, DynamicPose, Trajectory) are defined in kernels/common/pose.py, and return-type dataclasses are defined in functional/return_types.py.

  10. Initialize development environment and run tests

    main

    After cloning, run the bootstrap script to initialize the repository and set up the pre-commit formatting hooks.

    To ensure code quality before committing, run the formatting script and execute tests locally. Since GitHub workflows do not support GPUs, local testing is required to verify CUDA unit tests.

    Note: pytest runs all functions prefixed with test_*.

  12. Compress and decompress Gaussian parameters with PngCompression

    main

    Use gsplat.PngCompression to reduce the storage and streaming costs of Gaussian parameters. This method can significantly reduce file sizes (e.g., from 236 MB to 16.5 MB for 1 million Gaussians) with minimal PSNR loss.

    Input Requirements: The compression API expects Gaussian parameters to be provided as a Dict[str, torch.Tensor] containing at least the following keys:

    • means
    • scales
    • quats
    • opacities
    • sh0
    • shN

    Arbitrary extra attributes (e.g., additional feature tensors) are also supported and will be included in the compression process.

    from gsplat import PngCompression
from torch import Tensor
from typing import Dict

# Define Gaussian parameters as a dictionary of tensors
splats: Dict[str, Tensor] = {
    "means": Tensor(N, 3), 
    "scales": Tensor(N, 3), 
    "quats": Tensor(N, 4), 
    "opacities": Tensor(N),
    "sh0": Tensor(N, 1, 3), 
    "shN": Tensor(N, 24, 3), 
    "features1": Tensor(N, 128), 
    "features2": Tensor(N, 64),
}

compression_method = PngCompression()

# Run compression and save the compressed files to compress_dir
compression_method.compress(compress_dir, splats)

# Decompress the compressed files back into tensors
splats_c = compression_method.decompress(compress_dir)