3D gaussian splatting 代码阅读(三)：与pytorch的联合

一、gaussian-splatting/submodules/diff-gaussian-rasterization/ext.cpp

先学习下pytorch c++拓展以及pybind 11

#include <torch/extension.h>
#include "rasterize_points.h"

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("rasterize_gaussians", &RasterizeGaussiansCUDA);
  m.def("rasterize_gaussians_backward", &RasterizeGaussiansBackwardCUDA);
  m.def("mark_visible", &markVisible);
}

先看ext.cpp，这里定义了将RasterizeGaussiansCUDA、RasterizeGaussiansBackwardCUDA、markVisible重定向到python可以识别到的rasterize_gaussians、rasterize_gaussians_backward、mark_visible。

二、gaussian-splatting/submodules/diff-gaussian-rasterization/rasterize_points.cu

这里就是把cuda实现的forward和backward用包了一下，同时输入变成了torch::Tensor，同时也做了一些把torch::Tensor转换成cuda显存指针操作。

2.1 RasterizeGaussiansCUDA(forward)

1. 初始化一些变量：3D gaussian 个数、图像H、W
2. 初始化显存：out_color、radii、geomBuffer、binningBuffer、imgBuffer
   这里是把初始化显存写成了callback函数了：

std::function<char*(size_t)> geomFunc = resizeFunctional(geomBuffer);
std::function<char*(size_t)> binningFunc = resizeFunctional(binningBuffer);
std::function<char*(size_t)> imgFunc = resizeFunctional(imgBuffer);

std::function<char*(size_t N)> resizeFunctional(torch::Tensor& t) {
    auto lambda = [&t](size_t N) {
        t.resize_({(long long)N});
		return reinterpret_cast<char*>(t.contiguous().data_ptr());
    };
    return lambda;
}

写的非常巧妙。

3. 最后调用cuda实现：CudaRasterizer::Rasterizer::forward

RasterizeGaussiansBackwardCUDA(backward)

1. 同样初始化一些tensor
2. 再调用CudaRasterizer::Rasterizer::backward，只不过这里没用callback

三、gaussian-splatting/submodules/diff-gaussian-rasterization/diff_gaussian_rasterization/init.py

这里是对pybind调用的接口做了二次封装，同时继承pytorch自动求导模块，使得可以在pytorch框架中实现自动forward和backward。

3.1 _RasterizeGaussians

3.1.1 forward

在看代码前我们先学习下pytorch中的context
pytroch中ctx和self的区别

ctx可以简单理解为pytorch torch.autograd.Function内部的一个更高级、包含保存和恢复tensor的self。

# Keep relevant tensors for backward
ctx.raster_settings = raster_settings
ctx.num_rendered = num_rendered
ctx.save_for_backward(colors_precomp, means3D, scales, rotations, cov3Ds_precomp, radii, sh, geomBuffer, binningBuffer, imgBuffer)

forward的同时，保存了中间变量用于后面backward。

3.1.2 backward

# Restore necessary values from context
num_rendered = ctx.num_rendered
raster_settings = ctx.raster_settings
colors_precomp, means3D, scales, rotations, cov3Ds_precomp, radii, sh, geomBuffer, binningBuffer, imgBuffer = ctx.saved_tensors

先恢复forward保存的tensor
再调用_C.rasterize_gaussians_backward(*args) pybind生成的c++接口，最后返回梯度。

3.2 rasterize_gaussians

相当与对3.1的RasterizeGaussians又包了一层，直接调用torch.autograd.Function的apply接口，简单理解就是同时跑了forward和backward

3.3 GaussianRasterizer

这个才是最外层python代码会调用到的模块，继承了nn.Module，这里主要看一下forward函数
这里讲一下输入参数：
● means3D：输入3D 高斯的center
● means2D：cuda代码并没有用到
● opacities：不透明度
● shs = None,：球谐参数，如果下面colors_precomp不为空，则shs可以为空
● colors_precomp = None, 颜色如果上面shs不为空，则可以为空，会重新根据shs创建
● scales = None：3D gaussian的eigen value或者轴长，只有cov3D_precomp不为空，他才可以为空
● rotations = None：3D gaussian的旋转，只有cov3D_precomp不为空，他才可以为空
● cov3D_precomp：3D gaussian的协方差，如果为空，可以从scales和rotation中创建

四、gaussian-splatting/submodules/diff-gaussian-rasterization/setup.py

相对与python版本的cmake