Implement base Model and Head classes

[alckasoc]

Summary by CodeRabbit

• New Feature: Introduced a variety of model output head classes, enhancing the flexibility in defining the output structure of a neural network model.
• New Feature: Added a Model module that serves as a higher-level wrapper around nn.Module, simplifying the process of defining and using a trainable model.
• Improvement: Updated the UNet class with additional attributes and a new output_channels property for better control over the output channels.
• Improvement: Modified the InstanceCentroidFinder and KorniaAugmenter instantiations for more precise data augmentation.
• Test: Added new test cases for the Model class and updated existing tests for the UNet class, ensuring the reliability of these modules.
• Refactor: Restructured the augmentation_config dictionary in the test function for better organization of intensity and geometric augmentations.

[coderabbitai]

Rate Limit Exceeded

You have exceeded the limit for the number of files or commits that can be reviewed per hour. Please wait 29 minutes and 59 seconds before requesting another review.

To re-trigger a review, convert this PR to a draft and then mark it as ready for review again. Alternatively, you can push a new commit to this PR.

[coderabbitai]

Review Status

Actionable comments generated: 2

Commits

Files that changed from the base of the PR and between 537c14b and dce7294.

Files selected for processing (1)

• sleap_nn/architectures/heads.py (1 hunks)

Additional comments (Suppressed): 1

sleap_nn/architectures/heads.py (1)

• 498-512: The make_head method in the ClassVectorsHead class uses a softmax activation function at the end. This is fine if you are dealing with exclusive classes, but if your problem allows for multi-label classification (an instance can belong to more than one class), then you should use a sigmoid activation function instead. Please verify whether the classes are mutually exclusive or not.

Also, there seems to be no dropout layers added in between the fully connected layers. Dropout layers can help prevent overfitting by randomly setting a fraction rate of input units to 0 at each update during training time. Consider adding dropout layers if overfitting is a concern.

for i in range(self.num_fc_layers):
    if i == 0:
        module_list.append(nn.Linear(x_in, self.num_fc_units))
    else:
        module_list.append(nn.Linear(self.num_fc_units, self.num_fc_units))
    module_list.append(get_act_fn("relu"))
+   module_list.append(nn.Dropout(0.5))  # Add dropout layer with dropout rate of 0.5

[coderabbitai]

Review Status

Actionable comments generated: 2

Commits

Files that changed from the base of the PR and between dce7294 and 99413a1.

Files selected for processing (6)

• .gitignore (1 hunks)
• sleap_nn/architectures/encoder_decoder.py (4 hunks)
• sleap_nn/architectures/model.py (1 hunks)
• sleap_nn/architectures/unet.py (3 hunks)
• tests/architectures/test_model.py (1 hunks)
• tests/architectures/test_unet.py (3 hunks)

Files skipped from review due to trivial changes (2)

• .gitignore
• sleap_nn/architectures/model.py

Additional comments (Suppressed): 11

sleap_nn/architectures/unet.py (3)

• 3-11: The import of math and typing modules has been removed. Ensure that this does not affect any other parts of the code that might be using these imports.

• 43-52: The initialization of several attributes (in_channels, kernel_size, filters, filters_rate, down_blocks, up_blocks, convs_per_block) has been added to the UNet class. This is a good practice as it makes the code more readable and maintainable by clearly defining what attributes are expected in an instance of the class.

• 82-93: A new attribute current_strides and a new property output_channels have been added to the UNet class. The current_strides attribute seems to be derived from the Decoder instance, while the output_channels property is calculated based on the filters, filters_rate, down_blocks, and up_blocks attributes. Make sure that these new additions are used appropriately throughout the rest of the codebase.

sleap_nn/architectures/encoder_decoder.py (4)

• 28-38: The import statement for get_act_fn has been moved from sleap_nn.architectures.common to sleap_nn.architectures.utils. Ensure that this function exists in the new module and all references to it have been updated accordingly.

• 430-437: The attribute self.current_strides has been added. This seems to be a list that will store the stride values at each level of the decoder. Make sure that this attribute is used correctly throughout the codebase.

• 458-463: The current stride value is now being stored in self.current_strides after each iteration. This change appears to be part of the addition of the current_strides attribute. Ensure that this list is updated correctly and used as intended.

• 469-479: The return type of the forward method has changed from a single tensor to a list of tensors. This could potentially break any code that expects the output of this method to be a single tensor. Verify that all calls to this method have been updated to handle the new return type.

tests/architectures/test_unet.py (4)

