A version of rasterization() that utilies on PyTorch's autograd. .. note:: This function still relies on gsplat's CUDA backend for some computation, but the entire differentiable graph is on of PyTorch (and nerfacc) so could use Pytorch's autograd for backpropagation.
(
means: Tensor, # [..., N, 3]
quats: Tensor, # [..., N, 4]
scales: Tensor, # [..., N, 3]
opacities: Tensor, # [..., N]
colors: Tensor, # [..., (C,) N, D] or [..., (C,) N, K, 3]
viewmats: Tensor, # [..., C, 4, 4]
Ks: Tensor, # [..., C, 3, 3]
width: int,
height: int,
near_plane: float = 0.01,
far_plane: float = 1e10,
eps2d: float = 0.3,
sh_degree: Optional[int] = None,
tile_size: int = 16,
backgrounds: Optional[Tensor] = None,
render_mode: Literal["RGB", "D", "ED", "RGB+D", "RGB+ED"] = "RGB",
rasterize_mode: Literal["classic", "antialiased"] = "classic",
channel_chunk: int = 32,
batch_per_iter: int = 100,
)
| 771 | |
| 772 | |
| 773 | def _rasterization( |
| 774 | means: Tensor, # [..., N, 3] |
| 775 | quats: Tensor, # [..., N, 4] |
| 776 | scales: Tensor, # [..., N, 3] |
| 777 | opacities: Tensor, # [..., N] |
| 778 | colors: Tensor, # [..., (C,) N, D] or [..., (C,) N, K, 3] |
| 779 | viewmats: Tensor, # [..., C, 4, 4] |
| 780 | Ks: Tensor, # [..., C, 3, 3] |
| 781 | width: int, |
| 782 | height: int, |
| 783 | near_plane: float = 0.01, |
| 784 | far_plane: float = 1e10, |
| 785 | eps2d: float = 0.3, |
| 786 | sh_degree: Optional[int] = None, |
| 787 | tile_size: int = 16, |
| 788 | backgrounds: Optional[Tensor] = None, |
| 789 | render_mode: Literal["RGB", "D", "ED", "RGB+D", "RGB+ED"] = "RGB", |
| 790 | rasterize_mode: Literal["classic", "antialiased"] = "classic", |
| 791 | channel_chunk: int = 32, |
| 792 | batch_per_iter: int = 100, |
| 793 | ) -> Tuple[Tensor, Tensor, Dict]: |
| 794 | """A version of rasterization() that utilies on PyTorch's autograd. |
| 795 | |
| 796 | .. note:: |
| 797 | This function still relies on gsplat's CUDA backend for some computation, but the |
| 798 | entire differentiable graph is on of PyTorch (and nerfacc) so could use Pytorch's |
| 799 | autograd for backpropagation. |
| 800 | |
| 801 | .. note:: |
| 802 | This function relies on installing latest nerfacc, via: |
| 803 | pip install git+https://github.com/nerfstudio-project/nerfacc |
| 804 | |
| 805 | .. note:: |
| 806 | Compared to rasterization(), this function does not support some arguments such as |
| 807 | `packed`, `sparse_grad` and `absgrad`. |
| 808 | """ |
| 809 | from gsplat.cuda._torch_impl import ( |
| 810 | _fully_fused_projection, |
| 811 | _quat_scale_to_covar_preci, |
| 812 | _rasterize_to_pixels, |
| 813 | ) |
| 814 | |
| 815 | batch_dims = means.shape[:-2] |
| 816 | num_batch_dims = len(batch_dims) |
| 817 | B = math.prod(batch_dims) |
| 818 | N = means.shape[-2] |
| 819 | C = viewmats.shape[-3] |
| 820 | I = B * C |
| 821 | device = means.device |
| 822 | assert means.shape == batch_dims + (N, 3), means.shape |
| 823 | assert quats.shape == batch_dims + (N, 4), quats.shape |
| 824 | assert scales.shape == batch_dims + (N, 3), scales.shape |
| 825 | assert opacities.shape == batch_dims + (N,), opacities.shape |
| 826 | assert viewmats.shape == batch_dims + (C, 4, 4), viewmats.shape |
| 827 | assert Ks.shape == batch_dims + (C, 3, 3), Ks.shape |
| 828 | assert render_mode in ["RGB", "D", "ED", "RGB+D", "RGB+ED"], render_mode |
| 829 | |
| 830 | if sh_degree is None: |
searching dependent graphs…