tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

Remove pytorch hash table if tcnn implementation is used (#2175)

Browse Source 
* fixes div by 1 error

* removes pytorch hash_table if tcnn is used

* formatting

* test fixes

---------

Co-authored-by: Matthew Tancik <matt@lumalabs.ai>
This commit is contained in:
Georg Hess 2023-07-06 19:55:22 +02:00 committed by GitHub
parent 77cf4c88ed
commit 492fa0e67b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
nerfstudio/field_components/encodings.py
View File

@ -274,18 +274,18 @@ class HashEncoding(Encoding):
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))
growth_factor = np.exp((np.log(max_res) - np.log(min_res)) / (num_levels - 1)) if num_levels > 1 else 1
self.scalings = torch.floor(min_res * growth_factor**levels)
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
self.hash_table = torch.empty(0)
if implementation == "tcnn" and not TCNN_EXISTS:
print_tcnn_speed_warning("HashEncoding")
elif implementation == "tcnn":
implementation = "torch"
if implementation == "tcnn":
encoding_config = {
"otype": "HashGrid",
"n_levels": self.num_levels,
@ -301,6 +301,10 @@ class HashEncoding(Encoding):
n_input_dims=3,
encoding_config=encoding_config,
)
elif implementation == "torch":
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)
if self.tcnn_encoding is None:
assert (