diff --git a/docs/guides/model_components/visualize_encoders.ipynb b/docs/guides/model_components/visualize_encoders.ipynb
index 8ed2dbf97..6ef2c0950 100644
--- a/docs/guides/model_components/visualize_encoders.ipynb
+++ b/docs/guides/model_components/visualize_encoders.ipynb
@@ -208,7 +208,7 @@
{
"data": {
"text/html": [
- ""
+ ""
],
"text/plain": [
""
@@ -228,7 +228,7 @@
{
"data": {
"text/html": [
- ""
+ ""
],
"text/plain": [
""
@@ -247,7 +247,7 @@
{
"data": {
"text/html": [
- ""
+ ""
],
"text/plain": [
""
@@ -266,7 +266,7 @@
{
"data": {
"text/html": [
- ""
+ ""
],
"text/plain": [
""
@@ -318,7 +318,7 @@
},
{
"cell_type": "code",
- "execution_count": 3,
+ "execution_count": 18,
"metadata": {
"tags": [
"hide-input"
@@ -329,6 +329,8 @@
"name": "stdout",
"output_type": "stream",
"text": [
+ "tensor(0.0010, grad_fn=)\n",
+ "tensor(-0.0010, grad_fn=)\n",
"Input Values:\n"
]
},
@@ -354,7 +356,7 @@
{
"data": {
"text/html": [
- ""
+ ""
],
"text/plain": [
""
@@ -373,7 +375,7 @@
"num_levels = 8\n",
"min_res = 2\n",
"max_res = 128\n",
- "hash_table_size = 2**4 # Typically much larger tables are used\n",
+ "log2_hashmap_size = 4 # Typically much larger tables are used\n",
"\n",
"resolution = 128\n",
"slice = 0\n",
@@ -385,15 +387,18 @@
" num_levels=num_levels,\n",
" min_res=min_res,\n",
" max_res=max_res,\n",
- " hash_table_size=hash_table_size,\n",
+ " log2_hashmap_size=log2_hashmap_size,\n",
" features_per_level=features_per_level,\n",
" hash_init_scale=0.001,\n",
+ " implementation=\"torch\",\n",
")\n",
"\n",
"x_samples = torch.linspace(0, 1, resolution)\n",
"grid = torch.stack(torch.meshgrid([x_samples, x_samples, x_samples], indexing=\"ij\"), dim=-1)\n",
"\n",
"encoded_values = encoder(grid)\n",
+ "print(torch.max(encoded_values))\n",
+ "print(torch.min(encoded_values))\n",
"\n",
"grid_slice = grid[slice, ...]\n",
"encoded_values_slice = encoded_values[slice, ...]\n",
@@ -420,7 +425,7 @@
},
{
"cell_type": "code",
- "execution_count": 4,
+ "execution_count": null,
"metadata": {
"tags": [
"hide-input"
diff --git a/nerfstudio/utils/activations.py b/nerfstudio/field_components/activations.py
similarity index 100%
rename from nerfstudio/utils/activations.py
rename to nerfstudio/field_components/activations.py
diff --git a/nerfstudio/field_components/encoding.py b/nerfstudio/field_components/encoding.py
index 1e4fedf58..3e33cae30 100644
--- a/nerfstudio/field_components/encoding.py
+++ b/nerfstudio/field_components/encoding.py
@@ -17,7 +17,7 @@ Encoding functions
"""
from abc import abstractmethod
-from typing import Optional
+from typing import Literal, Optional
import numpy as np
import torch
@@ -26,8 +26,16 @@ from torch import nn
from torchtyping import TensorType
from nerfstudio.field_components.base import FieldModule
+from nerfstudio.utils.comments import print_tcnn_speed_warning
from nerfstudio.utils.math import components_from_spherical_harmonics, expected_sin
+try:
+ import tinycudann as tcnn
+
+ TCNN_EXISTS = True
+except ImportError:
+ TCNN_EXISTS = False
+
class Encoding(FieldModule):
"""Encode an input tensor. Intended to be subclassed
@@ -215,9 +223,10 @@ class HashEncoding(Encoding):
num_levels: Number of feature grids.
min_res: Resolution of smallest feature grid.
max_res: Resolution of largest feature grid.
- hash_table_size: Size of hash table.
+ log2_hashmap_size: Size of hash map is 2^log2_hashmap_size.
features_per_level: Number of features per level.
hash_init_scale: Value to initialize hash grid.
+ implementation: Implementation of hash encoding. Fallback to torch if tcnn not available.
"""
def __init__(
@@ -225,25 +234,43 @@ class HashEncoding(Encoding):
num_levels: int = 16,
min_res: int = 16,
max_res: int = 1024,
- hash_table_size: int = 2**19,
+ log2_hashmap_size: int = 19,
features_per_level: int = 2,
hash_init_scale: float = 0.001,
+ implementation: Literal["tcnn", "torch"] = "tcnn",
) -> None:
super().__init__(in_dim=3)
self.num_levels = num_levels
self.features_per_level = features_per_level
- self.hash_table_size = hash_table_size
+ self.log2_hashmap_size = log2_hashmap_size
+ self.hash_table_size = 2**log2_hashmap_size
levels = torch.arange(num_levels)
growth_factor = np.exp((np.log(max_res) - np.log(min_res)) / (num_levels - 1))
self.scalings = torch.floor(min_res * growth_factor**levels)
- self.hash_offset = levels * hash_table_size
- self.hash_table = torch.rand(size=(hash_table_size * num_levels, features_per_level)) * 2 - 1
+ self.hash_offset = levels * self.hash_table_size
+ self.hash_table = torch.rand(size=(self.hash_table_size * num_levels, features_per_level)) * 2 - 1
self.hash_table *= hash_init_scale
self.hash_table = nn.Parameter(self.hash_table)
+ self.tcnn_encoding = None
+ if not TCNN_EXISTS and implementation == "tcnn":
+ print_tcnn_speed_warning("HashEncoding")
+ elif implementation == "tcnn":
+ self.tcnn_encoding = tcnn.Encoding(
+ n_input_dims=3,
+ encoding_config={
+ "otype": "HashGrid",
+ "n_levels": self.num_levels,
+ "n_features_per_level": self.features_per_level,
+ "log2_hashmap_size": self.log2_hashmap_size,
+ "base_resolution": min_res,
+ "per_level_scale": growth_factor,
+ },
+ )
+
def get_out_dim(self) -> int:
return self.num_levels * self.features_per_level
@@ -266,9 +293,10 @@ class HashEncoding(Encoding):
x += self.hash_offset.to(x.device)
return x
- def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
- assert in_tensor.shape[-1] == 3
+ def pytorch_fwd(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
+ """Forward pass using pytorch. Significantly slower than TCNN implementation."""
+ assert in_tensor.shape[-1] == 3
in_tensor = in_tensor[..., None, :] # [..., 1, 3]
scaled = in_tensor * self.scalings.view(-1, 1).to(in_tensor.device) # [..., L, 3]
scaled_c = torch.ceil(scaled).type(torch.int32)
@@ -308,6 +336,11 @@ class HashEncoding(Encoding):
return torch.flatten(encoded_value, start_dim=-2, end_dim=-1) # [..., num_levels * features_per_level]
+ def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
+ if TCNN_EXISTS and self.tcnn_encoding is not None:
+ return self.tcnn_encoding(in_tensor)
+ return self.pytorch_fwd(in_tensor)
+
class TensorCPEncoding(Encoding):
"""Learned CANDECOMP/PARFAC (CP) decomposition encoding used in TensoRF
diff --git a/nerfstudio/fields/density_fields.py b/nerfstudio/fields/density_fields.py
index e807213ac..c2fad1e17 100644
--- a/nerfstudio/fields/density_fields.py
+++ b/nerfstudio/fields/density_fields.py
@@ -26,9 +26,9 @@ from torchtyping import TensorType
from nerfstudio.cameras.rays import RaySamples
from nerfstudio.data.scene_box import SceneBox
+from nerfstudio.field_components.activations import trunc_exp
from nerfstudio.field_components.spatial_distortions import SpatialDistortion
from nerfstudio.fields.base_field import Field
-from nerfstudio.utils.activations import trunc_exp
try:
import tinycudann as tcnn
diff --git a/nerfstudio/fields/instant_ngp_field.py b/nerfstudio/fields/instant_ngp_field.py
index dc0ec8c71..de88d2714 100644
--- a/nerfstudio/fields/instant_ngp_field.py
+++ b/nerfstudio/fields/instant_ngp_field.py
@@ -26,10 +26,10 @@ from torchtyping import TensorType
from nerfstudio.cameras.rays import RaySamples
from nerfstudio.data.scene_box import SceneBox
+from nerfstudio.field_components.activations import trunc_exp
from nerfstudio.field_components.embedding import Embedding
from nerfstudio.field_components.field_heads import FieldHeadNames
from nerfstudio.fields.base_field import Field
-from nerfstudio.utils.activations import trunc_exp
try:
import tinycudann as tcnn
diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py
index 3a5cb2682..b3c0f6cc4 100644
--- a/nerfstudio/fields/nerfacto_field.py
+++ b/nerfstudio/fields/nerfacto_field.py
@@ -27,6 +27,7 @@ from torchtyping import TensorType
from nerfstudio.cameras.rays import RaySamples
from nerfstudio.data.scene_box import SceneBox
+from nerfstudio.field_components.activations import trunc_exp
from nerfstudio.field_components.embedding import Embedding
from nerfstudio.field_components.encoding import Encoding, HashEncoding, SHEncoding
from nerfstudio.field_components.field_heads import (
@@ -41,7 +42,6 @@ from nerfstudio.field_components.spatial_distortions import (
SpatialDistortion,
)
from nerfstudio.fields.base_field import Field
-from nerfstudio.utils.activations import trunc_exp
try:
import tinycudann as tcnn
diff --git a/nerfstudio/utils/comments.py b/nerfstudio/utils/comments.py
new file mode 100644
index 000000000..8291de25c
--- /dev/null
+++ b/nerfstudio/utils/comments.py
@@ -0,0 +1,31 @@
+# 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.
+
+"""A collection of common strings and print statements used throughout the codebase."""
+
+from rich.console import Console
+
+console = Console(width=120)
+
+
+def print_tcnn_speed_warning(method_name: str):
+ """Prints a warning about the speed of the TCNN."""
+ console.line()
+ console.print(f"[bold yellow]WARNING: Using a slow implementation of {method_name}. ")
+ console.print(
+ "[bold yellow]:person_running: :person_running: "
+ + "Install tcnn for speedups :person_running: :person_running:"
+ )
+ console.print("[yellow]pip install git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch")
+ console.line()
diff --git a/nerfstudio/utils/rich.py b/nerfstudio/utils/rich_utils.py
similarity index 100%
rename from nerfstudio/utils/rich.py
rename to nerfstudio/utils/rich_utils.py
diff --git a/scripts/eval.py b/scripts/eval.py
index 9bfe8c338..098740556 100755
--- a/scripts/eval.py
+++ b/scripts/eval.py
@@ -30,7 +30,7 @@ from nerfstudio.cameras.camera_paths import get_path_from_json, get_spiral_path
from nerfstudio.cameras.cameras import Cameras
from nerfstudio.configs import base_config as cfg
from nerfstudio.pipelines.base import Pipeline
-from nerfstudio.utils.rich import ItersPerSecColumn
+from nerfstudio.utils.rich_utils import ItersPerSecColumn
console = Console(width=120)
diff --git a/tests/fields/modules/test_encodings.py b/tests/fields/modules/test_encodings.py
index 3b8b350d9..fe79f7799 100644
--- a/tests/fields/modules/test_encodings.py
+++ b/tests/fields/modules/test_encodings.py
@@ -139,6 +139,38 @@ def test_tensor_sh_encoder():
assert encoded.shape == (10, out_dim)
+def test_tensor_hash_encoder():
+ """Test Spherical Harmonic encoder"""
+
+ num_levels = 4
+ features_per_level = 4
+ out_dim = num_levels * features_per_level
+
+ encoder = encoding.HashEncoding(
+ num_levels=num_levels,
+ features_per_level=features_per_level,
+ log2_hashmap_size=5,
+ implementation="torch",
+ )
+ assert encoder.get_out_dim() == out_dim
+
+ in_tensor = torch.zeros((10, 3))
+ in_tensor[..., 1] = 1
+ encoded = encoder(in_tensor)
+ assert encoded.shape == (10, out_dim)
+
+ encoder = encoding.HashEncoding(
+ num_levels=num_levels,
+ features_per_level=features_per_level,
+ log2_hashmap_size=5,
+ implementation="tcnn",
+ )
+ in_tensor = torch.zeros((10, 3))
+ in_tensor[..., 1] = 1
+ encoded = encoder(in_tensor)
+ assert encoded.shape == (10, out_dim)
+
+
if __name__ == "__main__":
test_scaling_and_offset()
test_nerf_encoder()
@@ -146,3 +178,4 @@ if __name__ == "__main__":
test_tensor_vm_encoder()
test_tensor_cp_encoder()
test_tensor_sh_encoder()
+ test_tensor_hash_encoder()