Repository: nerfstudio-project/nerfstudio
Pull request: Ability to use vanilla pytorch instead of tcnn for Nerfacto 
- [x] `--pipeline.model.implementation tcnn` is default but `--pipeline.model.implementation torch` can be used for vanilla pytorch. This will only work for the Nerfacto method at the moment.
- [x] Removed "nerfacto" from the blacklisted github actions tests.

The vanilla pytorch implementation is maybe ~4x slower. Looking into numerical equivalence is a good next step.

Example commands are as follows,

```
ns-train nerfacto --vis viewer+wandb --experiment-name aspen-tcnn --data data/nerfstudio/aspen --pipeline.model.implementation tcnn
```

and

```
ns-train nerfacto --vis viewer+wandb --experiment-name aspen-torch --data data/nerfstudio/aspen --pipeline.model.implementation torch
```

The quality looks very similar.

The main branch and this PR (which is "tcnn" here and still the default for nerfacto) should be equivalent. Here is a small sanity check.

<img width="477" alt="image" src="https://github.com/nerfstudio-project/nerfstudio/assets/16965789/1a29dbf3-c33b-4e6f-9dd5-8756440b6ba7">


Problem statement:
Does anyone meet pretty low performance using TorchNerfactoField?
**Describe the bug**
I tried both TCNNNerfactoField and TorchNerfactoField on my captured set of images. TCNNNerfactoField does not have any issues, and perfect reconstruction, but TorchNerfactoField has very low evaluation performance (20psnr vs 40 psnr of TCNNNerfactoField).

In addition, the nn.softplus causes NAN gradients, I change to trunc_exp as used in TCNNNerfactoField.

**Additional context**
I also include _sample_locations and _density_before_activation for normal prediction. Without the normal loss, the performance are still unexpected.

Hints from issue discussion:

Commit message: removed/reverted files

Patch to review:
diff --git a/docs/quickstart/export_geometry.md b/docs/quickstart/export_geometry.md
index 2ffd51895f..490638422a 100644
--- a/docs/quickstart/export_geometry.md
+++ b/docs/quickstart/export_geometry.md
@@ -2,6 +2,14 @@
 
 Here we document how to export point clouds and meshes from nerfstudio. The main command you'll be working with is `ns-export`. Our point clouds are exported as `.ply` files and the textured meshes are exported as `.obj` files.
 
+## Install dependencies
+
+We require some extra dependencies for exporting geometry. Run the following command to install these packages.
+
+```bash
+pip install -e .[export]
+```
+
 ## Exporting a mesh
 
 ### 1. TSDF Fusion
diff --git a/nerfstudio/field_components/encodings.py b/nerfstudio/field_components/encodings.py
index b455468f25..8f840bfd63 100644
--- a/nerfstudio/field_components/encodings.py
+++ b/nerfstudio/field_components/encodings.py
@@ -109,7 +109,13 @@ class NeRFEncoding(Encoding):
     """
 
     def __init__(
-        self, in_dim: int, num_frequencies: int, min_freq_exp: float, max_freq_exp: float, include_input: bool = False
+        self,
+        in_dim: int,
+        num_frequencies: int,
+        min_freq_exp: float,
+        max_freq_exp: float,
+        include_input: bool = False,
+        implementation: Literal["tcnn", "torch"] = "torch",
     ) -> None:
         super().__init__(in_dim)
 
@@ -118,6 +124,18 @@ def __init__(
         self.max_freq = max_freq_exp
         self.include_input = include_input
 
+        self.tcnn_encoding = None
+        if not TCNN_EXISTS and implementation == "tcnn":
+            print_tcnn_speed_warning("NeRFEncoding")
+        elif implementation == "tcnn":
+            encoding_config = {"otype": "Frequency", "n_frequencies": num_frequencies}
+            assert min_freq_exp == 0, "tcnn only supports min_freq_exp = 0"
+            assert max_freq_exp == num_frequencies - 1, "tcnn only supports max_freq_exp = num_frequencies - 1"
+            self.tcnn_encoding = tcnn.Encoding(
+                n_input_dims=in_dim,
+                encoding_config=encoding_config,
+            )
+
     def get_out_dim(self) -> int:
         if self.in_dim is None:
             raise ValueError("Input dimension has not been set")
@@ -239,7 +257,7 @@ def __init__(
         log2_hashmap_size: int = 19,
         features_per_level: int = 2,
         hash_init_scale: float = 0.001,
-        implementation: Literal["tcnn", "torch"] = "tcnn",
+        implementation: Literal["tcnn", "torch"] = "torch",
         interpolation: Optional[Literal["Nearest", "Linear", "Smoothstep"]] = None,
     ) -> None:
         super().__init__(in_dim=3)
@@ -644,7 +662,7 @@ class SHEncoding(Encoding):
         levels: Number of spherical harmonic levels to encode.
     """
 
