From 8a53684131bb7588da3bb87f43f90f8d583f0162 Mon Sep 17 00:00:00 2001
From: Matt Tancik <tancik@berkeley.edu>
Date: Thu, 7 Apr 2022 16:32:22 -0700
Subject: [PATCH] add mlp model

---
 mattport/nerf/models.py | 58 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 58 insertions(+)
 create mode 100644 mattport/nerf/models.py

diff --git a/mattport/nerf/models.py b/mattport/nerf/models.py
new file mode 100644
index 000000000..f58ecc717
--- /dev/null
+++ b/mattport/nerf/models.py
@@ -0,0 +1,58 @@
+"""
+Collection of Models
+"""
+from typing import Optional
+
+from torch import nn
+from torchtyping import TensorType
+
+
+class MLP(nn.Module):
+    """Multilayer perceptron"""
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        num_layers: int,
+        layer_width: int,
+        activation: Optional[nn.Module] = nn.ReLU,
+        out_activation: Optional[nn.Module] = None,
+    ) -> None:
+        """Initialize mulilayer perceptron.
+
+        Args:
+            in_dim (int): Input layer dimension
+            out_dim (int): Ouput layer dimension
+            num_layers (int): Number of network layers
+            layer_width (int): Width of each MLP layer
+            activation (Optional[nn.Module], optional): intermediate layer activation function. Defaults to nn.ReLU.
+            out_activation (Optional[nn.Module], optional): output activation function. Defaults to None.
+        """
+        super().__init__()
+
+        layers = []
+        for i in range(num_layers):
+            if i == 0:
+                layers.append(nn.Linear(in_dim, layer_width))
+            else:
+                layers.append(nn.Linear(layer_width, layer_width))
+            if activation is not None:
+                layers.append(activation)
+        layers.append(nn.Linear(layer_width, out_dim))
+        if out_activation is not None:
+            layers.append(out_activation)
+
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType[..., "out_dim"]:
+        """Process input with a multilayer perceptron.
+
+        Args:
+            in_tensor (TensorType[..., "in_dim]): Network input
+
+        Returns:
+            TensorType[..., "out_dim"]: Network output
+        """
+
+        return self.net(in_tensor)