mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 10:55:57 +08:00
NeRFPlayer sliding window implementation (#1255)
* NeRFPlayer sliding window implementation * add comments
This commit is contained in:
LSong authored and 
GitHub committed
1 parent
a4b023b171
commit
b170e175ed
7 files changed
+1236
No files matched your search
@@ -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")
|
||||
@@ -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"]
|
||||
@@ -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 <stdint.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
// 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<at::Tensor> 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<at::Tensor> dy_dx,
|
||||
at::optional<at::Tensor> grad_inputs,
|
||||
const uint32_t gridtype,
|
||||
const bool align_corners
|
||||
);
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,14 @@
|
||||
/*
|
||||
* Code adapted from (@ashawkey) https://github.com/ashawkey/torch-ngp/
|
||||
* Author: @lsongx
|
||||
*/
|
||||
|
||||
|
||||
#include <torch/extension.h>
|
||||
|
||||
#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)");
|
||||
}
|
||||
@@ -0,0 +1,667 @@
|
||||
/*
|
||||
* Code adapted from (@ashawkey) https://github.com/ashawkey/torch-ngp/blob/main/gridencoder/src/gridencoder.cu
|
||||
* Author: @lsongx
|
||||
*/
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <stdexcept>
|
||||
|
||||
#include <stdint.h>
|
||||
#include <cstdio>
|
||||
|
||||
|
||||
#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 <typename T>
|
||||
static inline __host__ __device__ T div_round_up(T val, T divisor) {
|
||||
return (val + divisor - 1) / divisor;
|
||||
}
|
||||
|
||||
|
||||
template <uint32_t D>
|
||||
__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 <uint32_t D>
|
||||
__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<D>(pos_grid);
|
||||
}
|
||||
|
||||
return (index % hashmap_size) * grid_C + ch;
|
||||
}
|
||||
|
||||
// grid sampling kernel
|
||||
template <typename scalar_t, uint32_t D, uint32_t C>
|
||||
__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<D>(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<D>(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<D>(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<D>(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<D>(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<D>(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<D>(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<D>(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<D>(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 <typename scalar_t, uint32_t D, uint32_t C, uint32_t N_C>
|
||||
__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<D>(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<D>(gridtype, grid_C, align_corners,
|
||||
__float2uint_rn(temporal_row_index[1]), hashmap_size, resolution, pos_grid_local);
|
||||
uint32_t index_b = get_grid_index<D>(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 <typename scalar_t, uint32_t D, uint32_t C>
|
||||
__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 <typename scalar_t, uint32_t D>
|
||||
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<scalar_t, D, 1><<<blocks_hashgrid, N_THREAD>>>(
|
||||
inputs, temporal_row_index, embeddings, offsets, outputs,
|
||||
B, grid_C, L, S, H, dy_dx, gridtype, align_corners); break;
|
||||
case 2: kernel_grid<scalar_t, D, 2><<<blocks_hashgrid, N_THREAD>>>(
|
||||
inputs, temporal_row_index, embeddings, offsets, outputs,
|
||||
B, grid_C, L, S, H, dy_dx, gridtype, align_corners); break;
|
||||
case 4: kernel_grid<scalar_t, D, 4><<<blocks_hashgrid, N_THREAD>>>(
|
||||
inputs, temporal_row_index, embeddings, offsets, outputs,
|
||||
B, grid_C, L, S, H, dy_dx, gridtype, align_corners); break;
|
||||
case 8: kernel_grid<scalar_t, D, 8><<<blocks_hashgrid, N_THREAD>>>(
|
||||
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 <typename scalar_t>
|
||||
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<scalar_t, 1>(
|
||||
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<scalar_t, 2>(
|
||||
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<scalar_t, 3>(
|
||||
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<scalar_t, 4>(
|
||||
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<scalar_t, 5>(
|
||||
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 <typename scalar_t, uint32_t D>
|
||||
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<scalar_t, D, 1, 1><<<blocks_hashgrid, N_THREAD>>>(
|
||||
grad, inputs, temporal_row_index,
|
||||
embeddings, offsets, grad_embeddings, B, grid_C, L, S, H, gridtype, align_corners);
|
||||
if (dy_dx) {
|
||||
kernel_input_backward<scalar_t, D, 1><<<div_round_up(B * D, N_THREAD), N_THREAD>>>(
|
||||
grad, dy_dx, grad_inputs, B, L);
|
||||
}
|
||||
break;
|
||||
case 2:
|
||||
kernel_grid_backward<scalar_t, D, 2, 1><<<blocks_hashgrid, N_THREAD>>>(
|
||||
grad, inputs, temporal_row_index,
|
||||
embeddings, offsets, grad_embeddings, B, grid_C, L, S, H, gridtype, align_corners);
|
||||
if (dy_dx) {
|
||||
kernel_input_backward<scalar_t, D, 2><<<div_round_up(B * D, N_THREAD), N_THREAD>>>(
|
||||
grad, dy_dx, grad_inputs, B, L);
|
||||
}
|
||||
break;
|
||||
case 4:
|
||||
kernel_grid_backward<scalar_t, D, 4, 1><<<blocks_hashgrid, N_THREAD>>>(
|
||||
grad, inputs, temporal_row_index,
|
||||
embeddings, offsets, grad_embeddings, B, grid_C, L, S, H, gridtype, align_corners);
|
||||
if (dy_dx) {
|
||||
kernel_input_backward<scalar_t, D, 4><<<div_round_up(B * D, N_THREAD), N_THREAD>>>(
|
||||
grad, dy_dx, grad_inputs, B, L);
|
||||
}
|
||||
break;
|
||||
case 8:
|
||||
kernel_grid_backward<scalar_t, D, 8, 1><<<blocks_hashgrid, N_THREAD>>>(
|
||||
grad, inputs, temporal_row_index,
|
||||
embeddings, offsets, grad_embeddings, B, grid_C, L, S, H, gridtype, align_corners);
|
||||
if (dy_dx) {
|
||||
kernel_input_backward<scalar_t, D, 8><<<div_round_up(B * D, N_THREAD), N_THREAD>>>(
|
||||
grad, dy_dx, grad_inputs, B, L);
|
||||
}
|
||||
break;
|
||||
default: throw std::runtime_error{"GridEncoding: C must be 1, 2, 4, or 8."};
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
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<scalar_t, 1>(
|
||||
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<scalar_t, 2>(
|
||||
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<scalar_t, 3>(
|
||||
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<scalar_t, 4>(
|
||||
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<scalar_t, 5>(
|
||||
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<at::Tensor> 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<scalar_t>(
|
||||
inputs.data_ptr<float>(), temporal_row_index.data_ptr<float>(), embeddings.data_ptr<scalar_t>(),
|
||||
offsets.data_ptr<int>(), outputs.data_ptr<scalar_t>(), B, D, grid_C, C, L, S, H,
|
||||
dy_dx.has_value() ? dy_dx.value().data_ptr<scalar_t>() : 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<at::Tensor> dy_dx,
|
||||
at::optional<at::Tensor> 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<scalar_t>(
|
||||
grad.data_ptr<scalar_t>(), inputs.data_ptr<float>(), temporal_row_index.data_ptr<float>(),
|
||||
embeddings.data_ptr<scalar_t>(), offsets.data_ptr<int>(), grad_embeddings.data_ptr<scalar_t>(),
|
||||
B, D, grid_C, C, L, S, H,
|
||||
dy_dx.has_value() ? dy_dx.value().data_ptr<scalar_t>() : nullptr,
|
||||
grad_inputs.has_value() ? grad_inputs.value().data_ptr<scalar_t>() : nullptr,
|
||||
gridtype, align_corners);
|
||||
}));
|
||||
}
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in new issue
Block a user