-    def __init__(self, levels: int = 4) -> None:
+    def __init__(self, levels: int = 4, implementation: Literal["tcnn", "torch"] = "tcnn") -> None:
         super().__init__(in_dim=3)
 
         if levels <= 0 or levels > 4:
@@ -652,9 +670,28 @@ def __init__(self, levels: int = 4) -> None:
 
         self.levels = levels
 
+        self.tcnn_encoding = None
+        if not TCNN_EXISTS and implementation == "tcnn":
+            print_tcnn_speed_warning("SHEncoding")
+        elif implementation == "tcnn":
+            encoding_config = {
+                "otype": "SphericalHarmonics",
+                "degree": levels,
+            }
+            self.tcnn_encoding = tcnn.Encoding(
+                n_input_dims=3,
+                encoding_config=encoding_config,
+            )
+
     def get_out_dim(self) -> int:
         return self.levels**2
 
     @torch.no_grad()
-    def forward(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
+    def pytorch_fwd(self, in_tensor: TensorType["bs":..., "input_dim"]) -> TensorType["bs":..., "output_dim"]:
+        """Forward pass using pytorch. Significantly slower than TCNN implementation."""
         return components_from_spherical_harmonics(levels=self.levels, directions=in_tensor)
+
+    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)
diff --git a/nerfstudio/field_components/mlp.py b/nerfstudio/field_components/mlp.py
index 3ce3cdad95..2b9a53ed0b 100644
--- a/nerfstudio/field_components/mlp.py
+++ b/nerfstudio/field_components/mlp.py
@@ -15,13 +15,39 @@
 """
 Multi Layer Perceptron
 """
-from typing import Optional, Set, Tuple
+from typing import Literal, Optional, Set, Tuple
 
 import torch
 from torch import nn
 from torchtyping import TensorType
 
 from nerfstudio.field_components.base_field_component import FieldComponent
+from nerfstudio.utils.printing import print_tcnn_speed_warning
+
+try:
+    import tinycudann as tcnn
+
+    TCNN_EXISTS = True
+except ImportError:
+    TCNN_EXISTS = False
+
+
+def activation_to_tcnn_string(activation: nn.Module) -> str:
+    """Converts a torch.nn activation function to a string that can be used to
+    initialize a TCNN activation function.
+
+    Args:
+        activation: torch.nn activation function
+    Returns:
+        str: TCNN activation function string
+    """
+    if isinstance(activation, nn.ReLU):
+        return "ReLU"
+    if isinstance(activation, nn.Sigmoid):
+        return "Sigmoid"
+    if isinstance(activation, type(None)):
+        return "None"
+    raise ValueError(f"TCNN activation {activation} not supported for now.")
 
 
 class MLP(FieldComponent):
@@ -45,6 +71,7 @@ def __init__(
         skip_connections: Optional[Tuple[int]] = None,
         activation: Optional[nn.Module] = nn.ReLU(),
         out_activation: Optional[nn.Module] = None,
+        implementation: Literal["tcnn", "torch"] = "torch",
     ) -> None:
         super().__init__()
         self.in_dim = in_dim