• 6-6: The import path for get_children_layers has been changed. Ensure that the function is still correctly imported from the new location.

- from sleap_nn.architectures.common import get_children_layers
+ from sleap_nn.architectures.utils import get_children_layers

• 59-60: The variable name for the UNet instance has been changed from model to unet. Make sure this change is consistent throughout the codebase and does not affect any other parts of the code that reference this variable.

• 65-66: The output of the UNet model is now a list, and the final output shape is checked on the last element of this list. This is a significant change from the previous version where the output was directly a tensor. Ensure that all parts of the code that use the output of the UNet model are updated to handle this change.

-    assert y.shape == (1, 13, 192, 192)
+    assert type(y) is list
+    assert y[-1].shape == (1, 64, 192, 192)

• 68-75: A new Conv2d layer is added after the UNet model, and its output shape is checked. This seems to be a new addition to the test case. If this Conv2d layer is part of the actual model architecture, ensure that it is properly integrated into the model and not just in the test case.

[alckasoc]

99413a1

@talmo I set up a hacky way (above) to select the correct head based on output stride. I was thinking of creating some IntermediateFeature-like class, but that might not be necessary...? Maybe I could just return a dictionary with the output features and the current stride in the forward method for the decoder. What's the ideal way to do this?

[coderabbitai]

Review Status

Actionable comments generated: 2

Commits

Files that changed from the base of the PR and between 99413a1 and 2880849.

Files selected for processing (4)

• sleap_nn/architectures/encoder_decoder.py (3 hunks)
• sleap_nn/architectures/model.py (1 hunks)
• sleap_nn/architectures/unet.py (3 hunks)
• tests/architectures/test_unet.py (3 hunks)

Additional comments (Suppressed): 12

sleap_nn/architectures/model.py (2)

• 1-6: The introduction of a higher-level wrapper around nn.Module is a good design choice. It allows for easy querying of the model configuration without instantiating the model itself, which can be beneficial for debugging and understanding the model structure.

• 55-97: Similar to get_backbone, the get_head function provides a clean way to instantiate different head models based on a string identifier. This function supports a variety of head types, which is good for flexibility and extensibility.

sleap_nn/architectures/encoder_decoder.py (3)

• 28-38: The import statement for get_act_fn has been moved from sleap_nn.architectures.common to sleap_nn.architectures.utils. Ensure that the function is correctly implemented in the new module and all references to it have been updated.

• 430-437: The addition of self.current_strides = [] suggests that you're now tracking the strides at each block level. This could be useful for debugging or understanding the model's structure, but ensure it doesn't introduce unnecessary overhead if not used elsewhere.

• 458-482: The forward method signature and return type have changed. It now returns a tuple containing a list of output tensors and the current strides. Make sure all calls to this method are updated to handle the new return type. Also, verify that the change in the return type does not break any functionality where the output of this method is used.

-    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]:

tests/architectures/test_unet.py (3)

• 1-6: The import statement for get_children_layers has been updated to reflect its new location in the sleap_nn.architectures.utils module. Ensure that this function has been correctly moved and that all other references to it have been updated accordingly.

• 56-78: The test for the final output shape of the UNet model has been modified. Previously, the test was checking if the output shape of the model was (1, 13, 192, 192). Now, it checks if the output is a list, and if the shape of the last element of the list is (1, 64, 192, 192). It also checks if current_strides is a list and its length is 4. Then, it applies a convolutional layer with kernel size 1 to the last output of the UNet model and checks if the output shape is (1, 13, 192, 192). Make sure these changes are consistent with the modifications made to the UNet model.

• 88-94: The assignment of the evaluation mode to the encoder instance has been changed from _ = enc.eval() to enc.eval(). This change is fine as the return value of eval() (which is the module itself) is not used anywhere else in the code.

sleap_nn/architectures/unet.py (4)

• 3-12: The import of math and torch.nn.functional as F has been removed, and the Tuple type hint has been added. Ensure that these changes do not affect other parts of the code that may rely on these imports.

• 44-54: The initialization of several attributes (in_channels, kernel_size, filters, filters_rate, down_blocks, up_blocks, convs_per_block) has been added to the UNet class. Make sure these are used appropriately in the rest of the class methods.

• 83-92: A new property output_channels has been added to calculate and return the output channels of the UNet. This seems to be a useful addition for understanding the structure of the network.

• 94-106: The forward method now returns a tuple containing the output tensor and a list of current strides from the decoder, instead of just the output tensor. Ensure all calls to this method have been updated to handle the new return type.

