Implement base Model and Head classes

.gitignore

@@ -143,6 +143,9 @@ venv.bak/
 .dmypy.json
 dmypy.json
 
+# ruff
+.ruff_cache/
+
 # Pyre type checker
 .pyre/
 

sleap_nn/architectures/common.py

@@ -101,58 +101,3 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             ceil_mode=self.ceil_mode,
             return_indices=self.return_indices,
         )
-
-
-def get_act_fn(activation: str) -> nn.Module:
-    """Get an instance of an activation function module based on the provided name.
-
-    This function returns an instance of a PyTorch activation function module
-    corresponding to the given activation function name.
-
-    Args:
-        activation (str): Name of the activation function. Supported values are 'relu', 'sigmoid', and 'tanh'.
-
-    Returns:
-        nn.Module: An instance of the requested activation function module.
-
-    Raises:
-        KeyError: If the provided activation function name is not one of the supported values.
-
-    Example:
-        # Get an instance of the ReLU activation function
-        relu_fn = get_act_fn('relu')
-
-        # Apply the activation function to an input tensor
-        input_tensor = torch.randn(1, 64, 64)
-        output = relu_fn(input_tensor)
-    """
-    activations = {"relu": nn.ReLU(), "sigmoid": nn.Sigmoid(), "tanh": nn.Tanh()}
-
-    if activation not in activations:
-        raise KeyError(
-            f"Unsupported activation function: {activation}. Supported activations are: {', '.join(activations.keys())}"
-        )
-
-    return activations[activation]
-
-
-def get_children_layers(model: torch.nn.Module) -> List[nn.Module]:
-    """Recursively retrieves a flattened list of all children modules and submodules within the given model.
-
-    Args:
-        model: The PyTorch model to extract children from.
-
-    Returns:
-        list of nn.Module: A flattened list containing all children modules and submodules.
-    """
-    children = list(model.children())
-    flattened_children = []
-    if children == []:
-        return model
-    else:
-        for child in children:
-            try:
-                flattened_children.extend(get_children_layers(child))
-            except TypeError:
-                flattened_children.append(get_children_layers(child))
-    return flattened_children

sleap_nn/architectures/encoder_decoder.py

@@ -28,9 +28,11 @@
 from typing import List, Text, Tuple, Union
 
 import torch
-from sleap_nn.architectures.common import MaxPool2dWithSamePadding, get_act_fn
 from torch import nn
 
+from sleap_nn.architectures.common import MaxPool2dWithSamePadding
+from sleap_nn.architectures.utils import get_act_fn
+
 
 class SimpleConvBlock(nn.Module):
     """A simple convolutional block module.
@@ -428,6 +430,8 @@ def __init__(
         self.convs_per_block = convs_per_block
         self.kernel_size = kernel_size
 
+        self.current_strides = []
+
         self.decoder_stack = nn.ModuleList([])
         for block in range(up_blocks):
             prev_block_filters_in = -1 if block == 0 else block_filters_in
@@ -454,19 +458,25 @@ def __init__(
                 )
             )
 
+            self.current_strides.append(current_stride)
             current_stride = next_stride
 
-    def forward(self, x: torch.Tensor, features: List[torch.Tensor]) -> torch.Tensor:
+    def forward(
+        self, x: torch.Tensor, features: List[torch.Tensor]
+    ) -> Tuple[List[torch.Tensor], List]:
         """Forward pass through the Decoder module.
 
         Args:
             x: Input tensor for the decoder.
             features: List of feature tensors from different encoder levels.
 
         Returns:
-            torch.Tensor: Output tensor after applying the decoder operations.
+            outputs: List of output tensors after applying the decoder operations.
+            current_strides: the current strides from the decoder blocks.
         """
+        outputs = []
         for i in range(len(self.decoder_stack)):
             x = self.decoder_stack[i](x, features[i])
+            outputs.append(x)
 
-        return x
+        return outputs, self.current_strides