@@ -57,7 +84,27 @@ def __init__(
         self.activation = activation
         self.out_activation = out_activation
         self.net = None
-        self.build_nn_modules()
+
+        self.tcnn_encoding = None
+        if implementation == "torch":
+            self.build_nn_modules()
+        elif not TCNN_EXISTS and implementation == "tcnn":
+            print_tcnn_speed_warning("MLP")
+        elif implementation == "tcnn":
+            activation_str = activation_to_tcnn_string(activation)
+            output_activation_str = activation_to_tcnn_string(out_activation)
+            network_config = {
+                "otype": "FullyFusedMLP",
+                "activation": activation_str,
+                "output_activation": output_activation_str,
+                "n_neurons": layer_width,
+                "n_hidden_layers": num_layers - 1,
+            }
+            self.tcnn_encoding = tcnn.Network(
+                n_input_dims=in_dim,
+                n_output_dims=out_dim,
+                network_config=network_config,
+            )
 
     def build_nn_modules(self) -> None:
         """Initialize multi-layer perceptron."""
@@ -76,7 +123,7 @@ def build_nn_modules(self) -> None:
             layers.append(nn.Linear(self.layer_width, self.out_dim))
         self.layers = nn.ModuleList(layers)
 
-    def forward(self, in_tensor: TensorType["bs":..., "in_dim"]) -> TensorType["bs":..., "out_dim"]:
+    def pytorch_fwd(self, in_tensor: TensorType["bs":..., "in_dim"]) -> TensorType["bs":..., "out_dim"]:
         """Process input with a multilayer perceptron.
 
         Args:
@@ -96,3 +143,8 @@ def forward(self, in_tensor: TensorType["bs":..., "in_dim"]) -> TensorType["bs":
         if self.out_activation is not None:
             x = self.out_activation(x)
         return x
+
+    def forward(self, in_tensor: TensorType["bs":..., "in_dim"]) -> TensorType["bs":..., "out_dim"]:
+        if TCNN_EXISTS and self.tcnn_encoding is not None:
+            return self.tcnn_encoding(in_tensor)
+        return self.pytorch_fwd(in_tensor)
diff --git a/nerfstudio/fields/density_fields.py b/nerfstudio/fields/density_fields.py
index 0364db00d6..b3fd759cff 100644
--- a/nerfstudio/fields/density_fields.py
+++ b/nerfstudio/fields/density_fields.py
@@ -17,24 +17,20 @@
 """
 
 
-from typing import Optional, Tuple
+from typing import Literal, Optional, Tuple
 
-import numpy as np
 import torch
+from torch import nn
 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.encodings import HashEncoding
+from nerfstudio.field_components.mlp import MLP
 from nerfstudio.field_components.spatial_distortions import SpatialDistortion
 from nerfstudio.fields.base_field import Field
 
-try:
-    import tinycudann as tcnn
-except ImportError:
-    # tinycudann module doesn't exist
-    pass
-
 
 class HashMLPDensityField(Field):
     """A lightweight density field module.
@@ -59,45 +55,39 @@ def __init__(
         base_res: int = 16,
         log2_hashmap_size: int = 18,
         features_per_level: int = 2,
+        implementation: Literal["tcnn", "torch"] = "tcnn",
     ) -> None:
         super().__init__()
         self.register_buffer("aabb", aabb)
         self.spatial_distortion = spatial_distortion
         self.use_linear = use_linear
-        growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
 
         self.register_buffer("max_res", torch.tensor(max_res))
         self.register_buffer("num_levels", torch.tensor(num_levels))
         self.register_buffer("log2_hashmap_size", torch.tensor(log2_hashmap_size))
 
-        config = {
-            "encoding": {
-                "otype": "HashGrid",
-                "n_levels": num_levels,
-                "n_features_per_level": features_per_level,
-                "log2_hashmap_size": log2_hashmap_size,
-                "base_resolution": base_res,
-                "per_level_scale": growth_factor,
-            },
-            "network": {
-                "otype": "FullyFusedMLP",
-                "activation": "ReLU",
-                "output_activation": "None",
-                "n_neurons": hidden_dim,
-                "n_hidden_layers": num_layers - 1,
-            },
-        }
+        self.encoding = HashEncoding(
+            num_levels=num_levels,
+            min_res=base_res,
+            max_res=max_res,
+            log2_hashmap_size=log2_hashmap_size,
+            features_per_level=features_per_level,
+            implementation=implementation,
+        )
 
         if not self.use_linear:
-            self.mlp_base = tcnn.NetworkWithInputEncoding(
-                n_input_dims=3,
-                n_output_dims=1,
-                encoding_config=config["encoding"],
-                network_config=config["network"],
+            network = MLP(
+                in_dim=self.encoding.get_out_dim(),
+                num_layers=num_layers,
+                layer_width=hidden_dim,
+                out_dim=1,
+                activation=nn.ReLU(),
+                out_activation=None,
+                implementation=implementation,
             )
+            self.mlp_base = torch.nn.Sequential(self.encoding, network)
         else:
-            self.encoding = tcnn.Encoding(n_input_dims=3, encoding_config=config["encoding"])
-            self.linear = torch.nn.Linear(self.encoding.n_output_dims, 1)
+            self.linear = torch.nn.Linear(self.encoding.get_out_dim(), 1)
 
     def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, None]:
         if self.spatial_distortion is not None:
diff --git a/nerfstudio/fields/instant_ngp_field.py b/nerfstudio/fields/instant_ngp_field.py
index 87bfb778af..8b1d4a6c5d 100644
--- a/nerfstudio/fields/instant_ngp_field.py
+++ b/nerfstudio/fields/instant_ngp_field.py
@@ -59,6 +59,7 @@ class TCNNInstantNGPField(Field):
         num_levels: number of levels of the hashmap for the base mlp
         log2_hashmap_size: size of the hashmap for the base mlp
         max_res: maximum resolution of the hashmap for the base mlp
+        mode: which implementation to use for the field
     """
 
     def __init__(
diff --git a/nerfstudio/fields/nerfacto_field.py b/nerfstudio/fields/nerfacto_field.py
index 840268bfce..fbb8ded231 100644
--- a/nerfstudio/fields/nerfacto_field.py
+++ b/nerfstudio/fields/nerfacto_field.py
@@ -17,45 +17,31 @@
 """
 
 
-from typing import Dict, Optional, Tuple
+from typing import Dict, Literal, Optional, Tuple
 
-import numpy as np
 import torch
 from torch import nn
-from torch.nn.parameter import Parameter
 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.encodings import Encoding, HashEncoding, SHEncoding
+from nerfstudio.field_components.encodings import HashEncoding, NeRFEncoding, SHEncoding
 from nerfstudio.field_components.field_heads import (
-    DensityFieldHead,
-    FieldHead,
     FieldHeadNames,
     PredNormalsFieldHead,
-    RGBFieldHead,
     SemanticFieldHead,
     TransientDensityFieldHead,
     TransientRGBFieldHead,
     UncertaintyFieldHead,
 )
 from nerfstudio.field_components.mlp import MLP
-from nerfstudio.field_components.spatial_distortions import (
-    SceneContraction,
-    SpatialDistortion,
-)
+from nerfstudio.field_components.spatial_distortions import SpatialDistortion
 from nerfstudio.fields.base_field import Field, shift_directions_for_tcnn
 
-try:
-    import tinycudann as tcnn
-except ImportError:
-    # tinycudann module doesn't exist
-    pass
-
 
-class TCNNNerfactoField(Field):
+class NerfactoField(Field):
     """Compound Field that uses TCNN
 
     Args:
@@ -104,6 +90,7 @@ def __init__(
         use_pred_normals: bool = False,
         use_average_appearance_embedding: bool = False,
         spatial_distortion: Optional[SpatialDistortion] = None,
+        implementation: Literal["tcnn", "torch"] = "tcnn",
     ) -> None:
         super().__init__()
 
@@ -126,55 +113,47 @@ def __init__(
 
         base_res: int = 16
         features_per_level: int = 2
-        growth_factor = np.exp((np.log(max_res) - np.log(base_res)) / (num_levels - 1))
-
-        self.direction_encoding = tcnn.Encoding(
-            n_input_dims=3,
-            encoding_config={
-                "otype": "SphericalHarmonics",
-                "degree": 4,
-            },
+
+        self.direction_encoding = SHEncoding(
+            levels=4,
+            implementation=implementation,
         )
 
-        self.position_encoding = tcnn.Encoding(
-            n_input_dims=3,
-            encoding_config={"otype": "Frequency", "n_frequencies": 2},
+        self.position_encoding = NeRFEncoding(
+            in_dim=3, num_frequencies=2, min_freq_exp=0, max_freq_exp=2 - 1, implementation=implementation
         )
 
-        self.mlp_base = tcnn.NetworkWithInputEncoding(
-            n_input_dims=3,
-            n_output_dims=1 + self.geo_feat_dim,
-            encoding_config={
-                "otype": "HashGrid",
-                "n_levels": num_levels,
-                "n_features_per_level": features_per_level,
-                "log2_hashmap_size": log2_hashmap_size,
-                "base_resolution": base_res,
-                "per_level_scale": growth_factor,
-            },
-            network_config={
-                "otype": "FullyFusedMLP",
-                "activation": "ReLU",
-                "output_activation": "None",
-                "n_neurons": hidden_dim,
-                "n_hidden_layers": num_layers - 1,
-            },
+        encoder = HashEncoding(
+            num_levels=num_levels,
+            min_res=base_res,
+            max_res=max_res,
+            log2_hashmap_size=log2_hashmap_size,
+            features_per_level=features_per_level,
+            implementation=implementation,
+        )
+        network = MLP(
+            in_dim=encoder.get_out_dim(),
+            num_layers=num_layers,
+            layer_width=hidden_dim,
+            out_dim=1 + self.geo_feat_dim,
+            activation=nn.ReLU(),
+            out_activation=None,
+            implementation=implementation,
         )
+        self.mlp_base = torch.nn.Sequential(encoder, network)
 
         # transients
         if self.use_transient_embedding:
             self.transient_embedding_dim = transient_embedding_dim
             self.embedding_transient = Embedding(self.num_images, self.transient_embedding_dim)
-            self.mlp_transient = tcnn.Network(
-                n_input_dims=self.geo_feat_dim + self.transient_embedding_dim,
-                n_output_dims=hidden_dim_transient,
-                network_config={
-                    "otype": "FullyFusedMLP",
-                    "activation": "ReLU",
-                    "output_activation": "None",
-                    "n_neurons": hidden_dim_transient,
-                    "n_hidden_layers": num_layers_transient - 1,
-                },
+            self.mlp_transient = MLP(
+                in_dim=self.geo_feat_dim + self.transient_embedding_dim,
+                num_layers=num_layers_transient,
+                layer_width=hidden_dim_transient,
+                out_dim=hidden_dim_transient,
+                activation=nn.ReLU(),
+                out_activation=None,
+                implementation=implementation,
             )
             self.field_head_transient_uncertainty = UncertaintyFieldHead(in_dim=self.mlp_transient.n_output_dims)
             self.field_head_transient_rgb = TransientRGBFieldHead(in_dim=self.mlp_transient.n_output_dims)
@@ -182,16 +161,14 @@ def __init__(
 
         # semantics
         if self.use_semantics:
-            self.mlp_semantics = tcnn.Network(
-                n_input_dims=self.geo_feat_dim,
-                n_output_dims=hidden_dim_transient,
-                network_config={
-                    "otype": "FullyFusedMLP",
-                    "activation": "ReLU",
-                    "output_activation": "None",
-                    "n_neurons": 64,
-                    "n_hidden_layers": 1,
-                },
+            self.mlp_semantics = MLP(
+                in_dim=self.geo_feat_dim,
+                num_layers=2,
+                layer_width=64,
+                out_dim=hidden_dim_transient,
+                activation=nn.ReLU(),
+                out_activation=None,
+                implementation=implementation,
             )
             self.field_head_semantics = SemanticFieldHead(
                 in_dim=self.mlp_semantics.n_output_dims, num_classes=num_semantic_classes
@@ -199,29 +176,25 @@ def __init__(
 
         # predicted normals
         if self.use_pred_normals:
-            self.mlp_pred_normals = tcnn.Network(
-                n_input_dims=self.geo_feat_dim + self.position_encoding.n_output_dims,
-                n_output_dims=hidden_dim_transient,
-                network_config={
-                    "otype": "FullyFusedMLP",
-                    "activation": "ReLU",
-                    "output_activation": "None",
-                    "n_neurons": 64,
-                    "n_hidden_layers": 2,
-                },
+            self.mlp_pred_normals = MLP(
+                in_dim=self.geo_feat_dim + self.position_encoding.get_out_dim(),
+                num_layers=3,
+                layer_width=64,
+                out_dim=hidden_dim_transient,
+                activation=nn.ReLU(),
+                out_activation=None,
+                implementation=implementation,
             )
             self.field_head_pred_normals = PredNormalsFieldHead(in_dim=self.mlp_pred_normals.n_output_dims)
 
-        self.mlp_head = tcnn.Network(
-            n_input_dims=self.direction_encoding.n_output_dims + self.geo_feat_dim + self.appearance_embedding_dim,
-            n_output_dims=3,
-            network_config={
-                "otype": "FullyFusedMLP",
-                "activation": "ReLU",
-                "output_activation": "Sigmoid",
-                "n_neurons": hidden_dim_color,
-                "n_hidden_layers": num_layers_color - 1,
-            },
+        self.mlp_head = MLP(
+            in_dim=self.direction_encoding.get_out_dim() + self.geo_feat_dim + self.appearance_embedding_dim,
+            num_layers=num_layers_color,
+            layer_width=hidden_dim_color,
+            out_dim=3,
+            activation=nn.ReLU(),
+            out_activation=nn.Sigmoid(),
+            implementation=implementation,
         )
 
     def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
@@ -323,102 +296,3 @@ def get_outputs(
         outputs.update({FieldHeadNames.RGB: rgb})
 
         return outputs
-
-
-class TorchNerfactoField(Field):
-    """
-    PyTorch implementation of the compound field.
-    """
-
-    def __init__(
-        self,
-        aabb: TensorType,
-        num_images: int,
-        position_encoding: Encoding = HashEncoding(),
-        direction_encoding: Encoding = SHEncoding(),
-        base_mlp_num_layers: int = 3,
-        base_mlp_layer_width: int = 64,
-        head_mlp_num_layers: int = 2,
-        head_mlp_layer_width: int = 32,
-        appearance_embedding_dim: int = 40,
-        skip_connections: Tuple = (4,),
-        field_heads: Tuple[FieldHead] = (RGBFieldHead(),),
-        spatial_distortion: SpatialDistortion = SceneContraction(),
-    ) -> None:
-        super().__init__()
-        self.aabb = Parameter(aabb, requires_grad=False)
-        self.spatial_distortion = spatial_distortion
-        self.num_images = num_images
-        self.appearance_embedding_dim = appearance_embedding_dim
-        self.embedding_appearance = Embedding(self.num_images, self.appearance_embedding_dim)
-
-        self.position_encoding = position_encoding
-        self.direction_encoding = direction_encoding
-
-        self.mlp_base = MLP(
-            in_dim=self.position_encoding.get_out_dim(),
-            num_layers=base_mlp_num_layers,
-            layer_width=base_mlp_layer_width,
-            skip_connections=skip_connections,
-            out_activation=nn.ReLU(),
-        )
-
-        self.mlp_head = MLP(
-            in_dim=self.mlp_base.get_out_dim() + self.direction_encoding.get_out_dim() + self.appearance_embedding_dim,
-            num_layers=head_mlp_num_layers,
-            layer_width=head_mlp_layer_width,
-            out_activation=nn.ReLU(),
-        )
-
-        self.field_output_density = DensityFieldHead(in_dim=self.mlp_base.get_out_dim())
-        self.field_heads = nn.ModuleList(field_heads)
-        for field_head in self.field_heads:
-            field_head.set_in_dim(self.mlp_head.get_out_dim())  # type: ignore
-
-    def get_density(self, ray_samples: RaySamples) -> Tuple[TensorType, TensorType]:
-        if self.spatial_distortion is not None:
-            positions = ray_samples.frustums.get_positions()
-            positions = self.spatial_distortion(positions)
-        else:
-            positions = ray_samples.frustums.get_positions()
-        encoded_xyz = self.position_encoding(positions.view(-1, 3)).to(torch.float32)
-        base_mlp_out = self.mlp_base(encoded_xyz).view(*ray_samples.frustums.shape, -1)
-        density = self.field_output_density(base_mlp_out)
-        return density, base_mlp_out
-
-    def get_outputs(
-        self, ray_samples: RaySamples, density_embedding: Optional[TensorType] = None
-    ) -> Dict[FieldHeadNames, TensorType]:
-        outputs_shape = ray_samples.frustums.directions.shape[:-1]
-
-        if ray_samples.camera_indices is None:
-            raise AttributeError("Camera indices are not provided.")
-        camera_indices = ray_samples.camera_indices.squeeze()
-        if self.training:
-            embedded_appearance = self.embedding_appearance(camera_indices)
-        else:
-            embedded_appearance = torch.zeros(
-                (*outputs_shape, self.appearance_embedding_dim),
-                device=ray_samples.frustums.directions.device,
-            )
-
-        outputs = {}
-        for field_head in self.field_heads:
-            directions = shift_directions_for_tcnn(ray_samples.frustums.directions)
-            directions_flat = directions.view(-1, 3)
-            encoded_dir = self.direction_encoding(directions_flat)
-            mlp_out = self.mlp_head(
-                torch.cat(
-                    [
-                        encoded_dir.view(-1, self.direction_encoding.get_out_dim()),
-                        density_embedding.view(-1, self.mlp_base.get_out_dim()),  # type:ignore
-                        embedded_appearance.view(-1, self.appearance_embedding_dim),
-                    ],
-                    dim=-1,  # type:ignore
-                )
-            )
-            outputs[field_head.field_head_name] = field_head(mlp_out).view(*outputs_shape, -1).to(directions)
-        return outputs
-
-
-field_implementation_to_class: Dict[str, Field] = {"tcnn": TCNNNerfactoField, "torch": TorchNerfactoField}
diff --git a/nerfstudio/models/base_surface_model.py b/nerfstudio/models/base_surface_model.py
index 9c68cfaf42..8c46d46e07 100644
--- a/nerfstudio/models/base_surface_model.py
+++ b/nerfstudio/models/base_surface_model.py
@@ -33,7 +33,7 @@
 from nerfstudio.field_components.encodings import NeRFEncoding
 from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.field_components.spatial_distortions import SceneContraction
-from nerfstudio.fields.nerfacto_field import TCNNNerfactoField
+from nerfstudio.fields.nerfacto_field import NerfactoField
 from nerfstudio.fields.sdf_field import SDFFieldConfig
 from nerfstudio.fields.vanilla_nerf_field import NeRFField
 from nerfstudio.model_components.losses import (
@@ -124,7 +124,7 @@ def populate_modules(self):
 
         # background model
         if self.config.background_model == "grid":
-            self.field_background = TCNNNerfactoField(
+            self.field_background = NerfactoField(
                 self.scene_box.aabb,
                 spatial_distortion=self.scene_contraction,
                 num_images=self.num_train_data,
diff --git a/nerfstudio/models/instant_ngp.py b/nerfstudio/models/instant_ngp.py
index 1eff0c836c..e8ac52b96d 100644
--- a/nerfstudio/models/instant_ngp.py
+++ b/nerfstudio/models/instant_ngp.py
@@ -36,7 +36,7 @@
 )
 from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.field_components.spatial_distortions import SceneContraction
-from nerfstudio.fields.nerfacto_field import TCNNNerfactoField
+from nerfstudio.fields.nerfacto_field import NerfactoField
 from nerfstudio.model_components.losses import MSELoss
 from nerfstudio.model_components.ray_samplers import VolumetricSampler
 from nerfstudio.model_components.renderers import (
@@ -94,7 +94,7 @@ class NGPModel(Model):
     """
 
     config: InstantNGPModelConfig
-    field: TCNNNerfactoField
+    field: NerfactoField
 
     def __init__(self, config: InstantNGPModelConfig, **kwargs) -> None:
         super().__init__(config=config, **kwargs)
@@ -108,7 +108,7 @@ def populate_modules(self):
         else:
             scene_contraction = SceneContraction(order=float("inf"))
 
-        self.field = TCNNNerfactoField(
+        self.field = NerfactoField(
             aabb=self.scene_box.aabb,
             num_images=self.num_train_data,
             log2_hashmap_size=self.config.log2_hashmap_size,
diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py
index 3af9059d61..0c49c51970 100644
--- a/nerfstudio/models/nerfacto.py
+++ b/nerfstudio/models/nerfacto.py
@@ -37,7 +37,7 @@
 from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.field_components.spatial_distortions import SceneContraction
 from nerfstudio.fields.density_fields import HashMLPDensityField
-from nerfstudio.fields.nerfacto_field import TCNNNerfactoField
+from nerfstudio.fields.nerfacto_field import NerfactoField
 from nerfstudio.model_components.losses import (
     MSELoss,
     distortion_loss,
@@ -130,6 +130,8 @@ class NerfactoModelConfig(ModelConfig):
     """Whether to disable scene contraction or not."""
     use_gradient_scaling: bool = False
     """Use gradient scaler where the gradients are lower for points closer to the camera."""
+    implementation: Literal["tcnn", "torch"] = ("tcnn",)
+    """Which implementation to use for the model."""
 
 
 class NerfactoModel(Model):
@@ -151,7 +153,7 @@ def populate_modules(self):
             scene_contraction = SceneContraction(order=float("inf"))
 
         # Fields
-        self.field = TCNNNerfactoField(
+        self.field = NerfactoField(
             self.scene_box.aabb,
             hidden_dim=self.config.hidden_dim,
             num_levels=self.config.num_levels,
@@ -163,6 +165,7 @@ def populate_modules(self):
             num_images=self.num_train_data,
             use_pred_normals=self.config.predict_normals,
             use_average_appearance_embedding=self.config.use_average_appearance_embedding,
+            implementation=self.config.implementation,
         )
 
         self.density_fns = []
@@ -172,7 +175,12 @@ def populate_modules(self):
         if self.config.use_same_proposal_network:
             assert len(self.config.proposal_net_args_list) == 1, "Only one proposal network is allowed."
             prop_net_args = self.config.proposal_net_args_list[0]
-            network = HashMLPDensityField(self.scene_box.aabb, spatial_distortion=scene_contraction, **prop_net_args)
+            network = HashMLPDensityField(
+                self.scene_box.aabb,
+                spatial_distortion=scene_contraction,
+                **prop_net_args,
+                implementation=self.config.implementation,
+            )
             self.proposal_networks.append(network)
             self.density_fns.extend([network.density_fn for _ in range(num_prop_nets)])
         else:
@@ -182,6 +190,7 @@ def populate_modules(self):
                     self.scene_box.aabb,
                     spatial_distortion=scene_contraction,
                     **prop_net_args,
+                    implementation=self.config.implementation,
                 )
                 self.proposal_networks.append(network)
             self.density_fns.extend([network.density_fn for network in self.proposal_networks])
diff --git a/nerfstudio/models/semantic_nerfw.py b/nerfstudio/models/semantic_nerfw.py
index 35623cc996..bcab63f199 100644
--- a/nerfstudio/models/semantic_nerfw.py
+++ b/nerfstudio/models/semantic_nerfw.py
@@ -38,7 +38,7 @@
 from nerfstudio.field_components.field_heads import FieldHeadNames
 from nerfstudio.field_components.spatial_distortions import SceneContraction
 from nerfstudio.fields.density_fields import HashMLPDensityField
-from nerfstudio.fields.nerfacto_field import TCNNNerfactoField
+from nerfstudio.fields.nerfacto_field import NerfactoField
 from nerfstudio.model_components.losses import MSELoss, distortion_loss, interlevel_loss
 from nerfstudio.model_components.ray_samplers import ProposalNetworkSampler
 from nerfstudio.model_components.renderers import (
@@ -90,7 +90,7 @@ def populate_modules(self):
             raise ValueError("Transient embedding is not fully working for semantic nerf-w.")
 
         # Fields
-        self.field = TCNNNerfactoField(
+        self.field = NerfactoField(
             self.scene_box.aabb,
             num_levels=self.config.num_levels,
             max_res=self.config.max_res,
diff --git a/pyproject.toml b/pyproject.toml
index abe400912b..5168624282 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,7 +31,6 @@ dependencies = [
     "msgpack>=1.0.4",
     "msgpack_numpy>=0.4.8",
     "nerfacc==0.5.2",
-    "open3d>=0.16.0",
     "opencv-python==4.6.0.66",
     "Pillow>=9.3.0",
     "plotly>=5.7.0",
@@ -39,7 +38,6 @@ dependencies = [
     # TODO(1480) enable when pycolmap windows wheels are available
     # "pycolmap==0.3.0",
     "pyequilib>=0.5.6",
-    "pymeshlab>=2022.2.post2",
     "pyngrok>=5.1.0",
     "python-socketio>=5.7.1",
     "pyquaternion>=0.9.9",
@@ -54,9 +52,7 @@ dependencies = [
     "typing_extensions>=4.4.0",
     "viser>=0.0.5",
     "nuscenes-devkit>=1.1.1",
-    "wandb>=0.13.3",
-    "xatlas",
-    "trimesh>=3.20.2"
+    "wandb>=0.13.3"
 ]
 
 [project.urls]
@@ -71,6 +67,13 @@ gen = [
     "transformers==4.24.0",
 ]
 
+# Geometry export related dependencies
+export = [
+    "open3d>=0.16.0",
+    "pymeshlab>=2022.2.post2",
+    "trimesh>=3.20.2"
+    "xatlas",
+]
 
 # Development packages
 dev = [
diff --git a/tests/field_components/test_encodings.py b/tests/field_components/test_encodings.py
index dc59b79ff2..a7dbb7f2cb 100644
--- a/tests/field_components/test_encodings.py
+++ b/tests/field_components/test_encodings.py
@@ -154,10 +154,9 @@ def test_tensor_hash_encoder():
     )
     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)
+    in_tensor = torch.rand((10, 3))
+    encoded_torch = encoder(in_tensor)
+    assert encoded_torch.shape == (10, out_dim)
 
     encoder = encodings.HashEncoding(
         num_levels=num_levels,
@@ -165,10 +164,8 @@ def test_tensor_hash_encoder():
         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)
+    encoded_tcnn = encoder(in_tensor).float().cpu()
+    assert encoded_tcnn.shape == (10, out_dim)
 
 
 def test_kplane_encoder():


Code review comments:
diff_hunk: @@ -71,6 +67,13 @@ gen = [
     "transformers==4.24.0",
 ]
 
+# Geometry export related dependencies
+export = [
+    "open3d>=0.16.0",
+    "pymeshlab>=2022.2.post2",
+    "trimesh>=3.20.2"
path: pyproject.toml
text: @user1:
missing comma