-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor) -> Tuple[List[torch.Tensor], List]:

[coderabbitai]

Review Status

Actionable comments generated: 4

Commits

Files that changed from the base of the PR and between 2880849 and 9afdebb.

Files selected for processing (1)

• sleap_nn/architectures/model.py (1 hunks)

[coderabbitai]

Review Status

Actionable comments generated: 0

Commits

Files that changed from the base of the PR and between 9afdebb and 5cd8801.

Files selected for processing (1)

• sleap_nn/architectures/model.py (1 hunks)

Additional comments (Suppressed): 4

sleap_nn/architectures/model.py (4)

• 1-6: The introduction of a higher-level wrapper around nn.Module is a good design choice. It provides an abstraction layer that simplifies the process of defining and configuring models.

• 27-52: The get_backbone function is well implemented. It checks if the provided backbone name is supported, raises an error if not, and returns an instance of the requested backbone model. However, consider adding more backbones in the future to provide more flexibility.

• 55-97: The get_head function follows the same pattern as get_backbone, which maintains consistency. It's also extensible for additional head types in the future.

• 100-145: The Model class is well structured. The __init__ method initializes the backbone and heads based on the provided configurations. The from_config class method allows creating a model from a configuration dictionary, which can be useful when loading models from saved configurations. The forward method correctly performs a forward pass through the model, applying each head to the appropriate feature map from the backbone.

[coderabbitai]

Review Status

Actionable comments generated: 0

Commits

Files that changed from the base of the PR and between 5cd8801 and 4eaf268.

Files selected for processing (1)

• sleap_nn/data/instance_cropping.py (3 hunks)

Files skipped from review due to trivial changes (1)

• sleap_nn/data/instance_cropping.py

[coderabbitai]

Review Status

Actionable comments generated: 0

Commits

Files that changed from the base of the PR and between 4eaf268 and 702009a.

Files selected for processing (4)

• sleap_nn/architectures/encoder_decoder.py (3 hunks)
• sleap_nn/architectures/model.py (1 hunks)
• sleap_nn/architectures/unet.py (3 hunks)
• tests/architectures/test_unet.py (3 hunks)

Additional comments (Suppressed): 11

tests/architectures/test_unet.py (3)

• 1-9: The import statement for get_children_layers has been updated to reflect its new location in sleap_nn.architectures.utils. Ensure that the function has been correctly moved to this new module and that all other references to this function in the codebase have been updated accordingly.

• 56-80: The output of the UNet model is now a dictionary containing the final output tensor and a list of current strides from the decoder. This is a significant change from the previous version, where the model directly returned the output tensor. Ensure that all parts of the codebase that use the UNet model have been updated to handle this new output format. Also, verify that the strides are being correctly calculated and returned by the model.

• 90-96: The eval method of the Encoder instance is now being called without assigning its return value to a variable. This is a minor change from the previous version, where the return value was assigned to _. This change is acceptable as the eval method does not return any value and is used for setting the model to evaluation mode.

sleap_nn/architectures/model.py (1)

• 1-154: The new hunks in sleap_nn/architectures/model.py introduce a Model class that serves as a higher-level wrapper around torch.nn.Module. This class holds the configuration parameters required to construct the actual model and allows for easy querying of the model configuration. It initializes the backbone and head based on the provided model_config and performs a forward pass through the model.

The get_backbone and get_head functions are used to instantiate the appropriate backbone and head modules based on the provided names and configurations. These functions raise a KeyError if the provided name is not one of the supported values.

The Model class has a from_config class method that creates a model instance from a configuration dictionary. The forward method performs a forward pass through the model, applying the backbone and each head layer to the input tensor and returning a dictionary of outputs.

The code appears to be logically sound and well-structured, with clear separation of concerns and good use of Python's object-oriented features. The use of a dictionary to map names to classes in the get_backbone and get_head functions is a good approach for handling multiple types of backbones and heads.

However, there are some commented-out lines of code (lines 130-136) that should be removed if they are not needed. Leaving commented-out code in the codebase can lead to confusion and clutter.

-        # self.heads = nn.ModuleList()
-        # self.heads_metadata = []
-        # for head_config in head_configs:
-        #     head = get_head(head_config.head_type, head_config.head_config)
-        #     self.heads_metadata.append((head.output_stride, head.name))
-        #     head = head.make_head(x_in=in_channels)
-        #     self.heads.append(head)

sleap_nn/architectures/unet.py (4)

