Advanced Finetuning#358
Open
kozlov721 wants to merge 33 commits into
Open
Conversation
kozlov721
requested review from
conorsim,
klemen1999 and
tersekmatija
and removed request for
a team
March 17, 2026 08:55
github-actions
Bot
added
documentation
dtronmans
reviewed
Mar 23, 2026
dtronmans
left a comment
dtronmans
left a comment
Contributor
There was a problem hiding this comment.
These are my comments so far regarding potential breaking changes or unexpected behavior, I will just take a quick look at the changes to ReduceLROnPlateau to note any changes to the scheduler
dtronmans
reviewed
Mar 23, 2026
| for scheduler in schedulers: | ||
| if isinstance(scheduler, ReduceLROnPlateau): | ||
| if scheduler.mode == "min": | ||
| scheduler.step(loss) |
Contributor
There was a problem hiding this comment.
In both modes, the validation results are needed. If the mode is "max" then the chosen val metric is taken into account, and if it is "min" then the val/loss is taken into account. But at this point (on the last batch of a training_step) the val metrics/loss are not computed yet so what is being taken into account? I cannot find any reference to scheduler.step() in the main branch so I assume that this was automatically handled by PyTorch Lightning before.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds per-node finetuning config types and schema, wires finetuning into predefined model NodeConfig creation, centralizes parameter-pattern-based optimizer/scheduler extraction and building in Nodes, updates Lightning for manual optimization and ReduceLROnPlateau handling, removes legacy fc_decay, and adds comprehensive finetuning unit tests and docs/examples. ChangesPer-node Finetuning Implementation
Sequence Diagram(s)sequenceDiagram
participant Client
participant Trainer as LuxonisLightningModule
participant Nodes
participant Builder as build_optimizer_scheduler
Client->>Trainer: start training
Trainer->>Nodes: nodes.build_optimizers(base_configs)
Nodes->>Builder: build_optimizer_scheduler per group
Builder-->>Nodes: Optimizer + Scheduler
Nodes-->>Trainer: sequences of Optimizers/Schedulers
Trainer->>Trainer: toggle automatic_optimization, training_step/manual_backward
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
kozlov721
force-pushed
the
feature/parameter-groups
branch
from
June 3, 2026 19:30
fdf1ed4 to
83885f1
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
luxonis_train/lightning/utils.py (1)
999-1017:⚠️ Potential issue | 🟠 Major | ⚡ Quick winUse the formatted metric key for
ReduceLROnPlateau(mode="max").Lines 1007-1008 build
val/metric/{main_metric.node_name}/{main_metric.metric_name}, but validation metrics are logged underval/metric/{formatted_node_name}/{metric_name}. For any node whose formatted name differs from its identifier, this scheduler will monitor a key that is never logged.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/lightning/utils.py` around lines 999 - 1017, The ReduceLROnPlateau branch builds a monitor key with main_metric.node_name which can differ from the logged/formatted node name; update the monitor construction to use the metric's formatted node name (e.g., main_metric.formatted_node_name) or call the existing formatting helper (e.g., format_node_name(main_metric.node_name)) so the monitor string becomes f"val/metric/{formatted_node_name}/{main_metric.metric_name}" when cfg_scheduler.name == "ReduceLROnPlateau" and mode == "max".
♻️ Duplicate comments (1)
luxonis_train/lightning/luxonis_lightning.py (1)
485-501:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftMove
ReduceLROnPlateaustepping out oftraining_step.Line 489 feeds the last train-batch loss into
mode="min", and Line 491 reads the main metric before the current epoch's validation has run. In manual optimization this makes plateau decisions use the wrong signal and ignore the validation cadence.#!/bin/bash set -euo pipefail printf "\n-- scheduler stepping in training_step --\n" sed -n '480,502p' luxonis_train/lightning/luxonis_lightning.py printf "\n-- metric updates / resets happen in evaluation hooks --\n" rg -n -C2 'metric\.run_update|metric\.reset|main_metric_reference\.compute' luxonis_train/lightning/luxonis_lightning.py🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/lightning/luxonis_lightning.py` around lines 485 - 501, The ReduceLROnPlateau logic currently inside training_step (the block that checks self.trainer.is_last_batch and calls scheduler.step(loss) or scheduler.step(value) using nodes.main_metric_reference.compute()) must be removed from training_step and moved to a post-validation hook where metrics have been updated (e.g., on_validation_epoch_end or a dedicated end-of-epoch method run after metric.run_update/metric.reset). Specifically, stop calling ReduceLROnPlateau.step from training_step; instead iterate the same schedulers list in the chosen validation-end hook, compute the main metric via nodes.main_metric_reference.compute() there (ensuring it’s not called before validation completes), and call scheduler.step(value) for non-loss modes or scheduler.step(loss) only when you have the correct training loss signal; keep other schedulers' scheduler.step() behavior unchanged. Ensure code references: training_step (remove block), nodes.main_metric_reference.compute(), and the schedulers iteration (re-add in the validation-end hook).
🧹 Nitpick comments (1)
tests/integration/test_finetuning.py (1)
552-553: 💤 Low valueRemove leftover
# CHECKmarkers.The
# CHECKcomments abovetest_inherits_omitted_optimizer_and_schedulerandtest_inherits_name_and_merges_params_when_name_is_missinglook like authoring reminders that should be resolved or dropped before merge.Want me to confirm whether these tests need any follow-up before removing the markers?
Also applies to: 582-583
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/test_finetuning.py` around lines 552 - 553, Remove the leftover authoring markers by deleting the stray "# CHECK" comments immediately above the test functions test_inherits_omitted_optimizer_and_scheduler and test_inherits_name_and_merges_params_when_name_is_missing; ensure no other sentinel comments remain in that test file and run the test suite to confirm nothing relied on those markers.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@luxonis_train/lightning/utils.py`:
- Around line 356-363: The grouping currently keys buckets by
(cfg_optimizer.name, cfg_scheduler.name) which collapses different scheduler
configs that share a name; change the key to include scheduler parameters as
well (for example include a stable representation such as
frozenset(cfg_scheduler.params.items()) or a serialized params string) so each
distinct (optimizer, scheduler name, scheduler params) combination gets its own
bucket; update the lookup/creation at groups[...] and keep storing the same
tuple (list, cfg_scheduler) so existing uses of groups and the stored
cfg_scheduler remain correct.
In `@tests/integration/test_finetuning.py`:
- Line 8: Replace the private import "from torch._prims_common import Tensor"
with the public API "from torch import Tensor" and remove any leftover marker
comments "#!CHECK" and "# CHECK" that remain in the test file (they trigger Ruff
EXE003/EXE005); search for the exact string "from torch._prims_common import
Tensor" to locate the import to change and search for the literal marker
comments "#!CHECK" and "# CHECK" to delete them.
---
Outside diff comments:
In `@luxonis_train/lightning/utils.py`:
- Around line 999-1017: The ReduceLROnPlateau branch builds a monitor key with
main_metric.node_name which can differ from the logged/formatted node name;
update the monitor construction to use the metric's formatted node name (e.g.,
main_metric.formatted_node_name) or call the existing formatting helper (e.g.,
format_node_name(main_metric.node_name)) so the monitor string becomes
f"val/metric/{formatted_node_name}/{main_metric.metric_name}" when
cfg_scheduler.name == "ReduceLROnPlateau" and mode == "max".
---
Duplicate comments:
In `@luxonis_train/lightning/luxonis_lightning.py`:
- Around line 485-501: The ReduceLROnPlateau logic currently inside
training_step (the block that checks self.trainer.is_last_batch and calls
scheduler.step(loss) or scheduler.step(value) using
nodes.main_metric_reference.compute()) must be removed from training_step and
moved to a post-validation hook where metrics have been updated (e.g.,
on_validation_epoch_end or a dedicated end-of-epoch method run after
metric.run_update/metric.reset). Specifically, stop calling
ReduceLROnPlateau.step from training_step; instead iterate the same schedulers
list in the chosen validation-end hook, compute the main metric via
nodes.main_metric_reference.compute() there (ensuring it’s not called before
validation completes), and call scheduler.step(value) for non-loss modes or
scheduler.step(loss) only when you have the correct training loss signal; keep
other schedulers' scheduler.step() behavior unchanged. Ensure code references:
training_step (remove block), nodes.main_metric_reference.compute(), and the
schedulers iteration (re-add in the validation-end hook).
---
Nitpick comments:
In `@tests/integration/test_finetuning.py`:
- Around line 552-553: Remove the leftover authoring markers by deleting the
stray "# CHECK" comments immediately above the test functions
test_inherits_omitted_optimizer_and_scheduler and
test_inherits_name_and_merges_params_when_name_is_missing; ensure no other
sentinel comments remain in that test file and run the test suite to confirm
nothing relied on those markers.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b3e28619-c0c4-4caa-9147-a38605033fc9
📒 Files selected for processing (10)
configs/README.mdluxonis_train/config/config.pyluxonis_train/config/predefined_models/base_predefined_model.pyluxonis_train/lightning/luxonis_lightning.pyluxonis_train/lightning/utils.pyluxonis_train/nodes/heads/ocr_ctc_head.pyluxonis_train/strategies/base_strategy.pyluxonis_train/strategies/triple_lr_sgd.pytests/integration/test_finetuning.pytests/unittests/test_config.py
🚧 Files skipped from review as they are similar to previous changes (5)
- tests/unittests/test_config.py
- configs/README.md
- luxonis_train/nodes/heads/ocr_ctc_head.py
- luxonis_train/config/config.py
- luxonis_train/config/predefined_models/base_predefined_model.py
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
luxonis_train/lightning/utils.py (2)
1061-1067:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winConfig mutation could cause side effects; also fix warning message formatting.
Line 1063 mutates
cfg_scheduler.paramsin place. If the sameSchedulerConfiginstance is reused or inspected later, the mutation persists unexpectedly.Additionally, line 1065-1066 has a missing space: the string should be
"'CosineAnnealingLR'. Automatically".Proposed fix
if cfg_scheduler.name == "CosineAnnealingLR": if "T_max" not in cfg_scheduler.params: - cfg_scheduler.params["T_max"] = cfg.trainer.epochs + cfg_scheduler = SchedulerConfig( + name=cfg_scheduler.name, + params={**cfg_scheduler.params, "T_max": cfg.trainer.epochs}, + ) logger.warning( - "`T_max` was not set for 'CosineAnnealingLR'" - "Automatically setting `T_max` to number of epochs." + "`T_max` was not set for 'CosineAnnealingLR'. " + "Automatically setting `T_max` to number of epochs." )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/lightning/utils.py` around lines 1061 - 1067, The code mutates cfg_scheduler.params in place which can cause unexpected side effects; instead create a shallow copy of the params dict for the scheduler (e.g., new_params = dict(cfg_scheduler.params)), set new_params["T_max"] = cfg.trainer.epochs when missing, and assign the copy back to cfg_scheduler.params so the original object isn’t mutated; also fix the logger.warning message in the block to include the missing space/period so it reads "'CosineAnnealingLR'. Automatically setting `T_max` to number of epochs." and reference the cfg_scheduler, params, cfg.trainer.epochs, and logger.warning symbols when making the change.
561-637:⚠️ Potential issue | 🟠 Major | ⚡ Quick winMissing
AIMETCallbackintegration in centralized callback builder.The new
Nodes.build_callbacksmethod does not include theAIMETCallbacklogic that exists in the module-levelbuild_callbacksfunction (lines 738-739). SinceLuxonisLightningModule.configure_callbackscallsself.nodes.build_callbacks(...), theAIMETCallbackwill never be added whencfg.exporter.aimet.activeisTrue.Proposed fix - add AIMETCallback before returning
+ if self.cfg.exporter.aimet.active: + callbacks.append(AIMETCallback()) + return callbacks🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/lightning/utils.py` around lines 561 - 637, The Nodes.build_callbacks method omits the AIMETCallback logic, so when cfg.exporter.aimet.active is True the AIMETCallback never gets added; modify the build_callbacks method to mirror the module-level build_callbacks by checking self.cfg.exporter.aimet.active (or cfg.exporter.aimet.active) and appending an AIMETCallback instance before returning callbacks (use the existing AIMETCallback class name and same construction used in the module-level build_callbacks), ensuring this happens regardless of other callback logic so LuxonisLightningModule.configure_callbacks receives the AIMETCallback.
🧹 Nitpick comments (1)
luxonis_train/lightning/utils.py (1)
702-769: Consider deprecating/removing the unused module-levelbuild_callbacks
LuxonisLightningModule.configure_callbacksreturnsself.nodes.build_callbacks(...), and there are no internal imports/usages of the module-levelbuild_callbacksinluxonis_train/lightning/utils.py(only the definitions themselves appear). If it’s not a supported public API, remove it or turn it into a deprecated wrapper/alias to avoid future divergence.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@luxonis_train/lightning/utils.py` around lines 702 - 769, The module-level function build_callbacks appears unused externally and duplicates Nodes.build_callbacks (LuxonisLightningModule.configure_callbacks already calls self.nodes.build_callbacks); either remove this module-level build_callbacks or make it a tiny deprecated wrapper that forwards to Nodes.build_callbacks while emitting a deprecation warning. Locate the function build_callbacks in luxonis_train/lightning/utils.py and either delete it and update any external references, or replace its body with a call to nodes.build_callbacks (preserving the same signature) and add a clear DeprecationWarning/log message mentioning that callers should use Nodes.build_callbacks (refer to LuxonisLightningModule.configure_callbacks for the existing call pattern).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@luxonis_train/lightning/utils.py`:
- Around line 156-165: The override of NodeWrapper.train is not calling
super().train(mode), so self.training won't be updated; modify the train method
(NodeWrapper.train) to call super().train(mode) before/after propagating mode to
self.module, each loss in self.losses.values(), each metric in
self.metrics.values(), and each visualizer in self.visualizers.values(),
ensuring self.training reflects the new state while still setting training on
the wrapped module and subcomponents.
---
Outside diff comments:
In `@luxonis_train/lightning/utils.py`:
- Around line 1061-1067: The code mutates cfg_scheduler.params in place which
can cause unexpected side effects; instead create a shallow copy of the params
dict for the scheduler (e.g., new_params = dict(cfg_scheduler.params)), set
new_params["T_max"] = cfg.trainer.epochs when missing, and assign the copy back
to cfg_scheduler.params so the original object isn’t mutated; also fix the
logger.warning message in the block to include the missing space/period so it
reads "'CosineAnnealingLR'. Automatically setting `T_max` to number of epochs."
and reference the cfg_scheduler, params, cfg.trainer.epochs, and logger.warning
symbols when making the change.
- Around line 561-637: The Nodes.build_callbacks method omits the AIMETCallback
logic, so when cfg.exporter.aimet.active is True the AIMETCallback never gets
added; modify the build_callbacks method to mirror the module-level
build_callbacks by checking self.cfg.exporter.aimet.active (or
cfg.exporter.aimet.active) and appending an AIMETCallback instance before
returning callbacks (use the existing AIMETCallback class name and same
construction used in the module-level build_callbacks), ensuring this happens
regardless of other callback logic so LuxonisLightningModule.configure_callbacks
receives the AIMETCallback.
---
Nitpick comments:
In `@luxonis_train/lightning/utils.py`:
- Around line 702-769: The module-level function build_callbacks appears unused
externally and duplicates Nodes.build_callbacks
(LuxonisLightningModule.configure_callbacks already calls
self.nodes.build_callbacks); either remove this module-level build_callbacks or
make it a tiny deprecated wrapper that forwards to Nodes.build_callbacks while
emitting a deprecation warning. Locate the function build_callbacks in
luxonis_train/lightning/utils.py and either delete it and update any external
references, or replace its body with a call to nodes.build_callbacks (preserving
the same signature) and add a clear DeprecationWarning/log message mentioning
that callers should use Nodes.build_callbacks (refer to
LuxonisLightningModule.configure_callbacks for the existing call pattern).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 82bd1192-d9f0-46b6-a48b-925b38f0e0c0
📒 Files selected for processing (5)
configs/README.mdluxonis_train/config/config.pyluxonis_train/config/predefined_models/base_predefined_model.pyluxonis_train/lightning/luxonis_lightning.pyluxonis_train/lightning/utils.py
✅ Files skipped from review due to trivial changes (1)
- configs/README.md
🚧 Files skipped from review as they are similar to previous changes (3)
- luxonis_train/config/config.py
- luxonis_train/config/predefined_models/base_predefined_model.py
- luxonis_train/lightning/luxonis_lightning.py
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #358 +/- ##
===========================================
- Coverage 93.71% 65.24% -28.48%
===========================================
Files 269 275 +6
Lines 13175 13764 +589
===========================================
- Hits 12347 8980 -3367
- Misses 828 4784 +3956
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Improves the control over the training process.
Specification
finetuningto theNodeConfigfinetuningargument toSimplePredefinedModelocr_recognition_model.yamlto use differentweight_decayfornn.Linearmodules in theOCRCTCHeadas per the reference implenentationExample:
Dependencies & Potential Impact
None / not applicable
Deployment Plan
None / not applicable
Testing & Validation
Added new
test_finetunings.pytest.Summary by CodeRabbit
New Features
Bug Fixes
Configuration
Documentation
Tests