diff --git a/.vscode/c_cpp_properties.json b/.vscode/c_cpp_properties.json new file mode 100644 index 000000000..eeb5276d2 --- /dev/null +++ b/.vscode/c_cpp_properties.json @@ -0,0 +1,15 @@ +{ + "configurations": [ + { + "name": "Linux", + "includePath": [ + "${workspaceFolder}/**" + ], + "defines": [], + "compilerPath": "/usr/bin/gcc", + "cStandard": "c11", + "intelliSenseMode": "linux-gcc-x64" + } + ], + "version": 4 +} \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json index c973d2df0..d84fbc08b 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,33 +1,95 @@ { - "files.watcherExclude": { - "**/__pycache__/": true, - "**/.ipynb_checkpoints": true, - "**/*.ipynb": true, - "**/old*": true, - "**/runs*": true, - "**/temp*": true, - }, - "editor.defaultFormatter": "esbenp.prettier-vscode", - "[python]": { - "editor.defaultFormatter": "ms-python.python" - }, - "editor.formatOnSave": true, - "python.envFile": "${workspaceFolder}/.env", - "python.formatting.provider": "black", - "python.formatting.blackArgs": [ - "--line-length=120" - ], - "python.linting.pylintEnabled": true, - "python.linting.flake8Enabled": false, - "python.linting.enabled": true, - "python.sortImports.args": [ - "--src=${workspaceFolder}" - ], - "python.testing.unittestEnabled": false, - "python.testing.pytestEnabled": true, - "python.linting.ignorePatterns": [ - "**/site-packages/**/*.py", - ".vscode/*.py", - "external/**/*.py", - ] -} \ No newline at end of file + "files.watcherExclude": { + "**/__pycache__/": true, + "**/.ipynb_checkpoints": true, + "**/*.ipynb": true, + "**/old*": true, + "**/runs*": true, + "**/temp*": true + }, + "editor.defaultFormatter": "esbenp.prettier-vscode", + "[python]": { + "editor.defaultFormatter": "ms-python.python" + }, + "editor.formatOnSave": true, + "python.envFile": "${workspaceFolder}/.env", + "python.formatting.provider": "black", + "python.formatting.blackArgs": ["--line-length=120"], + "python.linting.pylintEnabled": true, + "python.linting.flake8Enabled": false, + "python.linting.enabled": true, + "python.sortImports.args": ["--src=${workspaceFolder}"], + "python.testing.unittestEnabled": false, + "python.testing.pytestEnabled": true, + "python.linting.ignorePatterns": [ + "**/site-packages/**/*.py", + ".vscode/*.py", + "external/**/*.py" + ], + "C_Cpp.errorSquiggles": "Disabled", + "files.associations": { + "array": "cpp", + "bitset": "cpp", + "string_view": "cpp", + "initializer_list": "cpp", + "utility": "cpp", + "__hash_table": "cpp", + "__split_buffer": "cpp", + "deque": "cpp", + "iterator": "cpp", + "queue": "cpp", + "stack": "cpp", + "string": "cpp", + "unordered_map": "cpp", + "vector": "cpp", + "atomic": "cpp", + "*.tcc": "cpp", + "cctype": "cpp", + "chrono": "cpp", + "clocale": "cpp", + "cmath": "cpp", + "codecvt": "cpp", + "condition_variable": "cpp", + "cstdarg": "cpp", + "cstddef": "cpp", + "cstdint": "cpp", + "cstdio": "cpp", + "cstdlib": "cpp", + "cstring": "cpp", + "ctime": "cpp", + "cwchar": "cpp", + "cwctype": "cpp", + "exception": "cpp", + "algorithm": "cpp", + "filesystem": "cpp", + "functional": "cpp", + "memory": "cpp", + "memory_resource": "cpp", + "optional": "cpp", + "ratio": "cpp", + "system_error": "cpp", + "tuple": "cpp", + "type_traits": "cpp", + "fstream": "cpp", + "iomanip": "cpp", + "iosfwd": "cpp", + "iostream": "cpp", + "istream": "cpp", + "limits": "cpp", + "mutex": "cpp", + "new": "cpp", + "ostream": "cpp", + "sstream": "cpp", + "stdexcept": "cpp", + "streambuf": "cpp", + "thread": "cpp", + "typeinfo": "cpp", + "__nullptr": "cpp", + "__config": "cpp", + "__locale": "cpp", + "__bit_reference": "cpp", + "ios": "cpp", + "__atomic": "cpp", + "__node_handle": "cpp" + } +} diff --git a/README.md b/README.md index 6632544e4..4579cca2e 100644 --- a/README.md +++ b/README.md @@ -60,7 +60,7 @@ pip install -e . # Install library with CUDA support. Change setup.py to `USE_CUDA = True` and then python setup.py develop -# Install tiny-cuda-nn (tcnn) to use with the graph_instant_ngp.yaml config +# Install tiny-cuda-nn (tcnn) and apex to use with the graph_instant_ngp.yaml config pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch # Run the test cases diff --git a/configs/graph_instant_ngp.yaml b/configs/graph_instant_ngp.yaml index e8b850c42..519360f99 100644 --- a/configs/graph_instant_ngp.yaml +++ b/configs/graph_instant_ngp.yaml @@ -26,8 +26,14 @@ graph: optimizers: fields: optimizer: - lr: 1e-2 - eps: 1e-7 # need to investigate further for float16 stability + lr: 3e-3 + eps: 1e-15 # NGP prefers small eps + # Note: If you want slightly faster coverange, you can use the following optimizer instead, which + # requires to install apex following https://github.com/NVIDIA/apex#linux + # optimizer: + # _target_: apex.optimizers.FusedAdam + # lr: 3e-3 + # eps: 1e-15 scheduler: null trainer: diff --git a/docs/tutorials/creating_graphs.md b/docs/tutorials/creating_graphs.md index 63e57a764..9cd1f9732 100644 --- a/docs/tutorials/creating_graphs.md +++ b/docs/tutorials/creating_graphs.md @@ -19,7 +19,7 @@ In general, modules can be thought of as individual component parts of the graph {ref}`fields` represents the "space", aka. the radiance field of the NeRF. Here, we define the field as the part of the network that takes in point samples and any other conditioning, and outputs any of the `FieldHeadNames` (`nerf/field_modules/field_heads.py`). -All fields are composed of modules continained within the `fields/` directory (e.g. `fields/modules/` or `fields/occupancy_fields/`). We can think of field modules as modules that actually define or interact with the field. +All fields are composed of modules continained within the `fields/` directory (e.g. `fields/modules/` or `fields/density_fields/`). We can think of field modules as modules that actually define or interact with the field. To build a NeRF field, we therefore follow these steps: diff --git a/environment/requirements.txt b/environment/requirements.txt index c165d97c2..3a6806495 100644 --- a/environment/requirements.txt +++ b/environment/requirements.txt @@ -3,6 +3,7 @@ aiortc av black[jupyter]==22.3.0 dcargs +ninja functorch==0.1.1 h5py>=2.9.0 hydra-colorlog==1.2.0 diff --git a/pyrad/cameras/rays.py b/pyrad/cameras/rays.py index 38f6991e0..210572505 100644 --- a/pyrad/cameras/rays.py +++ b/pyrad/cameras/rays.py @@ -94,6 +94,7 @@ class RaySamples(TensorDataclass): camera_indices (TensorType[..., 1]): Camera index. valid_mask (TensorType[..., 1]): Rays that are valid. deltas (TensorType[..., 1]): "width" of each sample. + ts (TensorType[..., 1]): sample-camera distance of each sample. """ frustums: Frustums diff --git a/pyrad/cuda/__init__.py b/pyrad/cuda/__init__.py index 02655775e..91ba5e016 100644 --- a/pyrad/cuda/__init__.py +++ b/pyrad/cuda/__init__.py @@ -11,3 +11,69 @@ # 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. + +"""CUDA Functions""" +import torch +from torch.cuda.amp import custom_bwd, custom_fwd + +from pyrad.cuda.backend import _C + +packbits = _C.packbits +ray_aabb_intersect = _C.ray_aabb_intersect +morton3D = _C.morton3D +morton3D_invert = _C.morton3D_invert +raymarching = _C.raymarching +volumetric_rendering_forward = _C.volumetric_rendering_forward +volumetric_rendering_backward = _C.volumetric_rendering_backward + +# pylint: disable=abstract-method,arguments-differ +class VolumeRenderer(torch.autograd.Function): + """CUDA Volumetirc Renderer""" + + @staticmethod + @custom_fwd(cast_inputs=torch.float32) + def forward(ctx, packed_info, starts, ends, sigmas, rgbs): + accumulated_weight, accumulated_depth, accumulated_color, mask = volumetric_rendering_forward( + packed_info, starts, ends, sigmas, rgbs + ) + # TODO(ruilongli): accelerate for torch.no_grad()? + ctx.save_for_backward( + accumulated_weight, + accumulated_depth, + accumulated_color, + packed_info, + starts, + ends, + sigmas, + rgbs, + ) + return accumulated_weight, accumulated_depth, accumulated_color, mask + + @staticmethod + @custom_bwd + def backward(ctx, grad_weight, grad_depth, grad_color, _grad_mask): + ( + accumulated_weight, + accumulated_depth, + accumulated_color, + packed_info, + starts, + ends, + sigmas, + rgbs, + ) = ctx.saved_tensors + grad_sigmas, grad_rgbs = volumetric_rendering_backward( + accumulated_weight, + accumulated_depth, + accumulated_color, + grad_weight, + grad_depth, + grad_color, + packed_info, + starts, + ends, + sigmas, + rgbs, + ) + # corresponds to the input argument list of forward() + return None, None, None, grad_sigmas, grad_rgbs diff --git a/pyrad/cuda/backend.py b/pyrad/cuda/backend.py new file mode 100644 index 000000000..76f9734dc --- /dev/null +++ b/pyrad/cuda/backend.py @@ -0,0 +1,47 @@ +# Copyright 2022 The Plenoptix 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. +import glob +import os +from subprocess import DEVNULL, call + +from torch.utils.cpp_extension import load + +PATH = os.path.dirname(os.path.abspath(__file__)) + + +def cuda_toolkit_available(): + """Check if the nvcc is avaiable on the machine.""" + # https://github.com/idiap/fast-transformers/blob/master/setup.py + try: + call(["nvcc"], stdout=DEVNULL, stderr=DEVNULL) + return True + except FileNotFoundError: + return False + + +if cuda_toolkit_available(): + sources = glob.glob(os.path.join(PATH, "csrc/*.cu")) +else: + sources = glob.glob(os.path.join(PATH, "csrc/*.cpp")) + +extra_cflags = ["-O2"] +extra_cuda_cflags = ["-O2"] +_C = load( + name="pyrad_extension", + sources=sources, + extra_cflags=extra_cflags, + extra_cuda_cflags=extra_cuda_cflags, +) + +__all__ = ["_C"] diff --git a/pyrad/cuda/csrc/include/helpers.h b/pyrad/cuda/csrc/include/helpers.h new file mode 100644 index 000000000..61ff8e3d5 --- /dev/null +++ b/pyrad/cuda/csrc/include/helpers.h @@ -0,0 +1,53 @@ +#pragma once + +#ifdef __CUDACC__ +#define CUDA_HOSTDEV __host__ __device__ +#else +#define CUDA_HOSTDEV +#endif + +#include + + +inline constexpr CUDA_HOSTDEV float __SQRT3() { return 1.73205080757f; } + +template +inline CUDA_HOSTDEV void __swap(scalar_t &a, scalar_t &b) +{ + scalar_t c = a; + a = b; + b = c; +} + +inline CUDA_HOSTDEV float __clamp(float f, float a, float b) { return fmaxf(a, fminf(f, b)); } + +inline CUDA_HOSTDEV int __clamp(int f, int a, int b) { return std::max(a, std::min(f, b)); } + +inline CUDA_HOSTDEV float __sign(float x) { return copysignf(1.0, x); } + +inline CUDA_HOSTDEV uint32_t __expand_bits(uint32_t v) +{ + v = (v * 0x00010001u) & 0xFF0000FFu; + v = (v * 0x00000101u) & 0x0F00F00Fu; + v = (v * 0x00000011u) & 0xC30C30C3u; + v = (v * 0x00000005u) & 0x49249249u; + return v; +} + +inline CUDA_HOSTDEV uint32_t __morton3D(uint32_t x, uint32_t y, uint32_t z) +{ + uint32_t xx = __expand_bits(x); + uint32_t yy = __expand_bits(y); + uint32_t zz = __expand_bits(z); + return xx | (yy << 1) | (zz << 2); +} + +inline CUDA_HOSTDEV uint32_t __morton3D_invert(uint32_t x) +{ + x = x & 0x49249249; + x = (x | (x >> 2)) & 0xc30c30c3; + x = (x | (x >> 4)) & 0x0f00f00f; + x = (x | (x >> 8)) & 0xff0000ff; + x = (x | (x >> 16)) & 0x0000ffff; + return x; +} diff --git a/pyrad/cuda/csrc/include/helpers_cuda.h b/pyrad/cuda/csrc/include/helpers_cuda.h new file mode 100644 index 000000000..7ca474309 --- /dev/null +++ b/pyrad/cuda/csrc/include/helpers_cuda.h @@ -0,0 +1,18 @@ +#pragma once + +#include "helpers.h" +#include + +#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor") +#define CHECK_CONTIGUOUS(x) \ + TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) \ + CHECK_CUDA(x); \ + CHECK_CONTIGUOUS(x) +#define CUDA_GET_THREAD_ID(tid, Q) \ + const int tid = blockIdx.x * blockDim.x + threadIdx.x; \ + if (tid >= Q) \ + return +#define CUDA_N_BLOCKS_NEEDED(Q, CUDA_N_THREADS) ((Q - 1) / CUDA_N_THREADS + 1) +#define DEVICE_GUARD(_ten) \ + const at::cuda::OptionalCUDAGuard device_guard(device_of(_ten)); \ No newline at end of file diff --git a/pyrad/cuda/csrc/intersection.cu b/pyrad/cuda/csrc/intersection.cu new file mode 100644 index 000000000..9b272c65f --- /dev/null +++ b/pyrad/cuda/csrc/intersection.cu @@ -0,0 +1,118 @@ +#include "include/helpers_cuda.h" + + +template +inline __host__ __device__ void _ray_aabb_intersect( + const scalar_t* rays_o, + const scalar_t* rays_d, + const scalar_t* aabb, + scalar_t* near, + scalar_t* far +) { + // aabb is [xmin, ymin, zmin, xmax, ymax, zmax] + scalar_t tmin = (aabb[0] - rays_o[0]) / rays_d[0]; + scalar_t tmax = (aabb[3] - rays_o[0]) / rays_d[0]; + if (tmin > tmax) __swap(tmin, tmax); + + scalar_t tymin = (aabb[1] - rays_o[1]) / rays_d[1]; + scalar_t tymax = (aabb[4] - rays_o[1]) / rays_d[1]; + if (tymin > tymax) __swap(tymin, tymax); + + if (tmin > tymax || tymin > tmax){ + *near = std::numeric_limits::max(); + *far = std::numeric_limits::max(); + return; + } + + if (tymin > tmin) tmin = tymin; + if (tymax < tmax) tmax = tymax; + + scalar_t tzmin = (aabb[2] - rays_o[2]) / rays_d[2]; + scalar_t tzmax = (aabb[5] - rays_o[2]) / rays_d[2]; + if (tzmin > tzmax) __swap(tzmin, tzmax); + + if (tmin > tzmax || tzmin > tmax){ + *near = std::numeric_limits::max(); + *far = std::numeric_limits::max(); + return; + } + + if (tzmin > tmin) tmin = tzmin; + if (tzmax < tmax) tmax = tzmax; + + *near = tmin; + *far = tmax; + return; +} + + +template +__global__ void kernel_ray_aabb_intersect( + const int N, + const scalar_t* rays_o, + const scalar_t* rays_d, + const scalar_t* aabb, + scalar_t* t_min, + scalar_t* t_max +){ + // aabb is [xmin, ymin, zmin, xmax, ymax, zmax] + CUDA_GET_THREAD_ID(thread_id, N); + + // locate + rays_o += thread_id * 3; + rays_d += thread_id * 3; + t_min += thread_id; + t_max += thread_id; + + _ray_aabb_intersect(rays_o, rays_d, aabb, t_min, t_max); + + scalar_t zero = static_cast(0.f); + *t_min = *t_min > zero ? *t_min : zero; + return; +} + +/** + * @brief Ray AABB Test + * + * @param rays_o Ray origins. Tensor with shape [N, 3]. + * @param rays_d Normalized ray directions. Tensor with shape [N, 3]. + * @param aabb Scene AABB [xmin, ymin, zmin, xmax, ymax, zmax]. Tensor with shape [6]. + * @return std::vector + * Ray AABB intersection {t_min, t_max} with shape [N] respectively. Note the t_min is + * clipped to minimum zero. + */ +std::vector ray_aabb_intersect( + const torch::Tensor rays_o, const torch::Tensor rays_d, const torch::Tensor aabb +) { + DEVICE_GUARD(rays_o); + CHECK_INPUT(rays_o); + CHECK_INPUT(rays_d); + CHECK_INPUT(aabb); + TORCH_CHECK(rays_o.ndimension() == 2 & rays_o.size(1) == 3) + TORCH_CHECK(rays_d.ndimension() == 2 & rays_d.size(1) == 3) + TORCH_CHECK(aabb.ndimension() == 1 & aabb.size(0) == 6) + + const int N = rays_o.size(0); + + const int threads = 256; + const int blocks = CUDA_N_BLOCKS_NEEDED(N, threads); + + torch::Tensor t_min = torch::empty({N}, rays_o.options()); + torch::Tensor t_max = torch::empty({N}, rays_o.options()); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF( + rays_o.scalar_type(), "ray_aabb_intersect", + ([&] { + kernel_ray_aabb_intersect<<>>( + N, + rays_o.data_ptr(), + rays_d.data_ptr(), + aabb.data_ptr(), + t_min.data_ptr(), + t_max.data_ptr() + ); + }) + ); + + return {t_min, t_max}; +} \ No newline at end of file diff --git a/pyrad/cuda/csrc/morton3d.cu b/pyrad/cuda/csrc/morton3d.cu new file mode 100644 index 000000000..5dd864f1f --- /dev/null +++ b/pyrad/cuda/csrc/morton3d.cu @@ -0,0 +1,107 @@ +#include "include/helpers_cuda.h" + + +template +__global__ void kernel_morton3D( + const int N, + const index_t* coords, + index_t* indices +){ + CUDA_GET_THREAD_ID(thread_id, N); + + // locate + coords += thread_id * 3; + indices += thread_id; + + uint32_t index = __morton3D( + static_cast(coords[0]), + static_cast(coords[1]), + static_cast(coords[2]) + ); + *indices = static_cast(index); + return; +} + + +template +__global__ void kernel_morton3D_invert( + const int N, + const index_t* indices, + index_t* coords +){ + CUDA_GET_THREAD_ID(thread_id, N); + + // locate + coords += thread_id * 3; + indices += thread_id; + + uint32_t index = static_cast(indices[0]); + coords[0] = static_cast(__morton3D_invert(index >> 0)); + coords[1] = static_cast(__morton3D_invert(index >> 1)); + coords[2] = static_cast(__morton3D_invert(index >> 2)); + return; +} + +/** + * @brief Convert 3D coords to indices in z-order + * + * @param coords Integer coords with shape [N, 3] + * @return torch::Tensor Integer indices with shape [N] + */ +torch::Tensor morton3D(const torch::Tensor coords){ + DEVICE_GUARD(coords); + CHECK_INPUT(coords); + TORCH_CHECK(coords.ndimension() == 2 & coords.size(1) == 3) + + const int N = coords.size(0); + + const int threads = 256; + const int blocks = CUDA_N_BLOCKS_NEEDED(N, threads); + + torch::Tensor indices = torch::empty({N}, coords.options()); + + AT_DISPATCH_INDEX_TYPES( + coords.scalar_type(), "morton3D", + ([&] { + kernel_morton3D<<>>( + N, + coords.data_ptr(), + indices.data_ptr() + ); + }) + ); + + return indices; +} + +/** + * @brief Convert 3D indices to coords in z-order + * + * @param indices Integer indices with shape [N] + * @return torch::Tensor Integer coords with shape [N, 3] + */ +torch::Tensor morton3D_invert(const torch::Tensor indices){ + DEVICE_GUARD(indices); + CHECK_INPUT(indices); + TORCH_CHECK(indices.ndimension() == 1) + + const int N = indices.size(0); + + const int threads = 256; + const int blocks = CUDA_N_BLOCKS_NEEDED(N, threads); + + torch::Tensor coords = torch::empty({N, 3}, indices.options()); + + AT_DISPATCH_INDEX_TYPES( + indices.scalar_type(), "morton3D_invert", + ([&] { + kernel_morton3D_invert<<>>( + N, + indices.data_ptr(), + coords.data_ptr() + ); + }) + ); + + return coords; +} \ No newline at end of file diff --git a/pyrad/cuda/csrc/packbits.cu b/pyrad/cuda/csrc/packbits.cu new file mode 100644 index 000000000..3b42f104e --- /dev/null +++ b/pyrad/cuda/csrc/packbits.cu @@ -0,0 +1,59 @@ +#include "include/helpers_cuda.h" + + +template +__global__ void kernel_packbits( + const int N, + const float threshold, + const scalar_t* data, + uint8_t* bitfield +){ + CUDA_GET_THREAD_ID(thread_id, N); + + data += thread_id * 8; + bitfield += thread_id; + + uint8_t bits = 0; + #pragma unroll 8 + for (uint8_t i = 0; i < 8; i++) { + bits |= (data[i] > threshold) ? ((uint8_t)1 << i) : 0; + } + bitfield[0] = bits; +} + +/** + * @brief Pack data into uint8 bits based on threshold. + * + * @param data: Tensor with any shape. Total elements must be N * 8. + * @param threshold: + * @return Tensor: bitfield that has shape of [N] + */ +torch::Tensor packbits( + const torch::Tensor data, const float threshold +) { + DEVICE_GUARD(data); + CHECK_INPUT(data); + TORCH_CHECK(data.numel() % 8 == 0) + + const int N = data.numel() / 8; + + const int threads = 256; + const int blocks = CUDA_N_BLOCKS_NEEDED(N, threads); + + torch::Tensor bitfield = torch::empty( + {N,}, data.options().dtype(torch::kUInt8)); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF( + data.scalar_type(), "packbits", + ([&] { + kernel_packbits<<>>( + N, + threshold, + data.data_ptr(), + bitfield.data_ptr() + ); + }) + ); + + return bitfield; +} \ No newline at end of file diff --git a/pyrad/cuda/csrc/pyrad.cpp b/pyrad/cuda/csrc/pyrad.cpp new file mode 100644 index 000000000..e16920fc3 --- /dev/null +++ b/pyrad/cuda/csrc/pyrad.cpp @@ -0,0 +1,68 @@ +#include "include/helpers.h" + + +torch::Tensor packbits( + const torch::Tensor data, const float threshold +) {}; + +std::vector ray_aabb_intersect( + const torch::Tensor rays_o, + const torch::Tensor rays_d, + const torch::Tensor aabb +) {}; + +torch::Tensor morton3D(const torch::Tensor coords) {}; + +torch::Tensor morton3D_invert(const torch::Tensor indices) {}; + +std::vector raymarching( + // rays + const torch::Tensor rays_o, + const torch::Tensor rays_d, + const torch::Tensor t_min, + const torch::Tensor t_max, + // density grid + const float grid_center, + const float grid_scale, + const int grid_cascades, + const int grid_size, + const torch::Tensor grid_bitfield, + // sampling args + const int max_total_samples, + const int num_steps, + const float cone_angle, + const float step_scale +) {}; + +std::vector volumetric_rendering_forward( + torch::Tensor packed_info, + torch::Tensor starts, + torch::Tensor ends, + torch::Tensor sigmas, + torch::Tensor rgbs +) {}; + +std::vector volumetric_rendering_backward( + torch::Tensor accumulated_weight, + torch::Tensor accumulated_depth, + torch::Tensor accumulated_color, + torch::Tensor grad_weight, + torch::Tensor grad_depth, + torch::Tensor grad_color, + torch::Tensor packed_info, + torch::Tensor starts, + torch::Tensor ends, + torch::Tensor sigmas, + torch::Tensor rgbs +) {}; + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("packbits", &packbits); + m.def("ray_aabb_intersect", &ray_aabb_intersect); + m.def("morton3D", &morton3D); + m.def("morton3D_invert", &morton3D_invert); + m.def("raymarching", &raymarching); + m.def("volumetric_rendering_forward", &volumetric_rendering_forward); + m.def("volumetric_rendering_backward", &volumetric_rendering_backward); +} \ No newline at end of file diff --git a/pyrad/cuda/csrc/pyrad.cu b/pyrad/cuda/csrc/pyrad.cu new file mode 100644 index 000000000..1a72b3d8a --- /dev/null +++ b/pyrad/cuda/csrc/pyrad.cu @@ -0,0 +1,68 @@ +#include "include/helpers_cuda.h" + + +torch::Tensor packbits( + const torch::Tensor data, const float threshold +); + +std::vector ray_aabb_intersect( + const torch::Tensor rays_o, + const torch::Tensor rays_d, + const torch::Tensor aabb +); + +torch::Tensor morton3D(const torch::Tensor coords); + +torch::Tensor morton3D_invert(const torch::Tensor indices); + +std::vector raymarching( + // rays + const torch::Tensor rays_o, + const torch::Tensor rays_d, + const torch::Tensor t_min, + const torch::Tensor t_max, + // density grid + const float grid_center, + const float grid_scale, + const int grid_cascades, + const int grid_size, + const torch::Tensor grid_bitfield, + // sampling args + const int max_total_samples, + const int num_steps, + const float cone_angle, + const float step_scale +); + +std::vector volumetric_rendering_forward( + torch::Tensor packed_info, + torch::Tensor starts, + torch::Tensor ends, + torch::Tensor sigmas, + torch::Tensor rgbs +); + +std::vector volumetric_rendering_backward( + torch::Tensor accumulated_weight, + torch::Tensor accumulated_depth, + torch::Tensor accumulated_color, + torch::Tensor grad_weight, + torch::Tensor grad_depth, + torch::Tensor grad_color, + torch::Tensor packed_info, + torch::Tensor starts, + torch::Tensor ends, + torch::Tensor sigmas, + torch::Tensor rgbs +); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("packbits", &packbits); + m.def("ray_aabb_intersect", &ray_aabb_intersect); + m.def("morton3D", &morton3D); + m.def("morton3D_invert", &morton3D_invert); + m.def("raymarching", &raymarching); + m.def("volumetric_rendering_forward", &volumetric_rendering_forward); + m.def("volumetric_rendering_backward", &volumetric_rendering_backward); +} \ No newline at end of file diff --git a/pyrad/cuda/csrc/raymarching.cu b/pyrad/cuda/csrc/raymarching.cu new file mode 100644 index 000000000..530e0dd1a --- /dev/null +++ b/pyrad/cuda/csrc/raymarching.cu @@ -0,0 +1,347 @@ +#include "include/helpers_cuda.h" + + +inline __device__ float min_step_size(uint32_t num_steps) { + return __SQRT3() / num_steps; +} + +// Maximum step size is the width of the coarsest gridsize cell. +inline __device__ float max_step_size(uint32_t grid_cascades, uint32_t grid_size) { + // Note(ruilongli): use mip=0 step size? + return __SQRT3() * (1 << (grid_cascades-1)) / grid_size; +} + +// Perform fixed-size stepping in unit-cube scenes (like original NeRF) and exponential +// stepping in larger scenes. +inline __device__ float calc_dt(float t, float cone_angle, float dt_min, float dt_max) { + // TODO(ruilongli): scene_scale related to cone_angle? + return __clamp(t * cone_angle, dt_min, dt_max); +} + +inline __device__ int mip_from_pos(float x, float y, float z, uint32_t grid_cascades, float grid_center, float grid_scale) { + float maxval = fmaxf(fmaxf(fabsf(x - grid_center), fabsf(y - grid_center)), fabsf(z - grid_center)) / grid_scale; + int exponent; frexpf(maxval, &exponent); + return min(grid_cascades-1, max(0, exponent+1)); +} + +inline __device__ int mip_from_dt( + float x, float y, float z, uint32_t grid_cascades, float dt, int grid_size, float grid_center, float grid_scale +) { + int mip = mip_from_pos(x, y, z, grid_cascades, grid_center, grid_scale); + dt *= 2 * grid_size; + if (dt<1.f) return mip; // exponent would be zero + int exponent; frexpf(dt, &exponent); + return min(grid_cascades-1, max(exponent, mip)); +} + +inline __device__ uint32_t grid_mip_offset(uint32_t mip, int grid_size) { + return (grid_size * grid_size * grid_size) * mip; +} + +inline __device__ uint32_t cascaded_grid_idx_at( + float x, float y, float z, uint32_t mip, int grid_size, float grid_center, float grid_scale +) { + float mip_scale = scalbnf(1.0f, -mip) / grid_scale; + int ix = (int)((mip_scale * (x - grid_center) + 0.5f) * grid_size); + int iy = (int)((mip_scale * (y - grid_center) + 0.5f) * grid_size); + int iz = (int)((mip_scale * (z - grid_center) + 0.5f) * grid_size); + // printf("[input] x %f, y %f, z %f, mip %d, grid_size %d, grid_center %f\n", x, y, z, mip, grid_size, grid_center); + // printf("[output] mip_scale %f, ix %d iy %d, iz %d\n", mip_scale, ix, iy, iz); + uint32_t idx = __morton3D( + __clamp(ix, 0, grid_size-1), + __clamp(iy, 0, grid_size-1), + __clamp(iz, 0, grid_size-1) + ); + return idx; +} + +inline __device__ bool grid_occupied_at( + float x, float y, float z, + const uint8_t* grid_bitfield, + uint32_t mip, int grid_size, float grid_center, float grid_scale +) { + uint32_t idx = ( + cascaded_grid_idx_at(x, y, z, mip, grid_size, grid_center, grid_scale) + + grid_mip_offset(mip, grid_size) + ); + return grid_bitfield[idx/8] & (1<<(idx%8)); +} + +inline __device__ float distance_to_next_voxel( + float x, float y, float z, + float dir_x, float dir_y, float dir_z, + float idir_x, float idir_y, float idir_z, + float res +) { // dda like step + // TODO: warning: expression has no effect? + x, y, z = res * x, res * y, res * z; + float tx = (floorf(x + 0.5f + 0.5f * __sign(dir_x)) - x) * idir_x; + float ty = (floorf(y + 0.5f + 0.5f * __sign(dir_y)) - y) * idir_y; + float tz = (floorf(z + 0.5f + 0.5f * __sign(dir_z)) - z) * idir_z; + float t = min(min(tx, ty), tz); + + return fmaxf(t / res, 0.0f); +} + +inline __device__ float advance_to_next_voxel( + float t, + float x, float y, float z, + float dir_x, float dir_y, float dir_z, + float idir_x, float idir_y, float idir_z, + float res, float dt_min) { + // Regular stepping (may be slower but matches non-empty space) + float t_target = t + distance_to_next_voxel( + x, y, z, dir_x, dir_y, dir_z, idir_x, idir_y, idir_z, res + ); + do { + t += dt_min; + } while (t < t_target); + return t; +} + +template +__global__ void kernel_raymarching( + // rays info + const uint32_t n_rays, + const scalar_t* rays_o, + const scalar_t* rays_d, + const scalar_t* t_min, + const scalar_t* t_max, + // density grid + const float grid_center, + const float grid_scale, + const int grid_cascades, + const int grid_size, + const uint8_t* grid_bitfield, + // sampling + const int max_total_samples, + const int num_steps, // default 1024 + const float cone_angle, // default 0. for nerf-syn and 1/256 for large scene + const float step_scale, + int* numsteps_counter, + int* rays_counter, + int* packed_info_out, + scalar_t* origins_out, + scalar_t* dirs_out, + scalar_t* starts_out, + scalar_t* ends_out +) { + CUDA_GET_THREAD_ID(i, n_rays); + + // locate + rays_o += i * 3; + rays_d += i * 3; + t_min += i; + t_max += i; + + const float ox = rays_o[0], oy = rays_o[1], oz = rays_o[2]; + const float dx = rays_d[0], dy = rays_d[1], dz = rays_d[2]; + const float rdx = 1 / dx, rdy = 1 / dy, rdz = 1 / dz; + const float near = t_min[0], far = t_max[0]; + + // TODO(ruilongli): perturb `near` as in ngp_pl? + // TODO(ruilongli): pre-compute `dt_min`, `dt_max` as it is a constant? + float dt_min = min_step_size(num_steps) * step_scale; + float dt_max = max_step_size(grid_cascades, grid_size) * grid_scale; + + // first pass to compute an accurate number of steps + uint32_t j = 0; + float t0 = near; + float dt = calc_dt(t0, cone_angle, dt_min, dt_max); + float t1 = t0 + dt; + float t_mid = (t0 + t1) * 0.5f; + + while (t_mid < far && j < num_steps) { + // current center + const float x = ox + t_mid * dx; + const float y = oy + t_mid * dy; + const float z = oz + t_mid * dz; + uint32_t mip = mip_from_dt(x, y, z, grid_cascades, dt, grid_size, grid_center, grid_scale); + + if (grid_occupied_at(x, y, z, grid_bitfield, mip, grid_size, grid_center, grid_scale)) { + ++j; + // march to next sample + t0 = t1; + dt = calc_dt(t0, cone_angle, dt_min, dt_max); + t1 = t0 + dt; + t_mid = (t0 + t1) * 0.5f; + } + else { + // march to next sample + float res = (grid_size >> mip) * grid_scale; + t_mid = advance_to_next_voxel( + t_mid, x, y, z, dx, dy, dz, rdx, rdy, rdz, res, dt_min + ); + dt = calc_dt(t_mid, cone_angle, dt_min, dt_max); + t0 = t_mid - dt * 0.5f; + t1 = t_mid + dt * 0.5f; + } + } + if (j == 0) return; + + uint32_t numsteps = j; + uint32_t base = atomicAdd(numsteps_counter, numsteps); + if (base + numsteps > max_total_samples) return; + + // locate + origins_out += base * 3; + dirs_out += base * 3; + starts_out += base; + ends_out += base; + + uint32_t ray_idx = atomicAdd(rays_counter, 1); + + packed_info_out[ray_idx * 3 + 0] = i; // ray idx in {rays_o, rays_d} + packed_info_out[ray_idx * 3 + 1] = base; // point idx start. + packed_info_out[ray_idx * 3 + 2] = numsteps; // point idx shift. + + // Second round + j = 0; + t0 = near; + dt = calc_dt(t0, cone_angle, dt_min, dt_max); + t1 = t0 + dt; + t_mid = (t0 + t1) / 2.; + + while (t_mid < far && j < num_steps) { + // current center + const float x = ox + t_mid * dx; + const float y = oy + t_mid * dy; + const float z = oz + t_mid * dz; + uint32_t mip = mip_from_dt(x, y, z, grid_cascades, dt, grid_size, grid_center, grid_scale); + + if (grid_occupied_at(x, y, z, grid_bitfield, mip, grid_size, grid_center, grid_scale)) { + origins_out[j * 3 + 0] = ox; + origins_out[j * 3 + 1] = oy; + origins_out[j * 3 + 2] = oz; + dirs_out[j * 3 + 0] = dx; + dirs_out[j * 3 + 1] = dy; + dirs_out[j * 3 + 2] = dz; + starts_out[j] = t0; + ends_out[j] = t1; + ++j; + // march to next sample + t0 = t1; + dt = calc_dt(t0, cone_angle, dt_min, dt_max); + t1 = t0 + dt; + t_mid = (t0 + t1) * 0.5f; + } + else { + // march to next sample + float res = (grid_size >> mip) * grid_scale; + t_mid = advance_to_next_voxel( + t_mid, x, y, z, dx, dy, dz, rdx, rdy, rdz, res, dt_min + ); + dt = calc_dt(t_mid, cone_angle, dt_min, dt_max); + t0 = t_mid - dt * 0.5f; + t1 = t_mid + dt * 0.5f; + } + } + return; +} + +/** + * @brief Sample points by ray marching. + * + * @param rays_o Ray origins Shape of [n_rays, 3]. + * @param rays_d Normalized ray directions. Shape of [n_rays, 3]. + * @param t_min Near planes of rays. Shape of [n_rays]. + * @param t_max Far planes of rays. Shape of [n_rays]. + * @param grid_center Density grid center. TODO: support 3-dims. + * @param grid_scale Density grid base level scale. TODO: support 3-dims. + * @param grid_cascades Density grid levels. + * @param grid_size Density grid resolution. + * @param grid_bitfield Density grid uint8 bit field. + * @param max_total_samples Maximum total number of samples in this batch. + * @param num_steps Used to define the minimal step size: SQRT3() / num_steps. + * @param cone_angle 0. for nerf-synthetic and 1./256 for real scenes. + * @param step_scale Scale up the step size by this much. Usually equals to scene scale. + * @return std::vector + * - packed_info: Stores how to index the ray samples from the returned values. + * Shape of [n_rays, 3]. First value is the ray index. Second value is the sample + * start index in the results for this ray. Third value is the number of samples for + * this ray. Note for rays that have zero samples, we simply skip them so the `packed_info` + * has some zero padding in the end. + * - origins: Ray origins for those samples. [max_total_samples, 3] + * - dirs: Ray directions for those samples. [max_total_samples, 3] + * - starts: Where the frustum-shape sample starts along a ray. [max_total_samples, 1] + * - ends: Where the frustum-shape sample ends along a ray. [max_total_samples, 1] + */ +std::vector raymarching( + // rays + const torch::Tensor rays_o, + const torch::Tensor rays_d, + const torch::Tensor t_min, + const torch::Tensor t_max, + // density grid + const float grid_center, + const float grid_scale, + const int grid_cascades, + const int grid_size, + const torch::Tensor grid_bitfield, + // sampling args + const int max_total_samples, + const int num_steps, + const float cone_angle, + const float step_scale +) { + DEVICE_GUARD(rays_o); + + CHECK_INPUT(rays_o); + CHECK_INPUT(rays_d); + CHECK_INPUT(t_min); + CHECK_INPUT(t_max); + CHECK_INPUT(grid_bitfield); + + const int n_rays = rays_o.size(0); + + const int threads = 256; + const int blocks = CUDA_N_BLOCKS_NEEDED(n_rays, threads); + + // helper counter + torch::Tensor numsteps_counter = torch::zeros( + {1}, rays_o.options().dtype(torch::kInt32)); + torch::Tensor rays_counter = torch::zeros( + {1}, rays_o.options().dtype(torch::kInt32)); + + // output samples + torch::Tensor packed_info = torch::zeros( + {n_rays, 3}, rays_o.options().dtype(torch::kInt32)); // ray_id, sample_id, num_samples + torch::Tensor origins = torch::zeros({max_total_samples, 3}, rays_o.options()); + torch::Tensor dirs = torch::zeros({max_total_samples, 3}, rays_o.options()); + torch::Tensor starts = torch::zeros({max_total_samples, 1}, rays_o.options()); + torch::Tensor ends = torch::zeros({max_total_samples, 1}, rays_o.options()); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF( + rays_o.scalar_type(), + "raymarching_train", + ([&] + { kernel_raymarching<<>>( + // rays + n_rays, + rays_o.data_ptr(), + rays_d.data_ptr(), + t_min.data_ptr(), + t_max.data_ptr(), + // density grid + grid_center, + grid_scale, + grid_cascades, + grid_size, + grid_bitfield.data_ptr(), + // sampling args + max_total_samples, + num_steps, + cone_angle, + step_scale, + numsteps_counter.data_ptr(), // total samples. + rays_counter.data_ptr(), // total rays. + packed_info.data_ptr(), + origins.data_ptr(), + dirs.data_ptr(), + starts.data_ptr(), + ends.data_ptr() + ); + })); + + return {packed_info, origins, dirs, starts, ends}; +} \ No newline at end of file diff --git a/pyrad/cuda/csrc/volume_rendering.cu b/pyrad/cuda/csrc/volume_rendering.cu new file mode 100644 index 000000000..8dc0b4635 --- /dev/null +++ b/pyrad/cuda/csrc/volume_rendering.cu @@ -0,0 +1,260 @@ +#include "include/helpers_cuda.h" + + +template +__global__ void volumetric_rendering_forward_kernel( + const uint32_t n_rays, + const int* packed_info, // input ray & point indices. + const scalar_t* starts, // input start t + const scalar_t* ends, // input end t + const scalar_t* sigmas, // input density after activation + const scalar_t* rgbs, // input rgb after activation + // should be all-zero initialized + scalar_t* accumulated_weight, // output + scalar_t* accumulated_depth, // output + scalar_t* accumulated_color, // output + bool* mask // output +) { + CUDA_GET_THREAD_ID(thread_id, n_rays); + + // locate + const int i = packed_info[thread_id * 3 + 0]; // ray idx in {rays_o, rays_d} + const int base = packed_info[thread_id * 3 + 1]; // point idx start. + const int numsteps = packed_info[thread_id * 3 + 2]; // point idx shift. + if (numsteps == 0) return; + + starts += base; + ends += base; + sigmas += base; + rgbs += base * 3; + + accumulated_weight += i; + accumulated_depth += i; + accumulated_color += i * 3; + mask += i; + + // accumulated rendering + scalar_t T = 1.f; + scalar_t EPSILON = 1e-4f; + int j = 0; + for (; j < numsteps; ++j) { + if (T < EPSILON) { + break; + } + const scalar_t delta = ends[j] - starts[j]; + const scalar_t t = (ends[j] + starts[j]) * 0.5f; + + const scalar_t alpha = 1.f - __expf(-sigmas[j] * delta); + const scalar_t weight = alpha * T; + accumulated_weight[0] += weight; + accumulated_depth[0] += weight * t; + accumulated_color[0] += weight * rgbs[j * 3 + 0]; + accumulated_color[1] += weight * rgbs[j * 3 + 1]; + accumulated_color[2] += weight * rgbs[j * 3 + 2]; + T *= (1.f - alpha); + } + mask[0] = true; +} + + +template +__global__ void volumetric_rendering_backward_kernel( + const uint32_t n_rays, + const int* packed_info, // input ray & point indices. + const scalar_t* starts, // input start t + const scalar_t* ends, // input end t + const scalar_t* sigmas, // input density after activation + const scalar_t* rgbs, // input rgb after activation + const scalar_t* accumulated_weight, // forward output + const scalar_t* accumulated_depth, // forward output + const scalar_t* accumulated_color, // forward output + const scalar_t* grad_weight, // input + const scalar_t* grad_depth, // input + const scalar_t* grad_color, // input + scalar_t* grad_sigmas, // output + scalar_t* grad_rgbs // output +) { + CUDA_GET_THREAD_ID(thread_id, n_rays); + + // locate + const int i = packed_info[thread_id * 3 + 0]; // ray idx in {rays_o, rays_d} + const int base = packed_info[thread_id * 3 + 1]; // point idx start. + const int numsteps = packed_info[thread_id * 3 + 2]; // point idx shift. + if (numsteps == 0) return; + + starts += base; + ends += base; + sigmas += base; + rgbs += base * 3; + + grad_sigmas += base; + grad_rgbs += base * 3; + + accumulated_weight += i; + accumulated_depth += i; + accumulated_color += i * 3; + + grad_weight += i; + grad_depth += i; + grad_color += i * 3; + + // backward of accumulated rendering + scalar_t T = 1.f; + scalar_t EPSILON = 1e-4f; + int j = 0; + scalar_t r = 0, g = 0, b = 0, d = 0; + for (; j < numsteps; ++j) { + if (T < EPSILON) { + break; + } + const scalar_t delta = ends[j] - starts[j]; + const scalar_t t = (ends[j] + starts[j]) * 0.5f; + + const scalar_t alpha = 1.f - __expf(-sigmas[j] * delta); + const scalar_t weight = alpha * T; + + r += weight * rgbs[j * 3 + 0]; + g += weight * rgbs[j * 3 + 1]; + b += weight * rgbs[j * 3 + 2]; + d += weight * t; + + T *= (1.f - alpha); + + grad_rgbs[j * 3 + 0] = grad_color[0] * weight; + grad_rgbs[j * 3 + 1] = grad_color[1] * weight; + grad_rgbs[j * 3 + 2] = grad_color[2] * weight; + + grad_sigmas[j] = delta * ( + grad_color[0] * (T * rgbs[j * 3 + 0] - (accumulated_color[0] - r)) + + grad_color[1] * (T * rgbs[j * 3 + 1] - (accumulated_color[1] - g)) + + grad_color[2] * (T * rgbs[j * 3 + 2] - (accumulated_color[2] - b)) + + grad_weight[0] * (1.f - accumulated_weight[0]) + + grad_depth[0] * (t * T - (accumulated_depth[0] - d)) + ); + } +} + +/** + * @brief Volumetric Rendering: Accumulating samples in the forward pass. + * The inputs, excepct for `sigmas` and `rgbs`, are the outputs of our + * cuda ray marching function in `raymarching.cu` + * + * @param packed_info Stores how to index the ray samples from the returned values. + * Shape of [n_rays, 3]. First value is the ray index. Second value is the sample + * start index in the results for this ray. Third value is the number of samples for + * this ray. Note for rays that have zero samples, we simply skip them so the `packed_info` + * has some zero padding in the end. + * @param starts: Where the frustum-shape sample starts along a ray. [total_samples, 1] + * @param ends: Where the frustum-shape sample ends along a ray. [total_samples, 1] + * @param sigmas Densities at those samples. [total_samples, 1] + * @param rgbs RGBs at those samples. [total_samples, 3] + * @return std::vector + * - accumulated_weight: Ray opacity. [n_rays, 1] + * - accumulated_depth: Ray depth. [n_rays, 1] + * - accumulated_color: Ray color. [n_rays, 3] + * - mask: Boolen value store if this ray has valid samples from packed_info. [n_rays] + */ +std::vector volumetric_rendering_forward( + torch::Tensor packed_info, + torch::Tensor starts, + torch::Tensor ends, + torch::Tensor sigmas, + torch::Tensor rgbs +) { + DEVICE_GUARD(packed_info); + CHECK_INPUT(packed_info); + CHECK_INPUT(starts); + CHECK_INPUT(ends); + CHECK_INPUT(sigmas); + CHECK_INPUT(rgbs); + TORCH_CHECK(packed_info.ndimension() == 2 & packed_info.size(1) == 3); + TORCH_CHECK(starts.ndimension() == 2 & starts.size(1) == 1); + TORCH_CHECK(ends.ndimension() == 2 & ends.size(1) == 1); + TORCH_CHECK(sigmas.ndimension() == 2 & sigmas.size(1) == 1); + TORCH_CHECK(rgbs.ndimension() == 2 & rgbs.size(1) == 3); + + const uint32_t n_rays = packed_info.size(0); + + const int threads = 256; + const int blocks = CUDA_N_BLOCKS_NEEDED(n_rays, threads); + + // outputs + torch::Tensor accumulated_weight = torch::zeros({n_rays, 1}, sigmas.options()); + torch::Tensor accumulated_depth = torch::zeros({n_rays, 1}, sigmas.options()); + torch::Tensor accumulated_color = torch::zeros({n_rays, 3}, sigmas.options()); + // The rays that are not skipped during sampling. + torch::Tensor mask = torch::zeros({n_rays}, sigmas.options().dtype(torch::kBool)); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF( + sigmas.scalar_type(), + "volumetric_rendering_forward", + ([&] + { volumetric_rendering_forward_kernel<<>>( + n_rays, + packed_info.data_ptr(), + starts.data_ptr(), + ends.data_ptr(), + sigmas.data_ptr(), + rgbs.data_ptr(), + accumulated_weight.data_ptr(), + accumulated_depth.data_ptr(), + accumulated_color.data_ptr(), + mask.data_ptr() + ); + })); + + return {accumulated_weight, accumulated_depth, accumulated_color, mask}; +} + + + +/** + * @brief Volumetric Rendering: Accumulating samples in the backward pass. + */ +std::vector volumetric_rendering_backward( + torch::Tensor accumulated_weight, + torch::Tensor accumulated_depth, + torch::Tensor accumulated_color, + torch::Tensor grad_weight, + torch::Tensor grad_depth, + torch::Tensor grad_color, + torch::Tensor packed_info, + torch::Tensor starts, + torch::Tensor ends, + torch::Tensor sigmas, + torch::Tensor rgbs +) { + DEVICE_GUARD(packed_info); + const uint32_t n_rays = packed_info.size(0); + + const int threads = 256; + const int blocks = CUDA_N_BLOCKS_NEEDED(n_rays, threads); + + // outputs + torch::Tensor grad_sigmas = torch::zeros(sigmas.sizes(), sigmas.options()); + torch::Tensor grad_rgbs = torch::zeros(rgbs.sizes(), rgbs.options()); + + AT_DISPATCH_FLOATING_TYPES_AND_HALF( + sigmas.scalar_type(), + "volumetric_rendering_backward", + ([&] + { volumetric_rendering_backward_kernel<<>>( + n_rays, + packed_info.data_ptr(), + starts.data_ptr(), + ends.data_ptr(), + sigmas.data_ptr(), + rgbs.data_ptr(), + accumulated_weight.data_ptr(), + accumulated_depth.data_ptr(), + accumulated_color.data_ptr(), + grad_weight.data_ptr(), + grad_depth.data_ptr(), + grad_color.data_ptr(), + grad_sigmas.data_ptr(), + grad_rgbs.data_ptr() + ); + })); + + return {grad_sigmas, grad_rgbs}; +} \ No newline at end of file diff --git a/pyrad/cuda/include/helpers.cuh b/pyrad/cuda/include/helpers.cuh deleted file mode 100644 index 040f61c9a..000000000 --- a/pyrad/cuda/include/helpers.cuh +++ /dev/null @@ -1,5 +0,0 @@ -/* - * Helper code. - */ - -#define CUDA_NUM_BLOCKS_NEEDED(Q, CUDA_N_THREADS) ((Q - 1) / CUDA_N_THREADS + 1) \ No newline at end of file diff --git a/pyrad/cuda/pyrad_cuda.cpp b/pyrad/cuda/pyrad_cuda.cpp deleted file mode 100644 index 50e9eb219..000000000 --- a/pyrad/cuda/pyrad_cuda.cpp +++ /dev/null @@ -1,30 +0,0 @@ -/* -This file uses pybind to make CUDA calls. -*/ - -#include - -#include - -// CUDA forward declarations - -std::vector sample_uniformly_along_ray_bundle( - torch::Tensor origins, - torch::Tensor directions, - torch::Tensor nears, - torch::Tensor fars, - torch::Tensor offsets, - int max_num_samples); - -// C++ interface - -#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor") -#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") -#define CHECK_INPUT(x) \ - CHECK_CUDA(x); \ - CHECK_CONTIGUOUS(x) - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) -{ - m.def("sample_uniformly_along_ray_bundle", &sample_uniformly_along_ray_bundle); -} \ No newline at end of file diff --git a/pyrad/cuda/pyrad_cuda_kernel.cu b/pyrad/cuda/pyrad_cuda_kernel.cu deleted file mode 100644 index fc0cb4009..000000000 --- a/pyrad/cuda/pyrad_cuda_kernel.cu +++ /dev/null @@ -1,133 +0,0 @@ -/* -Note that this file has the _kernel.cu extension because -setuptools cannot handle files with the same name but different extensions. -*/ - -#include - -#include -#include - -#include - -#define CUDA_N_BLOCKS_NEEDED(Q, CUDA_N_THREADS) ((Q - 1) / CUDA_N_THREADS + 1) - -namespace -{ - template - __device__ __forceinline__ scalar_t sigmoid(scalar_t z) - { - return 1.0 / (1.0 + exp(-z)); - } - - template - __device__ __forceinline__ scalar_t d_sigmoid(scalar_t z) - { - const auto s = sigmoid(z); - return (1.0 - s) * s; - } - - template - __device__ __forceinline__ scalar_t d_tanh(scalar_t z) - { - const auto t = tanh(z); - return 1 - (t * t); - } - - template - __device__ __forceinline__ scalar_t elu(scalar_t z, scalar_t alpha = 1.0) - { - return fmaxf(0.0, z) + fminf(0.0, alpha * (exp(z) - 1.0)); - } - - template - __device__ __forceinline__ scalar_t d_elu(scalar_t z, scalar_t alpha = 1.0) - { - const auto e = exp(z); - const auto d_relu = z < 0.0 ? 0.0 : 1.0; - return d_relu + (((alpha * (e - 1.0)) < 0.0) ? (alpha * e) : 0.0); - } - - template - __global__ void sample_uniformly_along_ray_bundle_kernel( - const torch::PackedTensorAccessor32 origins, - const torch::PackedTensorAccessor32 directions, - const torch::PackedTensorAccessor32 nears, - const torch::PackedTensorAccessor32 fars, - const torch::PackedTensorAccessor32 offsets, - torch::PackedTensorAccessor32 output_time_steps, - torch::PackedTensorAccessor32 output_samples, - torch::PackedTensorAccessor32 output_time_steps_mask) - { - - const int thread_id = blockIdx.x * blockDim.x + threadIdx.x; - const int max_num_samples = output_samples.size(1); - - scalar_t t_min = nears[thread_id]; - scalar_t t_max = fars[thread_id]; - scalar_t offset = offsets[thread_id]; - int sample_index = 0; - - // #pragma unroll - for (scalar_t t = t_min; t <= t_max; t += offset) - { - // #pragma unroll - for (int i = 0; i < 3; ++i) - { - if (sample_index < max_num_samples) - { - scalar_t pos_i = origins[thread_id][i] + t * directions[thread_id][i]; - output_samples[thread_id][sample_index][i] = pos_i; - } - } - if (sample_index < max_num_samples) - { - output_time_steps[thread_id][sample_index] = t; - output_time_steps_mask[thread_id][sample_index] = 1.0; - } - sample_index += 1; - } - } -} // namespace - -std::vector sample_uniformly_along_ray_bundle( - torch::Tensor origins, - torch::Tensor directions, - torch::Tensor nears, - torch::Tensor fars, - torch::Tensor offsets, - int max_num_samples) -{ - std::cout << "sample_uniformly_along_ray_bundle -- start" << std::endl; - - TORCH_CHECK(origins.is_floating_point()); // should be float32 - TORCH_CHECK(origins.ndimension() == 2); // and check that shape is (num_rays, 3) - TORCH_CHECK(origins.size(1) == 3); - const auto num_rays = origins.size(0); - - // TODO: change the number of threads based on GPU - const int threads = 1024; // number of threads in a thread block - // const int blocks = num_rays / threads; - const int blocks = CUDA_N_BLOCKS_NEEDED(num_rays, threads); - - // std::cout << origins.options() << std::endl; - torch::Tensor output_time_steps = torch::zeros({num_rays, max_num_samples}, origins.options()); - torch::Tensor output_samples = torch::zeros({num_rays, max_num_samples, 3}, origins.options()); - torch::Tensor output_time_steps_mask = torch::zeros({num_rays, max_num_samples}, origins.options()); - - // The second argument is used for error messages. - // __FUNCTION__ macro == "sample_uniformly_along_ray_bundle" in this case. - AT_DISPATCH_FLOATING_TYPES(origins.type(), __FUNCTION__, [&] - { sample_uniformly_along_ray_bundle_kernel<<>>( - origins.packed_accessor32(), - directions.packed_accessor32(), - nears.packed_accessor32(), - fars.packed_accessor32(), - offsets.packed_accessor32(), - output_time_steps.packed_accessor32(), - output_samples.packed_accessor32(), - output_time_steps_mask.packed_accessor32()); }); - - std::cout << "sample_uniformly_along_ray_bundle -- end" << std::endl; - return {output_time_steps, output_samples, output_time_steps_mask}; -} diff --git a/pyrad/cuda/sampler.py b/pyrad/cuda/sampler.py deleted file mode 100644 index b61cf51a2..000000000 --- a/pyrad/cuda/sampler.py +++ /dev/null @@ -1,59 +0,0 @@ -# Copyright 2022 The Plenoptix 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. - -""" -Sampler. -""" - -import math - -import torch -from torch import nn - -import pyrad_cuda - - -class LLTMFunction(torch.autograd.Function): - @staticmethod - def forward(ctx, input, weights, bias, old_h, old_cell): - outputs = pyrad_cuda.forward(input, weights, bias, old_h, old_cell) - new_h, new_cell = outputs[:2] - variables = outputs[1:] + [weights] - ctx.save_for_backward(*variables) - - return new_h, new_cell - - @staticmethod - def backward(ctx, grad_h, grad_cell): - outputs = pyrad_cuda.backward(grad_h.contiguous(), grad_cell.contiguous(), *ctx.saved_variables) - d_old_h, d_input, d_weights, d_bias, d_old_cell, d_gates = outputs - return d_input, d_weights, d_bias, d_old_h, d_old_cell - - -class Sampler(nn.Module): - def __init__(self, input_features, state_size): - super().__init__() - self.input_features = input_features - self.state_size = state_size - self.weights = nn.Parameter(torch.Tensor(3 * state_size, input_features + state_size)) - self.bias = nn.Parameter(torch.Tensor(1, 3 * state_size)) - self.reset_parameters() - - def reset_parameters(self): - stdv = 1.0 / math.sqrt(self.state_size) - for weight in self.parameters(): - weight.data.uniform_(-stdv, +stdv) - - def forward(self, input, state): - return LLTMFunction.apply(input, self.weights, self.bias, *state) diff --git a/pyrad/fields/base.py b/pyrad/fields/base.py index 43e11911a..ee69e7b10 100644 --- a/pyrad/fields/base.py +++ b/pyrad/fields/base.py @@ -19,12 +19,12 @@ Base class for the graphs. from abc import abstractmethod from typing import Dict, Optional, Tuple -from torch import nn import torch +from torch import nn from torchtyping import TensorType -from pyrad.fields.modules.field_heads import FieldHeadNames from pyrad.cameras.rays import Frustums, RaySamples +from pyrad.fields.modules.field_heads import FieldHeadNames from pyrad.utils.misc import is_not_none @@ -32,7 +32,7 @@ class Field(nn.Module): """Base class for fields.""" def density_fn(self, positions): - """Returns only the density. Used primarily with the occupancy grid.""" + """Returns only the density. Used primarily with the density grid.""" # Need to figure out a better way to descibe positions with a ray. ray_samples = RaySamples( frustums=Frustums( diff --git a/pyrad/fields/occupancy_fields/base.py b/pyrad/fields/density_fields/base.py similarity index 100% rename from pyrad/fields/occupancy_fields/base.py rename to pyrad/fields/density_fields/base.py diff --git a/pyrad/fields/density_fields/density_grid.py b/pyrad/fields/density_fields/density_grid.py new file mode 100644 index 000000000..03fe5d65a --- /dev/null +++ b/pyrad/fields/density_fields/density_grid.py @@ -0,0 +1,194 @@ +# Copyright 2022 The Plenoptix 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. + +""" +Code to implement the density grid. +""" +from typing import Callable, List, NoReturn, Tuple + +import torch +from torch import nn +from torchtyping import TensorType + +import pyrad.cuda as pyrad_cuda + + +def _create_grid_coords(resolution: int, device: torch.device = "cpu") -> TensorType["n_coords", 3]: + """Create 3D grid coordinates + + Args: + resolution (int): The 3D resolution of the grid. + device (torch.device): Which device you want the returned tensor to live int. + + Returns: + TensorType["n_coords", 3]: All grid coordinates with shape [res * res * res, 3] + """ + arrange = torch.arange(resolution, device=device) + grid_x, grid_y, grid_z = torch.meshgrid([arrange, arrange, arrange], indexing="ij") + coords = torch.stack([grid_x, grid_y, grid_z]) # [3, res, res, res] + coords = coords.reshape(3, -1).t().contiguous() # [N, 3] + return coords + + +class DensityGrid(nn.Module): + """Cascaded multi-res density grids. + + The multi-res grids are all centerd at [center, center, center] in the world space. + The first grid covers regions (center - base_scale / 2, center + base_scale / 2) + and the followings grows by 2x scale each. + + TODO(ruilongli): support `center` and `base_scale` with 3-dim. + + Args: + center (float, optional): Center of all the grids in the world space. Defaults to 0. + base_scale (float, optional): Center of the scale of the base level grid. Defaults to 1. + num_cascades (int, optional): Number of the cascaded multi-res levels. Defaults to 1. + resolution (int, optional): Resolution of the grid. Defaults to 128. + update_every_num_iters (int, optional): How frequently to update the grid values. Defaults to 16. + """ + + def __init__( + self, + center: float = 0.0, + base_scale: float = 1.0, + num_cascades: int = 1, + resolution: int = 128, + update_every_num_iters: int = 16, + ) -> None: + super().__init__() + self.center = center + self.base_scale = base_scale + self.num_cascades = num_cascades # the number of levels (i.e, cascades) + self.resolution = resolution + self.update_every_num_iters = update_every_num_iters + + density_grid = torch.zeros([self.num_cascades] + [self.resolution**3]) + self.register_buffer("density_grid", density_grid) + + density_bitfield = torch.zeros([self.num_cascades * self.resolution**3 // 8], dtype=torch.uint8) + self.register_buffer("density_bitfield", density_bitfield) + + # Integer grid coords / indices that do not related to cascades + grid_coords = _create_grid_coords(resolution) + # TODO(ruilongli): hacky way for now. support cpu version of morton3D? + grid_indices = pyrad_cuda.morton3D(grid_coords.to("cuda:0")).to("cpu") + self.register_buffer("grid_coords", grid_coords) + self.register_buffer("grid_indices", grid_indices) + + @torch.no_grad() + def get_all_cells(self) -> List[Tuple[TensorType["n_cells"], TensorType["n_cells", 3]]]: + """Returns all cells of the grid. + + Returns: + List[Tuple[TensorType["n_cells"], TensorType["n_cells", 3]]]: All grid indices and grid + coordinates of all cascade levels. + """ + return [(self.grid_indices, self.grid_coords)] * self.num_cascades + + @torch.no_grad() + def sample_uniform_and_occupied_cells(self, n: int) -> List[Tuple[TensorType["n"], TensorType["n", 3]]]: + """Samples both n uniform and occupied cells (per cascade) + + Args: + n: Number of cells to be sampled. + + Returns: + List[Tuple[TensorType["n"], TensorType["n", 3]]]: Sampled grid indices and grid + coordinates of all cascade levels. + """ + device = self.density_grid.device + + cells = [] + for c in range(self.num_cascades): + uniform_coords = torch.randint(self.resolution, (n, 3), device=device) + uniform_indices = pyrad_cuda.morton3D(uniform_coords) + + occupied_indices = torch.nonzero(self.density_grid[c] > 0)[:, 0] + if n < len(occupied_indices): + selector = torch.randint(len(occupied_indices), (n,), device=device) + occupied_indices = occupied_indices[selector] + occupied_coords = pyrad_cuda.morton3D_invert(occupied_indices) + + # concatenate + coords = torch.cat([uniform_coords, occupied_coords], dim=0) + indices = torch.cat([uniform_indices, occupied_indices], dim=0) + cells.append((indices, coords)) + return cells + + @torch.no_grad() + def update_density_grid( + self, + density_eval_func: Callable, + step: int, + density_threshold: float = 2, # 0.01 / (SQRT3 / 1024 * 3) + decay: float = 0.95, + ) -> NoReturn: + """Update the density grid in EMA way. + + Args: + density_eval_func (Callable): A Callable function that takes in sample positions (N, 3) and + returns densities (N, 1). + step (int): The current training step. Instant-NGP has a warmup stage where all cells are updated. + After certain amount of steps (256), it speeds up by only update sampled cells. + density_threshold (float): The threshold to prune cells. Recommand to calculate it using this rule: + `weight_threshold / dt`. For example for Instant-NGP on Lego, the minimum step size `dt` is + `3` (scale of the scene) * `sqrt(3)` / `1024` (number of samples). And if we want to the weight + threshold to be `0.01`, we will get `0.01 / (3 * sqrt(3) / 1024) ~= 2.` for density threshold. + decay (float): EMA decay for density updating. + """ + # create temporary grid + tmp_grid = -torch.ones_like(self.density_grid) + if step < 256: + cells = self.get_all_cells() + else: + N = self.resolution**3 // 4 + cells = self.sample_uniform_and_occupied_cells(N) + + # infer sigmas + for mip_level in range(self.num_cascades): + mip_scale = 2 ** (-mip_level) / self.base_scale + indices, coords = cells[mip_level] + # `coords` denotes the i-th cell. It's a poor naming here. + # the actually coordinates in world space x has mapping to + # the cell like this: + # x \in [0, 1/res] maps to 0-th cell. + # x \in [1 - 1/res, 1] maps to (res-1)-th cell. + x = (coords + torch.rand_like(coords.float())) / self.resolution + x = (x - 0.5) / mip_scale + self.center + tmp_grid[mip_level, indices] = density_eval_func(x).squeeze(-1) + + # ema update + valid_mask = (self.density_grid >= 0) & (tmp_grid >= 0) + self.density_grid[valid_mask] = torch.maximum(self.density_grid[valid_mask] * decay, tmp_grid[valid_mask]) + mean_density = self.density_grid.mean().item() + + # pack to bitfield + self.density_bitfield.data = pyrad_cuda.packbits(self.density_grid, min(mean_density, density_threshold)) + + # TODO(ruilongli): max pooling? https://github.com/NVlabs/instant-ngp/blob/master/src/testbed_nerf.cu#L578 + + @torch.no_grad() + def get_densities(self, xyzs: TensorType[..., 3]) -> TensorType[..., 1]: + """Trilinear interpolation to get the density values. + Args: + xyzs (TensorType[..., 3]): 3D querry coordinate + Returns: + TensorType[..., 1]: Density values + """ + # TODO(ruilongli) pass in RaySamples because we need dt to decide mip level. + raise NotImplementedError("We haven't implement the logic of querying densities from cascaded grid.") + + def forward(self, x): + """Not implemented.""" + raise RuntimeError("Shouldn't be called!") diff --git a/pyrad/fields/instant_ngp_field.py b/pyrad/fields/instant_ngp_field.py index 7cee20855..d3c69dbc8 100644 --- a/pyrad/fields/instant_ngp_field.py +++ b/pyrad/fields/instant_ngp_field.py @@ -22,11 +22,11 @@ from typing import Tuple import torch from torch.nn.parameter import Parameter +from pyrad.cameras.rays import RaySamples +from pyrad.fields.base import Field from pyrad.fields.modules.encoding import Encoding, HashEncoding, SHEncoding from pyrad.fields.modules.field_heads import FieldHeadNames -from pyrad.fields.base import Field from pyrad.fields.nerf_field import NeRFField -from pyrad.cameras.rays import RaySamples from pyrad.utils.activations import trunc_exp try: @@ -113,7 +113,7 @@ class TCNNInstantNGPField(Field): # Rectifying the density with an exponential is much more stable than a ReLU or # softplus, because it enables high post-activation (float32) density outputs # from smaller internal (float16) parameters. - density = trunc_exp(density_before_activation) + density = trunc_exp(density_before_activation.to(positions)) return density, base_mlp_out def get_outputs(self, ray_samples: RaySamples, density_embedding=None): @@ -123,7 +123,7 @@ class TCNNInstantNGPField(Field): directions_flat = directions.view(-1, 3) d = self.direction_encoding(directions_flat) h = torch.cat([d, density_embedding.view(-1, self.geo_feat_dim)], dim=-1) - rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1) + rgb = self.mlp_head(h).view(*ray_samples.frustums.directions.shape[:-1], -1).to(directions) return {FieldHeadNames.RGB: rgb} diff --git a/pyrad/fields/occupancy_fields/occupancy_grid.py b/pyrad/fields/occupancy_fields/occupancy_grid.py deleted file mode 100644 index 73b858ab0..000000000 --- a/pyrad/fields/occupancy_fields/occupancy_grid.py +++ /dev/null @@ -1,144 +0,0 @@ -# Copyright 2022 The Plenoptix 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. - -""" -Code to implement the density grid. -""" - -import torch -import torch.nn.functional as F -from torch import nn -from torchtyping import TensorType - -from pyrad.utils.misc import is_not_none - - -class OccupancyGrid(nn.Module): - """Module to keep track of the density and occupancy.""" - - def __init__( - self, num_cascades: int = 1, resolution: int = 128, aabb=None, update_every_num_iters: int = 16 - ) -> None: - super().__init__() - assert is_not_none(aabb), "The axis-aligned bounding box aabb is not defined!" - self.num_cascades = num_cascades # the number of levels (i.e, cascades) - self.resolution = resolution - self.register_buffer("aabb", aabb) # axis-aligned bounding box - occupancy_grid = torch.ones( - [self.num_cascades] + [self.resolution] * 3 - ) # (num_cascades, x_resolution, y_resolution, z_resolution) - # not a module parameter, but it still should be saved and part of `state_dict` - self.register_buffer("occupancy_grid", occupancy_grid) - self.mean_density = 0.0 - - self.update_every_num_iters = update_every_num_iters - - def reset(self): - """Zero out the occupancy grid.""" - self.occupancy_grid.zero_() - - @torch.no_grad() - def mark_untrained_grid(self): - """Marks the locations in the density grid that are never seen by training views.""" - # TODO: only query points that are covered by the training cameras viewing frustrum - # We'll need to pass in the camera poses and intrinsics to this function - raise NotImplementedError() - - @torch.no_grad() - def update_occupancy_grid(self, density_fn, decay=0.95, split_size=64): - """Update the occupancy grid with density values. - - Args: - density_fn: Function from takes points in the shape (..., 3) and - returns a tensor of densities of shape (..., 1). - """ - # TODO(ethan): remove this pylint disable - # pylint: disable=invalid-name - device = self.occupancy_grid.device - - assert self.num_cascades == 1, "Still need to implement code for multiple cascades." - temp_occupancy_grid = torch.zeros_like(self.occupancy_grid) - - X_length, Y_length, Z_length = self.aabb[1] - self.aabb[0] - X_half_voxel_length = 0.5 * X_length / self.resolution - Y_half_voxel_length = 0.5 * Y_length / self.resolution - Z_half_voxel_length = 0.5 * Z_length / self.resolution - half_voxel_lengths = torch.tensor( - [X_half_voxel_length, Y_half_voxel_length, Z_half_voxel_length], device=device - ) - X = torch.linspace( - self.aabb[0][0] + X_half_voxel_length, self.aabb[1][0] - X_half_voxel_length, self.resolution, device=device - ).split(split_size) - Y = torch.linspace( - self.aabb[0][1] + Y_half_voxel_length, self.aabb[1][1] - Y_half_voxel_length, self.resolution, device=device - ).split(split_size) - Z = torch.linspace( - self.aabb[0][2] + Z_half_voxel_length, self.aabb[1][2] - Z_half_voxel_length, self.resolution, device=device - ).split(split_size) - for x_index, xs in enumerate(X): - for y_index, ys in enumerate(Y): - for z_index, zs in enumerate(Z): - x_len, y_len, z_len = len(xs), len(ys), len(zs) - grid_x, grid_y, grid_z = torch.meshgrid(xs, ys, zs, indexing="ij") - xyzs = torch.stack([grid_x, grid_y, grid_z], dim=-1) # (resolution, resolution,, resolution, 3) - # add noise to the points - xyzs_rand = ( - torch.rand_like(xyzs) * half_voxel_lengths - ) # range [-half_voxel_length, half_voxel_length] for each dimension - xyzs += xyzs_rand - densities = density_fn(xyzs) - assert densities.shape[-1] == 1 - temp_occupancy_grid[ - 0, - x_index * split_size : x_index * split_size + x_len, - y_index * split_size : y_index * split_size + y_len, - z_index * split_size : z_index * split_size + z_len, - ] = densities[..., 0] - - # exponential moving average update - valid_mask = self.occupancy_grid >= 0.0 - self.occupancy_grid[valid_mask] = torch.maximum( - self.occupancy_grid[valid_mask] * decay, temp_occupancy_grid[valid_mask] - ) - self.mean_density = torch.mean(self.occupancy_grid[valid_mask]).item() - - @torch.no_grad() - def get_densities(self, xyzs: TensorType[..., 3]) -> TensorType[..., 1]: - """Trilinear interpolation to get the density values. - - Args: - xyzs (TensorType[..., 3]): 3D querry coordinate - - Returns: - TensorType[..., 1]: Density values - """ - occupancy_grid = self.occupancy_grid[None, ...] # shape (1, num_cascades, X_res, Y_res, Z_res) - xyzs_shape = xyzs.shape - xyzs_reshaped = xyzs.view(1, -1, 1, 1, 3) - voxel_lengths = self.aabb[1] - self.aabb[0] - xyzs_reshaped_normalized = ((xyzs_reshaped - self.aabb[0]) / voxel_lengths) * 2.0 - 1.0 - densities = F.grid_sample( - occupancy_grid.permute( - 0, 1, 4, 3, 2 - ), # (xyz to zyx) for grid sample useage because the xyzs have xyz ordering - xyzs_reshaped_normalized, - align_corners=True, - padding_mode="zeros", - ) - densities = densities.view(*xyzs_shape[:-1], 1) - return densities - - def forward(self, x): - """Not implemented.""" - raise NotImplementedError() diff --git a/pyrad/graphs/base.py b/pyrad/graphs/base.py index 5e8d30b7d..28c21b000 100644 --- a/pyrad/graphs/base.py +++ b/pyrad/graphs/base.py @@ -76,7 +76,7 @@ class Graph(AbstractGraph): intrinsics (torch.Tensor): Camera intrinsics. camera_to_world (torch.Tensor): Camera to world transformation. loss_coefficients (DictConfig): Loss specific weights. - steps_per_occupancy_grid_update (int): How often to update occupancy grid. + steps_per_density_grid_update (int): How often to update density grid. scene_bounds (SceneBounds): Bounds of target scene. collider_config (DictConfig): Configuration of scene collider. """ @@ -86,19 +86,19 @@ class Graph(AbstractGraph): intrinsics: torch.Tensor = None, camera_to_world: torch.Tensor = None, loss_coefficients: DictConfig = None, - steps_per_occupancy_grid_update: int = 16, + steps_per_density_grid_update: int = 16, scene_bounds: SceneBounds = None, collider_config: DictConfig = None, **kwargs, ) -> None: super().__init__() - assert is_not_none(scene_bounds), "scene_bounds is needed to use the occupancy grid" + assert is_not_none(scene_bounds), "scene_bounds is needed to use the density grid" self.intrinsics = intrinsics self.camera_to_world = camera_to_world self.scene_bounds = scene_bounds self.collider_config = collider_config self.loss_coefficients = loss_coefficients - self.steps_per_occupancy_grid_update = steps_per_occupancy_grid_update + self.steps_per_density_grid_update = steps_per_density_grid_update self.kwargs = kwargs self.collider = None self.ray_generator = RayGenerator(self.intrinsics, self.camera_to_world) diff --git a/pyrad/graphs/instant_ngp.py b/pyrad/graphs/instant_ngp.py index c074bde16..d47b5aaa2 100644 --- a/pyrad/graphs/instant_ngp.py +++ b/pyrad/graphs/instant_ngp.py @@ -21,19 +21,18 @@ from typing import Dict, List import torch from torch.nn import Parameter from torchmetrics import PeakSignalNoiseRatio -from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from torchmetrics.functional import structural_similarity_index_measure +from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity -from pyrad.fields.modules.field_heads import FieldHeadNames -from pyrad.fields.instant_ngp_field import field_implementation_to_class -from pyrad.graphs.base import Graph -from pyrad.optimizers.loss import MSELoss -from pyrad.fields.occupancy_fields.occupancy_grid import OccupancyGrid -from pyrad.renderers.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer -from pyrad.graphs.modules.ray_sampler import UniformSampler -from pyrad.utils import colors +import pyrad.cuda as pyrad_cuda from pyrad.cameras.rays import RayBundle -from pyrad.utils import visualization, writer +from pyrad.fields.density_fields.density_grid import DensityGrid +from pyrad.fields.instant_ngp_field import field_implementation_to_class +from pyrad.fields.modules.field_heads import FieldHeadNames +from pyrad.graphs.base import Graph +from pyrad.graphs.modules.ray_sampler import NGPSpacedSampler +from pyrad.optimizers.loss import MSELoss +from pyrad.utils import colors, visualization, writer from pyrad.utils.callbacks import Callback @@ -51,9 +50,9 @@ class NGPGraph(Graph): """defining callbacks to run after every training iteration""" self.callbacks = [ Callback( - self.occupancy_grid.update_every_num_iters, - self.occupancy_grid.update_occupancy_grid, - density_fn=self.field.density_fn, + self.density_grid.update_every_num_iters, + self.density_grid.update_density_grid, + density_eval_func=self.field.density_fn, ) ] @@ -63,19 +62,11 @@ class NGPGraph(Graph): self.field = field_implementation_to_class[self.field_implementation](self.scene_bounds.aabb) def populate_misc_modules(self): - # occupancy grid - self.occupancy_grid = OccupancyGrid(aabb=self.scene_bounds.aabb) + # density grid + self.density_grid = DensityGrid(center=0.0, base_scale=3, num_cascades=1) # samplers - self.sampler_occupancy_grid = UniformSampler(num_samples=128) - - # TODO stabalize occupancy grid. - # self.sampler_occupancy_grid = UniformSampler(num_samples=128, occupancy_field=self.occupancy_grid) - - # renderers - self.renderer_rgb = RGBRenderer(background_color=colors.WHITE) - self.renderer_accumulation = AccumulationRenderer() - self.renderer_depth = DepthRenderer() + self.sampler = NGPSpacedSampler(num_samples=1024, density_field=self.density_grid) # losses self.rgb_loss = MSELoss() @@ -90,37 +81,39 @@ class NGPGraph(Graph): param_groups["fields"] = list(self.field.parameters()) return param_groups + @torch.cuda.amp.autocast() def get_outputs(self, ray_bundle: RayBundle): + # TODO(ruilongli) + # - train test difference + # - visualize "depth_density_grid" + ray_samples, packed_info, t_min, t_max = self.sampler(ray_bundle, self.field.aabb) - # uniform sampling - ray_samples = self.sampler_occupancy_grid(ray_bundle) field_outputs = self.field.forward(ray_samples) - weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY]) + rgbs = field_outputs[FieldHeadNames.RGB] + sigmas = field_outputs[FieldHeadNames.DENSITY] - weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY]) - rgb = self.renderer_rgb( - rgb=field_outputs[FieldHeadNames.RGB], - weights=weights, + accumulated_weight, accumulated_depth, accumulated_color, mask = pyrad_cuda.VolumeRenderer.apply( + packed_info, ray_samples.frustums.starts, ray_samples.frustums.ends, sigmas, rgbs ) - accumulation = self.renderer_accumulation(weights) - depth = self.renderer_depth(weights, ray_samples) - - densities_occupancy_grid = self.occupancy_grid.get_densities(ray_samples.frustums.get_positions()) - weights_occupancy_grid = ray_samples.get_weights(densities_occupancy_grid) - depth_occupancy_grid = self.renderer_depth(weights_occupancy_grid, ray_samples) + accumulated_depth = torch.clip(accumulated_depth, t_min[:, None], t_max[:, None]) + accumulated_color = accumulated_color + colors.WHITE.to(accumulated_color) * (1.0 - accumulated_weight) outputs = { - "rgb": rgb, - "accumulation": accumulation, - "depth": depth, - "depth_occupancy_grid": depth_occupancy_grid, + "rgb": accumulated_color, + "accumulation": accumulated_weight, + "depth": accumulated_depth, + "mask": mask, # the rays we skipped from sampler } return outputs def get_loss_dict(self, outputs, batch): device = self.get_device() image = batch["image"].to(device) - rgb_loss = self.rgb_loss(image, outputs["rgb"]) + if "mask" in outputs: + mask = outputs["mask"] + rgb_loss = self.rgb_loss(image[mask], outputs["rgb"][mask]) + else: + rgb_loss = self.rgb_loss(image, outputs["rgb"]) loss_dict = {"rgb_loss": rgb_loss} return loss_dict @@ -137,8 +130,6 @@ class NGPGraph(Graph): outputs["accumulation"] = combined_acc combined_depth = torch.cat([depth], dim=1) outputs["depth"] = combined_depth - depth = visualization.apply_depth_colormap(outputs["depth_occupancy_grid"]) - outputs["depth_occupancy_grid"] = combined_depth def log_test_image_outputs(self, image_idx, step, batch, outputs): image = batch["image"] @@ -157,9 +148,6 @@ class NGPGraph(Graph): writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step) writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step) - occupancy_depth = visualization.apply_depth_colormap(outputs["depth_occupancy_grid"]) - writer.put_image(name=f"depth_occupancy_grid/image_idx_{image_idx}", image=occupancy_depth, step=step) - # Switch images from [H, W, C] to [1, C, H, W] for metrics computations image = torch.moveaxis(image, -1, 0)[None, ...] rgb = torch.moveaxis(rgb, -1, 0)[None, ...] diff --git a/pyrad/graphs/modules/ray_sampler.py b/pyrad/graphs/modules/ray_sampler.py index 34e5d5be7..2306cee70 100644 --- a/pyrad/graphs/modules/ray_sampler.py +++ b/pyrad/graphs/modules/ray_sampler.py @@ -17,25 +17,26 @@ Collection of sampling strategies """ from abc import abstractmethod -from typing import Callable, Optional +from typing import Callable, Optional, Tuple import torch from torch import nn from torchtyping import TensorType -from pyrad.cameras.rays import RayBundle, RaySamples -from pyrad.fields.occupancy_fields.occupancy_grid import OccupancyGrid +import pyrad.cuda as pyrad_cuda +from pyrad.cameras.rays import Frustums, RayBundle, RaySamples +from pyrad.fields.density_fields.density_grid import DensityGrid class Sampler(nn.Module): """Generate Samples""" def __init__( - self, num_samples: int, occupancy_field: Optional[OccupancyGrid] = None, weight_threshold: float = 1e-4 + self, num_samples: int, density_field: Optional[DensityGrid] = None, weight_threshold: float = 1e-4 ) -> None: super().__init__() self.num_samples = num_samples - self.occupancy_field = occupancy_field + self.density_field = density_field self.weight_threshold = weight_threshold @abstractmethod @@ -43,10 +44,10 @@ class Sampler(nn.Module): """Generate Ray Samples""" def forward(self, *args, **kwargs) -> RaySamples: - """Generate ray samples with optional occupancy filtering""" + """Generate ray samples with optional density filtering""" ray_samples = self.generate_ray_samples(*args, **kwargs) - if self.occupancy_field is not None: - densities = self.occupancy_field.get_densities(ray_samples.frustums.get_positions()) + if self.density_field is not None: + densities = self.density_field.get_densities(ray_samples.frustums.get_positions()) weights = ray_samples.get_weights(densities) valid_mask = weights >= self.weight_threshold @@ -63,7 +64,7 @@ class SpacedSampler(Sampler): spacing_fn: Callable, spacing_fn_inv: Callable, train_stratified=True, - occupancy_field: Optional[OccupancyGrid] = None, + density_field: Optional[DensityGrid] = None, weight_threshold: float = 1e-4, ) -> None: """ @@ -72,11 +73,11 @@ class SpacedSampler(Sampler): spacing_fn (Callable): Function that dictates sample spacing (ie `lambda x : x` is uniform). spacing_fn_inv (Callable): The inverse of spacing_fn. train_stratified (bool): Use stratified sampling during training. Defults to True - occupancy_field (OccupancyGrid, optional): Occupancy grid. If provides, + density_field (DensityGrid, optional): Density grid. If provides, samples below weight_threshold as set as invalid. - weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided. + weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided. """ - super().__init__(num_samples=num_samples, occupancy_field=occupancy_field, weight_threshold=weight_threshold) + super().__init__(num_samples=num_samples, density_field=density_field, weight_threshold=weight_threshold) self.train_stratified = train_stratified self.spacing_fn = spacing_fn self.spacing_fn_inv = spacing_fn_inv @@ -127,23 +128,23 @@ class UniformSampler(SpacedSampler): self, num_samples: int, train_stratified=True, - occupancy_field: Optional[OccupancyGrid] = None, + density_field: Optional[DensityGrid] = None, weight_threshold: float = 1e-4, ) -> None: """ Args: num_samples (int): Number of samples per ray train_stratified (bool): Use stratified sampling during training. Defults to True - occupancy_field (OccupancyGrid, optional): Occupancy grid. If provides, + density_field (DensityGrid, optional): Density grid. If provides, samples below weight_threshold as set as invalid. - weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided. + weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided. """ super().__init__( num_samples=num_samples, spacing_fn=lambda x: x, spacing_fn_inv=lambda x: x, train_stratified=train_stratified, - occupancy_field=occupancy_field, + density_field=density_field, weight_threshold=weight_threshold, ) @@ -155,23 +156,23 @@ class LinearDisparitySampler(SpacedSampler): self, num_samples: int, train_stratified=True, - occupancy_field: Optional[OccupancyGrid] = None, + density_field: Optional[DensityGrid] = None, weight_threshold: float = 1e-4, ) -> None: """ Args: num_samples (int): Number of samples per ray train_stratified (bool): Use stratified sampling during training. Defults to True - occupancy_field (OccupancyGrid, optional): Occupancy grid. If provides, + density_field (DensityGrid, optional): Density grid. If provides, samples below weight_threshold as set as invalid. - weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided. + weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided. """ super().__init__( num_samples=num_samples, spacing_fn=lambda x: 1 / x, spacing_fn_inv=lambda x: 1 / x, train_stratified=train_stratified, - occupancy_field=occupancy_field, + density_field=density_field, weight_threshold=weight_threshold, ) @@ -183,23 +184,23 @@ class SqrtSampler(SpacedSampler): self, num_samples: int, train_stratified=True, - occupancy_field: Optional[OccupancyGrid] = None, + density_field: Optional[DensityGrid] = None, weight_threshold: float = 1e-4, ) -> None: """ Args: num_samples (int): Number of samples per ray train_stratified (bool): Use stratified sampling during training. Defults to True - occupancy_field (OccupancyGrid, optional): Occupancy grid. If provides, + density_field (DensityGrid, optional): Density grid. If provides, samples below weight_threshold as set as invalid. - weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided. + weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided. """ super().__init__( num_samples=num_samples, spacing_fn=torch.sqrt, spacing_fn_inv=lambda x: x**2, train_stratified=train_stratified, - occupancy_field=occupancy_field, + density_field=density_field, weight_threshold=weight_threshold, ) @@ -211,23 +212,23 @@ class LogSampler(SpacedSampler): self, num_samples: int, train_stratified=True, - occupancy_field: Optional[OccupancyGrid] = None, + density_field: Optional[DensityGrid] = None, weight_threshold: float = 1e-4, ) -> None: """ Args: num_samples (int): Number of samples per ray train_stratified (bool): Use stratified sampling during training. Defults to True - occupancy_field (OccupancyGrid, optional): Occupancy grid. If provides, + density_field (DensityGrid, optional): Density grid. If provides, samples below weight_threshold as set as invalid. - weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided. + weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided. """ super().__init__( num_samples=num_samples, spacing_fn=torch.log, spacing_fn_inv=torch.exp, train_stratified=train_stratified, - occupancy_field=occupancy_field, + density_field=density_field, weight_threshold=weight_threshold, ) @@ -240,7 +241,7 @@ class PDFSampler(Sampler): num_samples: int, train_stratified: bool = True, include_original: bool = True, - occupancy_field: OccupancyGrid = None, + density_field: DensityGrid = None, weight_threshold: float = 1e-4, histogram_padding: float = 0.01, ) -> None: @@ -249,12 +250,12 @@ class PDFSampler(Sampler): num_samples (int): Number of samples per ray train_stratified: boolean: Randomize location within each bin during training. Defaults to True include_original: Add original samples to ray. Defaults to True - occupancy_field (OccupancyGrid, optional): Occupancy grid. If provides, + density_field (DensityGrid, optional): Density grid. If provides, samples below weight_threshold as set as invalid. - weight_thershold (float): Removes samples below threshold weight. Only used if occupancy field is provided. + weight_threshold (float): Removes samples below threshold weight. Only used if density field is provided. histogram_padding (float): Amount to weights prior to computing PDF. Defaults to 0.01. """ - super().__init__(num_samples=num_samples, occupancy_field=occupancy_field, weight_threshold=weight_threshold) + super().__init__(num_samples=num_samples, density_field=density_field, weight_threshold=weight_threshold) self.train_stratified = train_stratified self.include_original = include_original self.histogram_padding = histogram_padding @@ -336,3 +337,83 @@ class PDFSampler(Sampler): ray_samples = ray_bundle.get_ray_samples(bin_starts=bins[..., :-1, None], bin_ends=bins[..., 1:, None]) return ray_samples + + +class NGPSpacedSampler(Sampler): + """Sampler that matches Instant-NGP paper. + + This sampler does ray-box AABB test, ray samples generation and density check, all together. + + TODO(ruilongli): check whether fuse AABB test together with `raymarching` can speed things up. + Otherwise we can seperate it out and use collision detecoter that already exists in the repo. + """ + + def generate_ray_samples(self) -> RaySamples: + raise RuntimeError("For NGP we fused ray samples and density check together. Please call forward() directly.") + + # pylint: disable=arguments-differ + def forward( + self, + ray_bundle: RayBundle, + aabb: TensorType[2, 3], + num_samples: Optional[int] = None, + ) -> Tuple[RaySamples, TensorType["total_samples", 3], TensorType["total_samples"], TensorType["total_samples"]]: + """Generate ray samples in a bounding box. + + TODO(ruilongli): write a Packed[Ray_samples] class to ray_samples with packed_info. + TODO(ruilongli): maybe move aabb test to the collision detector? + + Args: + ray_bundle (RayBundle): Rays to generate samples for + aabb (TensorType[2, 3]): Bounding box of the scene. + num_samples (Optional[int]): Number of samples per ray + + Returns: + Tuple[RaySamples, TensorType["total_samples", 3], TensorType["total_samples"], TensorType["total_samples"]]: + First return is ray samples in a packed way where only the valid samples are kept. + Second return contains all the information to recover packed samples into unpacked mode for rendering. + The last two returns are t_min and t_max from ray-aabb test. + """ + num_samples = num_samples or self.num_samples + + aabb = aabb.flatten() + rays_o = ray_bundle.origins.contiguous() + rays_d = ray_bundle.directions.contiguous() + t_min, t_max = pyrad_cuda.ray_aabb_intersect(rays_o, rays_d, aabb) + scene_scale = (aabb[3:6] - aabb[0:3]).max() + + # TODO(ruilongli): * 16 is for original impl for training. Need to run + # some profiling test with this choice. + max_samples_per_batch = len(rays_o) * num_samples + + packed_info, origins, dirs, starts, ends = pyrad_cuda.raymarching( + # rays + rays_o, + rays_d, + t_min, + t_max, + # density grid + self.density_field.center, + self.density_field.base_scale, + self.density_field.num_cascades, + self.density_field.resolution, + self.density_field.density_bitfield, + # sampling args + max_samples_per_batch, + num_samples, + 0.0, + scene_scale, + ) + total_samples = max(packed_info[:, -1].sum(), 1) + packed_info = packed_info[:total_samples] + origins = origins[:total_samples] + dirs = dirs[:total_samples] + starts = starts[:total_samples] + ends = ends[:total_samples] + + zeros = torch.zeros_like(origins[:, :1]) + + ray_samples = RaySamples( + frustums=Frustums(origins=origins, directions=dirs, starts=starts, ends=ends, pixel_area=zeros), + ) + return ray_samples, packed_info, t_min, t_max diff --git a/pyrad/utils/callbacks.py b/pyrad/utils/callbacks.py index 691c4e0a5..c2de09f49 100644 --- a/pyrad/utils/callbacks.py +++ b/pyrad/utils/callbacks.py @@ -43,4 +43,4 @@ class Callback: step (int): current iteration step """ if step % self.update_every_num_iters == 0: - self.func(*self.args, **self.kwargs) + self.func(*self.args, **self.kwargs, step=step) diff --git a/pyrad/utils/visualization.py b/pyrad/utils/visualization.py index 94d103e7f..11d00193a 100644 --- a/pyrad/utils/visualization.py +++ b/pyrad/utils/visualization.py @@ -15,6 +15,7 @@ """ Helper functions for visualizing outputs """ from typing import Optional + import torch from matplotlib import cm from torchtyping import TensorType diff --git a/setup.py b/setup.py index e31e2ae57..2263a6ef8 100644 --- a/setup.py +++ b/setup.py @@ -1,36 +1,12 @@ import os + from setuptools import setup ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) -USE_CUDA = True ext_modules = [] cmdclass = {} setup_requires = [] -if USE_CUDA: - try: - from torch.utils.cpp_extension import BuildExtension, CUDAExtension - - ext_modules = [ - CUDAExtension( - name="pyrad_cuda", - sources=[ - "pyrad/cuda/pyrad_cuda.cpp", - "pyrad/cuda/pyrad_cuda_kernel.cu", - ], - include_dirs=[os.path.join("pyrad", "cuda", "include")], - extra_compile_args={ - "cxx": [], - "nvcc": [], - }, - verbose=True, - ), - ] - cmdclass = {"build_ext": BuildExtension} - setup_requires = ["pybind11>=2.5.0"] - except Exception as e: - print(e) - print("Failed to build the CUDA extension.") if __name__ == "__main__": setup( diff --git a/tests/cuda/test.py b/tests/cuda/test.py new file mode 100644 index 000000000..ec34acb52 --- /dev/null +++ b/tests/cuda/test.py @@ -0,0 +1,3 @@ +if __name__ == "__main__": + # TODO(ruilongli): write tests + pass diff --git a/tests/cuda/test_mattport_cuda.py b/tests/cuda/test_mattport_cuda.py deleted file mode 100644 index cb5efb5da..000000000 --- a/tests/cuda/test_mattport_cuda.py +++ /dev/null @@ -1,102 +0,0 @@ -""" -Test the CUDA bindings. -""" -import torch -from pyrad.cuda.sampler import Sampler -import time -import pyrad_cuda -import os - - -def test_gradients(): - """_summary_""" - pass - - -def test_outputs(): - """_summary_""" - pass - - -def pytorch_sample_uniformly_along_ray_bundle(origins, directions, nears, fars, range_offsets): - """Sample along every ray with a constant step size.""" - time_steps = nears[:, None] + range_offsets[None, :] # (num_rays, max_num_samples) - samples = origins[:, None, :] + directions[:, None, :] * time_steps[:, :, None] - time_steps_mask = (time_steps <= fars[:, None]).float() - return time_steps, samples, time_steps_mask - - -def get_rgb_density(samples, time_steps_mask): - num_rays, max_num_samples = time_steps_mask.shape - total_max_num_samples = num_rays * max_num_samples - - time_steps_mask_flat = time_steps_mask.view(total_max_num_samples) == 1.0 - samples_flat = samples.view(total_max_num_samples, 3) - samples_flat_masked = samples_flat[time_steps_mask_flat] - - # TODO: this goes through the forward pass of the MLP - - rgb = torch.zeros((num_rays, max_num_samples, 3)) - density = torch.zeros((num_rays, max_num_samples)) - rgb.view(total_max_num_samples, 3)[time_steps_mask_flat] = 1.0 - density.view(total_max_num_samples)[time_steps_mask_flat] = 1.0 - return rgb, density - - -def test_benchmark(): - """The test cases assume that rays can be at most of length 1.0.""" - - device = torch.device("cuda") - dtype = torch.float32 - - near_range = [0, 0.2] - far_range = [0.8, 1.0] - max_num_samples = 128 * 8 - offset = 1.0 / max_num_samples - - print("offset is", offset) - # import sys; sys.exit() - - print("Running ray sampling tests.") - # origins = torch.tensor() - # we want these to be dtype=torch.float32 and device="cuda" - num_rays = 1024 * 8 - print("num_rays:", num_rays) - print("max_num_samples:", max_num_samples) - origins = torch.zeros((num_rays, 3)).to(device) - directions = torch.ones((num_rays, 3)).to(device) - nears = torch.rand((num_rays)).to(device) * (near_range[1] - near_range[0]) + near_range[0] - fars = torch.rand((num_rays)).to(device) * (far_range[1] - far_range[0]) + far_range[0] - offsets = torch.ones(max_num_samples).to(device) * offset - range_offsets = torch.arange(0, max_num_samples).to(device) * offset - - start = time.time() - time_steps_pt, samples_pt, time_steps_mask_pt = pytorch_sample_uniformly_along_ray_bundle( - origins, directions, nears, fars, range_offsets - ) - PYTORCH_SAMPLE_TIME = time.time() - start - print("PYTORCH_SAMPLE_TIME:", PYTORCH_SAMPLE_TIME) - - start = time.time() - time_steps_cu, samples_cu, time_steps_mask_cu = pyrad_cuda.sample_uniformly_along_ray_bundle( - origins, directions, nears, fars, offsets, max_num_samples - ) - CUDA_SAMPLE_TIME = time.time() - start - print("CUDA_SAMPLE_TIME: ", CUDA_SAMPLE_TIME) - - # assert torch.allclose(time_steps_pt, time_steps_cu) - # assert torch.allclose(samples_pt, samples_cu) - # assert torch.allclose(time_steps_mask_pt, time_steps_mask_cu) - - # rgb_pt, density_pt = get_rgb_density(samples_pt, time_steps_mask_pt) - # rgb_cu, density_cu = get_rgb_density(samples_cu, time_steps_mask_cu) - - perc_samples_used = float(time_steps_mask_pt.sum() / (time_steps_mask_pt.shape[0] * time_steps_mask_pt.shape[1])) - perc_samples_used_str = str(round(perc_samples_used * 100.0, 2)) - print(f"Using field network for %{perc_samples_used_str} of samples in this ray bundle.") - - # TODO: rendering of rgb, density, and deltas with pytorch vs. cuda - - -if __name__ == "__main__": - test_benchmark() diff --git a/tests/graphs/modules/test_density_grid.py b/tests/graphs/modules/test_density_grid.py new file mode 100644 index 000000000..eac757139 --- /dev/null +++ b/tests/graphs/modules/test_density_grid.py @@ -0,0 +1,26 @@ +""" +Code to test the density grid. +""" + +# import torch + +# from pyrad.fields.density_fields.density_grid import DensityGrid + + +def test_density_grid(): + """For testing the density grid.""" + # TODO(ruilongli): currently the DensityGrid only supports GPU init. + # after we support for CPU init then we re-enable this test + + # device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + # num_cascades = 1 + # resolution = 128 + # aabb = torch.tensor([[-1, -1, -1], [1, 1, 1]], dtype=torch.float32).to(device) + # grid = DensityGrid(num_cascades, resolution, aabb) + # grid.to(device) + # density_eval_func = lambda x: x[..., 0:1] # xyz (..., 3) => density (...,) + # grid.update_density_grid(density_eval_func=density_eval_func, step=0) + + +if __name__ == "__main__": + test_density_grid() diff --git a/tests/graphs/modules/test_occupancy_grid.py b/tests/graphs/modules/test_occupancy_grid.py deleted file mode 100644 index 5110ee0bf..000000000 --- a/tests/graphs/modules/test_occupancy_grid.py +++ /dev/null @@ -1,22 +0,0 @@ -""" -Code to test the occupancy grid. -""" - -import torch -from pyrad.fields.occupancy_fields.occupancy_grid import OccupancyGrid - - -def test_occupancy_grid(): - """For testing the occupancy grid.""" - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - num_cascades = 1 - resolution = 128 - aabb = torch.tensor([[-1, -1, -1], [1, 1, 1]], dtype=torch.float32).to(device) - grid = OccupancyGrid(num_cascades, resolution, aabb) - grid.to(device) - density_fn = lambda x: x[..., 0:1] # xyz (..., 3) => density (...,) - grid.update_occupancy_grid(density_fn=density_fn) - - -if __name__ == "__main__": - test_occupancy_grid() diff --git a/tests/test_run_train.py b/tests/test_run_train.py index 40c87a140..29c357e66 100644 --- a/tests/test_run_train.py +++ b/tests/test_run_train.py @@ -10,7 +10,7 @@ from omegaconf import DictConfig from scripts.run_train import main -BLACKLIST = ["graph_semantic_nerf.yaml", "graph_mipnerf_360.yaml"] +BLACKLIST = ["graph_semantic_nerf.yaml", "graph_mipnerf_360.yaml", "graph_instant_ngp.yaml"] def set_reduced_config(config: DictConfig): @@ -19,7 +19,7 @@ def set_reduced_config(config: DictConfig): config.trainer.max_num_iterations = 2 # reduce dataset factors; set dataset to test # switch to using the vanilla ImageDataset class - config.data.image_dataset_train._target_ = ( # pylint: disable=protected-access + config.data.image_dataset_train._target_ = ( # pylint: disable=protected-access "pyrad.data.image_dataset.ImageDataset" ) config.data.dataset_inputs_train.data_directory = "tests/data/lego_test" @@ -27,7 +27,7 @@ def set_reduced_config(config: DictConfig): config.data.dataset_inputs_train.downscale_factor = 16 config.data.dataloader_train.image_sampler.num_images_to_sample_from = 1 config.data.dataloader_train.pixel_sampler.num_rays_per_batch = 4 - config.data.image_dataset_eval._target_ = ( # pylint: disable=protected-access + config.data.image_dataset_eval._target_ = ( # pylint: disable=protected-access "pyrad.data.image_dataset.ImageDataset" ) config.data.dataset_inputs_eval.data_directory = "tests/data/lego_test" diff --git a/tests/utils/test_tensor_dataclass.py b/tests/utils/test_tensor_dataclass.py index 63814955f..4b37d9d02 100644 --- a/tests/utils/test_tensor_dataclass.py +++ b/tests/utils/test_tensor_dataclass.py @@ -2,6 +2,7 @@ Test tensor dataclass """ from dataclasses import dataclass + import pytest import torch