Add option to load trained weights for Head layers #114
Conversation
WalkthroughThe pull request introduces significant enhancements to the configuration documentation and files related to the Changes
Possibly related PRs
Suggested reviewers
Poem
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (15)
🧰 Additional context used🪛 LanguageTooldocs/config.md[uncategorized] ~179-~179: Loose punctuation mark. (UNLIKELY_OPENING_PUNCTUATION) [uncategorized] ~182-~182: Loose punctuation mark. (UNLIKELY_OPENING_PUNCTUATION) 🪛 Markdownlintdocs/config.md178-178: Expected: 4; Actual: 8 (MD007, ul-indent) 179-179: Expected: 4; Actual: 8 (MD007, ul-indent) 180-180: Expected: 6; Actual: 12 (MD007, ul-indent) 181-181: Expected: 6; Actual: 12 (MD007, ul-indent) 182-182: Expected: 4; Actual: 8 (MD007, ul-indent) 183-183: Expected: 6; Actual: 12 (MD007, ul-indent) 184-184: Expected: 6; Actual: 12 (MD007, ul-indent) 185-185: Expected: 6; Actual: 12 (MD007, ul-indent) 186-186: Expected: 6; Actual: 12 (MD007, ul-indent) 187-187: Expected: 6; Actual: 12 (MD007, ul-indent) 188-188: Expected: 6; Actual: 12 (MD007, ul-indent) 184-184: null (MD037, no-space-in-emphasis) 184-184: null (MD037, no-space-in-emphasis) 🪛 yamllintinitial_config.yaml[error] 1-1: wrong new line character: expected \n (new-lines) 🪛 Ruffsleap_nn/training/model_trainer.py582-582: Use Remove (SIM118) 593-593: Use Remove (SIM118) 🔇 Additional comments (29)tests/assets/minimal_instance_centroid/initial_config.yaml (1)
The chunk_size of 100 might affect memory usage and training performance. Consider validating this value:
✅ Verification successfulThe chunk size value is consistent across all configurations The value of
This consistent usage across test cases and documentation suggests this is a validated default value. The review concern about validation can be dismissed. 🏁 Scripts executedThe following scripts were executed for the analysis: Script: #!/bin/bash
# Search for other chunk_size configurations to ensure consistency
rg "chunk_size:" -A 1
Length of output: 1310 tests/assets/minimal_instance/initial_config.yaml (3)
The ReduceLROnPlateau scheduler configuration is comprehensive and includes all necessary parameters with reasonable default values.
The chunk_size parameter should be validated to ensure it's a positive integer. #!/bin/bash
# Check if chunk_size validation exists in the codebase
rg "chunk_size.*validation|validate.*chunk_size" --type python
Specify a value for bin_files_path or remove if not needed. The #!/bin/bash
# Check if bin_files_path is used in the codebase
rg -l "bin_files_path" --type pythontests/assets/minimal_instance_bottomup/initial_config.yaml (2)
The PR objectives mention
#!/bin/bash
# Search for references to the weight loading parameters
rg "backbone_trained_ckpts_path|head_trained_ckpts_path" --type py
The new data loading parameters could significantly affect training:
Please ensure these values are optimal for your use case. ✅ Verification successfulParameters are well-defined with safe default values Based on the codebase analysis:
The values in the config (True, 100) match the defaults in the codebase, making them safe choices. The parameters are well-documented and consistently used across the data loading pipeline. 🏁 Scripts executedThe following scripts were executed for the analysis: Script: #!/bin/bash
# Search for any documentation or tests related to these parameters
rg -l "user_instances_only|chunk_size" --type py
Length of output: 4414 Script: #!/bin/bash
# Let's search for the actual parameters we're interested in
rg "user_instances_only|chunk_size" --type py -A 2
Length of output: 10463 tests/assets/minimal_instance_centroid/training_config.yaml (1)
The new data loading parameters need validation:
✅ Verification successfulParameters are well-documented and consistent with examples The values in your configuration align with both the documentation and example configurations:
🏁 Scripts executedThe following scripts were executed for the analysis: Script: #!/bin/bash
# Search for documentation about these parameters
rg -l "user_instances_only|chunk_size" docs/
Length of output: 156 Script: #!/bin/bash
# Check the documentation content for these parameters
rg -A 5 "user_instances_only|chunk_size" docs/config.md
# Also check usage in example configs
rg -A 2 "user_instances_only|chunk_size" docs/config_*.yaml
Length of output: 1404 tests/assets/minimal_instance_bottomup/training_config.yaml (2)
The new data configuration parameters have important implications:
Please confirm these values are appropriate for your use case, especially for:
The PR objective mentions loading
Consider adding a comment in the YAML to document this parameter's purpose and format requirements. tests/assets/minimal_instance/training_config.yaml (4)
The new
The ReduceLROnPlateau configuration includes some parameters that might need validation:
Consider:
#!/bin/bash
# Check if these parameters are documented
rg -A 5 "ReduceLROnPlateau" docs/
The new
#!/bin/bash
# Search for documentation or usage of bin_files_path
rg "bin_files_path" docs/
The new parameters
Please add validation for these values and document their recommended ranges. docs/config_bottomup.yaml (1)
The restructuring of the lr_scheduler configuration improves clarity by:
This change makes it easier to extend the configuration for other scheduler types in the future. docs/config_topdown_centered_instance.yaml (2)
The ReduceLROnPlateau configuration is technically sound with appropriate values. However, this change appears unrelated to the PR's stated objective of adding support for loading trained weights. Let's verify if this change is part of a larger refactoring: #!/bin/bash
# Search for similar scheduler changes in other config files
rg -l "ReduceLROnPlateau" --type yaml
The newly added
Let's check if this parameter is documented elsewhere: ✅ Verification successfulLet me search for more specific information about the chunk_size parameter in the codebase. The chunk_size parameter is correctly configured and consistent with the codebase The value of 100 for
🏁 Scripts executedThe following scripts were executed for the analysis: Script: #!/bin/bash
# Search for any documentation or usage of chunk_size
rg -i "chunk.?size" --type md --type rst
Length of output: 491 Script: #!/bin/bash
# Search for chunk_size in Python files to understand its usage
rg -i "chunk.?size" --type py -C 2
# Also check YAML files for any existing usage
rg -i "chunk.?size" --type yaml -C 2
Length of output: 10807 initial_config.yaml (2)
The configuration hardcodes #!/bin/bash
# Search for scheduler validation in the codebase
rg -A 5 'scheduler.*(?:ReduceLROnPlateau|StepLR)'
The docs/config_centroid.yaml (2)
While the learning rate scheduler configuration is well-structured and the values are reasonable, these changes appear to be outside the scope of the PR's objective of "Add option to load trained weights for Head layers". Consider moving these changes to a separate PR focused on learning rate scheduling improvements. Let's verify if there are related changes in other files: #!/bin/bash
# Search for other files with similar lr_scheduler changes
rg -l "reduce_lr_on_plateau" --type yaml
The newly added Let's check if this parameter is documented elsewhere: ✅ Verification successfulLet's search more broadly for any mentions or usage of chunk_size in the codebase to understand its purpose and implementation: Based on the search results, I can now provide a final response as the documentation for The The 🏁 Scripts executedThe following scripts were executed for the analysis: Script: #!/bin/bash
# Search for documentation of chunk_size parameter
rg -i "chunk.?size" docs/
Length of output: 494 Script: #!/bin/bash
# Search for chunk_size in all files
rg -i "chunk.?size" -A 3 -B 3
# Also check for any related documentation in README or docs
fd "README|\.md$" -t f -x cat {}
Length of output: 40633 tests/fixtures/datasets.py (3)
Changing
#!/bin/bash
# Search for tests that might be affected by this change
rg -A 5 "up_interpolate" --type python tests/Line range hint The changes in this file are focused on configuration updates for data processing and training, but I don't see any additions related to the PR's primary objective of loading trained weights for head layers. The PR description mentions new parameters Please ensure that:
#!/bin/bash
# Search for implementation of weight loading functionality
rg -A 5 "trained_ckpts_path" --type python
# Look for related test cases
fd -e py | xargs rg -l "test.*trained.*weights|test.*load.*weights"
The changes to the data provider configuration introduce breaking changes and new parameters:
Please ensure these changes are documented and verify their impact on existing tests. Also applies to: 54-55 sleap_nn/training/model_trainer.py (6)
The
When calling
Ensure that the learning rate scheduler configuration handles missing or incorrect parameters gracefully. Adding validation or default values can prevent runtime errors.
Verify that all necessary parameters are passed correctly to the superclass
Ensure that the updated parameters and initialization logic are consistently applied across all model subclasses ( Also applies to: 804-814
Double-check that Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## divya/add-config-wandb #114 +/- ##
=======================================================
Coverage 97.37% 97.37%
=======================================================
Files 38 38
Lines 3695 3702 +7
=======================================================
+ Hits 3598 3605 +7
Misses 97 97 ☔ View full report in Codecov by Sentry. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's limits when posting large numbers of comments.
Actionable comments posted: 15
🧹 Outside diff range and nitpick comments (19)
tests/assets/minimal_instance_centroid/initial_config.yaml (2)
64-64: Document the bin_files_path parameter.Given that this PR introduces functionality for loading pre-trained weights:
- Consider adding a comment explaining the expected format and purpose of bin_files_path
- Should there be validation to ensure the path exists when provided?
- How does this relate to the backbone_trained_ckpts_path and head_trained_ckpts_path mentioned in the PR description?
78-85: LGTM! Consider adding parameter descriptions.The ReduceLROnPlateau configuration looks well-structured with reasonable default values. To improve maintainability, consider adding inline comments explaining each parameter's purpose and impact on training.
Example documentation:
reduce_lr_on_plateau: # Minimum change in loss to be considered an improvement threshold: 1.0e-07 # How to measure the threshold (abs or rel) threshold_mode: abs # Number of epochs to wait after a reduction before resuming normal operation cooldown: 3 # Number of epochs with no improvement after which LR will be reduced patience: 5 # Factor by which the learning rate will be reduced factor: 0.5 # Lower bound on the learning rate min_lr: 1.0e-08tests/assets/minimal_instance/initial_config.yaml (1)
13-14: Document the significance of crop_hw values.The hardcoded values [160, 160] for crop_hw should be documented to explain their significance and any constraints.
crop_hw: - - 160 - - 160 + - 160 # Height of the crop window + - 160 # Width of the crop windowtests/assets/minimal_instance_bottomup/initial_config.yaml (2)
70-71: Document the purpose and format ofbin_files_path.The PR mentions loading
.ckptformat weights, butbin_files_pathsuggests binary files. Please clarify:
- The expected file format
- Whether this is related to the pre-trained weights loading feature
- Add documentation for this parameter's usage
84-91: LGTM! Consider adding parameter descriptions.The learning rate scheduler configuration is well-structured with appropriate values for ReduceLROnPlateau. Consider adding inline comments or documentation explaining the purpose of each parameter for future maintainability.
lr_scheduler: scheduler: ReduceLROnPlateau reduce_lr_on_plateau: threshold: 1.0e-07 # Minimum change in loss to be considered an improvement threshold_mode: abs # Use absolute change in loss cooldown: 3 # Epochs to wait before resuming normal operation patience: 5 # Epochs to wait before reducing LR factor: 0.5 # Factor to reduce LR by min_lr: 1.0e-08 # Minimum LR thresholdtests/assets/minimal_instance_centroid/training_config.yaml (1)
75-75: Document thebin_files_pathparameter.Since this parameter is related to loading trained weights (PR objective), please:
- Document the expected path format and structure
- Add validation to handle empty/invalid paths
- Consider providing an example path in the config
tests/assets/minimal_instance_bottomup/training_config.yaml (1)
99-106: LGTM: Well-configured learning rate scheduler.The ReduceLROnPlateau configuration is comprehensive and uses sensible values. The scheduler will:
- Reduce learning rate by half when no improvement is seen
- Wait 5 epochs before reducing (patience)
- Allow 3 epochs of cooldown between reductions
- Stop at min_lr of 1e-8
Consider adding comments in the YAML to explain the purpose of less obvious parameters like
threshold_modeandcooldown.docs/config_bottomup.yaml (2)
6-6: Document the purpose of chunk_size parameter.The newly added
chunk_sizeparameter would benefit from a comment explaining its purpose and impact on data processing. This helps users understand how to tune this value for their specific use cases.
Parameters should be added to trainer_config section
Based on the implementation in
model_trainer.py, the parametersbackbone_trained_ckpts_pathandhead_trained_ckpts_pathare used during model initialization and training. These parameters should be added to thetrainer_configsection of the configuration file, not themodel_configsection, as they are handled by the trainer.trainer_config: # ... existing config ... resume_ckpt_path: + backbone_trained_ckpts_path: # Path to trained checkpoint for backbone + head_trained_ckpts_path: # Path to trained checkpoint for head layers optimizer_name: Adam🔗 Analysis chain
Line range hint
1-115: Missing configuration parameters for trained weights loading.According to the PR objectives, this change should add support for loading pre-trained weights for head layers. However, I don't see the new parameters
backbone_trained_ckpts_pathandhead_trained_ckpts_pathin the configuration. These parameters should be added to enable the described functionality.Let's check if these parameters are defined elsewhere:
Consider adding these parameters to the configuration file:
trainer_config: # ... existing config ... resume_ckpt_path: + backbone_trained_ckpts_path: + head_trained_ckpts_path: optimizer_name: Adam🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Search for the new parameters in other configuration files rg -l "backbone_trained_ckpts_path|head_trained_ckpts_path"Length of output: 133
Script:
#!/bin/bash # Let's check the implementation to understand how these parameters are used rg -A 10 "backbone_trained_ckpts_path|head_trained_ckpts_path" tests/training/test_model_trainer.py sleap_nn/training/model_trainer.pyLength of output: 17142
docs/config_topdown_centered_instance.yaml (1)
Line range hint
1-114: Missing configuration for trained weights loading.Based on the PR objectives, this configuration file should include parameters for specifying trained weight paths:
backbone_trained_ckpts_pathhead_trained_ckpts_pathThese parameters are essential for the PR's core functionality of loading pre-trained weights for head layers.
Consider adding these parameters under the appropriate section (possibly under
model_configortrainer_config).docs/config_centroid.yaml (1)
Line range hint
89-93: Add configuration for head layer pre-trained weights.The PR aims to add support for loading pre-trained weights for head layers, but the configuration is missing the necessary parameters. Consider adding
head_trained_ckpts_pathunder thehead_configssection to align with the PR's objective.Example addition:
head_configs: single_instance: centered_instance: bottomup: centroid: + trained_weights_path: # Path to pre-trained weights for centroid head confmaps: anchor_part: 0🧰 Tools
🪛 yamllint
[error] 8-8: trailing spaces
(trailing-spaces)
[error] 9-9: trailing spaces
(trailing-spaces)
tests/fixtures/datasets.py (2)
62-62: Document the purpose of min_crop_size parameter.The new
min_crop_sizeparameter has been added to preprocessing configuration. Please add documentation explaining:
- The purpose of this parameter
- The impact when set to None
- Valid values and their effects
159-167: LGTM! Consider documenting scheduler parameters.The restructured LR scheduler configuration with ReduceLROnPlateau is well-organized and provides comprehensive control over learning rate adjustments. The parameters are sensibly configured with:
- Relative threshold mode with small threshold (1e-07)
- Conservative reduction factor (0.5)
- Reasonable patience (5) and cooldown (3) periods
- Safety floor for minimum learning rate (1e-08)
Consider adding inline comments or documentation explaining the purpose of each parameter and their recommended ranges.
tests/training/test_model_trainer.py (2)
108-110: Consider adding edge cases for StepLR parameters.While the basic StepLR configuration is tested, consider adding test cases for:
- Edge case values for
step_size(e.g., 0, negative)- Edge case values for
gamma(e.g., 0, negative, >1)- Verification of learning rate changes after each step
337-342: Enhance scheduler type validation tests.Consider improving the error handling tests:
- Test multiple invalid scheduler types
- Verify the error message content
- Test edge cases (empty string, None, etc.)
Example:
invalid_schedulers = ["ReduceLR", "", None, "CustomScheduler"] for scheduler in invalid_schedulers: OmegaConf.update(config, "trainer_config.lr_scheduler.scheduler", scheduler) with pytest.raises(ValueError, match=f"Unsupported scheduler type: {scheduler}"): trainer = ModelTrainer(config) trainer.train()docs/config.md (1)
178-188: Technical review of learning rate scheduler documentationThe learning rate scheduler documentation is well-structured but could benefit from additional clarity:
- For
StepLR, consider adding an example to illustrate the decay pattern.- For
ReduceLROnPlateau, thethreshold_modeparameter's explanation could be clearer.Consider updating the documentation with this improved explanation:
- `threshold_mode`: (str) One of "rel", "abs". In rel mode, dynamic_threshold = best * ( 1 + threshold ) in max mode or best * ( 1 - threshold ) in min mode. In abs mode, dynamic_threshold = best + threshold in max mode or best - threshold in min mode. *Default*: "rel". + - `threshold_mode`: (str) One of "rel" (relative) or "abs" (absolute). Determines how to interpret the threshold parameter: + - In "rel" mode: Triggers when improvement is less than (best * threshold) for min mode, or (best * (1 + threshold)) for max mode + - In "abs" mode: Triggers when improvement is less than (best - threshold) for min mode, or (best + threshold) for max mode + *Default*: "rel"🧰 Tools
🪛 LanguageTool
[uncategorized] ~179-~179: Loose punctuation mark.
Context: ...ReduceLROnPlateau". -step_lr: -step_size`: (int) Period...(UNLIKELY_OPENING_PUNCTUATION)
[uncategorized] ~182-~182: Loose punctuation mark.
Context: ...*: 0.1. -reduce_lr_on_plateau: -threshold: (float) Thre...(UNLIKELY_OPENING_PUNCTUATION)
🪛 Markdownlint
178-178: Expected: 4; Actual: 8
Unordered list indentation(MD007, ul-indent)
179-179: Expected: 4; Actual: 8
Unordered list indentation(MD007, ul-indent)
180-180: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
181-181: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
182-182: Expected: 4; Actual: 8
Unordered list indentation(MD007, ul-indent)
183-183: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
184-184: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
185-185: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
186-186: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
187-187: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
188-188: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
184-184: null
Spaces inside emphasis markers(MD037, no-space-in-emphasis)
184-184: null
Spaces inside emphasis markers(MD037, no-space-in-emphasis)
sleap_nn/training/model_trainer.py (3)
497-498: Include parameter types in docstringsIn the docstring for
TrainingModel, you have added descriptions forbackbone_trained_ckpts_pathandhead_trained_ckpts_path. Consider specifying the parameter types for better clarity.Apply this diff to enhance the docstring:
head_trained_ckpts_path: Path to trained ckpts for head layer. + + Args: + backbone_trained_ckpts_path (Optional[str]): Path to trained checkpoints for backbone. + head_trained_ckpts_path (Optional[str]): Path to trained checkpoints for head layers.
713-714: Add parameter types in docstringsFor consistency and clarity, include the parameter types in the docstrings for
SingleInstanceModel.Apply this diff to enhance the docstring:
head_trained_ckpts_path: Path to trained ckpts for head layer. + + Args: + backbone_trained_ckpts_path (Optional[str]): Path to trained checkpoints for backbone. + head_trained_ckpts_path (Optional[str]): Path to trained checkpoints for head layers.
882-895: Update docstrings and parameter types forCentroidModelInclude the parameter types in the docstrings to improve clarity.
Apply this diff:
head_trained_ckpts_path: Path to trained ckpts for head layer. + + Args: + backbone_trained_ckpts_path (Optional[str]): Path to trained checkpoints for backbone. + head_trained_ckpts_path (Optional[str]): Path to trained checkpoints for head layers.
🛑 Comments failed to post (15)
tests/assets/minimal_instance/initial_config.yaml (1)
15-15: 💡 Codebase verification
⚠️ Potential issuemin_crop_size parameter should have an integer value
The
min_crop_sizeparameter is used to specify the minimum crop size whencrop_hwis not set. Based on the implementation insleap_nn/data/instance_cropping.py, it expects an integer value or None. Please either:
- Set an appropriate integer value for
min_crop_size- Remove the parameter if you want to use the default behavior (equivalent to setting it to 0)
🔗 Analysis chain
Specify a value for min_crop_size or remove if not needed.
The
min_crop_sizeparameter is added but left empty, which could lead to validation issues or undefined behavior.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Check if min_crop_size is used in the codebase rg -l "min_crop_size" --type pythonLength of output: 161
Script:
#!/bin/bash # Check if min_crop_size is used in the codebase with ripgrep rg -l "min_crop_size" # Also search for any YAML files that might reference this parameter fd -e yaml -x grep -l "min_crop_size" {}Length of output: 586
Script:
#!/bin/bash # Check the implementation and usage of min_crop_size in the core code rg "min_crop_size" sleap_nn/data/instance_cropping.py -C 2 # Check documentation for min_crop_size rg "min_crop_size" docs/config.md -C 2Length of output: 2095
tests/assets/minimal_instance_centroid/training_config.yaml (2)
89-96: 🛠️ Refactor suggestion
Review learning rate scheduler parameters.
The ReduceLROnPlateau configuration needs attention:
threshold: 1e-7is very small - consider a larger value like 1e-4patience: 5epochs might be too short for complex modelsmin_lr: 1e-8is extremely small - typical values are around 1e-6Consider these adjustments:
reduce_lr_on_plateau: - threshold: 1.0e-07 + threshold: 1.0e-04 - patience: 5 + patience: 10 - min_lr: 1.0e-08 + min_lr: 1.0e-06📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.scheduler: ReduceLROnPlateau reduce_lr_on_plateau: threshold: 1.0e-04 threshold_mode: abs cooldown: 3 patience: 10 factor: 0.5 min_lr: 1.0e-06
75-77:
⚠️ Potential issueImplementation incomplete for PR objective.
The configuration is missing the parameters mentioned in the PR objective:
backbone_trained_ckpts_pathhead_trained_ckpts_pathThese parameters are essential for loading pre-trained weights for head layers.
Add the missing parameters:
bin_files_path: + backbone_trained_ckpts_path: + head_trained_ckpts_path: resume_ckpt_path:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.bin_files_path: backbone_trained_ckpts_path: head_trained_ckpts_path: resume_ckpt_path: wandb:docs/config_topdown_centered_instance.yaml (1)
15-15:
⚠️ Potential issueSpecify a value for min_crop_size or document if it's optional.
The
min_crop_sizeparameter is added without a value. This could lead to:
- Configuration validation errors
- Unexpected runtime behavior
Please either:
- Set a default value
- Document if this is an optional parameter
- Add validation to handle empty values
initial_config.yaml (4)
1-104:
⚠️ Potential issueFix line endings for cross-platform compatibility
The file uses incorrect line endings. Ensure consistent use of Unix-style line endings (\n) across the codebase.
Run the following command to fix line endings:
#!/bin/bash # Convert CRLF to LF sed -i 's/\r$//' initial_config.yaml🧰 Tools
🪛 yamllint
[error] 1-1: wrong new line character: expected \n
(new-lines)
84-84:
⚠️ Potential issueSecurity: Remove empty API key field
For security reasons, sensitive information like API keys should not be included in configuration files, even if empty. Consider removing this field and using environment variables instead.
- api_key: '' + # API key should be provided via environment variable WANDB_API_KEY📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.# API key should be provided via environment variable WANDB_API_KEY
3-4:
⚠️ Potential issueReplace hardcoded paths with relative paths
The configuration contains Windows-specific absolute paths, which will break when used on different machines or operating systems.
Replace absolute paths with relative paths:
- train_labels_path: C:\Users\TalmoLab\Desktop\Divya\sleap-nn\tests\assets/minimal_instance.pkg.slp - val_labels_path: C:\Users\TalmoLab\Desktop\Divya\sleap-nn\tests\assets/minimal_instance.pkg.slp + train_labels_path: tests/assets/minimal_instance.pkg.slp + val_labels_path: tests/assets/minimal_instance.pkg.slp📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.train_labels_path: tests/assets/minimal_instance.pkg.slp val_labels_path: tests/assets/minimal_instance.pkg.slp
27-28:
⚠️ Potential issueAdd configuration for backbone and head trained weights
According to the PR objectives, the configuration should support
backbone_trained_ckpts_pathandhead_trained_ckpts_pathfor loading pre-trained weights. However, these parameters are missing from the configuration.Add the missing parameters:
init_weights: default pre_trained_weights: null +backbone_trained_ckpts_path: null +head_trained_ckpts_path: null backbone_type: unet📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.init_weights: default pre_trained_weights: null backbone_trained_ckpts_path: null head_trained_ckpts_path: nulltests/training/test_model_trainer.py (1)
346-384: 🛠️ Refactor suggestion
Enhance test coverage for weight loading functionality.
The current test verifies basic weight loading, but consider adding tests for:
- Loading only backbone weights
- Loading only head weights
- Invalid checkpoint paths
- Incompatible model architectures
- Verifying weights are maintained after training
Example additions:
# Test loading only backbone weights trainer = ModelTrainer(load_weights_config) trainer._initialize_model( backbone_trained_ckpts_path=(Path(minimal_instance_ckpt) / "best.ckpt").as_posix(), head_trained_ckpts_path=None ) model_ckpt = next(trainer.model.parameters())[0, 0, :].detach().numpy() assert np.all(np.abs(first_layer_ckpt - model_ckpt) < 1e-6) # Test invalid checkpoint path with pytest.raises(FileNotFoundError): trainer._initialize_model( backbone_trained_ckpts_path="invalid_path.ckpt", head_trained_ckpts_path=None ) # Test incompatible architecture incompatible_ckpt = {...} # Create checkpoint with different architecture with pytest.raises(ValueError, match="Incompatible model architecture"): trainer._initialize_model( backbone_trained_ckpts_path=incompatible_ckpt, head_trained_ckpts_path=None )docs/config.md (1)
178-188:
⚠️ Potential issueMarkdown formatting improvements needed
The indentation of list items is inconsistent with the rest of the document. The static analysis tool suggests using 4-space indentation for unordered lists.
Apply these formatting changes:
- Adjust indentation to match the rest of the document (4 spaces for top-level items, 6 spaces for nested items)
- Remove extra spaces around emphasis markers in the
threshold_modedescription- Fix loose punctuation marks after the scheduler options
- - `scheduler`: (str) Name of the scheduler to use. Valid schedulers: `"StepLR"`, `"ReduceLROnPlateau"`. - - `step_lr`: - - `step_size`: (int) Period of learning rate decay. If `step_size`=10, then every 10 epochs, learning rate will be reduced by a factor of `gamma`. + - `scheduler`: (str) Name of the scheduler to use. Valid schedulers: `"StepLR"`, `"ReduceLROnPlateau"` + - `step_lr`: + - `step_size`: (int) Period of learning rate decay. If `step_size`=10, then every 10 epochs, learning rate will be reduced by a factor of `gamma`.Committable suggestion was skipped due to low confidence.
🧰 Tools
🪛 LanguageTool
[uncategorized] ~179-~179: Loose punctuation mark.
Context: ...ReduceLROnPlateau". -step_lr: -step_size`: (int) Period...(UNLIKELY_OPENING_PUNCTUATION)
[uncategorized] ~182-~182: Loose punctuation mark.
Context: ...*: 0.1. -reduce_lr_on_plateau: -threshold: (float) Thre...(UNLIKELY_OPENING_PUNCTUATION)
🪛 Markdownlint
178-178: Expected: 4; Actual: 8
Unordered list indentation(MD007, ul-indent)
179-179: Expected: 4; Actual: 8
Unordered list indentation(MD007, ul-indent)
180-180: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
181-181: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
182-182: Expected: 4; Actual: 8
Unordered list indentation(MD007, ul-indent)
183-183: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
184-184: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
185-185: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
186-186: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
187-187: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
188-188: Expected: 6; Actual: 12
Unordered list indentation(MD007, ul-indent)
184-184: null
Spaces inside emphasis markers(MD037, no-space-in-emphasis)
184-184: null
Spaces inside emphasis markers(MD037, no-space-in-emphasis)
sleap_nn/training/model_trainer.py (5)
332-336: 🛠️ Refactor suggestion
Update docstrings to reflect new parameters
The
_initialize_modelmethod now accepts two new parameters:backbone_trained_ckpts_pathandhead_trained_ckpts_path. Please update the method's docstring to include descriptions for these parameters to maintain clarity and consistency in documentation.
455-455:
⚠️ Potential issueAvoid logging sensitive configuration data
When updating the WandB experiment configuration, ensure that no sensitive information (such as API keys or personal data) is being inadvertently logged.
517-518: 🛠️ Refactor suggestion
Remove redundancy in attribute assignment
You are already passing
backbone_trained_ckpts_pathandhead_trained_ckpts_pathto the superclass. Assigning them again here may be unnecessary.Apply this diff to remove redundant assignments:
- self.backbone_trained_ckpts_path = backbone_trained_ckpts_path - self.head_trained_ckpts_path = head_trained_ckpts_pathCommittable suggestion was skipped due to low confidence.
576-585:
⚠️ Potential issueSimplify dictionary key iteration
As per the static analysis hint, you can simplify the dictionary iteration by removing the
.keys()method.Apply this diff to address the issue:
- for k in ckpt["state_dict"].keys() + for k in ckpt["state_dict"] if ".backbone" in k📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.# Initializing backbone (encoder + decoder) with trained ckpts if backbone_trained_ckpts_path is not None: print(f"Loading backbone weights from `{backbone_trained_ckpts_path}` ...") ckpt = torch.load(backbone_trained_ckpts_path) ckpt["state_dict"] = { k: ckpt["state_dict"][k] for k in ckpt["state_dict"] if ".backbone" in k } self.load_state_dict(ckpt["state_dict"], strict=False)🧰 Tools
🪛 Ruff
582-582: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
587-598:
⚠️ Potential issueRemove debugging
- The
- Simplify the dictionary key iteration by removing
.keys()as per the static analysis hint.Apply this diff to address both issues:
} - for k in ckpt["state_dict"].keys() + for k in ckpt["state_dict"] if ".head_layers" in k } - print(f"from main code: {ckpt['state_dict'].keys()}")Committable suggestion was skipped due to low confidence.
🧰 Tools
🪛 Ruff
593-593: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
This PR adds an option to load trained weights for head layers right before training (instead of random initialization) for transfer learning. The
.ckptpaths can be passed toModelTrainer.train(backbone_trained_ckpts_path=<your_path>, head_trained_ckpts_path=<your_path>).Summary by CodeRabbit
New Features
schedulertypes and parameters.chunk_sizeanduser_instances_onlyparameters for improved data handling.bin_files_pathfor specifying binary file paths in training configurations.Bug Fixes
Tests