From 1dd50cbe4d41ae7d3564d0e27d54345957a1020b Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Thu, 2 Jul 2026 20:19:58 +0000
Subject: [PATCH 1/2] feat: new cnn model

---
 .../models/geometric_gnn_models.py            | 585 ++++++++++++++++++
 src/neuron_proofreader/models/gnn_models.py   | 189 ++++++
 .../models/new_vision_models.py               | 231 +++++++
 .../models/point_cloud_models.py              | 347 +++++++++++
 .../models/vision_models.py                   | 490 +++++++++++++++
 5 files changed, 1842 insertions(+)
 create mode 100644 src/neuron_proofreader/models/geometric_gnn_models.py
 create mode 100644 src/neuron_proofreader/models/gnn_models.py
 create mode 100644 src/neuron_proofreader/models/new_vision_models.py
 create mode 100644 src/neuron_proofreader/models/point_cloud_models.py
 create mode 100644 src/neuron_proofreader/models/vision_models.py

diff --git a/src/neuron_proofreader/models/geometric_gnn_models.py b/src/neuron_proofreader/models/geometric_gnn_models.py
new file mode 100644
index 00000000..0c68d7f1
--- /dev/null
+++ b/src/neuron_proofreader/models/geometric_gnn_models.py
@@ -0,0 +1,585 @@
+"""
+Created on Sat July 15 12:00:00 2025
+
+@author: Anna Grim
+@email: anna.grim@alleninstitute.org
+
+Code for graph neural network models that perform machine learning tasks
+within NeuronProofreader pipelines.
+
+"""
+
+from torch import nn
+from torch_geometric.nn import GATv2Conv
+
+import torch
+import torch.nn.functional as F
+
+from neuron_proofreader.machine_learning.vision_models import CNN3D
+from neuron_proofreader.utils.ml_util import FeedForwardNet
+
+
+# --- Multimodal GNN Architectures ---
+class VisionSkeleton(nn.Module):
+
+    def __init__(self, ggnn_name, patch_shape, output_dim=64):
+        # Call parent class
+        super().__init__()
+        assert ggnn_name in ["egnn"]
+
+        # Architecture
+        self.skeleton_model = SkeletonGNN(ggnn_name, output_dim=32)
+        self.vision_model = CNN3D(
+            patch_shape,
+            n_conv_layers=6,
+            n_feat_channels=20,
+            output_dim=output_dim,
+            use_double_conv=True,
+        )
+        self.output = FeedForwardNet(output_dim + 35, 1, 3)
+
+    def forward(self, x):
+        """
+        Passes the given input through this neural network.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input vector of features.
+
+        Returns
+        -------
+        x : torch.Tensor
+            Output of the neural network.
+        """
+        # Modality-based embeddings
+        x_img = self.vision_model(x["img"])
+        x_skel = self.skeleton_model(*x["graph"])
+        x = torch.cat((x_img, x_skel), dim=1)
+
+        # Output layer
+        x = self.output(x)
+        return x
+
+
+class SkeletonGNN(nn.Module):
+
+    def __init__(self, ggnn_name, output_dim=64):
+        # Call parent class
+        super().__init__()
+
+        # Instance attributes
+        self.gnn_h = GAT(output_dim, 2 * output_dim, output_dim)
+        self.gnn_x = GAT(3, 16, 3)
+
+        # Set geometric gnn
+        if ggnn_name == "egnn":
+            self.geometric_gnn = EGNN(
+                in_node_dim=1, hidden_dim=32, out_node_dim=output_dim
+            )
+
+    # --- Core Routines ---
+    def forward(self, h, x, edge_index, batch):
+        # Node-level embeddings
+        h, x = self.geometric_gnn(h, x, edge_index)
+
+        # Graph-level embedddings
+        h_skels = list()
+        edge_index = edge_index
+        num_graphs = int(batch.max().item()) + 1
+        for graph_id in range(num_graphs):
+            # Extract subgraph
+            node_mask = batch == graph_id
+            h_g, x_g, edge_index_g = self.extract_subgraph(
+                h, x, edge_index, node_mask
+            )
+
+            # Pool node embeddings
+            h_g, x_g, edge_index_g = self.pool_nonbranching_paths(
+                h_g, x_g, edge_index_g
+            )
+
+            # Encode pooled graph
+            h_g = self.encode_pooled_graph(h_g, x_g, edge_index_g)
+            h_skels.append(h_g)
+        return torch.cat(h_skels, dim=0)
+
+    def pool_nonbranching_paths(self, h, x, edge_index):
+        # Extract adjacency matrix and degrees
+        num_nodes = h.size(0)
+        adj, deg = self.get_adj_and_deg(edge_index, num_nodes)
+
+        # Search graph
+        node_to_path = torch.full((num_nodes,), -1, device=h.device)
+        path_idx = 0
+        h_pooled = list()
+        x_pooled = list()
+        visited = set()
+        for start in range(num_nodes):
+            # Check whether to visit
+            if start in visited:
+                continue
+
+            # Case 1: Branch points are singleton paths
+            if deg[start] > 2:
+                node_to_path[start] = path_idx
+                h_pooled.append(h[start])
+                x_pooled.append(x[start])
+                path_idx += 1
+                visited.add(start)
+                continue
+
+            # Case 2: Non-branching path traversal
+            path = [start]
+            visited.add(start)
+            prev = None
+            cur = start
+            while True:
+                nbs = [n for n in adj[cur] if n != prev]
+                if len(nbs) != 1:
+                    break
+
+                nxt = nbs[0]
+                if nxt in visited or deg[nxt] > 2:
+                    break
+
+                path.append(nxt)
+                visited.add(nxt)
+                prev, cur = cur, nxt
+
+            h_pooled.append(h[path].mean(dim=0))
+            x_pooled.append(x[path].mean(dim=0))
+
+            for n in path:
+                node_to_path[n] = path_idx
+
+            path_idx += 1
+
+        # Finish
+        h_pooled = torch.stack(h_pooled, dim=0)
+        x_pooled = torch.stack(x_pooled, dim=0)
+        edge_index_pooled = self.get_edge_index_pooled(
+            edge_index, node_to_path
+        )
+        return h_pooled, x_pooled, edge_index_pooled
+
+    def get_adj_and_deg(self, edge_index, num_nodes):
+        # Compute node degrees
+        deg = torch.zeros(num_nodes, dtype=torch.long)
+        ones = torch.ones(edge_index.shape[1], dtype=torch.long)
+        deg.scatter_add_(0, edge_index[0], ones)
+        deg.scatter_add_(0, edge_index[1], ones)
+
+        # Build adjacency list
+        adj = [[] for _ in range(num_nodes)]
+        for u, v in edge_index.t().tolist():
+            adj[u].append(v)
+            adj[v].append(u)
+        return adj, deg
+
+    def encode_pooled_graph(self, h, x, edge_index):
+        # Message passing over pooled graph
+        h = self.gnn_h(h, edge_index)
+        x = self.gnn_x(x, edge_index)
+
+        # temp
+        if h.size(0) == 0 or x.size(0) == 0:
+            print(h.size(0), x.size(0))
+            raise RuntimeError("Empty tensor passed to graph pooling")
+
+        # Graph-level pooling
+        h = h.mean(dim=0, keepdim=True)
+        x = x.mean(dim=0, keepdim=True)
+        return torch.cat((h, x), dim=1)
+
+    # --- Helpers ---
+    def extract_subgraph(self, h, x, edge_index, node_mask):
+        # Build subgraph
+        node_ids = (node_mask).nonzero(as_tuple=True)[0]
+        h_g = h[node_ids]
+        x_g = x[node_ids]
+
+        # Remap nodes and edges
+        id_map = {int(n): i for i, n in enumerate(node_ids.tolist())}
+        edge_mask = node_mask[edge_index[0]] & node_mask[edge_index[1]]
+        edge_index_g = edge_index[:, edge_mask]
+        edge_index_g = torch.stack(
+            [
+                torch.tensor([id_map[int(u)] for u in edge_index_g[0]]),
+                torch.tensor([id_map[int(v)] for v in edge_index_g[1]]),
+            ],
+            dim=0,
+        )
+        return h_g, x_g, edge_index_g
+
+    @staticmethod
+    def get_edge_index_pooled(edge_index, node_to_path):
+        # Extract edges in pooled graph
+        src, dst = edge_index
+        src_p = node_to_path[src]
+        dst_p = node_to_path[dst]
+
+        # Remove intra-path edges
+        mask = src_p != dst_p
+        edge_index_pooled = torch.stack([src_p[mask], dst_p[mask]], dim=0)
+        return torch.unique(edge_index_pooled, dim=1)
+
+
+# --- Geometric GNN Architectures ---
+class EGNN(nn.Module):
+
+    def __init__(
+        self,
+        in_node_dim,
+        hidden_dim,
+        out_node_dim,
+        in_edge_dim=0,
+        device="cuda",
+        act_fn=nn.SiLU(),
+        n_layers=4,
+        residual=True,
+        attention=False,
+        normalize=False,
+        tanh=False,
+    ):
+        """
+        Instantiates an EGNN object.
+
+        Parameters
+        ----------
+        in_node_dim : int
+            Number of features for 'h' at the input.
+        hidden_dim : int
+            Number of hidden features.
+        out_node_dim : int
+            Number of features for 'h' at the output.
+        in_edge_dim : int, optional
+            Number of features for the edge features.
+        device : str
+            Device to load model and inputs. Default is "cuda".
+        act_fn : ...
+            Non-linearity
+        n_layers : int
+            Number of layer for the EGNN.
+        residual : bool
+            Indication of whether to use residual connections.
+        attention : bool
+            Indication of whether using attention mechanism.
+        normalize : bool
+            Normalizes the coordinates messages such that:
+                x^{l+1}_i = x^{l}_i + Σ(x_i - x_j)phi_x(m_ij)
+        tanh : ...
+            Sets a tanh activation function at the output of phi_x(m_ij).
+        """
+        # Call parent class
+        super(EGNN, self).__init__()
+
+        # Instance attributes
+        self.hidden_dim = hidden_dim
+        self.device = device
+        self.n_layers = n_layers
+        self.embedding_in = nn.Linear(in_node_dim, self.hidden_dim)
+        self.embedding_out = nn.Linear(self.hidden_dim, out_node_dim)
+
+        # Build architecture
+        for i in range(0, n_layers):
+            self.add_module(
+                "gcl_%d" % i,
+                E_GCL(
+                    self.hidden_dim,
+                    self.hidden_dim,
+                    self.hidden_dim,
+                    edges_in_dim=in_edge_dim,
+                    act_fn=act_fn,
+                    residual=residual,
+                    attention=attention,
+                    normalize=normalize,
+                    tanh=tanh,
+                ),
+            )
+        self.to(self.device)
+
+    # --- Core Routines ---
+    def forward(self, h, x, edge_index):
+        h = self.embedding_in(h)
+        for i in range(0, self.n_layers):
+            h, x, _ = self._modules["gcl_%d" % i](h, edge_index, x)
+        h = self.embedding_out(h)
+        return h, x
+
+
+class E_GCL(nn.Module):
+    """
+    Class that implements an equivariant convolutional layer (i.e. E(n)).
+    """
+
+    def __init__(
+        self,
+        input_dim,
+        output_dim,
+        hidden_dim,
+        edges_in_dim=0,
+        act_fn=nn.SiLU(),
+        residual=True,
+        attention=False,
+        normalize=False,
+        coords_agg="mean",
+        tanh=False,
+    ):
+        # Call parent class
+        super(E_GCL, self).__init__()
+
+        # Instance attributes
+        input_edge = input_dim * 2
+        self.residual = residual
+        self.attention = attention
+        self.normalize = normalize
+        self.coords_agg = coords_agg
+        self.tanh = tanh
+        self.epsilon = 1e-8
+        edge_coords_dim = 1
+
+        # Architecture
+        self.node_mlp = nn.Sequential(
+            nn.Linear(hidden_dim + input_dim, hidden_dim),
+            act_fn,
+            nn.Linear(hidden_dim, output_dim),
+        )
+        self.coord_mlp = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            act_fn,
+            nn.Linear(hidden_dim, 1, bias=False),
+        )
+        self.edge_mlp = nn.Sequential(
+            nn.Linear(input_edge + edge_coords_dim + edges_in_dim, hidden_dim),
+            act_fn,
+            nn.Linear(hidden_dim, hidden_dim),
+            act_fn,
+        )
+        if self.attention:
+            self.att_mlp = nn.Sequential(
+                nn.Linear(hidden_dim, 1), nn.Sigmoid()
+            )
+
+    def edge_model(self, source, target, radial, edge_attr):
+        if edge_attr is None:
+            out = torch.cat([source, target, radial], dim=1)
+        else:
+            out = torch.cat([source, target, radial, edge_attr], dim=1)
+        out = self.edge_mlp(out)
+        if self.attention:
+            att_val = self.att_mlp(out)
+            out = out * att_val
+        return out
+
+    def node_model(self, x, edge_index, edge_attr, node_attr):
+        row, col = edge_index
+        agg = unsorted_segment_sum(edge_attr, row, num_segments=x.size(0))
+        if node_attr is not None:
+            agg = torch.cat([x, agg, node_attr], dim=1)
+        else:
+            agg = torch.cat([x, agg], dim=1)
+        out = self.node_mlp(agg)
+        if self.residual:
+            out = x + out
+        return out, agg
+
+    def coord_model(self, coord, edge_index, coord_diff, edge_feat):
+        row, col = edge_index
+        trans = coord_diff * self.coord_mlp(edge_feat)
+        if self.coords_agg == "sum":
+            agg = unsorted_segment_sum(trans, row, num_segments=coord.size(0))
+        elif self.coords_agg == "mean":
+            agg = unsorted_segment_mean(trans, row, num_segments=coord.size(0))
+        else:
+            raise Exception("Wrong coords_agg parameter" % self.coords_agg)
+        coord += agg
+        return coord
+
+    def coord2radial(self, edge_index, coord):
+        row, col = edge_index
+        coord_diff = coord[row] - coord[col]
+        radial = torch.sum(coord_diff**2, 1).unsqueeze(1)
+
+        if self.normalize:
+            norm = torch.sqrt(radial).detach() + self.epsilon
+            coord_diff = coord_diff / norm
+
+        radial = torch.zeros_like(radial, device="cuda")
+        return radial, coord_diff
+
+    def forward(self, h, edge_index, coord, edge_attr=None, node_attr=None):
+        row, col = edge_index
+        radial, coord_diff = self.coord2radial(edge_index, coord)
+
+        edge_feat = self.edge_model(h[row], h[col], radial, edge_attr)
+        coord = self.coord_model(coord, edge_index, coord_diff, edge_feat)
+        h, agg = self.node_model(h, edge_index, edge_feat, node_attr)
+
+        return h, coord, edge_attr
+
+
+def unsorted_segment_sum(data, segment_ids, num_segments):
+    result_shape = (num_segments, data.size(1))
+    result = data.new_full(result_shape, 0)  # Init empty result tensor.
+    segment_ids = segment_ids.unsqueeze(-1).expand(-1, data.size(1))
+    result.scatter_add_(0, segment_ids, data)
+    return result
+
+
+def unsorted_segment_mean(data, segment_ids, num_segments):
+    result_shape = (num_segments, data.size(1))
+    segment_ids = segment_ids.unsqueeze(-1).expand(-1, data.size(1))
+    result = data.new_full(result_shape, 0)  # Init empty result tensor.
+    count = data.new_full(result_shape, 0)
+    result.scatter_add_(0, segment_ids, data)
+    count.scatter_add_(0, segment_ids, torch.ones_like(data))
+    return result / count.clamp(min=1)
+
+
+# --- GNN Architectures ---
+class GAT(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        out_channels,
+        num_layers=2,
+        heads=4,
+        dropout=0.1,
+    ):
+        # Call parent class
+        super().__init__()
+
+        # Instance attributes
+        self.convs = nn.ModuleList()
+        self.dropout = dropout
+
+        # First layer
+        self.convs.append(
+            GATv2Conv(
+                in_channels,
+                hidden_channels,
+                heads=heads,
+                concat=True,
+                dropout=dropout,
+            )
+        )
+
+        # Hidden layers
+        for _ in range(num_layers - 2):
+            self.convs.append(
+                GATv2Conv(
+                    hidden_channels * heads,
+                    hidden_channels,
+                    heads=heads,
+                    concat=True,
+                    dropout=dropout,
+                )
+            )
+
+        # Output layer
+        self.convs.append(
+            GATv2Conv(
+                hidden_channels * heads,
+                out_channels,
+                heads=1,
+                concat=False,
+                dropout=dropout,
+            )
+        )
+
+    def forward(self, x, edge_index):
+        for conv in self.convs[:-1]:
+            x = conv(x, edge_index)
+            x = F.elu(x)
+        x = self.convs[-1](x, edge_index)
+        return x
+
+
+class GATGraphEncoder(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        hidden_channels,
+        out_channels,
+        heads=4,
+        num_layers=2,
+        dropout=0.2,
+    ):
+        # Call parent class
+        super().__init__()
+
+        # Instance attributes
+        self.gnn = GAT(
+            in_channels=in_channels,
+            hidden_channels=hidden_channels,
+            out_channels=hidden_channels,
+            num_layers=num_layers,
+            heads=heads,
+            dropout=dropout,
+        )
+        self.readout = nn.Linear(hidden_channels, out_channels)
+
+    def forward(self, x, edge_index):
+        x = self.gnn(x, edge_index)
+        x = x.mean(dim=0, keepdim=True)
+        return self.readout(x)
+
+
+# --- Helpers ---
+def get_edges(n_nodes):
+    rows, cols = [], []
+    for i in range(n_nodes):
+        for j in range(n_nodes):
+            if i != j:
+                rows.append(i)
+                cols.append(j)
+
+    edges = [rows, cols]
+    return edges
+
+
+def get_edges_batch(n_nodes, batch_size):
+    edges = get_edges(n_nodes)
+    edge_attr = torch.ones(len(edges[0]) * batch_size, 1)
+    edges = [torch.LongTensor(edges[0]), torch.LongTensor(edges[1])]
+    if batch_size == 1:
+        return edges, edge_attr
+    elif batch_size > 1:
+        rows, cols = [], []
+        for i in range(batch_size):
+            rows.append(edges[0] + n_nodes * i)
+            cols.append(edges[1] + n_nodes * i)
+        edges = [torch.cat(rows), torch.cat(cols)]
+    return edges, edge_attr
+
+
+def subgraph_to_data(subgraph):
+    h = torch.tensor(subgraph.node_radius[:, None], dtype=torch.float32)
+    x = torch.tensor(subgraph.node_xyz, dtype=torch.float32)
+    edges = torch.tensor(list(subgraph.edges), dtype=torch.long).T
+    return h, x, edges
+
+
+if __name__ == "__main__":
+    # Dummy parameters
+    batch_size = 8
+    n_nodes = 4
+    n_feat = 1
+    x_dim = 3
+
+    # Dummy variables h, x and fully connected edges
+    h = torch.ones(batch_size * n_nodes, n_feat)
+    x = torch.ones(batch_size * n_nodes, x_dim)
+    edges, edge_attr = get_edges_batch(n_nodes, batch_size)
+
+    # Initialize EGNN
+    egnn = EGNN(
+        in_node_dim=n_feat, hidden_dim=32, out_node_dim=1, in_edge_dim=1
+    )
+
+    # Run EGNN
+    h, x = egnn(h, x, edges)
diff --git a/src/neuron_proofreader/models/gnn_models.py b/src/neuron_proofreader/models/gnn_models.py
new file mode 100644
index 00000000..95f13873
--- /dev/null
+++ b/src/neuron_proofreader/models/gnn_models.py
@@ -0,0 +1,189 @@
+"""
+Created on Fri April 11 11:00:00 2024
+
+@author: Anna Grim
+@email: anna.grim@alleninstitute.org
+
+Graph neural network architectures that classify edge proposals.
+
+"""
+
+from torch import nn
+from torch_geometric import nn as nn_geometric
+
+import ast
+import torch
+import torch.nn.init as init
+
+from neuron_proofreader.machine_learning.vision_models import CNN3D
+from neuron_proofreader.split_proofreading import split_feature_extraction
+from neuron_proofreader.utils.ml_util import FeedForwardNet
+
+
+# --- Models ---
+class VisionHGAT(torch.nn.Module):
+    """
+    Heterogeneous graph attention network that processes multimodal features
+    such as image patches and feature vectors.
+    """
+
+    # Class attributes
+    relations = [
+        str(("branch", "to", "branch")),
+        str(("proposal", "to", "proposal")),
+        str(("branch", "to", "proposal")),
+        str(("proposal", "to", "branch")),
+    ]
+
+    def __init__(
+        self,
+        patch_shape,
+        disable_msg_passing=False,
+        heads=4,
+        hidden_dim=128,
+        n_layers=2,
+    ):
+        # Call parent class
+        super().__init__()
+
+        # Initial embeddings
+        self.node_embedding = init_node_embedding(hidden_dim)
+        self.patch_embedding = init_patch_embedding(
+            patch_shape, hidden_dim // 2
+        )
+
+        # Message passing layers
+        self.disable_msg_passing = disable_msg_passing
+        if self.disable_msg_passing:
+            self.gat1 = self.init_mlp_layers(hidden_dim, n_layers)
+            self.gat2 = self.init_mlp_layers(hidden_dim, n_layers)
+            self.output = nn.Linear(hidden_dim, 1)
+        else:
+            self.gat1 = self.init_gat(hidden_dim, hidden_dim, heads)
+            self.gat2 = self.init_gat(hidden_dim * heads, hidden_dim, heads)
+            self.output = nn.Linear(hidden_dim * heads**2, 1)
+
+        # Initialize weights
+        self.init_weights()
+
+    def init_gat(self, hidden_dim, edge_dim, heads):
+        gat_dict = dict()
+        for relation in VisionHGAT.relations:
+            # Parse relation string
+            relation = ast.literal_eval(relation)
+            node_type_1, _, node_type_2 = relation
+            is_same = node_type_1 == node_type_2
+
+            # Initialize layer
+            init_gat = init_gat_same if is_same else init_gat_mixed
+            gat_dict[relation] = init_gat(hidden_dim, edge_dim, heads)
+        return nn_geometric.HeteroConv(gat_dict)
+
+    def init_mlp_layers(self, hidden_dim, n_layers=2):
+        layers = nn.ModuleList()
+        for _ in range(n_layers):
+            layers.append(
+                nn_geometric.HeteroDictLinear(
+                    hidden_dim, hidden_dim, types=("branch", "proposal")
+                )
+            )
+        return layers
+
+    def init_weights(self):
+        """
+        Initializes linear layers.
+        """
+        for layer in [self.node_embedding, self.patch_embedding, self.output]:
+            for param in layer.parameters():
+                if len(param.shape) > 1:
+                    init.kaiming_normal_(param)
+                else:
+                    init.zeros_(param)
+
+    def forward(self, input_dict):
+        x_dict = input_dict["x_dict"]
+        x_img = input_dict["img"]
+        edge_index_dict = input_dict["edge_index_dict"]
+
+        # Initial embedding
+        x_img = self.patch_embedding(x_img)
+        for key, f in self.node_embedding.items():
+            x_dict[key] = f(x_dict[key])
+        x_dict["proposal"] = torch.cat((x_dict["proposal"], x_img), dim=1)
+
+        # Message passing
+        if self.disable_msg_passing:
+            for layer in self.gat1:
+                x_dict = layer(x_dict)
+            for layer in self.gat2:
+                x_dict = layer(x_dict)
+        else:
+            x_dict = self.gat1(x_dict, edge_index_dict)
+            x_dict = self.gat2(x_dict, edge_index_dict)
+        return self.output(x_dict["proposal"])
+
+
+# --- Helpers ---
+def init_gat_same(hidden_dim, edge_dim, heads):
+    gat = nn_geometric.GATv2Conv(
+        -1, hidden_dim, dropout=0.1, edge_dim=edge_dim, heads=heads
+    )
+    return gat
+
+
+def init_gat_mixed(hidden_dim, edge_dim, heads):
+    gat = nn_geometric.GATv2Conv(
+        (hidden_dim, hidden_dim),
+        hidden_dim,
+        add_self_loops=False,
+        dropout=0.1,
+        edge_dim=edge_dim,
+        heads=heads,
+    )
+    return gat
+
+
+def init_node_embedding(output_dim):
+    """
+    Builds a node embedding layer using a feed forward network for each
+    node type.
+
+    Parameters
+    ----------
+    output_dim : int
+        Output dimension for the embeddings. Note that the proposal output
+        dimension must be divided by 2 to account for the image patch
+        features.
+    """
+    # Get feature dimensions
+    node_input_dims = split_feature_extraction.get_feature_dict()
+    dim_b = node_input_dims["branch"]
+    dim_p = node_input_dims["proposal"]
+
+    # Set node embedding layer
+    node_embedding = nn.ModuleDict(
+        {
+            "branch": FeedForwardNet(dim_b, output_dim, 3),
+            "proposal": FeedForwardNet(dim_p, output_dim // 2, 3),
+        }
+    )
+    return node_embedding
+
+
+def init_patch_embedding(patch_shape, output_dim):
+    """
+    Builds the initial image patch embedding layer using a Convolutional
+    Neural Network (CNN).
+
+    Parameters
+    ----------
+    output_dim : int
+        Output dimension of the embedding.
+    """
+    patch_embedding = CNN3D(
+        patch_shape,
+        output_dim=output_dim,
+        n_conv_layers=6,
+        n_feat_channels=24,
+    )
+    return patch_embedding
diff --git a/src/neuron_proofreader/models/new_vision_models.py b/src/neuron_proofreader/models/new_vision_models.py
new file mode 100644
index 00000000..9a413770
--- /dev/null
+++ b/src/neuron_proofreader/models/new_vision_models.py
@@ -0,0 +1,231 @@
+"""
+Created on Thu July 2 13:00:00 2026
+
+@author: Anna Grim
+@email: anna.grim@alleninstitute.org
+
+Code for vision models that perform image classification tasks within
+NeuronProofreader pipelines.
+
+"""
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+from neuron_proofreader.utils.ml_util import FeedForwardNet
+
+
+class NewCNN3D(nn.Module):
+    """
+    Convolutional neural network for 3D images.
+    """
+
+    def __init__(
+        self,
+        patch_shape,
+        in_channels=2,
+        first_out_channels=16,
+        output_dim=1,
+        channel_growth=2,
+        dropout=0.1,
+        max_channels=256,
+        num_blocks=5,
+        num_single_blocks=2,
+        use_double=True,
+    ):
+        """
+        Instantiates a CNN3D object.
+
+        Parameters
+        ----------
+        patch_shape : Tuple[int]
+            Shape of input image patch.
+        output_dim : int, optional
+            Dimension of output. Default is 1.
+        dropout : float, optional
+            Fraction of values to randomly drop during training. Default is
+            0.1.
+        num_blocks : int, optional
+            Number of convolutional blocks. Default is 5.
+        use_double : bool, optional
+            Indication of whether to use double convolution. Default is True.
+        """
+        # Call parent class
+        nn.Module.__init__(self)
+
+        # Instance attributes
+        self.dropout = dropout
+        self.patch_shape = patch_shape
+
+        self.encode = Encoder3D(
+            in_channels,
+            first_out_channels,
+            num_blocks,
+            channel_growth=channel_growth,
+            use_double=use_double,
+            max_channels=max_channels,
+            num_single_blocks=num_single_blocks,
+        )
+        self.output = FeedForwardNet(self.encode.out_channels, output_dim, 3)
+
+        # Initialize weights
+        self.apply(self.init_weights)
+
+    @staticmethod
+    def init_weights(m):
+        """
+        Initializes the weights and biases of a given PyTorch layer.
+
+        Parameters
+        ----------
+        m : nn.Module
+            PyTorch layer or module.
+        """
+        if isinstance(m, nn.Conv3d):
+            nn.init.xavier_normal_(m.weight)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.Linear):
+            nn.init.xavier_normal_(m.weight)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.GroupNorm):
+            nn.init.constant_(m.weight, 1)
+            nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        """
+        Passes the given input through this neural network.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input vector of features.
+
+        Returns
+        -------
+        x : torch.Tensor
+            Output of the neural network.
+        """
+        x = self.encode(x)
+        x = F.adaptive_avg_pool3d(x, 1).flatten(1)
+        x = self.output(x)
+        return x
+
+
+class Encoder3D(nn.Module):
+    """
+    Sequence of ConvBlock3D blocks with growing (capped) channel width.
+    """
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        num_blocks,
+        channel_growth=2,
+        use_double=True,
+        max_channels=256,
+        num_single_blocks=2,
+    ):
+        """
+        Instantiates an Encoder3D object.
+
+        Parameters
+        ----------
+        in_channels : int
+            Number of channels input to the first block.
+        out_channels : int
+            Number of channels output by the first block.
+        num_blocks : int
+            Number of conv blocks in the encoder.
+        channel_growth : float, optional
+            Multiplicative channel growth factor per layer. Default is 2.
+        use_double : bool, optional
+            Indication of whether blocks use double convolution. Default is True.
+        max_channels : int, optional
+            Cap on channel growth across layers. Default is 128.
+        """
+        # Call parent class
+        super().__init__()
+
+        # Create encoding blocks
+        blocks = list()
+        for i in range(num_blocks):
+            use_double = i > num_single_blocks
+            block = ConvBlock3D(
+                in_channels, out_channels, use_double=use_double
+            )
+            blocks.append(block)
+            in_channels = block.out_channels
+            out_channels = int(min(out_channels * channel_growth, max_channels))
+
+        # Instance attributes
+        self.blocks = nn.ModuleList(blocks)
+        self.out_channels = self.blocks[-1].out_channels
+
+    def forward(self, x):
+        for block in self.blocks:
+            x = block(x)
+        return x
+
+
+class ConvBlock3D(nn.Module):
+
+    def __init__(
+        self, in_channels, out_channels, kernel_size=3, use_double=True
+    ):
+        """
+        Instantiates a ConvBlock3D object.
+ 
+        Parameters
+        ----------
+        in_channels : int
+            Number of input channels to this block.
+        out_channels : int
+            Number of output channels from this block.
+        kernel_size : int, optional
+            Size of kernel used on convolutional layers. Default is 3.
+        use_double : bool, optional
+            Indication of whether to apply a second conv+norm+act before
+            pooling. Default is True.
+        """
+        # Call parent class
+        super().__init__()
+
+        # Instance attributes
+        self.out_channels = out_channels
+
+        # Create encoding layers
+        layers = self.create_unit(in_channels, out_channels, kernel_size)
+        if use_double:
+            layers.extend(
+                self.create_unit(out_channels, out_channels, kernel_size)
+            )
+
+        self.net = nn.Sequential(*layers, nn.MaxPool3d(kernel_size=2))
+
+    def forward(self, x):
+        return self.net(x)
+
+    # --- Helpers ---
+    def create_unit(self, in_channels, out_channels, kernel_size):
+        padding = kernel_size // 2
+        n_groups = self.get_num_groups(out_channels, 8)
+        unit = [
+            nn.Conv3d(
+                in_channels,
+                out_channels,
+                kernel_size,
+                padding=padding,
+            ),
+            nn.GroupNorm(n_groups, out_channels),
+            nn.GELU()
+        ]
+        return unit
+
+    @staticmethod
+    def get_num_groups(num_channels, max_groups=8):
+        for g in reversed(range(1, max_groups + 1)):
+            if num_channels % g == 0:
+                return g
+        return 1
diff --git a/src/neuron_proofreader/models/point_cloud_models.py b/src/neuron_proofreader/models/point_cloud_models.py
new file mode 100644
index 00000000..23d46cab
--- /dev/null
+++ b/src/neuron_proofreader/models/point_cloud_models.py
@@ -0,0 +1,347 @@
+"""
+Created on Thu Nov 20 5:00:00 2025
+
+@author: Anna Grim
+@email: anna.grim@alleninstitute.org
+
+Code for point cloud models that perform machine learning tasks within
+NeuronProofreader pipelines.
+
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from neuron_proofreader.machine_learning.vision_models import CNN3D
+from neuron_proofreader.utils.ml_util import FeedForwardNet
+
+
+# --- Architectures ---
+class PointNet2(nn.Module):
+
+    def __init__(self, output_dim=1):
+        super().__init__()
+        self.sa1 = PointNetSetAbstraction(
+            n_points=512, mlp_channels=[64, 64, 128]
+        )
+        self.sa2 = PointNetSetAbstraction(
+            n_points=128, mlp_channels=[128, 128, 256]
+        )
+        self.sa3 = PointNetSetAbstraction(
+            n_points=None, mlp_channels=[256, 512, 1024]
+        )  # global
+        self.fc1 = nn.Linear(1024, 512)
+        self.bn1 = nn.BatchNorm1d(512)
+        self.drop1 = nn.Dropout(0.4)
+        self.fc2 = nn.Linear(512, 256)
+        self.bn2 = nn.BatchNorm1d(256)
+        self.drop2 = nn.Dropout(0.4)
+        self.fc3 = nn.Linear(256, output_dim)
+
+    def forward(self, x):
+        """
+        x: [B, N, 3]
+        """
+        xyz, f1 = self.sa1(x)
+        xyz, f2 = self.sa2(xyz)
+        xyz, f3 = self.sa3(xyz)
+        x = F.relu(self.bn1(self.fc1(f3)))
+        x = self.drop1(x)
+        x = F.relu(self.bn2(self.fc2(x)))
+        x = self.drop2(x)
+        x = self.fc3(x)
+        return x
+
+
+class VisionPointNet(nn.Module):
+
+    def __init__(self, patch_shape, output_dim=128):
+        super().__init__()
+
+        self.point_net = PointNet2(output_dim=output_dim)
+        self.vision_model = CNN3D(
+            patch_shape,
+            n_conv_layers=6,
+            n_feat_channels=20,
+            output_dim=output_dim,
+            use_double_conv=True,
+        )
+        self.output = FeedForwardNet(2 * output_dim, 1, 3)
+
+    def forward(self, x):
+        """
+        Passes the given input through this neural network.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input vector of features.
+
+        Returns
+        -------
+        x : torch.Tensor
+            Output of the neural network.
+        """
+        x_img = self.vision_model(x["img"])
+        x_pc = self.point_net(x["point_cloud"])
+        x = torch.cat((x_img, x_pc), dim=1)
+        x = self.output(x)
+        return x
+
+
+class PointNetSetAbstraction(nn.Module):
+    """
+    PointNet++ Set Abstraction (SA) layer
+    """
+
+    def __init__(self, n_points, mlp_channels):
+        super().__init__()
+        self.n_points = n_points
+        layers = []
+        last_channel = 3
+        for out_ch in mlp_channels:
+            layers.append(nn.Conv1d(last_channel, out_ch, 1))
+            layers.append(nn.BatchNorm1d(out_ch))
+            layers.append(nn.ReLU())
+            last_channel = out_ch
+        self.mlp = nn.Sequential(*layers)
+
+    def forward(self, xyz):
+        """
+        xyz: [B, N, 3]
+        """
+        B, N, _ = xyz.shape
+        if self.n_points is not None:
+            idx = farthest_point_sample(xyz, self.n_points)
+            new_xyz = index_points(xyz, idx)
+        else:
+            new_xyz = xyz
+        # MLP expects [B, C, N]
+        features = self.mlp(new_xyz.transpose(1, 2))
+        features = torch.max(features, 2)[0]  # global pooling
+        return new_xyz, features
+
+
+class DGCNN(nn.Module):
+
+    def __init__(self, input_dim=3, output_dim=64, k=16):
+        super().__init__()
+        self.k = k  # number of neighbors
+
+        # BatchNorm layers
+        self.bn1 = nn.BatchNorm2d(32)
+        self.bn2 = nn.BatchNorm2d(32)
+        self.bn3 = nn.BatchNorm2d(64)
+
+        # Conv layers (Conv2d on edge features)
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(input_dim * 2, 32, kernel_size=1, bias=False),
+            self.bn1,
+            nn.LeakyReLU(0.2),
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(32 * 2, 32, kernel_size=1, bias=False),
+            self.bn2,
+            nn.LeakyReLU(0.2),
+        )
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(32 * 2, 64, kernel_size=1, bias=False),
+            self.bn3,
+            nn.LeakyReLU(0.2),
+        )
+
+        # Final fully-connected layer for embedding
+        self.fc = nn.Linear(128, output_dim)
+
+    def get_graph_feature(self, x, k):
+        # x: (B, C, N)
+        # Compute pairwise distance
+        B, C, N = x.size()
+        x = x.transpose(2, 1)  # (B, N, C)
+        idx = self.knn(x, k)  # (B, N, k)
+
+        # Gather neighbors
+        neighbors = self.index_points(x, idx)  # (B, N, k, C)
+        x = x.unsqueeze(2).expand_as(neighbors)  # (B, N, k, C)
+        edge_feature = torch.cat((x, neighbors - x), dim=-1)  # (B, N, k, 2C)
+        return edge_feature.permute(0, 3, 1, 2).contiguous()  # (B, 2C, N, k)
+
+    @staticmethod
+    def knn(x, k):
+        # Compute pairwise distance
+        inner = -2 * torch.matmul(x, x.transpose(2, 1))
+        xx = torch.sum(x**2, dim=2, keepdim=True)
+        pairwise_distance = -xx - inner - xx.transpose(2, 1)
+        _, idx = pairwise_distance.topk(k=k, dim=-1)  # (B, N, k)
+        return idx
+
+    @staticmethod
+    def index_points(x, idx):
+        # x: (B, N, C), idx: (B, N, k)
+        B = x.size(0)
+        batch_indices = (
+            torch.arange(B, device=x.device)
+            .view(B, 1, 1)
+            .repeat(1, idx.size(1), idx.size(2))
+        )
+        return x[batch_indices, idx, :]
+
+    def forward(self, x):
+        """
+        Passes the given input through this neural network.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input vector of features.
+
+        Returns
+        -------
+        x : torch.Tensor
+            Output of the neural network.
+        """
+        x1 = F.relu(self.conv1(self.get_graph_feature(x, self.k)))
+        x1 = x1.max(dim=-1)[0]
+
+        x2 = F.relu(self.conv2(self.get_graph_feature(x1, self.k)))
+        x2 = x2.max(dim=-1)[0]
+
+        x3 = F.relu(self.conv3(self.get_graph_feature(x2, self.k)))
+        x3 = x3.max(dim=-1)[0]
+
+        x_cat = torch.cat((x1, x2, x3), dim=1)
+        x_global = x_cat.max(dim=-1)[0]
+        return self.fc(x_global)
+
+
+class VisionDGCNN(nn.Module):
+
+    def __init__(self, patch_shape, output_dim=128):
+        super().__init__()
+
+        self.dgcnn = DGCNN(output_dim=output_dim)
+        self.vision_model = CNN3D(
+            patch_shape,
+            n_conv_layers=6,
+            n_feat_channels=20,
+            output_dim=output_dim,
+            use_double_conv=True,
+        )
+        self.output = FeedForwardNet(2 * output_dim, 1, 3)
+
+    def forward(self, x):
+        """
+        Passes the given input through this neural network.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input vector of features.
+
+        Returns
+        -------
+        x : torch.Tensor
+            Output of the neural network.
+        """
+        x_img = self.vision_model(x["img"])
+        x_pc = self.dgcnn(x["point_cloud"])
+        x = torch.cat((x_img, x_pc), dim=1)
+        x = self.output(x)
+        return x
+
+
+# --- Point Cloud Generation ---
+def farthest_point_sample(xyz, n_points):
+    """
+    Farthest Point Sampling (FPS)
+    xyz: [B, N, 3]
+    return: [B, n_points] indices
+    """
+    B, N, _ = xyz.shape
+    centroids = torch.zeros(B, n_points, dtype=torch.long, device=xyz.device)
+    distance = torch.ones(B, N, device=xyz.device) * 1e10
+    farthest = torch.randint(0, N, (B,), dtype=torch.long, device=xyz.device)
+    batch_indices = torch.arange(B, dtype=torch.long, device=xyz.device)
+    for i in range(n_points):
+        centroids[:, i] = farthest
+        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
+        dist = torch.sum((xyz - centroid) ** 2, -1)
+        mask = dist < distance
+        distance[mask] = dist[mask]
+        farthest = torch.max(distance, -1)[1]
+    return centroids
+
+
+def index_points(points, idx):
+    """
+    Index points from xyz/features
+    points: [B, N, C]
+    idx: [B, S]
+    return: [B, S, C]
+    """
+    B = points.shape[0]
+    S = idx.shape[1]
+    batch_indices = (
+        torch.arange(B, dtype=torch.long, device=points.device)
+        .view(B, 1)
+        .repeat(1, S)
+    )
+    return points[batch_indices, idx, :]
+
+
+def subgraph_to_point_cloud(graph, n_points=3600):
+    point_cloud = list()
+    for n1, n2 in graph.edges:
+        # Use average radius
+        r1 = graph.node_radius[n1]
+        r2 = graph.node_radius[n2]
+        r = (r1 + r2) / 2
+
+        pts = sample_cylinder_between_points(
+            graph.node_xyz[n1], graph.node_xyz[n2], r
+        )
+        point_cloud.append(pts)
+    point_cloud = np.vstack(point_cloud)
+
+    total_points = point_cloud.shape[0]
+    idxs = np.arange(len(point_cloud))
+    if total_points >= n_points:
+        # Downsample randomly
+        sampled_idxs = np.random.choice(idxs, n_points, replace=False)
+        point_cloud = point_cloud[sampled_idxs]
+    else:
+        # Upsample with replacement
+        sampled_idxs = np.random.choice(idxs, n_points, replace=True)
+        point_cloud = point_cloud[sampled_idxs]
+    return point_cloud.T
+
+
+def sample_cylinder_between_points(p1, p2, r, n_samples=25):
+    p1 = np.array(p1)
+    p2 = np.array(p2)
+    axis = p2 - p1
+    axis_length = np.linalg.norm(axis)
+    axis_dir = axis / axis_length
+
+    # Build orthogonal basis
+    if np.allclose(axis_dir, [0, 0, 1]):
+        ortho1 = np.array([1, 0, 0])
+    else:
+        ortho1 = np.cross(axis_dir, [0, 0, 1])
+        ortho1 /= np.linalg.norm(ortho1)
+    ortho2 = np.cross(axis_dir, ortho1)
+
+    # Random samples along axis
+    t_values = np.random.rand(n_samples)
+    centers = p1 + np.outer(t_values, axis)
+
+    # Random samples in circular cross-section
+    angles = np.random.rand(n_samples) * 2 * np.pi
+    radii = r * np.sqrt(np.random.rand(n_samples))
+    offsets = np.outer(radii * np.cos(angles), ortho1) + np.outer(
+        radii * np.sin(angles), ortho2
+    )
+    points = centers + offsets
+    return points
diff --git a/src/neuron_proofreader/models/vision_models.py b/src/neuron_proofreader/models/vision_models.py
new file mode 100644
index 00000000..d514df69
--- /dev/null
+++ b/src/neuron_proofreader/models/vision_models.py
@@ -0,0 +1,490 @@
+"""
+Created on Sat July 15 12:00:00 2025
+
+@author: Anna Grim
+@email: anna.grim@alleninstitute.org
+
+Code for vision models that perform image classification tasks within
+NeuronProofreader pipelines.
+
+"""
+
+from einops import rearrange
+
+import torch
+import torch.nn as nn
+
+from neuron_proofreader.utils.ml_util import FeedForwardNet, init_mlp
+
+
+# --- CNNs ---
+class CNN3D(nn.Module):
+    """
+    Convolutional neural network for 3D images.
+    """
+
+    def __init__(
+        self,
+        patch_shape,
+        output_dim=1,
+        dropout=0.1,
+        n_conv_layers=5,
+        n_feat_channels=16,
+        use_double_conv=True,
+    ):
+        """
+        Instantiates a CNN3D object.
+
+        Parameters
+        ----------
+        patch_shape : Tuple[int]
+            Shape of input image patch.
+        output_dim : int, optional
+            Dimension of output. Default is 1.
+        dropout : float, optional
+            Fraction of values to randomly drop during training. Default is
+            0.1.
+        n_conv_layers : int, optional
+            Number of convolutional layers. Default is 5.
+        use_double_conv : bool, optional
+            Indication of whether to use double convolution. Default is True.
+        """
+        # Call parent class
+        nn.Module.__init__(self)
+
+        # Class attributes
+        self.dropout = dropout
+        self.patch_shape = patch_shape
+
+        # Convolutional layers
+        self.conv_layers = init_cnn3d(
+            2, n_feat_channels, n_conv_layers, use_double_conv=use_double_conv
+        )
+
+        # Output layer
+        flat_size = self._get_flattened_size()
+        self.output = FeedForwardNet(flat_size, output_dim, 3)
+
+        # Initialize weights
+        self.apply(self.init_weights)
+
+    def _get_flattened_size(self):
+        """
+        Compute the flattened feature vector size after applying a sequence
+        of convolutional and pooling layers on an input tensor with the given
+        shape.
+
+        Returns
+        -------
+        int
+            Length of the flattened feature vector after the convolutions and
+            pooling.
+        """
+        with torch.no_grad():
+            x = torch.zeros(1, 2, *self.patch_shape)
+            x = self.conv_layers(x)
+            return x.view(1, -1).size(1)
+
+    @staticmethod
+    def init_weights(m):
+        """
+        Initializes the weights and biases of a given PyTorch layer.
+
+        Parameters
+        ----------
+        m : nn.Module
+            PyTorch layer or module.
+        """
+        if isinstance(m, nn.Conv3d):
+            nn.init.kaiming_normal_(
+                m.weight, mode="fan_in", nonlinearity="leaky_relu"
+            )
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.Linear):
+            nn.init.xavier_normal_(m.weight)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.BatchNorm3d):
+            nn.init.constant_(m.weight, 1)
+            nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        """
+        Passes the given input through this neural network.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input vector of features.
+
+        Returns
+        -------
+        x : torch.Tensor
+            Output of the neural network.
+        """
+        x = self.conv_layers(x)
+        x = x.view(x.size(0), -1)
+        x = self.output(x)
+        return x
+
+
+# --- Transformers ---
+class ViT3D(nn.Module):
+    """
+    A class that implements a 3D Vision transformer.
+
+    Attributes
+    ----------
+    """
+
+    def __init__(
+        self,
+        img_shape=(128, 128, 128),
+        emb_dim=512,
+        depth=6,
+        heads=8,
+        mlp_dim=1024,
+        output_dim=1,
+    ):
+        """
+        Instantiates a ViT3D object.
+
+        Parameters
+        ----------
+        img_shape : Tuple[int], optional
+            Shape of the input image. Default is (128, 128, 128).
+        emb_dim : int, optional
+            Dimension of the embedding space. Default is 512.
+        depth : int, optional
+            Number of transformer blocks. Default is 6.
+        heads : int, optional
+            Number of attention heads in each transformer block. Default 8.
+        mlp_dim : int, optional
+            Dimension of MLP embedding space. Default is 1024.
+        output_dim : int, optional
+            Dimension of output. Default is 1.
+        """
+        # Call parent class
+        super().__init__()
+
+        # Token embedding
+        self.cls_token = nn.Parameter(torch.empty(1, 1, emb_dim))
+        self.img_tokenizer = ImageTokenizer3D(emb_dim, img_shape)
+
+        # Position embedding
+        n_tokens = self.img_tokenizer.count_tokens() + 1
+        self.pos_embedding = nn.Parameter(torch.empty(1, n_tokens, emb_dim))
+        print("# Tokens:", n_tokens)
+
+        # Transformer Blocks
+        self.transformer = nn.Sequential(
+            *[
+                TransformerEncoderBlock(emb_dim, heads, mlp_dim)
+                for _ in range(depth)
+            ]
+        )
+        self.norm = nn.LayerNorm(emb_dim)
+
+        # Output layer
+        self.output = FeedForwardNet(emb_dim, output_dim, 2)
+
+        # Initialize weights
+        self._init_weights()
+
+    def _init_weights(self):
+        """
+        Initializes the model's weights.
+        """
+        # Initialize token embedding
+        nn.init.trunc_normal_(self.cls_token, std=0.02)
+        nn.init.trunc_normal_(self.pos_embedding, std=0.02)
+
+        # Initialize Transformer and output layers
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.LayerNorm):
+                nn.init.ones_(module.weight)
+                nn.init.zeros_(module.bias)
+
+    def forward(self, x):
+        """
+        Passes the given input through this neural network.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input vector of features.
+
+        Returns
+        -------
+        x : torch.Tensor
+            Output of the neural network.
+        """
+        # Tokenize input -> (b, n_tokens, emb_dim)
+        x = self.img_tokenizer(x)
+        cls_tokens = self.cls_token.expand(x.size(0), -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # Transformer
+        x = x + self.pos_embedding[:, : x.size(1)]
+        x = self.transformer(x)
+        x = self.norm(x[:, 0])
+
+        # Output layer
+        x = self.output(x)
+        return x
+
+
+class ImageTokenizer3D(nn.Module):
+    """
+    A class for learning image token embeddings for transformer-based
+    architectures.
+
+    Attributes
+    ----------
+    dropout : nn.Dropout
+        Dropout layer.
+    emb_dim : int
+        Dimension of the embedding space.
+    img_shape : Tuple[int]
+        Shape of the input image (D, H, W).
+    pos_embedding : nn.Parameter
+        Learnable position encoding.
+    proj : nn.Conv3d
+        Convolutional layer that generates a learnable projection of the
+        tokens.
+    """
+
+    def __init__(
+        self,
+        emb_dim,
+        img_shape,
+        dropout=0.05,
+        n_cnn_layers=3,
+        n_cnn_channels=32,
+    ):
+        """
+        Instantiates a ImageTokenizer3D object.
+
+        Parameters
+        ----------
+        emb_dim : int
+            Dimension of the embedding space.
+        img_shape : Tuple[int]
+            Shape of the input image (D, H, W).
+        dropout : float, optional
+            Dropout probability applied after adding positional embeddings.
+            Default is 0.05.
+        n_cnn_layers : int, optional
+            Number of layers in the CNN that generates the initial token
+            embedding. Default is 3.
+        """
+        # Call parent class
+        super().__init__()
+
+        # Class attributes
+        self.emb_dim = emb_dim
+        self.img_shape = img_shape
+
+        # Image embedding
+        cnn_out_channels = n_cnn_channels * (2 ** (n_cnn_layers - 1))
+        self.proj = nn.Conv3d(cnn_out_channels, emb_dim, kernel_size=1)
+        self.tokenizer = init_cnn3d(
+            2, n_cnn_channels, n_cnn_layers, use_double_conv=False
+        )
+
+        # Positional embedding
+        n_tokens = self.count_tokens()
+        self.pos_embedding = nn.Parameter(torch.randn(1, n_tokens, emb_dim))
+        self.dropout = nn.Dropout(p=dropout)
+
+    def count_tokens(self):
+        """
+        Counts the number of tokens that are generated given the patch shape,
+        CCN3D architecture, and embedding dimension.
+
+        Returns
+        -------
+        int
+            Number of tokens generated by tokenizer.
+        """
+        with torch.no_grad():
+            dummy = torch.zeros(1, 2, *self.img_shape)
+            feats = self.tokenizer(dummy)
+            feats = self.proj(feats)
+            return feats.flatten(2).shape[-1]
+
+    def forward(self, x):
+        """
+        Forward pass that converts an input image into a sequence of token
+        embeddings.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, C, D, H, W).
+
+        Returns
+        -------
+        torch.Tensor
+            Token embeddings of shape (B, N, E).
+        """
+        # Embed image
+        x = self.tokenizer(x)
+        x = self.proj(x)
+
+        # Tokenize
+        x = rearrange(x, "b c d h w -> b (d h w) c")
+        x = x + self.pos_embedding
+        return self.dropout(x)
+
+
+class TransformerEncoderBlock(nn.Module):
+    """
+    Single transformer encoder block.
+
+    Attributes
+    ----------
+    attn : nn.MultiheadAttention
+        Multihead attention block.
+    dropout : nn.Dropout
+        Dropout layer.
+    mlp : nn.Sequential
+        Multi-layer perceptron.
+    norm1 : nn.LayerNorm
+        Applies layer normalization over a mini-batch of the inputs.
+    norm2 : nn.LayerNorm
+        Applied layer normalization over a mini-batch of the outputs of the
+        multihead attention block.
+    """
+
+    def __init__(self, emb_dim, heads, mlp_dim, dropout=0):
+        """
+        Instantiates a TransformerEncoderBlock object.
+
+        Parameters
+        ----------
+        emb_dim : int
+            Dimension of the embedding space.
+        heads : int
+            Number of attention heads.
+        mlp_dim : int
+            Dimensionality of the hidden layer in the MLP.
+        dropout : float, optional
+            Dropout probability applied after attention and MLP layers.
+            Default is 0.
+        """
+        # Call parent class
+        super().__init__()
+
+        # Attention head
+        self.norm1 = nn.LayerNorm(emb_dim)
+        self.attn = nn.MultiheadAttention(
+            emb_dim, heads, dropout=dropout, batch_first=True
+        )
+        self.norm2 = nn.LayerNorm(emb_dim)
+        self.mlp = init_mlp(emb_dim, mlp_dim, emb_dim, dropout)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x):
+        """
+        Forward pass of the encoder block.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (B, N, E).
+
+        Returns
+        -------
+        x : torch.Tensor
+            Output tensor of shape (B, N, E), same as input.
+        """
+        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
+        x = x + self.mlp(self.norm2(x))
+        x = self.dropout(x)
+        return x
+
+
+# --- Build Simple Neural Networks ---
+def init_cnn3d(in_channels, n_feat_channels, n_layers, use_double_conv=True):
+    """
+    Initializes a convolutional neural network.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of channels that are input to this convolutional layer.
+    out_channels : int
+        Number of channels that are output from this convolutional layer.
+    n_layers : int
+        Number of layers in the network.
+    use_double_conv : bool, optional
+        Indication of whether to use double convolution. Default is True.
+
+    Returns
+    -------
+    layers : torch.nn.Sequential
+        Sequence of operations that define the network.
+    """
+    layers = list()
+    in_channels = in_channels
+    out_channels = n_feat_channels
+    for i in range(n_layers):
+        # Build layer
+        layers.append(
+            init_conv_layer(in_channels, out_channels, 3, use_double_conv)
+        )
+
+        # Update channel sizes
+        in_channels = out_channels
+        out_channels *= 2
+    return nn.Sequential(*layers)
+
+
+def init_conv_layer(in_channels, out_channels, kernel_size, use_double_conv):
+    """
+    Initializes a convolutional layer.
+
+    Parameters
+    ----------
+    in_channels : int
+        Number of channels that are input to this convolutional layer.
+    out_channels : int
+        Number of channels that are output from this convolutional layer.
+    kernel_size : int
+        Size of kernel used on convolutional layers.
+    use_double_conv : bool
+        Indication of whether to use double convolution.
+
+    Returns
+    -------
+    layers : torch.nn.Sequential
+        Sequence of operations that define this convolutional layer.
+    """
+    # Convolution
+    layers = [
+        nn.Conv3d(
+            in_channels, out_channels, kernel_size, padding=kernel_size // 2
+        ),
+        nn.BatchNorm3d(out_channels),
+        nn.GELU(),
+    ]
+    if use_double_conv:
+        layers += [
+            nn.Conv3d(
+                out_channels,
+                out_channels,
+                kernel_size,
+                padding=kernel_size // 2,
+            ),
+            nn.BatchNorm3d(out_channels),
+            nn.GELU(),
+        ]
+    # Pooling
+    layers.append(nn.MaxPool3d(kernel_size=2))
+    return nn.Sequential(*layers)

From 2a009c1e3dacf427b2d9a2f714b056897601501a Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Thu, 2 Jul 2026 20:20:58 +0000
Subject: [PATCH 2/2] refactor: rm duplicates

---
 .../machine_learning/geometric_gnn_models.py  | 585 ------------------
 .../machine_learning/gnn_models.py            | 189 ------
 .../machine_learning/point_cloud_models.py    | 347 -----------
 .../machine_learning/vision_models.py         | 490 ---------------
 4 files changed, 1611 deletions(-)
 delete mode 100644 src/neuron_proofreader/machine_learning/geometric_gnn_models.py
 delete mode 100644 src/neuron_proofreader/machine_learning/gnn_models.py
 delete mode 100644 src/neuron_proofreader/machine_learning/point_cloud_models.py
 delete mode 100644 src/neuron_proofreader/machine_learning/vision_models.py

diff --git a/src/neuron_proofreader/machine_learning/geometric_gnn_models.py b/src/neuron_proofreader/machine_learning/geometric_gnn_models.py
deleted file mode 100644
index 0c68d7f1..00000000
--- a/src/neuron_proofreader/machine_learning/geometric_gnn_models.py
+++ /dev/null
@@ -1,585 +0,0 @@
-"""
-Created on Sat July 15 12:00:00 2025
-
-@author: Anna Grim
-@email: anna.grim@alleninstitute.org
-
-Code for graph neural network models that perform machine learning tasks
-within NeuronProofreader pipelines.
-
-"""
-
-from torch import nn
-from torch_geometric.nn import GATv2Conv
-
-import torch
-import torch.nn.functional as F
-
-from neuron_proofreader.machine_learning.vision_models import CNN3D
-from neuron_proofreader.utils.ml_util import FeedForwardNet
-
-
-# --- Multimodal GNN Architectures ---
-class VisionSkeleton(nn.Module):
-
-    def __init__(self, ggnn_name, patch_shape, output_dim=64):
-        # Call parent class
-        super().__init__()
-        assert ggnn_name in ["egnn"]
-
-        # Architecture
-        self.skeleton_model = SkeletonGNN(ggnn_name, output_dim=32)
-        self.vision_model = CNN3D(
-            patch_shape,
-            n_conv_layers=6,
-            n_feat_channels=20,
-            output_dim=output_dim,
-            use_double_conv=True,
-        )
-        self.output = FeedForwardNet(output_dim + 35, 1, 3)
-
-    def forward(self, x):
-        """
-        Passes the given input through this neural network.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Input vector of features.
-
-        Returns
-        -------
-        x : torch.Tensor
-            Output of the neural network.
-        """
-        # Modality-based embeddings
-        x_img = self.vision_model(x["img"])
-        x_skel = self.skeleton_model(*x["graph"])
-        x = torch.cat((x_img, x_skel), dim=1)
-
-        # Output layer
-        x = self.output(x)
-        return x
-
-
-class SkeletonGNN(nn.Module):
-
-    def __init__(self, ggnn_name, output_dim=64):
-        # Call parent class
-        super().__init__()
-
-        # Instance attributes
-        self.gnn_h = GAT(output_dim, 2 * output_dim, output_dim)
-        self.gnn_x = GAT(3, 16, 3)
-
-        # Set geometric gnn
-        if ggnn_name == "egnn":
-            self.geometric_gnn = EGNN(
-                in_node_dim=1, hidden_dim=32, out_node_dim=output_dim
-            )
-
-    # --- Core Routines ---
-    def forward(self, h, x, edge_index, batch):
-        # Node-level embeddings
-        h, x = self.geometric_gnn(h, x, edge_index)
-
-        # Graph-level embedddings
-        h_skels = list()
-        edge_index = edge_index
-        num_graphs = int(batch.max().item()) + 1
-        for graph_id in range(num_graphs):
-            # Extract subgraph
-            node_mask = batch == graph_id
-            h_g, x_g, edge_index_g = self.extract_subgraph(
-                h, x, edge_index, node_mask
-            )
-
-            # Pool node embeddings
-            h_g, x_g, edge_index_g = self.pool_nonbranching_paths(
-                h_g, x_g, edge_index_g
-            )
-
-            # Encode pooled graph
-            h_g = self.encode_pooled_graph(h_g, x_g, edge_index_g)
-            h_skels.append(h_g)
-        return torch.cat(h_skels, dim=0)
-
-    def pool_nonbranching_paths(self, h, x, edge_index):
-        # Extract adjacency matrix and degrees
-        num_nodes = h.size(0)
-        adj, deg = self.get_adj_and_deg(edge_index, num_nodes)
-
-        # Search graph
-        node_to_path = torch.full((num_nodes,), -1, device=h.device)
-        path_idx = 0
-        h_pooled = list()
-        x_pooled = list()
-        visited = set()
-        for start in range(num_nodes):
-            # Check whether to visit
-            if start in visited:
-                continue
-
-            # Case 1: Branch points are singleton paths
-            if deg[start] > 2:
-                node_to_path[start] = path_idx
-                h_pooled.append(h[start])
-                x_pooled.append(x[start])
-                path_idx += 1
-                visited.add(start)
-                continue
-
-            # Case 2: Non-branching path traversal
-            path = [start]
-            visited.add(start)
-            prev = None
-            cur = start
-            while True:
-                nbs = [n for n in adj[cur] if n != prev]
-                if len(nbs) != 1:
-                    break
-
-                nxt = nbs[0]
-                if nxt in visited or deg[nxt] > 2:
-                    break
-
-                path.append(nxt)
-                visited.add(nxt)
-                prev, cur = cur, nxt
-
-            h_pooled.append(h[path].mean(dim=0))
-            x_pooled.append(x[path].mean(dim=0))
-
-            for n in path:
-                node_to_path[n] = path_idx
-
-            path_idx += 1
-
-        # Finish
-        h_pooled = torch.stack(h_pooled, dim=0)
-        x_pooled = torch.stack(x_pooled, dim=0)
-        edge_index_pooled = self.get_edge_index_pooled(
-            edge_index, node_to_path
-        )
-        return h_pooled, x_pooled, edge_index_pooled
-
-    def get_adj_and_deg(self, edge_index, num_nodes):
-        # Compute node degrees
-        deg = torch.zeros(num_nodes, dtype=torch.long)
-        ones = torch.ones(edge_index.shape[1], dtype=torch.long)
-        deg.scatter_add_(0, edge_index[0], ones)
-        deg.scatter_add_(0, edge_index[1], ones)
-
-        # Build adjacency list
-        adj = [[] for _ in range(num_nodes)]
-        for u, v in edge_index.t().tolist():
-            adj[u].append(v)
-            adj[v].append(u)
-        return adj, deg
-
-    def encode_pooled_graph(self, h, x, edge_index):
-        # Message passing over pooled graph
-        h = self.gnn_h(h, edge_index)
-        x = self.gnn_x(x, edge_index)
-
-        # temp
-        if h.size(0) == 0 or x.size(0) == 0:
-            print(h.size(0), x.size(0))
-            raise RuntimeError("Empty tensor passed to graph pooling")
-
-        # Graph-level pooling
-        h = h.mean(dim=0, keepdim=True)
-        x = x.mean(dim=0, keepdim=True)
-        return torch.cat((h, x), dim=1)
-
-    # --- Helpers ---
-    def extract_subgraph(self, h, x, edge_index, node_mask):
-        # Build subgraph
-        node_ids = (node_mask).nonzero(as_tuple=True)[0]
-        h_g = h[node_ids]
-        x_g = x[node_ids]
-
-        # Remap nodes and edges
-        id_map = {int(n): i for i, n in enumerate(node_ids.tolist())}
-        edge_mask = node_mask[edge_index[0]] & node_mask[edge_index[1]]
-        edge_index_g = edge_index[:, edge_mask]
-        edge_index_g = torch.stack(
-            [
-                torch.tensor([id_map[int(u)] for u in edge_index_g[0]]),
-                torch.tensor([id_map[int(v)] for v in edge_index_g[1]]),
-            ],
-            dim=0,
-        )
-        return h_g, x_g, edge_index_g
-
-    @staticmethod
-    def get_edge_index_pooled(edge_index, node_to_path):
-        # Extract edges in pooled graph
-        src, dst = edge_index
-        src_p = node_to_path[src]
-        dst_p = node_to_path[dst]
-
-        # Remove intra-path edges
-        mask = src_p != dst_p
-        edge_index_pooled = torch.stack([src_p[mask], dst_p[mask]], dim=0)
-        return torch.unique(edge_index_pooled, dim=1)
-
-
-# --- Geometric GNN Architectures ---
-class EGNN(nn.Module):
-
-    def __init__(
-        self,
-        in_node_dim,
-        hidden_dim,
-        out_node_dim,
-        in_edge_dim=0,
-        device="cuda",
-        act_fn=nn.SiLU(),
-        n_layers=4,
-        residual=True,
-        attention=False,
-        normalize=False,
-        tanh=False,
-    ):
-        """
-        Instantiates an EGNN object.
-
-        Parameters
-        ----------
-        in_node_dim : int
-            Number of features for 'h' at the input.
-        hidden_dim : int
-            Number of hidden features.
-        out_node_dim : int
-            Number of features for 'h' at the output.
-        in_edge_dim : int, optional
-            Number of features for the edge features.
-        device : str
-            Device to load model and inputs. Default is "cuda".
-        act_fn : ...
-            Non-linearity
-        n_layers : int
-            Number of layer for the EGNN.
-        residual : bool
-            Indication of whether to use residual connections.
-        attention : bool
-            Indication of whether using attention mechanism.
-        normalize : bool
-            Normalizes the coordinates messages such that:
-                x^{l+1}_i = x^{l}_i + Σ(x_i - x_j)phi_x(m_ij)
-        tanh : ...
-            Sets a tanh activation function at the output of phi_x(m_ij).
-        """
-        # Call parent class
-        super(EGNN, self).__init__()
-
-        # Instance attributes
-        self.hidden_dim = hidden_dim
-        self.device = device
-        self.n_layers = n_layers
-        self.embedding_in = nn.Linear(in_node_dim, self.hidden_dim)
-        self.embedding_out = nn.Linear(self.hidden_dim, out_node_dim)
-
-        # Build architecture
-        for i in range(0, n_layers):
-            self.add_module(
-                "gcl_%d" % i,
-                E_GCL(
-                    self.hidden_dim,
-                    self.hidden_dim,
-                    self.hidden_dim,
-                    edges_in_dim=in_edge_dim,
-                    act_fn=act_fn,
-                    residual=residual,
-                    attention=attention,
-                    normalize=normalize,
-                    tanh=tanh,
-                ),
-            )
-        self.to(self.device)
-
-    # --- Core Routines ---
-    def forward(self, h, x, edge_index):
-        h = self.embedding_in(h)
-        for i in range(0, self.n_layers):
-            h, x, _ = self._modules["gcl_%d" % i](h, edge_index, x)
-        h = self.embedding_out(h)
-        return h, x
-
-
-class E_GCL(nn.Module):
-    """
-    Class that implements an equivariant convolutional layer (i.e. E(n)).
-    """
-
-    def __init__(
-        self,
-        input_dim,
-        output_dim,
-        hidden_dim,
-        edges_in_dim=0,
-        act_fn=nn.SiLU(),
-        residual=True,
-        attention=False,
-        normalize=False,
-        coords_agg="mean",
-        tanh=False,
-    ):
-        # Call parent class
-        super(E_GCL, self).__init__()
-
-        # Instance attributes
-        input_edge = input_dim * 2
-        self.residual = residual
-        self.attention = attention
-        self.normalize = normalize
-        self.coords_agg = coords_agg
-        self.tanh = tanh
-        self.epsilon = 1e-8
-        edge_coords_dim = 1
-
-        # Architecture
-        self.node_mlp = nn.Sequential(
-            nn.Linear(hidden_dim + input_dim, hidden_dim),
-            act_fn,
-            nn.Linear(hidden_dim, output_dim),
-        )
-        self.coord_mlp = nn.Sequential(
-            nn.Linear(hidden_dim, hidden_dim),
-            act_fn,
-            nn.Linear(hidden_dim, 1, bias=False),
-        )
-        self.edge_mlp = nn.Sequential(
-            nn.Linear(input_edge + edge_coords_dim + edges_in_dim, hidden_dim),
-            act_fn,
-            nn.Linear(hidden_dim, hidden_dim),
-            act_fn,
-        )
-        if self.attention:
-            self.att_mlp = nn.Sequential(
-                nn.Linear(hidden_dim, 1), nn.Sigmoid()
-            )
-
-    def edge_model(self, source, target, radial, edge_attr):
-        if edge_attr is None:
-            out = torch.cat([source, target, radial], dim=1)
-        else:
-            out = torch.cat([source, target, radial, edge_attr], dim=1)
-        out = self.edge_mlp(out)
-        if self.attention:
-            att_val = self.att_mlp(out)
-            out = out * att_val
-        return out
-
-    def node_model(self, x, edge_index, edge_attr, node_attr):
-        row, col = edge_index
-        agg = unsorted_segment_sum(edge_attr, row, num_segments=x.size(0))
-        if node_attr is not None:
-            agg = torch.cat([x, agg, node_attr], dim=1)
-        else:
-            agg = torch.cat([x, agg], dim=1)
-        out = self.node_mlp(agg)
-        if self.residual:
-            out = x + out
-        return out, agg
-
-    def coord_model(self, coord, edge_index, coord_diff, edge_feat):
-        row, col = edge_index
-        trans = coord_diff * self.coord_mlp(edge_feat)
-        if self.coords_agg == "sum":
-            agg = unsorted_segment_sum(trans, row, num_segments=coord.size(0))
-        elif self.coords_agg == "mean":
-            agg = unsorted_segment_mean(trans, row, num_segments=coord.size(0))
-        else:
-            raise Exception("Wrong coords_agg parameter" % self.coords_agg)
-        coord += agg
-        return coord
-
-    def coord2radial(self, edge_index, coord):
-        row, col = edge_index
-        coord_diff = coord[row] - coord[col]
-        radial = torch.sum(coord_diff**2, 1).unsqueeze(1)
-
-        if self.normalize:
-            norm = torch.sqrt(radial).detach() + self.epsilon
-            coord_diff = coord_diff / norm
-
-        radial = torch.zeros_like(radial, device="cuda")
-        return radial, coord_diff
-
-    def forward(self, h, edge_index, coord, edge_attr=None, node_attr=None):
-        row, col = edge_index
-        radial, coord_diff = self.coord2radial(edge_index, coord)
-
-        edge_feat = self.edge_model(h[row], h[col], radial, edge_attr)
-        coord = self.coord_model(coord, edge_index, coord_diff, edge_feat)
-        h, agg = self.node_model(h, edge_index, edge_feat, node_attr)
-
-        return h, coord, edge_attr
-
-
-def unsorted_segment_sum(data, segment_ids, num_segments):
-    result_shape = (num_segments, data.size(1))
-    result = data.new_full(result_shape, 0)  # Init empty result tensor.
-    segment_ids = segment_ids.unsqueeze(-1).expand(-1, data.size(1))
-    result.scatter_add_(0, segment_ids, data)
-    return result
-
-
-def unsorted_segment_mean(data, segment_ids, num_segments):
-    result_shape = (num_segments, data.size(1))
-    segment_ids = segment_ids.unsqueeze(-1).expand(-1, data.size(1))
-    result = data.new_full(result_shape, 0)  # Init empty result tensor.
-    count = data.new_full(result_shape, 0)
-    result.scatter_add_(0, segment_ids, data)
-    count.scatter_add_(0, segment_ids, torch.ones_like(data))
-    return result / count.clamp(min=1)
-
-
-# --- GNN Architectures ---
-class GAT(nn.Module):
-
-    def __init__(
-        self,
-        in_channels,
-        hidden_channels,
-        out_channels,
-        num_layers=2,
-        heads=4,
-        dropout=0.1,
-    ):
-        # Call parent class
-        super().__init__()
-
-        # Instance attributes
-        self.convs = nn.ModuleList()
-        self.dropout = dropout
-
-        # First layer
-        self.convs.append(
-            GATv2Conv(
-                in_channels,
-                hidden_channels,
-                heads=heads,
-                concat=True,
-                dropout=dropout,
-            )
-        )
-
-        # Hidden layers
-        for _ in range(num_layers - 2):
-            self.convs.append(
-                GATv2Conv(
-                    hidden_channels * heads,
-                    hidden_channels,
-                    heads=heads,
-                    concat=True,
-                    dropout=dropout,
-                )
-            )
-
-        # Output layer
-        self.convs.append(
-            GATv2Conv(
-                hidden_channels * heads,
-                out_channels,
-                heads=1,
-                concat=False,
-                dropout=dropout,
-            )
-        )
-
-    def forward(self, x, edge_index):
-        for conv in self.convs[:-1]:
-            x = conv(x, edge_index)
-            x = F.elu(x)
-        x = self.convs[-1](x, edge_index)
-        return x
-
-
-class GATGraphEncoder(nn.Module):
-
-    def __init__(
-        self,
-        in_channels,
-        hidden_channels,
-        out_channels,
-        heads=4,
-        num_layers=2,
-        dropout=0.2,
-    ):
-        # Call parent class
-        super().__init__()
-
-        # Instance attributes
-        self.gnn = GAT(
-            in_channels=in_channels,
-            hidden_channels=hidden_channels,
-            out_channels=hidden_channels,
-            num_layers=num_layers,
-            heads=heads,
-            dropout=dropout,
-        )
-        self.readout = nn.Linear(hidden_channels, out_channels)
-
-    def forward(self, x, edge_index):
-        x = self.gnn(x, edge_index)
-        x = x.mean(dim=0, keepdim=True)
-        return self.readout(x)
-
-
-# --- Helpers ---
-def get_edges(n_nodes):
-    rows, cols = [], []
-    for i in range(n_nodes):
-        for j in range(n_nodes):
-            if i != j:
-                rows.append(i)
-                cols.append(j)
-
-    edges = [rows, cols]
-    return edges
-
-
-def get_edges_batch(n_nodes, batch_size):
-    edges = get_edges(n_nodes)
-    edge_attr = torch.ones(len(edges[0]) * batch_size, 1)
-    edges = [torch.LongTensor(edges[0]), torch.LongTensor(edges[1])]
-    if batch_size == 1:
-        return edges, edge_attr
-    elif batch_size > 1:
-        rows, cols = [], []
-        for i in range(batch_size):
-            rows.append(edges[0] + n_nodes * i)
-            cols.append(edges[1] + n_nodes * i)
-        edges = [torch.cat(rows), torch.cat(cols)]
-    return edges, edge_attr
-
-
-def subgraph_to_data(subgraph):
-    h = torch.tensor(subgraph.node_radius[:, None], dtype=torch.float32)
-    x = torch.tensor(subgraph.node_xyz, dtype=torch.float32)
-    edges = torch.tensor(list(subgraph.edges), dtype=torch.long).T
-    return h, x, edges
-
-
-if __name__ == "__main__":
-    # Dummy parameters
-    batch_size = 8
-    n_nodes = 4
-    n_feat = 1
-    x_dim = 3
-
-    # Dummy variables h, x and fully connected edges
-    h = torch.ones(batch_size * n_nodes, n_feat)
-    x = torch.ones(batch_size * n_nodes, x_dim)
-    edges, edge_attr = get_edges_batch(n_nodes, batch_size)
-
-    # Initialize EGNN
-    egnn = EGNN(
-        in_node_dim=n_feat, hidden_dim=32, out_node_dim=1, in_edge_dim=1
-    )
-
-    # Run EGNN
-    h, x = egnn(h, x, edges)
diff --git a/src/neuron_proofreader/machine_learning/gnn_models.py b/src/neuron_proofreader/machine_learning/gnn_models.py
deleted file mode 100644
index 95f13873..00000000
--- a/src/neuron_proofreader/machine_learning/gnn_models.py
+++ /dev/null
@@ -1,189 +0,0 @@
-"""
-Created on Fri April 11 11:00:00 2024
-
-@author: Anna Grim
-@email: anna.grim@alleninstitute.org
-
-Graph neural network architectures that classify edge proposals.
-
-"""
-
-from torch import nn
-from torch_geometric import nn as nn_geometric
-
-import ast
-import torch
-import torch.nn.init as init
-
-from neuron_proofreader.machine_learning.vision_models import CNN3D
-from neuron_proofreader.split_proofreading import split_feature_extraction
-from neuron_proofreader.utils.ml_util import FeedForwardNet
-
-
-# --- Models ---
-class VisionHGAT(torch.nn.Module):
-    """
-    Heterogeneous graph attention network that processes multimodal features
-    such as image patches and feature vectors.
-    """
-
-    # Class attributes
-    relations = [
-        str(("branch", "to", "branch")),
-        str(("proposal", "to", "proposal")),
-        str(("branch", "to", "proposal")),
-        str(("proposal", "to", "branch")),
-    ]
-
-    def __init__(
-        self,
-        patch_shape,
-        disable_msg_passing=False,
-        heads=4,
-        hidden_dim=128,
-        n_layers=2,
-    ):
-        # Call parent class
-        super().__init__()
-
-        # Initial embeddings
-        self.node_embedding = init_node_embedding(hidden_dim)
-        self.patch_embedding = init_patch_embedding(
-            patch_shape, hidden_dim // 2
-        )
-
-        # Message passing layers
-        self.disable_msg_passing = disable_msg_passing
-        if self.disable_msg_passing:
-            self.gat1 = self.init_mlp_layers(hidden_dim, n_layers)
-            self.gat2 = self.init_mlp_layers(hidden_dim, n_layers)
-            self.output = nn.Linear(hidden_dim, 1)
-        else:
-            self.gat1 = self.init_gat(hidden_dim, hidden_dim, heads)
-            self.gat2 = self.init_gat(hidden_dim * heads, hidden_dim, heads)
-            self.output = nn.Linear(hidden_dim * heads**2, 1)
-
-        # Initialize weights
-        self.init_weights()
-
-    def init_gat(self, hidden_dim, edge_dim, heads):
-        gat_dict = dict()
-        for relation in VisionHGAT.relations:
-            # Parse relation string
-            relation = ast.literal_eval(relation)
-            node_type_1, _, node_type_2 = relation
-            is_same = node_type_1 == node_type_2
-
-            # Initialize layer
-            init_gat = init_gat_same if is_same else init_gat_mixed
-            gat_dict[relation] = init_gat(hidden_dim, edge_dim, heads)
-        return nn_geometric.HeteroConv(gat_dict)
-
-    def init_mlp_layers(self, hidden_dim, n_layers=2):
-        layers = nn.ModuleList()
-        for _ in range(n_layers):
-            layers.append(
-                nn_geometric.HeteroDictLinear(
-                    hidden_dim, hidden_dim, types=("branch", "proposal")
-                )
-            )
-        return layers
-
-    def init_weights(self):
-        """
-        Initializes linear layers.
-        """
-        for layer in [self.node_embedding, self.patch_embedding, self.output]:
-            for param in layer.parameters():
-                if len(param.shape) > 1:
-                    init.kaiming_normal_(param)
-                else:
-                    init.zeros_(param)
-
-    def forward(self, input_dict):
-        x_dict = input_dict["x_dict"]
-        x_img = input_dict["img"]
-        edge_index_dict = input_dict["edge_index_dict"]
-
-        # Initial embedding
-        x_img = self.patch_embedding(x_img)
-        for key, f in self.node_embedding.items():
-            x_dict[key] = f(x_dict[key])
-        x_dict["proposal"] = torch.cat((x_dict["proposal"], x_img), dim=1)
-
-        # Message passing
-        if self.disable_msg_passing:
-            for layer in self.gat1:
-                x_dict = layer(x_dict)
-            for layer in self.gat2:
-                x_dict = layer(x_dict)
-        else:
-            x_dict = self.gat1(x_dict, edge_index_dict)
-            x_dict = self.gat2(x_dict, edge_index_dict)
-        return self.output(x_dict["proposal"])
-
-
-# --- Helpers ---
-def init_gat_same(hidden_dim, edge_dim, heads):
-    gat = nn_geometric.GATv2Conv(
-        -1, hidden_dim, dropout=0.1, edge_dim=edge_dim, heads=heads
-    )
-    return gat
-
-
-def init_gat_mixed(hidden_dim, edge_dim, heads):
-    gat = nn_geometric.GATv2Conv(
-        (hidden_dim, hidden_dim),
-        hidden_dim,
-        add_self_loops=False,
-        dropout=0.1,
-        edge_dim=edge_dim,
-        heads=heads,
-    )
-    return gat
-
-
-def init_node_embedding(output_dim):
-    """
-    Builds a node embedding layer using a feed forward network for each
-    node type.
-
-    Parameters
-    ----------
-    output_dim : int
-        Output dimension for the embeddings. Note that the proposal output
-        dimension must be divided by 2 to account for the image patch
-        features.
-    """
-    # Get feature dimensions
-    node_input_dims = split_feature_extraction.get_feature_dict()
-    dim_b = node_input_dims["branch"]
-    dim_p = node_input_dims["proposal"]
-
-    # Set node embedding layer
-    node_embedding = nn.ModuleDict(
-        {
-            "branch": FeedForwardNet(dim_b, output_dim, 3),
-            "proposal": FeedForwardNet(dim_p, output_dim // 2, 3),
-        }
-    )
-    return node_embedding
-
-
-def init_patch_embedding(patch_shape, output_dim):
-    """
-    Builds the initial image patch embedding layer using a Convolutional
-    Neural Network (CNN).
-
-    Parameters
-    ----------
-    output_dim : int
-        Output dimension of the embedding.
-    """
-    patch_embedding = CNN3D(
-        patch_shape,
-        output_dim=output_dim,
-        n_conv_layers=6,
-        n_feat_channels=24,
-    )
-    return patch_embedding
diff --git a/src/neuron_proofreader/machine_learning/point_cloud_models.py b/src/neuron_proofreader/machine_learning/point_cloud_models.py
deleted file mode 100644
index 23d46cab..00000000
--- a/src/neuron_proofreader/machine_learning/point_cloud_models.py
+++ /dev/null
@@ -1,347 +0,0 @@
-"""
-Created on Thu Nov 20 5:00:00 2025
-
-@author: Anna Grim
-@email: anna.grim@alleninstitute.org
-
-Code for point cloud models that perform machine learning tasks within
-NeuronProofreader pipelines.
-
-"""
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from neuron_proofreader.machine_learning.vision_models import CNN3D
-from neuron_proofreader.utils.ml_util import FeedForwardNet
-
-
-# --- Architectures ---
-class PointNet2(nn.Module):
-
-    def __init__(self, output_dim=1):
-        super().__init__()
-        self.sa1 = PointNetSetAbstraction(
-            n_points=512, mlp_channels=[64, 64, 128]
-        )
-        self.sa2 = PointNetSetAbstraction(
-            n_points=128, mlp_channels=[128, 128, 256]
-        )
-        self.sa3 = PointNetSetAbstraction(
-            n_points=None, mlp_channels=[256, 512, 1024]
-        )  # global
-        self.fc1 = nn.Linear(1024, 512)
-        self.bn1 = nn.BatchNorm1d(512)
-        self.drop1 = nn.Dropout(0.4)
-        self.fc2 = nn.Linear(512, 256)
-        self.bn2 = nn.BatchNorm1d(256)
-        self.drop2 = nn.Dropout(0.4)
-        self.fc3 = nn.Linear(256, output_dim)
-
-    def forward(self, x):
-        """
-        x: [B, N, 3]
-        """
-        xyz, f1 = self.sa1(x)
-        xyz, f2 = self.sa2(xyz)
-        xyz, f3 = self.sa3(xyz)
-        x = F.relu(self.bn1(self.fc1(f3)))
-        x = self.drop1(x)
-        x = F.relu(self.bn2(self.fc2(x)))
-        x = self.drop2(x)
-        x = self.fc3(x)
-        return x
-
-
-class VisionPointNet(nn.Module):
-
-    def __init__(self, patch_shape, output_dim=128):
-        super().__init__()
-
-        self.point_net = PointNet2(output_dim=output_dim)
-        self.vision_model = CNN3D(
-            patch_shape,
-            n_conv_layers=6,
-            n_feat_channels=20,
-            output_dim=output_dim,
-            use_double_conv=True,
-        )
-        self.output = FeedForwardNet(2 * output_dim, 1, 3)
-
-    def forward(self, x):
-        """
-        Passes the given input through this neural network.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Input vector of features.
-
-        Returns
-        -------
-        x : torch.Tensor
-            Output of the neural network.
-        """
-        x_img = self.vision_model(x["img"])
-        x_pc = self.point_net(x["point_cloud"])
-        x = torch.cat((x_img, x_pc), dim=1)
-        x = self.output(x)
-        return x
-
-
-class PointNetSetAbstraction(nn.Module):
-    """
-    PointNet++ Set Abstraction (SA) layer
-    """
-
-    def __init__(self, n_points, mlp_channels):
-        super().__init__()
-        self.n_points = n_points
-        layers = []
-        last_channel = 3
-        for out_ch in mlp_channels:
-            layers.append(nn.Conv1d(last_channel, out_ch, 1))
-            layers.append(nn.BatchNorm1d(out_ch))
-            layers.append(nn.ReLU())
-            last_channel = out_ch
-        self.mlp = nn.Sequential(*layers)
-
-    def forward(self, xyz):
-        """
-        xyz: [B, N, 3]
-        """
-        B, N, _ = xyz.shape
-        if self.n_points is not None:
-            idx = farthest_point_sample(xyz, self.n_points)
-            new_xyz = index_points(xyz, idx)
-        else:
-            new_xyz = xyz
-        # MLP expects [B, C, N]
-        features = self.mlp(new_xyz.transpose(1, 2))
-        features = torch.max(features, 2)[0]  # global pooling
-        return new_xyz, features
-
-
-class DGCNN(nn.Module):
-
-    def __init__(self, input_dim=3, output_dim=64, k=16):
-        super().__init__()
-        self.k = k  # number of neighbors
-
-        # BatchNorm layers
-        self.bn1 = nn.BatchNorm2d(32)
-        self.bn2 = nn.BatchNorm2d(32)
-        self.bn3 = nn.BatchNorm2d(64)
-
-        # Conv layers (Conv2d on edge features)
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(input_dim * 2, 32, kernel_size=1, bias=False),
-            self.bn1,
-            nn.LeakyReLU(0.2),
-        )
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(32 * 2, 32, kernel_size=1, bias=False),
-            self.bn2,
-            nn.LeakyReLU(0.2),
-        )
-        self.conv3 = nn.Sequential(
-            nn.Conv2d(32 * 2, 64, kernel_size=1, bias=False),
-            self.bn3,
-            nn.LeakyReLU(0.2),
-        )
-
-        # Final fully-connected layer for embedding
-        self.fc = nn.Linear(128, output_dim)
-
-    def get_graph_feature(self, x, k):
-        # x: (B, C, N)
-        # Compute pairwise distance
-        B, C, N = x.size()
-        x = x.transpose(2, 1)  # (B, N, C)
-        idx = self.knn(x, k)  # (B, N, k)
-
-        # Gather neighbors
-        neighbors = self.index_points(x, idx)  # (B, N, k, C)
-        x = x.unsqueeze(2).expand_as(neighbors)  # (B, N, k, C)
-        edge_feature = torch.cat((x, neighbors - x), dim=-1)  # (B, N, k, 2C)
-        return edge_feature.permute(0, 3, 1, 2).contiguous()  # (B, 2C, N, k)
-
-    @staticmethod
-    def knn(x, k):
-        # Compute pairwise distance
-        inner = -2 * torch.matmul(x, x.transpose(2, 1))
-        xx = torch.sum(x**2, dim=2, keepdim=True)
-        pairwise_distance = -xx - inner - xx.transpose(2, 1)
-        _, idx = pairwise_distance.topk(k=k, dim=-1)  # (B, N, k)
-        return idx
-
-    @staticmethod
-    def index_points(x, idx):
-        # x: (B, N, C), idx: (B, N, k)
-        B = x.size(0)
-        batch_indices = (
-            torch.arange(B, device=x.device)
-            .view(B, 1, 1)
-            .repeat(1, idx.size(1), idx.size(2))
-        )
-        return x[batch_indices, idx, :]
-
-    def forward(self, x):
-        """
-        Passes the given input through this neural network.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Input vector of features.
-
-        Returns
-        -------
-        x : torch.Tensor
-            Output of the neural network.
-        """
-        x1 = F.relu(self.conv1(self.get_graph_feature(x, self.k)))
-        x1 = x1.max(dim=-1)[0]
-
-        x2 = F.relu(self.conv2(self.get_graph_feature(x1, self.k)))
-        x2 = x2.max(dim=-1)[0]
-
-        x3 = F.relu(self.conv3(self.get_graph_feature(x2, self.k)))
-        x3 = x3.max(dim=-1)[0]
-
-        x_cat = torch.cat((x1, x2, x3), dim=1)
-        x_global = x_cat.max(dim=-1)[0]
-        return self.fc(x_global)
-
-
-class VisionDGCNN(nn.Module):
-
-    def __init__(self, patch_shape, output_dim=128):
-        super().__init__()
-
-        self.dgcnn = DGCNN(output_dim=output_dim)
-        self.vision_model = CNN3D(
-            patch_shape,
-            n_conv_layers=6,
-            n_feat_channels=20,
-            output_dim=output_dim,
-            use_double_conv=True,
-        )
-        self.output = FeedForwardNet(2 * output_dim, 1, 3)
-
-    def forward(self, x):
-        """
-        Passes the given input through this neural network.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Input vector of features.
-
-        Returns
-        -------
-        x : torch.Tensor
-            Output of the neural network.
-        """
-        x_img = self.vision_model(x["img"])
-        x_pc = self.dgcnn(x["point_cloud"])
-        x = torch.cat((x_img, x_pc), dim=1)
-        x = self.output(x)
-        return x
-
-
-# --- Point Cloud Generation ---
-def farthest_point_sample(xyz, n_points):
-    """
-    Farthest Point Sampling (FPS)
-    xyz: [B, N, 3]
-    return: [B, n_points] indices
-    """
-    B, N, _ = xyz.shape
-    centroids = torch.zeros(B, n_points, dtype=torch.long, device=xyz.device)
-    distance = torch.ones(B, N, device=xyz.device) * 1e10
-    farthest = torch.randint(0, N, (B,), dtype=torch.long, device=xyz.device)
-    batch_indices = torch.arange(B, dtype=torch.long, device=xyz.device)
-    for i in range(n_points):
-        centroids[:, i] = farthest
-        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
-        dist = torch.sum((xyz - centroid) ** 2, -1)
-        mask = dist < distance
-        distance[mask] = dist[mask]
-        farthest = torch.max(distance, -1)[1]
-    return centroids
-
-
-def index_points(points, idx):
-    """
-    Index points from xyz/features
-    points: [B, N, C]
-    idx: [B, S]
-    return: [B, S, C]
-    """
-    B = points.shape[0]
-    S = idx.shape[1]
-    batch_indices = (
-        torch.arange(B, dtype=torch.long, device=points.device)
-        .view(B, 1)
-        .repeat(1, S)
-    )
-    return points[batch_indices, idx, :]
-
-
-def subgraph_to_point_cloud(graph, n_points=3600):
-    point_cloud = list()
-    for n1, n2 in graph.edges:
-        # Use average radius
-        r1 = graph.node_radius[n1]
-        r2 = graph.node_radius[n2]
-        r = (r1 + r2) / 2
-
-        pts = sample_cylinder_between_points(
-            graph.node_xyz[n1], graph.node_xyz[n2], r
-        )
-        point_cloud.append(pts)
-    point_cloud = np.vstack(point_cloud)
-
-    total_points = point_cloud.shape[0]
-    idxs = np.arange(len(point_cloud))
-    if total_points >= n_points:
-        # Downsample randomly
-        sampled_idxs = np.random.choice(idxs, n_points, replace=False)
-        point_cloud = point_cloud[sampled_idxs]
-    else:
-        # Upsample with replacement
-        sampled_idxs = np.random.choice(idxs, n_points, replace=True)
-        point_cloud = point_cloud[sampled_idxs]
-    return point_cloud.T
-
-
-def sample_cylinder_between_points(p1, p2, r, n_samples=25):
-    p1 = np.array(p1)
-    p2 = np.array(p2)
-    axis = p2 - p1
-    axis_length = np.linalg.norm(axis)
-    axis_dir = axis / axis_length
-
-    # Build orthogonal basis
-    if np.allclose(axis_dir, [0, 0, 1]):
-        ortho1 = np.array([1, 0, 0])
-    else:
-        ortho1 = np.cross(axis_dir, [0, 0, 1])
-        ortho1 /= np.linalg.norm(ortho1)
-    ortho2 = np.cross(axis_dir, ortho1)
-
-    # Random samples along axis
-    t_values = np.random.rand(n_samples)
-    centers = p1 + np.outer(t_values, axis)
-
-    # Random samples in circular cross-section
-    angles = np.random.rand(n_samples) * 2 * np.pi
-    radii = r * np.sqrt(np.random.rand(n_samples))
-    offsets = np.outer(radii * np.cos(angles), ortho1) + np.outer(
-        radii * np.sin(angles), ortho2
-    )
-    points = centers + offsets
-    return points
diff --git a/src/neuron_proofreader/machine_learning/vision_models.py b/src/neuron_proofreader/machine_learning/vision_models.py
deleted file mode 100644
index d514df69..00000000
--- a/src/neuron_proofreader/machine_learning/vision_models.py
+++ /dev/null
@@ -1,490 +0,0 @@
-"""
-Created on Sat July 15 12:00:00 2025
-
-@author: Anna Grim
-@email: anna.grim@alleninstitute.org
-
-Code for vision models that perform image classification tasks within
-NeuronProofreader pipelines.
-
-"""
-
-from einops import rearrange
-
-import torch
-import torch.nn as nn
-
-from neuron_proofreader.utils.ml_util import FeedForwardNet, init_mlp
-
-
-# --- CNNs ---
-class CNN3D(nn.Module):
-    """
-    Convolutional neural network for 3D images.
-    """
-
-    def __init__(
-        self,
-        patch_shape,
-        output_dim=1,
-        dropout=0.1,
-        n_conv_layers=5,
-        n_feat_channels=16,
-        use_double_conv=True,
-    ):
-        """
-        Instantiates a CNN3D object.
-
-        Parameters
-        ----------
-        patch_shape : Tuple[int]
-            Shape of input image patch.
-        output_dim : int, optional
-            Dimension of output. Default is 1.
-        dropout : float, optional
-            Fraction of values to randomly drop during training. Default is
-            0.1.
-        n_conv_layers : int, optional
-            Number of convolutional layers. Default is 5.
-        use_double_conv : bool, optional
-            Indication of whether to use double convolution. Default is True.
-        """
-        # Call parent class
-        nn.Module.__init__(self)
-
-        # Class attributes
-        self.dropout = dropout
-        self.patch_shape = patch_shape
-
-        # Convolutional layers
-        self.conv_layers = init_cnn3d(
-            2, n_feat_channels, n_conv_layers, use_double_conv=use_double_conv
-        )
-
-        # Output layer
-        flat_size = self._get_flattened_size()
-        self.output = FeedForwardNet(flat_size, output_dim, 3)
-
-        # Initialize weights
-        self.apply(self.init_weights)
-
-    def _get_flattened_size(self):
-        """
-        Compute the flattened feature vector size after applying a sequence
-        of convolutional and pooling layers on an input tensor with the given
-        shape.
-
-        Returns
-        -------
-        int
-            Length of the flattened feature vector after the convolutions and
-            pooling.
-        """
-        with torch.no_grad():
-            x = torch.zeros(1, 2, *self.patch_shape)
-            x = self.conv_layers(x)
-            return x.view(1, -1).size(1)
-
-    @staticmethod
-    def init_weights(m):
-        """
-        Initializes the weights and biases of a given PyTorch layer.
-
-        Parameters
-        ----------
-        m : nn.Module
-            PyTorch layer or module.
-        """
-        if isinstance(m, nn.Conv3d):
-            nn.init.kaiming_normal_(
-                m.weight, mode="fan_in", nonlinearity="leaky_relu"
-            )
-            if m.bias is not None:
-                nn.init.constant_(m.bias, 0)
-        elif isinstance(m, nn.Linear):
-            nn.init.xavier_normal_(m.weight)
-            if m.bias is not None:
-                nn.init.constant_(m.bias, 0)
-        elif isinstance(m, nn.BatchNorm3d):
-            nn.init.constant_(m.weight, 1)
-            nn.init.constant_(m.bias, 0)
-
-    def forward(self, x):
-        """
-        Passes the given input through this neural network.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Input vector of features.
-
-        Returns
-        -------
-        x : torch.Tensor
-            Output of the neural network.
-        """
-        x = self.conv_layers(x)
-        x = x.view(x.size(0), -1)
-        x = self.output(x)
-        return x
-
-
-# --- Transformers ---
-class ViT3D(nn.Module):
-    """
-    A class that implements a 3D Vision transformer.
-
-    Attributes
-    ----------
-    """
-
-    def __init__(
-        self,
-        img_shape=(128, 128, 128),
-        emb_dim=512,
-        depth=6,
-        heads=8,
-        mlp_dim=1024,
-        output_dim=1,
-    ):
-        """
-        Instantiates a ViT3D object.
-
-        Parameters
-        ----------
-        img_shape : Tuple[int], optional
-            Shape of the input image. Default is (128, 128, 128).
-        emb_dim : int, optional
-            Dimension of the embedding space. Default is 512.
-        depth : int, optional
-            Number of transformer blocks. Default is 6.
-        heads : int, optional
-            Number of attention heads in each transformer block. Default 8.
-        mlp_dim : int, optional
-            Dimension of MLP embedding space. Default is 1024.
-        output_dim : int, optional
-            Dimension of output. Default is 1.
-        """
-        # Call parent class
-        super().__init__()
-
-        # Token embedding
-        self.cls_token = nn.Parameter(torch.empty(1, 1, emb_dim))
-        self.img_tokenizer = ImageTokenizer3D(emb_dim, img_shape)
-
-        # Position embedding
-        n_tokens = self.img_tokenizer.count_tokens() + 1
-        self.pos_embedding = nn.Parameter(torch.empty(1, n_tokens, emb_dim))
-        print("# Tokens:", n_tokens)
-
-        # Transformer Blocks
-        self.transformer = nn.Sequential(
-            *[
-                TransformerEncoderBlock(emb_dim, heads, mlp_dim)
-                for _ in range(depth)
-            ]
-        )
-        self.norm = nn.LayerNorm(emb_dim)
-
-        # Output layer
-        self.output = FeedForwardNet(emb_dim, output_dim, 2)
-
-        # Initialize weights
-        self._init_weights()
-
-    def _init_weights(self):
-        """
-        Initializes the model's weights.
-        """
-        # Initialize token embedding
-        nn.init.trunc_normal_(self.cls_token, std=0.02)
-        nn.init.trunc_normal_(self.pos_embedding, std=0.02)
-
-        # Initialize Transformer and output layers
-        for module in self.modules():
-            if isinstance(module, nn.Linear):
-                nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.zeros_(module.bias)
-            elif isinstance(module, nn.LayerNorm):
-                nn.init.ones_(module.weight)
-                nn.init.zeros_(module.bias)
-
-    def forward(self, x):
-        """
-        Passes the given input through this neural network.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Input vector of features.
-
-        Returns
-        -------
-        x : torch.Tensor
-            Output of the neural network.
-        """
-        # Tokenize input -> (b, n_tokens, emb_dim)
-        x = self.img_tokenizer(x)
-        cls_tokens = self.cls_token.expand(x.size(0), -1, -1)
-        x = torch.cat((cls_tokens, x), dim=1)
-
-        # Transformer
-        x = x + self.pos_embedding[:, : x.size(1)]
-        x = self.transformer(x)
-        x = self.norm(x[:, 0])
-
-        # Output layer
-        x = self.output(x)
-        return x
-
-
-class ImageTokenizer3D(nn.Module):
-    """
-    A class for learning image token embeddings for transformer-based
-    architectures.
-
-    Attributes
-    ----------
-    dropout : nn.Dropout
-        Dropout layer.
-    emb_dim : int
-        Dimension of the embedding space.
-    img_shape : Tuple[int]
-        Shape of the input image (D, H, W).
-    pos_embedding : nn.Parameter
-        Learnable position encoding.
-    proj : nn.Conv3d
-        Convolutional layer that generates a learnable projection of the
-        tokens.
-    """
-
-    def __init__(
-        self,
-        emb_dim,
-        img_shape,
-        dropout=0.05,
-        n_cnn_layers=3,
-        n_cnn_channels=32,
-    ):
-        """
-        Instantiates a ImageTokenizer3D object.
-
-        Parameters
-        ----------
-        emb_dim : int
-            Dimension of the embedding space.
-        img_shape : Tuple[int]
-            Shape of the input image (D, H, W).
-        dropout : float, optional
-            Dropout probability applied after adding positional embeddings.
-            Default is 0.05.
-        n_cnn_layers : int, optional
-            Number of layers in the CNN that generates the initial token
-            embedding. Default is 3.
-        """
-        # Call parent class
-        super().__init__()
-
-        # Class attributes
-        self.emb_dim = emb_dim
-        self.img_shape = img_shape
-
-        # Image embedding
-        cnn_out_channels = n_cnn_channels * (2 ** (n_cnn_layers - 1))
-        self.proj = nn.Conv3d(cnn_out_channels, emb_dim, kernel_size=1)
-        self.tokenizer = init_cnn3d(
-            2, n_cnn_channels, n_cnn_layers, use_double_conv=False
-        )
-
-        # Positional embedding
-        n_tokens = self.count_tokens()
-        self.pos_embedding = nn.Parameter(torch.randn(1, n_tokens, emb_dim))
-        self.dropout = nn.Dropout(p=dropout)
-
-    def count_tokens(self):
-        """
-        Counts the number of tokens that are generated given the patch shape,
-        CCN3D architecture, and embedding dimension.
-
-        Returns
-        -------
-        int
-            Number of tokens generated by tokenizer.
-        """
-        with torch.no_grad():
-            dummy = torch.zeros(1, 2, *self.img_shape)
-            feats = self.tokenizer(dummy)
-            feats = self.proj(feats)
-            return feats.flatten(2).shape[-1]
-
-    def forward(self, x):
-        """
-        Forward pass that converts an input image into a sequence of token
-        embeddings.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Input tensor of shape (B, C, D, H, W).
-
-        Returns
-        -------
-        torch.Tensor
-            Token embeddings of shape (B, N, E).
-        """
-        # Embed image
-        x = self.tokenizer(x)
-        x = self.proj(x)
-
-        # Tokenize
-        x = rearrange(x, "b c d h w -> b (d h w) c")
-        x = x + self.pos_embedding
-        return self.dropout(x)
-
-
-class TransformerEncoderBlock(nn.Module):
-    """
-    Single transformer encoder block.
-
-    Attributes
-    ----------
-    attn : nn.MultiheadAttention
-        Multihead attention block.
-    dropout : nn.Dropout
-        Dropout layer.
-    mlp : nn.Sequential
-        Multi-layer perceptron.
-    norm1 : nn.LayerNorm
-        Applies layer normalization over a mini-batch of the inputs.
-    norm2 : nn.LayerNorm
-        Applied layer normalization over a mini-batch of the outputs of the
-        multihead attention block.
-    """
-
-    def __init__(self, emb_dim, heads, mlp_dim, dropout=0):
-        """
-        Instantiates a TransformerEncoderBlock object.
-
-        Parameters
-        ----------
-        emb_dim : int
-            Dimension of the embedding space.
-        heads : int
-            Number of attention heads.
-        mlp_dim : int
-            Dimensionality of the hidden layer in the MLP.
-        dropout : float, optional
-            Dropout probability applied after attention and MLP layers.
-            Default is 0.
-        """
-        # Call parent class
-        super().__init__()
-
-        # Attention head
-        self.norm1 = nn.LayerNorm(emb_dim)
-        self.attn = nn.MultiheadAttention(
-            emb_dim, heads, dropout=dropout, batch_first=True
-        )
-        self.norm2 = nn.LayerNorm(emb_dim)
-        self.mlp = init_mlp(emb_dim, mlp_dim, emb_dim, dropout)
-        self.dropout = nn.Dropout(p=dropout)
-
-    def forward(self, x):
-        """
-        Forward pass of the encoder block.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Input tensor of shape (B, N, E).
-
-        Returns
-        -------
-        x : torch.Tensor
-            Output tensor of shape (B, N, E), same as input.
-        """
-        x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
-        x = x + self.mlp(self.norm2(x))
-        x = self.dropout(x)
-        return x
-
-
-# --- Build Simple Neural Networks ---
-def init_cnn3d(in_channels, n_feat_channels, n_layers, use_double_conv=True):
-    """
-    Initializes a convolutional neural network.
-
-    Parameters
-    ----------
-    in_channels : int
-        Number of channels that are input to this convolutional layer.
-    out_channels : int
-        Number of channels that are output from this convolutional layer.
-    n_layers : int
-        Number of layers in the network.
-    use_double_conv : bool, optional
-        Indication of whether to use double convolution. Default is True.
-
-    Returns
-    -------
-    layers : torch.nn.Sequential
-        Sequence of operations that define the network.
-    """
-    layers = list()
-    in_channels = in_channels
-    out_channels = n_feat_channels
-    for i in range(n_layers):
-        # Build layer
-        layers.append(
-            init_conv_layer(in_channels, out_channels, 3, use_double_conv)
-        )
-
-        # Update channel sizes
-        in_channels = out_channels
-        out_channels *= 2
-    return nn.Sequential(*layers)
-
-
-def init_conv_layer(in_channels, out_channels, kernel_size, use_double_conv):
-    """
-    Initializes a convolutional layer.
-
-    Parameters
-    ----------
-    in_channels : int
-        Number of channels that are input to this convolutional layer.
-    out_channels : int
-        Number of channels that are output from this convolutional layer.
-    kernel_size : int
-        Size of kernel used on convolutional layers.
-    use_double_conv : bool
-        Indication of whether to use double convolution.
-
-    Returns
-    -------
-    layers : torch.nn.Sequential
-        Sequence of operations that define this convolutional layer.
-    """
-    # Convolution
-    layers = [
-        nn.Conv3d(
-            in_channels, out_channels, kernel_size, padding=kernel_size // 2
-        ),
-        nn.BatchNorm3d(out_channels),
-        nn.GELU(),
-    ]
-    if use_double_conv:
-        layers += [
-            nn.Conv3d(
-                out_channels,
-                out_channels,
-                kernel_size,
-                padding=kernel_size // 2,
-            ),
-            nn.BatchNorm3d(out_channels),
-            nn.GELU(),
-        ]
-    # Pooling
-    layers.append(nn.MaxPool3d(kernel_size=2))
-    return nn.Sequential(*layers)