【计算机视觉】Gaussian Splatting源码解读补充（二）

本文是对学习笔记之——3D Gaussian Splatting源码解读的补充，并订正了一些错误。

目录

• 三、相机相关
• • `scene/cameras.py`：`class Camera`
  • 四、前向传播（渲染）：`submodules/diff-gaussian-rasterization/cuda_rasterizer/forward.cu`
  • • 预备知识：CUDA编程基础
    • `def render(viewpoint_camera, pc : GaussianModel, pipe, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None)`
    • CUDA文件`rasterizer_impl.cu`: `submodules/diff-gaussian-rasterization/cuda_rasterizer`
    • • 1. 引用头文件
      • • (1) Cooperative Groups
        • (2) CUB
        • (3) GLM
        • 2. `getHigherMsb`
        • 3. `checkFrustum`和`markVisible`
        • 4. 几个`fromChunk`
        • 5. 为排序做准备：`duplicateWithKeys`
        • 6. 排序之后的处理：`identifyTileRanges`
        • 7. 重头戏：前向传播`forward`
        • CUDA文件`forward.cu`: `submodules/diff-gaussian-rasterization/cuda_rasterizer`
        • • 1. `computeColorFromSH`
          • 2. `computeCov2D`
          • 3. `computeCov3D`
          • 4. `preprocessCUDA`
          • 5. `render`和`renderCUDA`

            ---

            三、相机相关

            scene/cameras.py：class Camera

            class Camera(nn.Module):    def __init__(self, colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask,                 image_name, uid,                 trans=np.array([0.0, 0.0, 0.0]), scale=1.0, data_device = "cuda"                 ):        super(Camera, self).__init__()        self.uid = uid        self.colmap_id = colmap_id        self.R = R # 旋转矩阵        self.T = T # 平移向量        self.FoVx = FoVx # x方向视场角        self.FoVy = FoVy # y方向视场角        self.image_name = image_name        try:            self.data_device = torch.device(data_device)        except Exception as e:            print(e)            print(f"[Warning] Custom device {data_device} failed, fallback to default cuda device" )            self.data_device = torch.device("cuda")        self.original_image = image.clamp(0.0, 1.0).to(self.data_device) # 原始图像        self.image_width = self.original_image.shape[2] # 图像宽度        self.image_height = self.original_image.shape[1] # 图像高度        if gt_alpha_mask is not None:            self.original_image *= gt_alpha_mask.to(self.data_device)        else:            self.original_image *= torch.ones((1, self.image_height, self.image_width), device=self.data_device)# 距离相机平面znear和zfar之间且在视锥内的物体才会被渲染        self.zfar = 100.0 # 最远能看到多远        self.znear = 0.01 # 最近能看到多近        self.trans = trans # 相机中心的平移        self.scale = scale # 相机中心坐标的缩放        self.world_view_transform = torch.tensor(getWorld2View2(R, T, trans, scale)).transpose(0, 1).cuda() # 世界到相机坐标系的变换矩阵，4×4        self.projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy).transpose(0,1).cuda() # 投影矩阵        self.full_proj_transform = (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0) # 从世界坐标系到图像的变换矩阵        self.camera_center = self.world_view_transform.inverse()[3, :3] # 相机在世界坐标系下的坐标

            其中出现的辅助函数：

            # utils/graphic_utils.pydef getWorld2View2(R, t, translate=np.array([.0, .0, .0]), scale=1.0):    Rt = np.zeros((4, 4)) # 按理来说是世界到相机的变换矩阵，但没有加平移和缩放    Rt[:3, :3] = R.transpose()    Rt[:3, 3] = t    Rt[3, 3] = 1.0    C2W = np.linalg.inv(Rt) # 相机到世界的变换矩阵    cam_center = C2W[:3, 3] # 相机中心在世界中的坐标，即C2W矩阵第四列的三维平移向量    cam_center = (cam_center + translate) * scale # 相机中心坐标需要平移和缩放处理    C2W[:3, 3] = cam_center # 重新填入C2W矩阵    Rt = np.linalg.inv(C2W) # 再取逆获得W2C矩阵    return np.float32(Rt)

            四、前向传播（渲染）：submodules/diff-gaussian-rasterization/cuda_rasterizer/forward.cu

            预备知识：CUDA编程基础

            这部分的参考资料：

            [1] CUDA Tutorial

            [2] An Even Easier Introduction to CUDA

            [3] Introduction to CUDA Programming

            CUDA是一个为支持CUDA的GPU提供的平台和编程模型。该平台使GPU能够进行通用计算。CUDA提供了C/C++语言扩展和API，以便用户利用GPU进行高效计算。一般称CPU为host，GPU为device。

            CUDA C++语言中有一个加在函数前面的关键字__global__，用于表明该函数是运行在GPU上的，并且可以被CPU调用。这种函数称为kernel。

            当我们调用kernel的时候，需要在函数名和括号之间加上，其中M是block的个数，T是每个block中线程的个数。这些线程都是并行执行的，每个线程中都在执行该函数。

            根据参考资料[3]，GPU在同一时刻运行一个kernel（也就是一组任务），kernel运行在grid上，每个grid由多个block组成（他们是独立的ALU组），每个block有多个线程。

            同一block中的线程一般合作完成任务，它们可以共享内存（这部分内存速度极快，用__shared__关键字声明）。每个线程“知道”它在哪个block（通过访问内置变量blockIdx.x）和它是第几个线程（通过访问变量threadIdx.x），以及每个block有多少个线程（blockDim.x），从而确定它应该完成什么任务。实际上线程和block的索引是三维的，这里只举了一个一维的例子。

            注意GPU和CPU的内存是隔离的，想要操作显存或者在显存和CPU内存之间进行交流必须显示的声明这些操作。指针也是不一样的，有可能虽然都是int*，但表示的含义却不同：device指针指向显存，host指针指向CPU内存。CUDA提供了操作内存的内置函数：cudaMalloc、cudaFree、cudaMemcpy等，它们分别类似于C函数malloc、free和memcpy。

            关于同步方面，内置函数 __syncthreads()可以同步一个block中的所有线程。在CPU中调用内置函数cudaDeviceSynchronize()可以可以阻塞CPU，直到所有先前发出的CUDA调用都完成为止。

            另外还有__host__关键字和__device__关键字，前者表示该函数只编译成CPU版本（这是默认状态），后者表示只编译成GPU版本。二者同时使用表示编译CPU和GPU两个版本。从CPU调用__device__函数和从GPU调用__host__函数都会报错。

            以下是一个CUDA并行计算向量加法的示例：

            #include int a[10] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};int b[10] = {0, 10, 20, 30, 40, 50, 60, 70, 80, 90};int c[10];__global__ void vec_add(int* A, int* B, int* C){    int i = blockIdx.x * blockDim.x + threadIdx.x;    C[i] = A[i] + B[i];}int main(){    int* gpua, *gpub, *gpuc;    int sz = 10 * sizeof(int);    cudaMalloc((void**)&gpua, sz); // 在GPU中分配内存    cudaMalloc((void**)&gpub, sz);    cudaMalloc((void**)&gpuc, sz);    cudaMemcpy(gpua, a, sz, cudaMemcpyHostToDevice); // 拷贝CPU中的数组到GPU    cudaMemcpy(gpub, b, sz, cudaMemcpyHostToDevice);    vec_add(gpua, gpub, gpuc); // 调用kernel    cudaDeviceSynchronize(); // 等GPU执行完（这里其实不是很有必要）    cudaMemcpy(c, gpuc, sz, cudaMemcpyDeviceToHost); // 把GPU的计算结果拷贝到CPU    for(int i = 0; i  
            有了这些预备知识之后，我们就可以开始看代码了。在看CUDA代码之前，我们先看看gaussian_renderer/__init__.py中的render函数。
             
            def render(viewpoint_camera, pc : GaussianModel, pipe, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None)
             
            def render(viewpoint_camera, pc : GaussianModel, pipe, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None):'''viewpoint_camera: scene.cameras.Camera类的实例pipe: 流水线，规定了要干什么'''    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0    try:        screenspace_points.retain_grad() # 让screenspace_points在反向传播时接受梯度    except:        pass    # Set up rasterization configuration    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)    raster_settings = GaussianRasterizationSettings(        image_height=int(viewpoint_camera.image_height),        image_width=int(viewpoint_camera.image_width),        tanfovx=tanfovx,        tanfovy=tanfovy,        bg=bg_color, # 背景颜色        scale_modifier=scaling_modifier,        viewmatrix=viewpoint_camera.world_view_transform, # W2C矩阵        projmatrix=viewpoint_camera.full_proj_transform, # 世界到图像的投影矩阵        sh_degree=pc.active_sh_degree, # 目前的球谐阶数        campos=viewpoint_camera.camera_center, # CAMera center POSition，相机中心文职        prefiltered=False,        debug=pipe.debug    )    rasterizer = GaussianRasterizer(raster_settings=raster_settings)    means3D = pc.get_xyz    means2D = screenspace_points # 疑似各个Gaussian的中心投影到在图像中的坐标    opacity = pc.get_opacity # 不透明度    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from    # scaling / rotation by the rasterizer.    scales = None    rotations = None    cov3D_precomp = None    if pipe.compute_cov3D_python:        cov3D_precomp = pc.get_covariance(scaling_modifier)    else:        scales = pc.get_scaling        rotations = pc.get_rotation    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.    shs = None    colors_precomp = None    if override_color is None:        if pipe.convert_SHs_python:            shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2)            dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1))            dir_pp_normalized = dir_pp/dir_pp.norm(dim=1, keepdim=True)            # 找到高斯中心到相机的光线在单位球体上的坐标            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)            # 球谐的傅里叶系数转成RGB颜色            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)        else:            shs = pc.get_features    else:        colors_precomp = override_color    # Rasterize visible Gaussians to image, obtain their radii (on screen).     rendered_image, radii = rasterizer(        means3D = means3D,        means2D = means2D,        shs = shs,        colors_precomp = colors_precomp,        opacities = opacity,        scales = scales,        rotations = rotations,        cov3D_precomp = cov3D_precomp)    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.    # They will be excluded from value updates used in the splitting criteria.    return {"render": rendered_image,            "viewspace_points": screenspace_points,            "visibility_filter" : radii > 0,            "radii": radii}
             
            CUDA文件rasterizer_impl.cu: submodules/diff-gaussian-rasterization/cuda_rasterizer
             
            1. 引用头文件
             
            #include "rasterizer_impl.h"#include #include #include #include #include #include "cuda_runtime.h"#include "device_launch_parameters.h"#include  // CUDA的CUB库#include #define GLM_FORCE_CUDA#include  // GLM (OpenGL Mathematics)库#include  // CUDA 9引入的Cooperative Groups库#include namespace cg = cooperative_groups;#include "auxiliary.h"#include "forward.h"#include "backward.h"
             
            有一些库需要我们讲解一下。
             
            (1) Cooperative Groups
             
            参考资料：
             

            • Cooperative Groups: Flexible CUDA Thread Programming
            • 这个库的官方文档

              我们直到 __syncthreads()函数提供了在一个block内同步各线程的方法，但有时要同步block内的一部分线程或者多个block的线程，这时候就需要Cooperative Groups库了。这个库定义了划分和同步一组线程的方法。

              在Gaussian Splatting的所有CUDA代码中，这个库仅以两种方式被调用：

              (a)

              auto idx = cg::this_grid().thread_rank();

              cg::this_grid()返回一个cg::grid_group实例，表示当前线程所处的grid。他有一个方法thread_rank()返回当前线程在该grid中排第几。

              (b)

              auto block = cg::this_thread_block();

              cg::this_thread_block返回一个cg::thread_block实例。代码中用到的成员函数有：

              • block.sync()：同步该block中的所有线程（等价于__syncthreads()）。
              • block.thread_rank()：返回非负整数，表示当前线程在该block中排第几。
              • block.group_index()：返回一个cg::dim3实例，表示该block在grid中的三维索引。
              • block.thread_index()：返回一个cg::dim3实例，表示当前线程在block中的三维索引。

                (2) CUB

                参考资料：CUB API

                CUB provides layered algorithms that correspond to the thread/warp/block/device hierarchy of threads in CUDA. There are distinct algorithms for each layer and higher-level layers build on top of those below.

                换言之，CUB就是针对不同的计算等级：线程、wap、block、device等设计了并行算法。例如，reduce函数有四个版本：ThreadReduce、WarpReduce、BlockReduce、DeviceReduce。

                diff-gaussian-rasterization模块调用了CUB库的两个函数：

                ⭐️ (a) cub::DeviceScan::InclusiveSum

                这个函数是算前缀和的。所谓"Inclusive"就是第 i i i个数被计入第 i i i个和中。函数定义如下：

                templatestatic inline cudaError_t InclusiveSum(void *d_temp_storage, // 额外需要的临时显存空间size_t &temp_storage_bytes, // 临时显存空间的大小InputIteratorT d_in, // 输入指针OutputIteratorT d_out, // 输出指针int num_items, // 元素个数cudaStream_t stream = 0)

                ⭐️ (b) cub::DeviceRadixSort::SortPairs：device级别的并行基数排序

                该函数根据key将(key, value)对进行升序排序。这是一种稳定排序。

                函数定义如下：

                templatestatic inline cudaError_t SortPairs(void *d_temp_storage, // 排序时用到的临时显存空间size_t &temp_storage_bytes, // 临时显存空间的大小const KeyT *d_keys_in,         KeyT *d_keys_out, // key的输入和输出指针const ValueT *d_values_in,     ValueT *d_values_out, // value的输入和输出指针NumItemsT num_items, // 对多少个条目进行排序int begin_bit = 0, // 低位int end_bit = sizeof(KeyT) * 8, // 高位cudaStream_t stream = 0)// 按照[begin_bit, end_bit)内的位进行排序

                示例代码：

                #include #include // or equivalently int main(){    // Declare, allocate, and initialize device-accessible pointers    // for sorting data    int  num_items = 7;    int  keys_in[7] = {8, 6, 7, 5, 3, 0, 9};    int* d_keys_in, *d_keys_out;    int  values_in[7] = {0, 1, 2, 3, 4, 5, 6};    int* d_values_in, *d_values_out;    int keys_out[7], values_out[7];    // 下面把数据搬到显存中    int sz = 7 * sizeof(int);    cudaMalloc((void**)&d_values_in, sz);    cudaMalloc((void**)&d_values_out, sz);    cudaMalloc((void**)&d_keys_in, sz);    cudaMalloc((void**)&d_keys_out, sz);    cudaMemcpy(d_keys_in, keys_in, sz, cudaMemcpyHostToDevice);    cudaMemcpy(d_values_in, values_in, sz, cudaMemcpyHostToDevice);    // Determine temporary device storage requirements    void     *d_temp_storage = NULL;    size_t   temp_storage_bytes = 0;    cub::DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes,        d_keys_in, d_keys_out, d_values_in, d_values_out, num_items);    // Allocate temporary storage    cudaMalloc(&d_temp_storage, temp_storage_bytes);    // Run sorting operation    cub::DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes,        d_keys_in, d_keys_out, d_values_in, d_values_out, num_items);    // 结果：    // d_keys_out                    printf("%d ", keys_out[i]);    }    puts("\nvalues");    for(int i = 0; i  msb)msb += step;elsemsb -= step;}if (n >> msb)msb++;return msb;}

                这个函数似乎是用二分法检测n表示成二进制数除去前导零有几位（n=0时返回1）。我写了一个测试程序推测它的功能：

                #include uint32_t getHigherMsb(uint32_t n);int main(){    uint32_t n[] = {0b0, 0b1, 0b10, 0b11, 0b111, 0b10101, 0b1110001};    for(int i = 0; i  
                4. 几个fromChunk
                 
                这些函数的作用是从以char数组形式存储的二进制块中读取GeometryState、ImageState、BinningState等类的信息。
                 
                CudaRasterizer::GeometryState CudaRasterizer::GeometryState::fromChunk(char*& chunk, size_t P){GeometryState geom;obtain(chunk, geom.depths, P, 128);obtain(chunk, geom.clamped, P * 3, 128);obtain(chunk, geom.internal_radii, P, 128);obtain(chunk, geom.means2D, P, 128);obtain(chunk, geom.cov3D, P * 6, 128);obtain(chunk, geom.conic_opacity, P, 128);obtain(chunk, geom.rgb, P * 3, 128);obtain(chunk, geom.tiles_touched, P, 128);cub::DeviceScan::InclusiveSum(nullptr, geom.scan_size, geom.tiles_touched, geom.tiles_touched, P);obtain(chunk, geom.scanning_space, geom.scan_size, 128);obtain(chunk, geom.point_offsets, P, 128);return geom;}CudaRasterizer::ImageState CudaRasterizer::ImageState::fromChunk(char*& chunk, size_t N){ImageState img;obtain(chunk, img.accum_alpha, N, 128);obtain(chunk, img.n_contrib, N, 128);obtain(chunk, img.ranges, N, 128);return img;}CudaRasterizer::BinningState CudaRasterizer::BinningState::fromChunk(char*& chunk, size_t P){BinningState binning;obtain(chunk, binning.point_list, P, 128);obtain(chunk, binning.point_list_unsorted, P, 128);obtain(chunk, binning.point_list_keys, P, 128);obtain(chunk, binning.point_list_keys_unsorted, P, 128);cub::DeviceRadixSort::SortPairs(nullptr, binning.sorting_size,binning.point_list_keys_unsorted, binning.point_list_keys,binning.point_list_unsorted, binning.point_list, P);obtain(chunk, binning.list_sorting_space, binning.sorting_size, 128);return binning;}
                 
                5. 为排序做准备：duplicateWithKeys
                 
                // Generates one key/value pair for all Gaussian / tile overlaps. // Run once per Gaussian (1:N mapping).__global__ void duplicateWithKeys(int P,const float2* points_xy,const float* depths,const uint32_t* offsets,uint64_t* gaussian_keys_unsorted,uint32_t* gaussian_values_unsorted,int* radii,dim3 grid){auto idx = cg::this_grid().thread_rank(); // 线程索引，该显线程处理第idx个Gaussianif (idx >= P)return;// Generate no key/value pair for invisible Gaussiansif (radii[idx] > 0){// Find this Gaussian's offset in buffer for writing keys/values.uint32_t off = (idx == 0) ? 0 : offsets[idx - 1];uint2 rect_min, rect_max;getRect(points_xy[idx], radii[idx], rect_min, rect_max, grid);// 因为要给Gaussian覆盖的每个tile生成一个(key, value)对，// 所以先获取它占了哪些tile// For each tile that the bounding rect overlaps, emit a // key/value pair. The key is |  tile ID  |      depth      |,// and the value is the ID of the Gaussian. Sorting the values // with this key yields Gaussian IDs in a list, such that they// are first sorted by tile and then by depth. for (int y = rect_min.y; y