From b170e175edea410fc7ce75733bf36135eb64503c Mon Sep 17 00:00:00 2001 From: LSong <12553253+lsongx@users.noreply.github.com> Date: Wed, 18 Jan 2023 18:57:55 -0800 Subject: [PATCH] NeRFPlayer sliding window implementation (#1255) * NeRFPlayer sliding window implementation * add comments --- nerfstudio/field_components/cuda/__init__.py | 32 + nerfstudio/field_components/cuda/_backend.py | 57 ++ .../cuda/csrc/include/temporal_gridencoder.h | 62 ++ .../field_components/cuda/csrc/pybind.cu | 14 + .../cuda/csrc/temporal_gridencoder.cu | 667 ++++++++++++++++++ nerfstudio/field_components/temporal_grid.py | 357 ++++++++++ tests/field_components/test_temporal_grid.py | 47 ++ 7 files changed, 1236 insertions(+) create mode 100644 nerfstudio/field_components/cuda/__init__.py create mode 100644 nerfstudio/field_components/cuda/_backend.py create mode 100644 nerfstudio/field_components/cuda/csrc/include/temporal_gridencoder.h create mode 100644 nerfstudio/field_components/cuda/csrc/pybind.cu create mode 100644 nerfstudio/field_components/cuda/csrc/temporal_gridencoder.cu create mode 100644 nerfstudio/field_components/temporal_grid.py create mode 100644 tests/field_components/test_temporal_grid.py diff --git a/nerfstudio/field_components/cuda/__init__.py b/nerfstudio/field_components/cuda/__init__.py new file mode 100644 index 000000000..f92662b85 --- /dev/null +++ b/nerfstudio/field_components/cuda/__init__.py @@ -0,0 +1,32 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""init cuda functions""" +from typing import Callable + + +def _make_lazy_cuda_func(name: str) -> Callable: + """_make_lazy_cuda_func from nerfacc.cuda""" + + def call_cuda(*args, **kwargs): + # pylint: disable=import-outside-toplevel + from ._backend import _C + + return getattr(_C, name)(*args, **kwargs) + + return call_cuda + + +temporal_grid_encode_forward = _make_lazy_cuda_func("temporal_grid_encode_forward") +temporal_grid_encode_backward = _make_lazy_cuda_func("temporal_grid_encode_backward") diff --git a/nerfstudio/field_components/cuda/_backend.py b/nerfstudio/field_components/cuda/_backend.py new file mode 100644 index 000000000..d4269b9ee --- /dev/null +++ b/nerfstudio/field_components/cuda/_backend.py @@ -0,0 +1,57 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""compiling the cuda kernels""" +import glob +import os +import shutil + +from nerfacc.cuda._backend import cuda_toolkit_available +from torch.utils.cpp_extension import _get_build_directory, load + +PATH = os.path.dirname(os.path.abspath(__file__)) +NAME = "nerfstudio_field_components_cuda" +BUILD_DIR = _get_build_directory(NAME, verbose=False) + + +_C = None +if cuda_toolkit_available(): + if os.path.exists(os.path.join(BUILD_DIR, f"{NAME}.so")): + # If the build exists, we assume the extension has been built + # and we can load it. + _C = load( + name=NAME, + sources=glob.glob(os.path.join(PATH, "csrc/*.cu")), + extra_cflags=["-O3", "-std=c++14"], + extra_cuda_cflags=["-O3", "-std=c++14"], + extra_include_paths=[], + ) + else: + # Build from scratch. Remove the build directory just to be safe: pytorch jit might stuck + # if the build directory exists. + shutil.rmtree(BUILD_DIR) + print("nerfstudio field components: Setting up CUDA (This may take a few minutes the first time)") + _C = load( + name=NAME, + sources=glob.glob(os.path.join(PATH, "csrc/*.cu")), + extra_cflags=["-O3", "-std=c++14"], + extra_cuda_cflags=["-O3", "-std=c++14"], + extra_include_paths=[], + ) + print("nerfstudio field components: Setting up CUDA finished") +else: + print("nerfstudio field components: No CUDA toolkit found. Some models may fail.") + + +__all__ = ["_C"] diff --git a/nerfstudio/field_components/cuda/csrc/include/temporal_gridencoder.h b/nerfstudio/field_components/cuda/csrc/include/temporal_gridencoder.h new file mode 100644 index 000000000..7b5ad3d28 --- /dev/null +++ b/nerfstudio/field_components/cuda/csrc/include/temporal_gridencoder.h @@ -0,0 +1,62 @@ +/* + * Code adapted from (@ashawkey) https://github.com/ashawkey/torch-ngp/ + * Author: @lsongx + */ + +#ifndef _TEMPORAL_HASH_ENCODE_H +#define _TEMPORAL_HASH_ENCODE_H + +#include +#include + +// inputs: input coordinates, [B, D], float, in [0, 1] +// temporal_row_index: row index for sampling from channels, [B, 4*num_of_channels], uint32_t +// embeddings: the grid embedding, [sO, grid_C], float +// offsets: offsets for different levels used in NGP, [L + 1], uint32_t +// outputs: interpolated outputs, [B, L * C], float +// grid_C: number of channels for the grid embedding +// B: batch size +// D: coord dim +// L: number of levels +// S: resolution multiplier at each level +// H: base resolution + +void temporal_grid_encode_forward( + const at::Tensor inputs, + const at::Tensor temporal_row_index, + const at::Tensor embeddings, + const at::Tensor offsets, + at::Tensor outputs, + const uint32_t B, + const uint32_t D, + const uint32_t grid_C, + const uint32_t C, + const uint32_t L, + const float S, + const uint32_t H, + at::optional dy_dx, + const uint32_t gridtype, + const bool align_corners +); + +void temporal_grid_encode_backward( + const at::Tensor grad, + const at::Tensor inputs, + const at::Tensor temporal_row_index, + const at::Tensor embeddings, + const at::Tensor offsets, + at::Tensor grad_embeddings, + const uint32_t B, + const uint32_t D, + const uint32_t grid_C, + const uint32_t C, + const uint32_t L, + const float S, + const uint32_t H, + const at::optional dy_dx, + at::optional grad_inputs, + const uint32_t gridtype, + const bool align_corners +); + +#endif \ No newline at end of file diff --git a/nerfstudio/field_components/cuda/csrc/pybind.cu b/nerfstudio/field_components/cuda/csrc/pybind.cu new file mode 100644 index 000000000..877b386ca --- /dev/null +++ b/nerfstudio/field_components/cuda/csrc/pybind.cu @@ -0,0 +1,14 @@ +/* + * Code adapted from (@ashawkey) https://github.com/ashawkey/torch-ngp/ + * Author: @lsongx + */ + + +#include + +#include "include/temporal_gridencoder.h" + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("temporal_grid_encode_forward", &temporal_grid_encode_forward, "temporal_grid_encode_forward (CUDA)"); + m.def("temporal_grid_encode_backward", &temporal_grid_encode_backward, "temporal_grid_encode_backward (CUDA)"); +} \ No newline at end of file diff --git a/nerfstudio/field_components/cuda/csrc/temporal_gridencoder.cu b/nerfstudio/field_components/cuda/csrc/temporal_gridencoder.cu new file mode 100644 index 000000000..b0ca666eb --- /dev/null +++ b/nerfstudio/field_components/cuda/csrc/temporal_gridencoder.cu @@ -0,0 +1,667 @@ +/* + * Code adapted from (@ashawkey) https://github.com/ashawkey/torch-ngp/blob/main/gridencoder/src/gridencoder.cu + * Author: @lsongx + */ + +#include +#include +#include + +#include +#include + +#include +#include + +#include +#include + + +#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor") +#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x) +#define CHECK_IS_INT(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, #x " must be an int tensor") +#define CHECK_IS_FLOATING(x) TORCH_CHECK( \ + x.scalar_type() == at::ScalarType::Float || \ + x.scalar_type() == at::ScalarType::Half || \ + x.scalar_type() == at::ScalarType::Double, #x " must be a floating tensor" \ +) + + +// just for compatability of half precision in AT_DISPATCH_FLOATING_TYPES_AND_HALF... +static inline __device__ at::Half atomicAdd(at::Half *address, at::Half val) { + // requires CUDA >= 10 and ARCH >= 70 + // this is very slow compared to float or __half2, and never used. + // The following line was commented in torch-ngp; uncomment it in case someone want to try fp16... + return atomicAdd(reinterpret_cast<__half*>(address), val); +} + + +template +static inline __host__ __device__ T div_round_up(T val, T divisor) { + return (val + divisor - 1) / divisor; +} + + +template +__device__ uint32_t fast_hash(const uint32_t pos_grid[D]) { + + // coherent type of hashing + constexpr uint32_t primes[7] = { 1u, 2654435761u, 805459861u, 3674653429u, 2097192037u, 1434869437u, 2165219737u }; + + uint32_t result = 0; + #pragma unroll + for (uint32_t i = 0; i < D; ++i) { + result ^= pos_grid[i] * primes[i]; + } + + return result; +} + +// locate the index of grid with channel (ch) +template +__device__ uint32_t get_grid_index( + const uint32_t gridtype, // gridtype: 0 == hash, 1 == tiled + const uint32_t grid_C, // number of total channels for the embedding + const bool align_corners, + const uint32_t ch, // number of output channels + const uint32_t hashmap_size, + const uint32_t resolution, + const uint32_t pos_grid[D] // location of the grid +) { + uint32_t stride = 1; + uint32_t index = 0; + + #pragma unroll + for (uint32_t d = 0; d < D && stride <= hashmap_size; d++) { + index += pos_grid[d] * stride; + stride *= align_corners ? resolution: (resolution + 1); + } + + // NOTE: for NeRF, the hash is in fact not necessary. Check https://github.com/NVlabs/instant-ngp/issues/97. + // gridtype: 0 == hash, 1 == tiled + if (gridtype == 0 && stride > hashmap_size) { + index = fast_hash(pos_grid); + } + + return (index % hashmap_size) * grid_C + ch; +} + +// grid sampling kernel +template +__global__ void kernel_grid( + const float * __restrict__ inputs, // input coordinates + const float * __restrict__ temporal_row_index, // row index for sampling from channels + const scalar_t * __restrict__ grid, // the grid embedding + const int * __restrict__ offsets, // offsets for different levels used in NGP + scalar_t * __restrict__ outputs, + const uint32_t B, // batch size + const uint32_t grid_C, // number of channels for the grid embedding + const uint32_t L, // number of levels + const float S, // resolution multiplier at each level + const uint32_t H, // base resolution + scalar_t * __restrict__ dy_dx, // gradient to the input coords + const uint32_t gridtype, // gridtype: 0 == hash, 1 == tiled + const bool align_corners +) { + const uint32_t b = blockIdx.x * blockDim.x + threadIdx.x; + + if (b >= B) return; + + const uint32_t level = blockIdx.y; + + // locate + grid += (uint32_t)offsets[level] * grid_C; + inputs += b * D; + outputs += level * B * C + b * C; + temporal_row_index += b*C*4; + + // check input range (should be in [0, 1]) + bool flag_oob = false; + #pragma unroll + for (uint32_t d = 0; d < D; d++) { + if (inputs[d] < 0 || inputs[d] > 1) { + flag_oob = true; + } + } + // if input out of bound, just set output to 0 + if (flag_oob) { + #pragma unroll + for (uint32_t ch = 0; ch < C; ch++) { + outputs[ch] = 0; + } + if (dy_dx) { + dy_dx += b * D * L * C + level * D * C; // B L D C + #pragma unroll + for (uint32_t d = 0; d < D; d++) { + #pragma unroll + for (uint32_t ch = 0; ch < C; ch++) { + dy_dx[d * C + ch] = 0; + } + } + } + return; + } + + const uint32_t hashmap_size = offsets[level + 1] - offsets[level]; + const float scale = exp2f(level * S) * H - 1.0f; + const uint32_t resolution = (uint32_t)ceil(scale) + 1; + + // calculate coordinate + float pos[D]; + uint32_t pos_grid[D]; + + #pragma unroll + for (uint32_t d = 0; d < D; d++) { + pos[d] = inputs[d] * scale + (align_corners ? 0.0f : 0.5f); + pos_grid[d] = floorf(pos[d]); + pos[d] -= (float)pos_grid[d]; + } + + // interpolate + scalar_t results[C] = {0}; // temp results in register + + #pragma unroll + for (uint32_t idx = 0; idx < (1 << D); idx++) { + float w = 1; + uint32_t pos_grid_local[D]; + + #pragma unroll + for (uint32_t d = 0; d < D; d++) { + if ((idx & (1 << d)) == 0) { + w *= 1 - pos[d]; + pos_grid_local[d] = pos_grid[d]; + } else { + w *= pos[d]; + pos_grid_local[d] = pos_grid[d] + 1; + } + } + + // writing to register (fast) + #pragma unroll + for (uint32_t ch = 0; ch < C; ch++) { + if (temporal_row_index[ch*4] == 1) { + // if a temporal_row_index == 1, then no interpolation between channels are used + uint32_t index = get_grid_index(gridtype, grid_C, align_corners, + __float2uint_rn(temporal_row_index[ch*4+1]), hashmap_size, resolution, pos_grid_local); + results[ch] += w * grid[index]; + } else { + // else, we need to interpolate between channels + uint32_t index_a = get_grid_index(gridtype, grid_C, align_corners, + __float2uint_rn(temporal_row_index[ch*4+1]), hashmap_size, resolution, pos_grid_local); + uint32_t index_b = get_grid_index(gridtype, grid_C, align_corners, + __float2uint_rn(temporal_row_index[ch*4+3]), hashmap_size, resolution, pos_grid_local); + results[ch] += w * (grid[index_a]*temporal_row_index[ch*4]+grid[index_b]*temporal_row_index[ch*4+2]); + } + } + } + + // writing to global memory (slow) + #pragma unroll + for (uint32_t ch = 0; ch < C; ch++) { + outputs[ch] = results[ch]; + } + + // prepare dy_dx + // differentiable (soft) indexing: https://discuss.pytorch.org/t/differentiable-indexing/17647/9 + if (dy_dx) { + + dy_dx += b * D * L * C + level * D * C; // B L D C + + #pragma unroll + for (uint32_t gd = 0; gd < D; gd++) { + + scalar_t results_grad[C] = {0}; + + #pragma unroll + for (uint32_t idx = 0; idx < (1 << (D - 1)); idx++) { + float w = scale; + uint32_t pos_grid_local[D]; + + #pragma unroll + for (uint32_t nd = 0; nd < D - 1; nd++) { + const uint32_t d = (nd >= gd) ? (nd + 1) : nd; + + if ((idx & (1 << nd)) == 0) { + w *= 1 - pos[d]; + pos_grid_local[d] = pos_grid[d]; + } else { + w *= pos[d]; + pos_grid_local[d] = pos_grid[d] + 1; + } + } + + scalar_t left[C] = {0}; + scalar_t right[C] = {0}; + pos_grid_local[gd] = pos_grid[gd]; + #pragma unroll + for (uint32_t ch = 0; ch < C; ch++) { + if (temporal_row_index[ch*4] == 1) { + uint32_t index = get_grid_index(gridtype, grid_C, align_corners, + __float2uint_rn(temporal_row_index[ch*4+1]), hashmap_size, resolution, pos_grid_local); + left[ch] += grid[index]; + } else { + uint32_t index_a = get_grid_index(gridtype, grid_C, align_corners, + __float2uint_rn(temporal_row_index[ch*4+1]), hashmap_size, resolution, pos_grid_local); + uint32_t index_b = get_grid_index(gridtype, grid_C, align_corners, + __float2uint_rn(temporal_row_index[ch*4+3]), hashmap_size, resolution, pos_grid_local); + left[ch] += grid[index_a]*temporal_row_index[ch*4]+grid[index_b]*temporal_row_index[ch*4+2]; + } + } + + pos_grid_local[gd] = pos_grid[gd] + 1; + #pragma unroll + for (uint32_t ch = 0; ch < C; ch++) { + if (temporal_row_index[ch*4] == 1) { + uint32_t index = get_grid_index(gridtype, grid_C, align_corners, + __float2uint_rn(temporal_row_index[ch*4+1]), hashmap_size, resolution, pos_grid_local); + right[ch] += grid[index]; + } else { + uint32_t index_a = get_grid_index(gridtype, grid_C, align_corners, + __float2uint_rn(temporal_row_index[ch*4+1]), hashmap_size, resolution, pos_grid_local); + uint32_t index_b = get_grid_index(gridtype, grid_C, align_corners, + __float2uint_rn(temporal_row_index[ch*4+3]), hashmap_size, resolution, pos_grid_local); + right[ch] += grid[index_a]*temporal_row_index[ch*4]+grid[index_b]*temporal_row_index[ch*4+2]; + } + } + + #pragma unroll + for (uint32_t ch = 0; ch < C; ch++) { + results_grad[ch] += w * (right[ch] - left[ch]); + } + } + + #pragma unroll + for (uint32_t ch = 0; ch < C; ch++) { + dy_dx[gd * C + ch] = results_grad[ch]; + } + } + } +} + +// grid backward kernel; N_C is kept here though always set to 1 (maybe useful for future updates) +template +__global__ void kernel_grid_backward( + const scalar_t * __restrict__ grad, + const float * __restrict__ inputs, + const float * __restrict__ temporal_row_index, + const scalar_t * __restrict__ grid, + const int * __restrict__ offsets, + scalar_t * __restrict__ grad_grid, + const uint32_t B, + const uint32_t grid_C, + const uint32_t L, + const float S, + const uint32_t H, + const uint32_t gridtype, + const bool align_corners +) { + const uint32_t b = (blockIdx.x * blockDim.x + threadIdx.x) * N_C / C; + if (b >= B) return; + + const uint32_t level = blockIdx.y; + // embed channel (from the outputs) + const uint32_t ch = (blockIdx.x * blockDim.x + threadIdx.x) * N_C - b * C; + + // locate + grad_grid += offsets[level] * grid_C; + inputs += b * D; + grad += level * B * C + b * C + ch; // L, B, C + temporal_row_index += b*C*4 + ch*4; + + const uint32_t hashmap_size = offsets[level + 1] - offsets[level]; + const float scale = exp2f(level * S) * H - 1.0f; + const uint32_t resolution = (uint32_t)ceil(scale) + 1; + + // check input range (should be in [0, 1]) + #pragma unroll + for (uint32_t d = 0; d < D; d++) { + if (inputs[d] < 0 || inputs[d] > 1) { + return; // grad is init as 0, so we simply return. + } + } + + // calculate coordinate + float pos[D]; + uint32_t pos_grid[D]; + + #pragma unroll + for (uint32_t d = 0; d < D; d++) { + pos[d] = inputs[d] * scale + (align_corners ? 0.0f : 0.5f); + pos_grid[d] = floorf(pos[d]); + pos[d] -= (float)pos_grid[d]; + } + + scalar_t grad_cur = grad[0]; + + // interpolate + #pragma unroll + for (uint32_t idx = 0; idx < (1 << D); idx++) { + float w = 1; + uint32_t pos_grid_local[D]; + + #pragma unroll + for (uint32_t d = 0; d < D; d++) { + if ((idx & (1 << d)) == 0) { + w *= 1 - pos[d]; + pos_grid_local[d] = pos_grid[d]; + } else { + w *= pos[d]; + pos_grid_local[d] = pos_grid[d] + 1; + } + } + + // get the weights used for calculating the interpolation + if (temporal_row_index[0] == 1) { + // if no interpolation, directly return grad + uint32_t index = get_grid_index(gridtype, grid_C, align_corners, + __float2uint_rn(temporal_row_index[1]), hashmap_size, resolution, pos_grid_local); + atomicAdd(&grad_grid[index], w*grad_cur); + } else { + // if interpolation used, return with weights multiplied + uint32_t index_a = get_grid_index(gridtype, grid_C, align_corners, + __float2uint_rn(temporal_row_index[1]), hashmap_size, resolution, pos_grid_local); + uint32_t index_b = get_grid_index(gridtype, grid_C, align_corners, + __float2uint_rn(temporal_row_index[3]), hashmap_size, resolution, pos_grid_local); + atomicAdd(&grad_grid[index_a], temporal_row_index[0]*w*grad_cur); + atomicAdd(&grad_grid[index_b], temporal_row_index[2]*w*grad_cur); + } + } +} + +// coord backward kernel +template +__global__ void kernel_input_backward( + const scalar_t * __restrict__ grad, + const scalar_t * __restrict__ dy_dx, + scalar_t * __restrict__ grad_inputs, + uint32_t B, + uint32_t L +) { + const uint32_t t = threadIdx.x + blockIdx.x * blockDim.x; + if (t >= B * D) return; + + const uint32_t b = t / D; + const uint32_t d = t - b * D; + + dy_dx += b * L * D * C; + + scalar_t result = 0; + + # pragma unroll + for (int l = 0; l < L; l++) { + # pragma unroll + for (int ch = 0; ch < C; ch++) { + result += grad[l * B * C + b * C + ch] * dy_dx[l * D * C + d * C + ch]; + } + } + + grad_inputs[t] = result; +} + + +template +void kernel_grid_wrapper( + const float *inputs, + const float *temporal_row_index, + const scalar_t *embeddings, + const int *offsets, + scalar_t *outputs, + const uint32_t B, + const uint32_t grid_C, + const uint32_t C, + const uint32_t L, + const float S, + const uint32_t H, + scalar_t *dy_dx, + const uint32_t gridtype, + const bool align_corners +) { + static constexpr uint32_t N_THREAD = 512; + const dim3 blocks_hashgrid = { div_round_up(B, N_THREAD), L, 1 }; + switch (C) { + case 1: kernel_grid<<>>( + inputs, temporal_row_index, embeddings, offsets, outputs, + B, grid_C, L, S, H, dy_dx, gridtype, align_corners); break; + case 2: kernel_grid<<>>( + inputs, temporal_row_index, embeddings, offsets, outputs, + B, grid_C, L, S, H, dy_dx, gridtype, align_corners); break; + case 4: kernel_grid<<>>( + inputs, temporal_row_index, embeddings, offsets, outputs, + B, grid_C, L, S, H, dy_dx, gridtype, align_corners); break; + case 8: kernel_grid<<>>( + inputs, temporal_row_index, embeddings, offsets, outputs, + B, grid_C, L, S, H, dy_dx, gridtype, align_corners); break; + default: throw std::runtime_error{"GridEncoding: C must be 1, 2, 4, or 8."}; + } +} + +template +void temporal_grid_encode_forward_cuda( + const float *inputs, + const float *temporal_row_index, + const scalar_t *embeddings, + const int *offsets, + scalar_t *outputs, + const uint32_t B, + const uint32_t D, + const uint32_t grid_C, + const uint32_t C, + const uint32_t L, + const float S, + const uint32_t H, + scalar_t *dy_dx, + const uint32_t gridtype, + const bool align_corners +) { + switch (D) { + case 1: kernel_grid_wrapper( + inputs, temporal_row_index, embeddings, offsets, outputs, + B, grid_C, C, L, S, H, dy_dx, gridtype, align_corners); break; + case 2: kernel_grid_wrapper( + inputs, temporal_row_index, embeddings, offsets, outputs, + B, grid_C, C, L, S, H, dy_dx, gridtype, align_corners); break; + case 3: kernel_grid_wrapper( + inputs, temporal_row_index, embeddings, offsets, outputs, + B, grid_C, C, L, S, H, dy_dx, gridtype, align_corners); break; + case 4: kernel_grid_wrapper( + inputs, temporal_row_index, embeddings, offsets, outputs, + B, grid_C, C, L, S, H, dy_dx, gridtype, align_corners); break; + case 5: kernel_grid_wrapper( + inputs, temporal_row_index, embeddings, offsets, outputs, + B, grid_C, C, L, S, H, dy_dx, gridtype, align_corners); break; + default: throw std::runtime_error{"GridEncoding: D must be 1, 2, 3, 4, or 5."}; + } +} + +template +void kernel_grid_backward_wrapper( + const scalar_t *grad, + const float *inputs, + const float *temporal_row_index, + const scalar_t *embeddings, + const int *offsets, + scalar_t *grad_embeddings, + const uint32_t B, + const uint32_t grid_C, + const uint32_t C, + const uint32_t L, + const float S, + const uint32_t H, + scalar_t *dy_dx, + scalar_t *grad_inputs, + const uint32_t gridtype, + const bool align_corners +) { + static constexpr uint32_t N_THREAD = 256; + const dim3 blocks_hashgrid = { div_round_up(B * C, N_THREAD), L, 1 }; + switch (C) { + case 1: + kernel_grid_backward<<>>( + grad, inputs, temporal_row_index, + embeddings, offsets, grad_embeddings, B, grid_C, L, S, H, gridtype, align_corners); + if (dy_dx) { + kernel_input_backward<<>>( + grad, dy_dx, grad_inputs, B, L); + } + break; + case 2: + kernel_grid_backward<<>>( + grad, inputs, temporal_row_index, + embeddings, offsets, grad_embeddings, B, grid_C, L, S, H, gridtype, align_corners); + if (dy_dx) { + kernel_input_backward<<>>( + grad, dy_dx, grad_inputs, B, L); + } + break; + case 4: + kernel_grid_backward<<>>( + grad, inputs, temporal_row_index, + embeddings, offsets, grad_embeddings, B, grid_C, L, S, H, gridtype, align_corners); + if (dy_dx) { + kernel_input_backward<<>>( + grad, dy_dx, grad_inputs, B, L); + } + break; + case 8: + kernel_grid_backward<<>>( + grad, inputs, temporal_row_index, + embeddings, offsets, grad_embeddings, B, grid_C, L, S, H, gridtype, align_corners); + if (dy_dx) { + kernel_input_backward<<>>( + grad, dy_dx, grad_inputs, B, L); + } + break; + default: throw std::runtime_error{"GridEncoding: C must be 1, 2, 4, or 8."}; + } +} + +template +void temporal_grid_encode_backward_cuda( + const scalar_t *grad, + const float *inputs, + const float *temporal_row_index, + const scalar_t *embeddings, + const int *offsets, + scalar_t *grad_embeddings, + const uint32_t B, + const uint32_t D, + const uint32_t grid_C, + const uint32_t C, + const uint32_t L, + const float S, + const uint32_t H, + scalar_t *dy_dx, + scalar_t *grad_inputs, + const uint32_t gridtype, + const bool align_corners +) { + switch (D) { + case 1: kernel_grid_backward_wrapper( + grad, inputs, temporal_row_index, embeddings, offsets, + grad_embeddings, B, grid_C, C, L, S, H, dy_dx, grad_inputs, gridtype, align_corners); break; + case 2: kernel_grid_backward_wrapper( + grad, inputs, temporal_row_index, embeddings, offsets, + grad_embeddings, B, grid_C, C, L, S, H, dy_dx, grad_inputs, gridtype, align_corners); break; + case 3: kernel_grid_backward_wrapper( + grad, inputs, temporal_row_index, embeddings, offsets, + grad_embeddings, B, grid_C, C, L, S, H, dy_dx, grad_inputs, gridtype, align_corners); break; + case 4: kernel_grid_backward_wrapper( + grad, inputs, temporal_row_index, embeddings, offsets, + grad_embeddings, B, grid_C, C, L, S, H, dy_dx, grad_inputs, gridtype, align_corners); break; + case 5: kernel_grid_backward_wrapper( + grad, inputs, temporal_row_index, embeddings, offsets, + grad_embeddings, B, grid_C, C, L, S, H, dy_dx, grad_inputs, gridtype, align_corners); break; + default: throw std::runtime_error{"GridEncoding: D must be 1, 2, 3, 4, or 5."}; + } +} + +void temporal_grid_encode_forward( + const at::Tensor inputs, + const at::Tensor temporal_row_index, + const at::Tensor embeddings, + const at::Tensor offsets, + at::Tensor outputs, + const uint32_t B, + const uint32_t D, + const uint32_t grid_C, + const uint32_t C, + const uint32_t L, + const float S, + const uint32_t H, + at::optional dy_dx, + const uint32_t gridtype, + const bool align_corners +) { + CHECK_INPUT(inputs); + CHECK_INPUT(temporal_row_index); + CHECK_INPUT(embeddings); + CHECK_INPUT(offsets); + CHECK_INPUT(outputs); + // CHECK_INPUT(dy_dx); + + CHECK_IS_FLOATING(inputs); + CHECK_IS_FLOATING(temporal_row_index); + CHECK_IS_FLOATING(embeddings); + CHECK_IS_INT(offsets); + CHECK_IS_FLOATING(outputs); + // CHECK_IS_FLOATING(dy_dx); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF( + embeddings.scalar_type(), "temporal_grid_encode_forward", ([&] { + temporal_grid_encode_forward_cuda( + inputs.data_ptr(), temporal_row_index.data_ptr(), embeddings.data_ptr(), + offsets.data_ptr(), outputs.data_ptr(), B, D, grid_C, C, L, S, H, + dy_dx.has_value() ? dy_dx.value().data_ptr() : nullptr, + gridtype, align_corners); + })); +} + +void temporal_grid_encode_backward( + const at::Tensor grad, + const at::Tensor inputs, + const at::Tensor temporal_row_index, + const at::Tensor embeddings, + const at::Tensor offsets, + at::Tensor grad_embeddings, + const uint32_t B, + const uint32_t D, + const uint32_t grid_C, + const uint32_t C, + const uint32_t L, + const float S, + const uint32_t H, + const at::optional dy_dx, + at::optional grad_inputs, + const uint32_t gridtype, + const bool align_corners +) { + CHECK_INPUT(grad); + CHECK_INPUT(inputs); + CHECK_INPUT(temporal_row_index); + CHECK_INPUT(embeddings); + CHECK_INPUT(offsets); + CHECK_INPUT(grad_embeddings); + // CHECK_INPUT(dy_dx); + // CHECK_INPUT(grad_inputs); + + CHECK_IS_FLOATING(grad); + CHECK_IS_FLOATING(inputs); + CHECK_IS_FLOATING(temporal_row_index); + CHECK_IS_FLOATING(embeddings); + CHECK_IS_INT(offsets); + CHECK_IS_FLOATING(grad_embeddings); + // CHECK_IS_FLOATING(dy_dx); + // CHECK_IS_FLOATING(grad_inputs); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF( + grad.scalar_type(), "temporal_grid_encode_backward", ([&] { + temporal_grid_encode_backward_cuda( + grad.data_ptr(), inputs.data_ptr(), temporal_row_index.data_ptr(), + embeddings.data_ptr(), offsets.data_ptr(), grad_embeddings.data_ptr(), + B, D, grid_C, C, L, S, H, + dy_dx.has_value() ? dy_dx.value().data_ptr() : nullptr, + grad_inputs.has_value() ? grad_inputs.value().data_ptr() : nullptr, + gridtype, align_corners); + })); +} diff --git a/nerfstudio/field_components/temporal_grid.py b/nerfstudio/field_components/temporal_grid.py new file mode 100644 index 000000000..c4f5e20c5 --- /dev/null +++ b/nerfstudio/field_components/temporal_grid.py @@ -0,0 +1,357 @@ +# Copyright 2022 The Nerfstudio Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Implements the temporal grid used by NeRFPlayer (https://arxiv.org/abs/2210.15947). +A time conditioned sliding window is applied on the feature channels, so +that the feature vectors become time-aware. +(A large) Part of the code are adapted from (@ashawkey) https://github.com/ashawkey/torch-ngp/ +""" +from typing import Optional + +import numpy as np +import torch +from torch import nn +from torch.autograd import Function +from torch.cuda.amp import custom_bwd, custom_fwd +from torchtyping import TensorType + +import nerfstudio.field_components.cuda as _C + + +# pylint: disable=abstract-method, arguments-differ +class TemporalGridEncodeFunc(Function): + """Class for autograd in pytorch.""" + + @staticmethod + @custom_fwd + def forward( + ctx, + inputs: TensorType["bs", "input_dim"], + temporal_row_index: TensorType["bs", "temporal_index_dim"], + embeddings: TensorType["table_size", "embed_dim"], + offsets: TensorType["num_levels+1"], + per_level_scale: float, + base_resolution: int, + calc_grad_inputs: bool = False, + gridtype: int = 0, + align_corners: bool = False, + ) -> TensorType["bs", "output_dim"]: + """Call forward and interpolate the feature from embeddings + + Args: + inputs: the input coords + temporal_row_index: the input index of channels for doing the interpolation + embeddings: the saved (hashing) table for the feature grid (of the full sequence) + offsets: offsets for each level in the multilevel table, used for locating in cuda kernels + per_level_scale: scale parameter for the table; same as InstantNGP + base_resolution: base resolution for the table; same as InstantNGP + calc_grad_inputs: bool indicator for calculating gradients on the inputs + gridtype: 0 == hash, 1 == tiled; tiled is a baseline in InstantNGP (not random collision) + align_corners: same as other interpolation operators + """ + + inputs = inputs.contiguous() + temporal_row_index = temporal_row_index.contiguous() + + B, D = inputs.shape # batch size, coord dim + L = offsets.shape[0] - 1 # level + grid_channel = embeddings.shape[1] # embedding dim for each level + C = temporal_row_index.shape[1] // 4 # output embedding dim for each level + S = np.log2(per_level_scale) # resolution multiplier at each level, apply log2 for later CUDA exp2f + H = base_resolution # base resolution + + # torch.half used by torch-ngp, but we disable it + # (could be of negative impact on the performance? not sure, but feel free to inform me and help improve it!) + # # manually handle autocast (only use half precision embeddings, inputs must be float for enough precision) + # # if C % 2 != 0, force float, since half for atomicAdd is very slow. + # if torch.is_autocast_enabled() and C % 2 == 0: + # embeddings = embeddings.to(torch.half) + + # L first, optimize cache for cuda kernel, but needs an extra permute later + outputs = torch.empty(L, B, C, device=inputs.device, dtype=embeddings.dtype) + + if calc_grad_inputs: + dy_dx = torch.empty(B, L * D * C, device=inputs.device, dtype=embeddings.dtype) + else: + dy_dx = None + + _C.temporal_grid_encode_forward( + inputs, + temporal_row_index, + embeddings, + offsets, + outputs, + B, + D, + grid_channel, + C, + L, + S, + H, + dy_dx, + gridtype, + align_corners, + ) + + # permute back to [B, L * C] + outputs = outputs.permute(1, 0, 2).reshape(B, L * C) + + ctx.save_for_backward(inputs, temporal_row_index, embeddings, offsets, dy_dx) + ctx.dims = [B, D, grid_channel, C, L, S, H, gridtype] + ctx.align_corners = align_corners + + return outputs + + @staticmethod + @custom_bwd + def backward(ctx, grad): + + inputs, temporal_row_index, embeddings, offsets, dy_dx = ctx.saved_tensors + B, D, grid_channel, C, L, S, H, gridtype = ctx.dims + align_corners = ctx.align_corners + + # grad: [B, L * C] --> [L, B, C] + grad = grad.view(B, L, C).permute(1, 0, 2).contiguous() + + grad_embeddings = torch.zeros_like(embeddings).contiguous() + + if dy_dx is not None: + grad_inputs = torch.zeros_like(inputs, dtype=embeddings.dtype) + else: + grad_inputs = None + + _C.temporal_grid_encode_backward( + grad, + inputs, + temporal_row_index, + embeddings, + offsets, + grad_embeddings, + B, + D, + grid_channel, + C, + L, + S, + H, + dy_dx, + grad_inputs, + gridtype, + align_corners, + ) + + if grad_inputs is not None: + grad_inputs = grad_inputs.to(inputs.dtype) + + return grad_inputs, None, grad_embeddings, None, None, None, None, None, None + + +class TemporalGridEncoder(nn.Module): + """Class for temporal grid encoding. + This class extends the grid encoding (from InstantNGP) by allowing the output time-dependent feature channels. + For example, for time 0 the interpolation uses channels [0,1], then for time 1 channels [2,1] are used. + This operation can be viewed as applying a time-dependent sliding window on the feature channels. + + Args: + temporal_dim: the dimension of temporal modeling; a higher dim indicates a higher freq on the time axis + input_dim: the dimension of input coords + num_levels: number of levels for multi-scale hashing; same as InstantNGP + level_dim: the dim of output feature vector for each level; same as InstantNGP + per_level_scale: scale factor; same as InstantNGP + base_resolution: base resolution for the table; same as InstantNGP + log2_hashmap_size: the size of the table; same as InstantNGP + desired_resolution: desired resolution at the last level; same as InstantNGP + gridtype: "tiled" or "hash" + align_corners: same as other interpolation operators + """ + + def __init__( + self, + temporal_dim: int = 64, + input_dim: int = 3, + num_levels: int = 16, + level_dim: int = 2, + per_level_scale: int = 2, + base_resolution: int = 16, + log2_hashmap_size: int = 19, + desired_resolution: Optional[int] = None, + gridtype: str = "hash", + align_corners: bool = False, + ) -> None: + super().__init__() + + # the finest resolution desired at the last level, if provided, overridee per_level_scale + if desired_resolution is not None: + per_level_scale = np.exp2(np.log2(desired_resolution / base_resolution) / (num_levels - 1)) + + self.temporal_dim = temporal_dim + self.input_dim = input_dim # coord dims, 2 or 3 + self.num_levels = num_levels # num levels, each level multiply resolution by 2 + self.level_dim = level_dim # encode channels per level + self.per_level_scale = per_level_scale # multiply resolution by this scale at each level. + self.log2_hashmap_size = log2_hashmap_size + self.base_resolution = base_resolution + self.output_dim = num_levels * level_dim + self.gridtype = gridtype + _gridtype_to_id = {"hash": 0, "tiled": 1} + self.gridtype_id = _gridtype_to_id[gridtype] # "tiled" or "hash" + self.align_corners = align_corners + + # allocate parameters + offsets = [] + offset = 0 + self.max_params = 2**log2_hashmap_size + for i in range(num_levels): + resolution = int(np.ceil(base_resolution * per_level_scale**i)) + params_in_level = min( + self.max_params, (resolution if align_corners else resolution + 1) ** input_dim + ) # limit max number + params_in_level = int(np.ceil(params_in_level / 8) * 8) # make divisible + offsets.append(offset) + offset += params_in_level + offsets.append(offset) + offsets = torch.from_numpy(np.array(offsets, dtype=np.int32)) + self.register_buffer("offsets", offsets) + self.n_params = offsets[-1] * level_dim + self.embeddings = nn.Parameter(torch.empty(offset, level_dim + temporal_dim)) + self.init_parameters() + + def init_parameters(self) -> None: + """Initialize the parameters: + 1. Uniform initialization of the embeddings + 2. Temporal interpolation index initialization: + For each temporal dim, we initialize a interpolation candidate. + For example, if temporal dim 0, we use channels [0,1,2,3], then for temporal dim 1, + we use channels [4,1,2,3]. After that, temporal dim 2, we use channels [4,5,2,3]. + This is for the alignment of the channels. I.e., each temporal dim should differ + on only one channel, otherwise moving from one temporal dim to the next one is not + that consistent. + To associate time w.r.t. temporal dim, we evenly distribute time into the temporal dims. + That is, if we have 16 temporal dims, then the 16th channel combinations is the time 1. + (Time should be within 0 and 1.) Given a time, we first look up which temporal dim should + be used. And then compute the linear combination weights. + For implementing it, a table for all possible channel combination are used. Each row in + the table is the candidate feature channels, and means we move from one temporal dim to + the next one. For example, the first row will use feature channels [0,1,2,3,4]. Each row + is of length `num_of_output_channel*4`. The expanding param 4 is for saving the combination + weights and channels. The first row will be [?,0,?,1, 1,2,0,0, 1,3,0,0, 1,4,0,0]. Each + 4 tuple means + `[weight_for_channel_A, index_for_channel_A, weight_for_channel_B, index_for_channel_B]` + If `weight_for_channel_A` is 1, then there is no interpolation on this channel. + """ + std = 1e-4 + self.embeddings.data.uniform_(-std, std) + # generate sampling index + temporal_grid_rows = self.temporal_dim + index_init = [0, self.level_dim] + list(range(1, self.level_dim)) + permute_base = list(range(2, self.level_dim + 1)) + last_entry = 0 # insert into ith place + permute_init = permute_base[:last_entry] + [0] + permute_base[last_entry:] + index_list = [torch.as_tensor(index_init, dtype=torch.long)] + permute_list = [torch.as_tensor(permute_init, dtype=torch.long)] + + # converts a list of channel candidates into sampling row + def to_sampling_index(index, permute, last_entry): + row = index[permute] + row = torch.stack([torch.ones_like(row), row, torch.zeros_like(row), torch.zeros_like(row)], 1) + row = row.reshape([-1]) + mask_a = torch.zeros_like(row).bool() + mask_b = torch.zeros_like(row).bool() + row[last_entry * 4 + 3] = index[1] + mask_a[last_entry * 4] = 1 + mask_b[last_entry * 4 + 2] = 1 + return row, mask_a, mask_b + + row, mask_a, mask_b = to_sampling_index(index_list[0], permute_list[0], last_entry) + sampling_index = [row] + index_a_mask, index_b_mask = [mask_a], [mask_b] + # iterate on all temporal grid to get all rows + for _ in range(1, temporal_grid_rows - 1): + # the following lines are a little confusing... + # the basic idea is to keep a buffer and then move to the next channel + last_entry += 1 + if last_entry >= self.level_dim: + last_entry = 0 + last_index_max = index_list[-1].max().item() + last_index_min = index_list[-1].min().item() + tem_permute_list = permute_list[-1].clone() # for rearrange + tem_permute_list[tem_permute_list == 0] += 1 + prev = index_list[-1][1:][tem_permute_list - 1].tolist() + prev.pop(last_entry) + new_index = [last_index_min + 1, last_index_max + 1] + prev + new_index = torch.as_tensor(new_index, dtype=torch.long) + new_permute = permute_base[:last_entry] + [0] + permute_base[last_entry:] + new_permute = torch.as_tensor(new_permute, dtype=torch.long) + index_list.append(torch.as_tensor(new_index, dtype=torch.long)) + permute_list.append(torch.as_tensor(new_permute, dtype=torch.long)) + row, mask_a, mask_b = to_sampling_index(index_list[-1], permute_list[-1], last_entry) + sampling_index.append(row) + index_a_mask.append(mask_a) + index_b_mask.append(mask_b) + self.register_buffer("index_list", torch.stack(index_list)) + self.register_buffer("sampling_index", torch.stack(sampling_index)) + # index_a_mask and index_b_mask are for inserting the combination weights + self.register_buffer("index_a_mask", torch.stack(index_a_mask)) + self.register_buffer("index_b_mask", torch.stack(index_b_mask)) + + def __repr__(self) -> str: + """For debug and logging purpose.""" + return ( + f"GridEncoder: input_dim={self.input_dim} num_levels={self.num_levels} " + f"level_dim={self.level_dim} resolution={self.base_resolution} -> " + f"{int(round(self.base_resolution * self.per_level_scale ** (self.num_levels - 1)))} " + f"per_level_scale={self.per_level_scale:.4f} params={tuple(self.embeddings.shape)} " + f"gridtype={self.gridtype} align_corners={self.align_corners}" + ) + + def get_temporal_index(self, time: TensorType["bs"]) -> TensorType["bs", "temporal_index_dim"]: + """Convert the time into sampling index lists.""" + row_idx_value = time * (len(self.sampling_index) - 1) + row_idx = row_idx_value.long() + row_idx[time == 1] = len(self.sampling_index) - 1 + temporal_row_index = self.sampling_index[row_idx].float() + mask_a = self.index_a_mask[row_idx] + mask_b = self.index_b_mask[row_idx] + temporal_row_index[mask_a] = row_idx + 1 - row_idx_value + temporal_row_index[mask_b] = row_idx_value - row_idx + return temporal_row_index + + def forward(self, xyz: TensorType["bs", "input_dim"], time: TensorType["bs", 1]) -> TensorType["bs", "output_dim"]: + """Forward and sampling feature vectors from the embedding. + + Args: + xyz: input coords, should be in [0,1] + time: input time, should be in [0,1] with shape [bs, 1] + """ + outputs = TemporalGridEncodeFunc.apply( + xyz, + self.get_temporal_index(time[:, 0]), + self.embeddings, + self.offsets, + self.per_level_scale, + self.base_resolution, + xyz.requires_grad, + self.gridtype_id, + self.align_corners, + ) + return outputs + + def get_temporal_tv_loss(self) -> TensorType[()]: + """Apply TV loss on the temporal channels. + Sample a random channel combination (i.e., row for the combination table), + and then compute loss on it. + """ + row_idx = torch.randint(0, len(self.index_list), [1]).item() + feat_idx = self.index_list[row_idx] + return (self.embeddings[:, feat_idx[0]] - self.embeddings[:, feat_idx[1]]).abs().mean() diff --git a/tests/field_components/test_temporal_grid.py b/tests/field_components/test_temporal_grid.py new file mode 100644 index 000000000..a83450c38 --- /dev/null +++ b/tests/field_components/test_temporal_grid.py @@ -0,0 +1,47 @@ +""" +Test if temporal grid run properly (forward and backward) +""" +import torch + +from nerfstudio.field_components import temporal_grid + + +def test_temporal_grid(): + """Test temporal grid""" + if not torch.cuda.is_available(): + print("Unable to test temporal grid without GPU, since CUDA kernels involved.") + return + + params = dict( + temporal_dim=2, + input_dim=1, + num_levels=1, + level_dim=1, + per_level_scale=1, + base_resolution=1, + log2_hashmap_size=2, + desired_resolution=None, + gridtype="tiled", + align_corners=False, + ) + model = temporal_grid.TemporalGridEncoder(**params).cuda() + random_embedding = torch.rand_like(model.embeddings) + random_embedding[:, 0] = torch.arange(8).to(model.embeddings) + model.embeddings = torch.nn.Parameter(random_embedding, requires_grad=True) + + x = torch.zeros([1024, 1]).cuda() + t = torch.zeros([1024, 1]).cuda() + out = model(x, t) + weight = torch.randn_like(out) + (out * weight).sum().backward() + assert torch.all(out == 0.5) + assert model.embeddings.grad.sum() - weight.sum() < 0.01 + assert torch.all(model.embeddings.grad[2:, :] == 0) + assert torch.all(model.embeddings.grad[:, 1:] == 0) + model.get_temporal_tv_loss() + # TODO: Any better way to numerically test it? Especially for the gradients. + # (maybe add some randomness and multiple cases for the testing?) + + +if __name__ == "__main__": + test_temporal_grid()