emmi_inference.models.abupt¶

Classes¶

AnchoredBranchedUPT

Base class for all neural network modules.

Functions¶

main()

Module Contents¶

class emmi_inference.models.abupt.AnchoredBranchedUPT(ndim=3, input_dim=3, output_dim_surface=4, output_dim_volume=7, dim=192, geometry_depth=1, num_heads=3, blocks='pscscs', num_volume_blocks=6, num_surface_blocks=6, radius=0.25, **kwargs)¶

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing to nest them in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will have their parameters converted too when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

Parameters:

ndim (int)
input_dim (int)
output_dim_surface (int)
output_dim_volume (int)
dim (int)
geometry_depth (int)
num_heads (int)
blocks (str)
num_volume_blocks (int)
num_surface_blocks (int)
radius (float)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

rope¶

encoder¶

geometry_blocks¶

pos_embed¶

surface_bias¶

volume_bias¶

blocks¶

surface_blocks¶

volume_blocks¶

surface_decoder¶

volume_decoder¶

forward(geometry_position, geometry_supernode_idx, geometry_batch_idx, surface_anchor_position, volume_anchor_position, surface_query_position=None, volume_query_position=None)¶

Parameters:

geometry_position (torch.Tensor)
geometry_supernode_idx (torch.Tensor)
geometry_batch_idx (torch.Tensor | None)
surface_anchor_position (torch.Tensor)
volume_anchor_position (torch.Tensor)
surface_query_position (torch.Tensor | None)
volume_query_position (torch.Tensor | None)

Return type:

dict[str, torch.Tensor]

emmi_inference.models.abupt.main()¶