• 2-12: The import of math and F (functional) from torch.nn has been removed. Ensure that these are not used elsewhere in the code or in the functions that are not shown in the hunks. If they are used, they should be imported back.

• 44-54: The initialization of the UNet class has been updated to include more attributes. This is a good practice as it makes the class more flexible and customizable. However, ensure that these new attributes are used correctly throughout the class and that their default values (if any) are set appropriately.

• 83-92: The output_channels property has been added to the UNet class. This property calculates the number of output channels based on the filters, filters_rate, down_blocks, and up_blocks attributes. This is a good addition as it provides a way to easily get the number of output channels without having to calculate it manually each time.

• 94-105: The forward method has been updated to return a tuple containing the output tensor and a list of current strides from the decoder. This is a significant change and could potentially break existing code that calls this method and expects a single output tensor. Ensure that all calls to this method throughout the codebase have been updated to handle the new return type.

sleap_nn/architectures/encoder_decoder.py (3)

• 28-38: The import statement for get_act_fn has been moved from sleap_nn.architectures.common to sleap_nn.architectures.utils. Ensure that the function is correctly defined in the new module and that all references to it have been updated.

• 430-437: The current_strides attribute has been added to the Decoder class. This attribute is used to keep track of the current stride at each level of the decoder. This is a good practice as it allows for dynamic tracking of the stride, which can be useful for understanding the scale of the features at each level.

• 458-485: The forward method of the Decoder class has been updated to return a dictionary containing the output tensors and the current strides. This change provides more information about the state of the decoder and can be useful for debugging and understanding the model's behavior. However, ensure that all calls to this method throughout the codebase have been updated to handle the new return type.

-    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]:

[talmo]

Minor changes

[coderabbitai]

Review Status

Actionable comments generated: 1

Commits

Files that changed from the base of the PR and between 702009a and 521fc9c.

Files selected for processing (1)

• sleap_nn/architectures/utils.py (1 hunks)

Additional comments (Suppressed): 1

sleap_nn/architectures/utils.py (1)

• 6-42: The get_act_fn function is well implemented. It provides a clear and concise way to get an instance of a PyTorch activation function module based on a string input. The error handling for unsupported activation functions is also appropriate.

[coderabbitai]

Review Status

Actionable comments generated: 3

Commits

Files that changed from the base of the PR and between 521fc9c and 4eff9be.

Files selected for processing (1)

• sleap_nn/architectures/model.py (1 hunks)

Additional comments (Suppressed): 1

sleap_nn/architectures/model.py (1)

• 108-130: The Model class is initialized with a list of head configurations but it's not clear if all heads are compatible with the backbone. Consider adding a compatibility check between the backbone and the heads to avoid potential runtime errors.

[coderabbitai]

Review Status

Actionable comments generated: 1

Commits

Files that changed from the base of the PR and between 4eff9be and 0325596.

Files ignored due to filter (1)

• pyproject.toml

Files selected for processing (2)

• sleap_nn/data/pipelines.py (3 hunks)
• tests/data/test_pipelines.py (3 hunks)

Additional comments (Suppressed): 6

tests/data/test_pipelines.py (2)

• 42-77: The augmentation configuration has been restructured into two categories: intensity and geometric. Ensure that the code that uses this configuration has been updated to handle the new structure. Also, a new key anchor_ind has been added to the preprocessing section. Verify its usage and impact on the system.

• 96-131: Similar to the previous comment, the augmentation configuration has been restructured and a new key anchor_ind has been added. Ensure that the code that uses this configuration has been updated to handle the new structure and the new key.

sleap_nn/data/pipelines.py (4)

• 6-6: The import statement for DictConfig has been changed. Ensure that the new import path is correct and does not break the code.

- from omegaconf.dictconfig import DictConfig
+ from omegaconf.omegaconf import DictConfig

• 49-51: The InstanceCentroidFinder instantiation now includes an additional parameter anchor_ind. Ensure that this parameter is correctly defined in self.data_config.preprocessing and that it is correctly used within the InstanceCentroidFinder class.

• 44-47: The KorniaAugmenter instantiation now includes a specific subset of augmentations (intensity). Ensure that this subset is correctly defined in self.data_config.augmentation_config.augmentations and that it is correctly used within the KorniaAugmenter class.

• 62-65: The KorniaAugmenter instantiation now includes a specific subset of augmentations (geometric). Ensure that this subset is correctly defined in self.data_config.augmentation_config.augmentations and that it is correctly used within the KorniaAugmenter class.