fix(training): forward optimizer_kwargs and respect a custom regression loss - #430
Open
ChrisW09 wants to merge 2 commits into
Open
fix(training): forward optimizer_kwargs and respect a custom regression loss#430ChrisW09 wants to merge 2 commits into
ChrisW09 wants to merge 2 commits into
Conversation
…on loss
normalize_optimizer_kwargs kept only keys prefixed with 'optimizer_', so
the documented TrainerConfig.optimizer_kwargs={'eps': ..., 'betas': ...}
was discarded and the optimizer silently trained with torch defaults.
Unprefixed keys now pass through; lr and weight_decay are dropped in
both spellings because build_optimizer receives them explicitly.
TaskModel's num_classes==1 branch also replaced a user-supplied loss_fct
with MSELoss unconditionally, unlike the binary and multiclass branches,
contradicting the documented 'custom loss function overriding the
automatic selection'.
Fixes #415
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
model._task_model is typed ITaskModel | None, which has no .trainer; cast to Any for the end-to-end optimizer inspection. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
1 participant
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.
Fixes #415
Problems
TrainerConfig.optimizer_kwargsare silently discarded._FitMixin.fitpasses them intoTaskModel(optimizer_args=...), which routes them throughnormalize_optimizer_kwargs— a filter that keeps only keys prefixedoptimizer_. Modern unprefixed kwargs are all dropped, and the optimizer trains with torch defaults. Reproduced live:optimizer_kwargs={"eps": 1e-1, "betas": (0.5, 0.6)}→ optimizer runs witheps=1e-08, betas=(0.9, 0.999).num_classes == 1branch ofTaskModel.__init__doesself.loss_fct = nn.MSELoss()unconditionally, while the binary and multiclass branches guard withif not self.loss_fct. The docstring documentsloss_fctas "Custom loss function overriding the automatic selection".Fix
normalize_optimizer_kwargsstrips the legacy prefix when present and passes unprefixed keys through;lr/weight_decayare reserved in both spellings becausebuild_optimizerreceives them explicitly (forwarding them would raise a duplicate-keyword error).if not self.loss_fct, matching the other two.Tests
test_unprefixed_keys_pass_through,test_reserved_keys_dropped_regardless_of_prefixtest_trainer_config_optimizer_kwargs_reach_the_optimizer— end-to-end: fits a model and assertseps/betason the live optimizer param group.tests/test_lightning_module_loss.pycovering custom/default loss selection for all three task types.🤖 Generated with Claude Code