Skip to content

ENH: early stopping in models & example with RUMnet - #294

Merged
VincentAuriau merged 2 commits into
mainfrom
enh-early-stop
Dec 13, 2025
Merged

ENH: early stopping in models & example with RUMnet#294
VincentAuriau merged 2 commits into
mainfrom
enh-early-stop

Conversation

@VincentAuriau

Copy link
Copy Markdown
Collaborator

No description provided.

@gemini-code-assist

Copy link
Copy Markdown
Contributor

Summary of Changes

Hello @VincentAuriau, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the model training capabilities by integrating an early stopping mechanism into the base model, which helps prevent overfitting and optimizes training duration. It also refines the logging of validation metrics for clarity and provides a practical example of early stopping in action within the RUMnet model's notebook.

Highlights

  • Early Stopping Implementation: Introduced early stopping logic within the fit method of the base model, allowing training to halt prematurely if self.stop_training is activated.
  • Validation Loss Metric Renaming: Renamed the test_loss key to val_loss in the losses_history dictionary to standardize the naming convention for validation metrics.
  • RUMnet Early Stopping Example: Added a new section to the RUMnet Jupyter notebook demonstrating how to configure and use tf.keras.callbacks.EarlyStopping with the model.
  • Notebook Metadata Updates: Updated the Jupyter notebook's kernel display name to 'choice_learn' and the Python version to '3.12.11'.
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces early stopping functionality to the base model and adds a corresponding example in the RUMnet notebook. The implementation in base_model.py correctly uses the self.stop_training flag set by Keras callbacks.

My review identifies a couple of issues:

  • There's a bug where the early stopping check is missing for the training path where sample_weight is not used, which would cause early stopping to fail.
  • The messages printed when early stopping is triggered are unconditional and should respect the verbose parameter.

The notebook changes are good, with a clear example of how to use the new early stopping feature and a fix to use val_loss consistently. After addressing the feedback, this will be a solid enhancement.

Comment on lines +349 to +351
if self.stop_training:
print("Training stopped with early stopping taking effect")
break

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

This early stopping check is great, but it's missing from the training loop in the else block below (when sample_weight is None), which will prevent early stopping from working in that case. Please add a similar check there.

Also, this print statement is unconditional. It should only be executed if verbose > 0 to avoid unexpected output when verbose is set to 0.

Suggested change
if self.stop_training:
print("Training stopped with early stopping taking effect")
break
if self.stop_training:
if verbose > 0:
print("Training stopped with early stopping taking effect")
break

Comment on lines 444 to 446
if self.stop_training:
print("Early Stopping taking effect")
print("Training stopped with early stopping taking effect")
break

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Similar to the check within the batch loop, this print statement is unconditional. It should be wrapped in an if verbose > 0: check to respect the verbosity level set for the fit method.

Suggested change
if self.stop_training:
print("Early Stopping taking effect")
print("Training stopped with early stopping taking effect")
break
if self.stop_training:
if verbose > 0:
print("Training stopped with early stopping taking effect")
break

@github-actions

github-actions Bot commented Dec 12, 2025

Copy link
Copy Markdown
Contributor

Coverage

Coverage Report for Python 3.9
FileStmtsMissCoverMissing
choice_learn
   __init__.py20100% 
   tf_ops.py62198%283
choice_learn/basket_models
   __init__.py40100% 
   alea_carta.py1482285%86–90, 92–96, 98–102, 106, 109, 131, 159, 308, 431–455
   base_basket_model.py2352789%111–112, 123, 141, 185, 255, 377, 485, 585–587, 676, 762, 772, 822–830, 891–894, 934–935
   basic_attention_model.py89496%424, 427, 433, 440
   self_attention_model.py133993%71, 73, 75, 450–454, 651
   shopper.py184995%130, 159, 325, 345, 360, 363, 377, 489, 618
choice_learn/basket_models/data
   __init__.py20100% 
   basket_dataset.py1903084%74–77, 295–297, 407, 540–576, 636, 658–661, 700–705, 790–801, 849
   preprocessing.py947817%43–45, 128–364
choice_learn/basket_models/datasets
   __init__.py30100% 
   bakery.py38392%47, 51, 61
   synthetic_dataset.py81693%62, 194–199, 247
choice_learn/basket_models/utils
   __init__.py00100% 
   permutation.py22195%37
choice_learn/data
   __init__.py30100% 
   choice_dataset.py6493395%198, 250, 283, 421, 463–464, 589, 724, 738, 840, 842, 937, 957–961, 1140, 1159–1161, 1179–1181, 1209, 1214, 1223, 1240, 1281, 1293, 1307, 1346, 1361, 1366, 1395, 1408, 1443–1444
   indexer.py2412390%20, 31, 45, 60–67, 202–204, 219–230, 265, 291, 582
   storage.py161696%22, 33, 51, 56, 61, 71
   store.py72720%3–275
choice_learn/datasets
   __init__.py40100% 
   base.py400599%42–43, 153–154, 714
   expedia.py1028319%37–301
   tafeng.py490100% 
choice_learn/datasets/data
   __init__.py00100% 
choice_learn/models
   __init__.py14286%15–16
   base_model.py3092592%144, 186, 283, 291, 297, 306, 346, 350–351, 356, 385, 389–390, 407, 438–439, 448–449, 550, 552, 568, 572, 574, 697–698
   baseline_models.py490100% 
   conditional_logit.py2692690%49, 52, 54, 85, 88, 91–95, 98–102, 136, 206, 212–216, 351, 388, 445, 520–526, 651, 685, 822, 826
   halo_mnl.py124298%186, 374
   latent_class_base_model.py2863986%55–61, 273–279, 288, 325–330, 497–500, 605, 624, 665–701, 715, 720, 751–752, 774–775, 869–870, 974
   latent_class_mnl.py62690%257–261, 296
   learning_mnl.py67396%157, 182, 188
   nested_logit.py2911296%55, 77, 160, 269, 351, 484, 530, 600, 679, 848, 900, 904
   reslogit.py132695%285, 360, 369, 374, 382, 432
   rumnet.py236399%748–751, 982
   simple_mnl.py139696%167, 275, 347, 355, 357, 359
   tastenet.py94397%142, 180, 188
choice_learn/toolbox
   __init__.py00100% 
   assortment_optimizer.py27678%28–30, 93–95, 160–162
   gurobi_opt.py2362360%3–675
   or_tools_opt.py2301195%103, 107, 296–305, 315, 319, 607, 611
choice_learn/utils
   metrics.py854349%74, 126–130, 147–166, 176, 190–199, 211–232, 242
TOTAL561884185% 

Tests Skipped Failures Errors Time
221 0 💤 0 ❌ 0 🔥 7m 56s ⏱️

@github-actions

github-actions Bot commented Dec 12, 2025

Copy link
Copy Markdown
Contributor

Coverage

Coverage Report for Python 3.10
FileStmtsMissCoverMissing
choice_learn
   __init__.py20100% 
   tf_ops.py62198%283
choice_learn/basket_models
   __init__.py40100% 
   alea_carta.py1482285%86–90, 92–96, 98–102, 106, 109, 131, 159, 308, 431–455
   base_basket_model.py2352789%111–112, 123, 141, 185, 255, 377, 485, 585–587, 676, 762, 772, 822–830, 891–894, 934–935
   basic_attention_model.py89496%424, 427, 433, 440
   self_attention_model.py133993%71, 73, 75, 450–454, 651
   shopper.py184995%130, 159, 325, 345, 360, 363, 377, 489, 618
choice_learn/basket_models/data
   __init__.py20100% 
   basket_dataset.py1903084%74–77, 295–297, 407, 540–576, 636, 658–661, 700–705, 790–801, 849
   preprocessing.py947817%43–45, 128–364
choice_learn/basket_models/datasets
   __init__.py30100% 
   bakery.py38392%47, 51, 61
   synthetic_dataset.py81693%62, 194–199, 247
choice_learn/basket_models/utils
   __init__.py00100% 
   permutation.py22195%37
choice_learn/data
   __init__.py30100% 
   choice_dataset.py6493395%198, 250, 283, 421, 463–464, 589, 724, 738, 840, 842, 937, 957–961, 1140, 1159–1161, 1179–1181, 1209, 1214, 1223, 1240, 1281, 1293, 1307, 1346, 1361, 1366, 1395, 1408, 1443–1444
   indexer.py2412390%20, 31, 45, 60–67, 202–204, 219–230, 265, 291, 582
   storage.py161696%22, 33, 51, 56, 61, 71
   store.py72720%3–275
choice_learn/datasets
   __init__.py40100% 
   base.py400599%42–43, 153–154, 714
   expedia.py1028319%37–301
   tafeng.py490100% 
choice_learn/datasets/data
   __init__.py00100% 
choice_learn/models
   __init__.py14286%15–16
   base_model.py3092592%144, 186, 283, 291, 297, 306, 346, 350–351, 356, 385, 389–390, 407, 438–439, 448–449, 550, 552, 568, 572, 574, 697–698
   baseline_models.py490100% 
   conditional_logit.py2692690%49, 52, 54, 85, 88, 91–95, 98–102, 136, 206, 212–216, 351, 388, 445, 520–526, 651, 685, 822, 826
   halo_mnl.py124298%186, 374
   latent_class_base_model.py2863986%55–61, 273–279, 288, 325–330, 497–500, 605, 624, 665–701, 715, 720, 751–752, 774–775, 869–870, 974
   latent_class_mnl.py62690%257–261, 296
   learning_mnl.py67396%157, 182, 188
   nested_logit.py2911296%55, 77, 160, 269, 351, 484, 530, 600, 679, 848, 900, 904
   reslogit.py132695%285, 360, 369, 374, 382, 432
   rumnet.py236399%748–751, 982
   simple_mnl.py139696%167, 275, 347, 355, 357, 359
   tastenet.py94397%142, 180, 188
choice_learn/toolbox
   __init__.py00100% 
   assortment_optimizer.py27678%28–30, 93–95, 160–162
   gurobi_opt.py2382380%3–675
   or_tools_opt.py2301195%103, 107, 296–305, 315, 319, 607, 611
choice_learn/utils
   metrics.py854349%74, 126–130, 147–166, 176, 190–199, 211–232, 242
TOTAL562084385% 

Tests Skipped Failures Errors Time
221 0 💤 0 ❌ 0 🔥 8m 36s ⏱️

@github-actions

github-actions Bot commented Dec 12, 2025

Copy link
Copy Markdown
Contributor

Coverage

Coverage Report for Python 3.11
FileStmtsMissCoverMissing
choice_learn
   __init__.py20100% 
   tf_ops.py62198%283
choice_learn/basket_models
   __init__.py40100% 
   alea_carta.py1482285%86–90, 92–96, 98–102, 106, 109, 131, 159, 308, 431–455
   base_basket_model.py2352789%111–112, 123, 141, 185, 255, 377, 485, 585–587, 676, 762, 772, 822–830, 891–894, 934–935
   basic_attention_model.py89496%424, 427, 433, 440
   self_attention_model.py133993%71, 73, 75, 450–454, 651
   shopper.py184995%130, 159, 325, 345, 360, 363, 377, 489, 618
choice_learn/basket_models/data
   __init__.py20100% 
   basket_dataset.py1903084%74–77, 295–297, 407, 540–576, 636, 658–661, 700–705, 790–801, 849
   preprocessing.py947817%43–45, 128–364
choice_learn/basket_models/datasets
   __init__.py30100% 
   bakery.py38392%47, 51, 61
   synthetic_dataset.py81693%62, 194–199, 247
choice_learn/basket_models/utils
   __init__.py00100% 
   permutation.py22195%37
choice_learn/data
   __init__.py30100% 
   choice_dataset.py6493395%198, 250, 283, 421, 463–464, 589, 724, 738, 840, 842, 937, 957–961, 1140, 1159–1161, 1179–1181, 1209, 1214, 1223, 1240, 1281, 1293, 1307, 1346, 1361, 1366, 1395, 1408, 1443–1444
   indexer.py2412390%20, 31, 45, 60–67, 202–204, 219–230, 265, 291, 582
   storage.py161696%22, 33, 51, 56, 61, 71
   store.py72720%3–275
choice_learn/datasets
   __init__.py40100% 
   base.py400599%42–43, 153–154, 714
   expedia.py1028319%37–301
   tafeng.py490100% 
choice_learn/datasets/data
   __init__.py00100% 
choice_learn/models
   __init__.py14286%15–16
   base_model.py3092592%144, 186, 283, 291, 297, 306, 346, 350–351, 356, 385, 389–390, 407, 438–439, 448–449, 550, 552, 568, 572, 574, 697–698
   baseline_models.py490100% 
   conditional_logit.py2692690%49, 52, 54, 85, 88, 91–95, 98–102, 136, 206, 212–216, 351, 388, 445, 520–526, 651, 685, 822, 826
   halo_mnl.py124298%186, 374
   latent_class_base_model.py2863986%55–61, 273–279, 288, 325–330, 497–500, 605, 624, 665–701, 715, 720, 751–752, 774–775, 869–870, 974
   latent_class_mnl.py62690%257–261, 296
   learning_mnl.py67396%157, 182, 188
   nested_logit.py2911296%55, 77, 160, 269, 351, 484, 530, 600, 679, 848, 900, 904
   reslogit.py132695%285, 360, 369, 374, 382, 432
   rumnet.py236399%748–751, 982
   simple_mnl.py139696%167, 275, 347, 355, 357, 359
   tastenet.py94397%142, 180, 188
choice_learn/toolbox
   __init__.py00100% 
   assortment_optimizer.py27678%28–30, 93–95, 160–162
   gurobi_opt.py2382380%3–675
   or_tools_opt.py2301195%103, 107, 296–305, 315, 319, 607, 611
choice_learn/utils
   metrics.py854349%74, 126–130, 147–166, 176, 190–199, 211–232, 242
TOTAL562084385% 

Tests Skipped Failures Errors Time
221 0 💤 0 ❌ 0 🔥 9m 38s ⏱️

@github-actions

github-actions Bot commented Dec 12, 2025

Copy link
Copy Markdown
Contributor

Coverage

Coverage Report for Python 3.12
FileStmtsMissCoverMissing
choice_learn
   __init__.py20100% 
   tf_ops.py62198%283
choice_learn/basket_models
   __init__.py40100% 
   alea_carta.py1482285%86–90, 92–96, 98–102, 106, 109, 131, 159, 308, 431–455
   base_basket_model.py2352789%111–112, 123, 141, 185, 255, 377, 485, 585–587, 676, 762, 772, 822–830, 891–894, 934–935
   basic_attention_model.py89496%424, 427, 433, 440
   self_attention_model.py133993%71, 73, 75, 450–454, 651
   shopper.py184995%130, 159, 325, 345, 360, 363, 377, 489, 618
choice_learn/basket_models/data
   __init__.py20100% 
   basket_dataset.py1903084%74–77, 295–297, 407, 540–576, 636, 658–661, 700–705, 790–801, 849
   preprocessing.py947817%43–45, 128–364
choice_learn/basket_models/datasets
   __init__.py30100% 
   bakery.py38392%47, 53, 61
   synthetic_dataset.py81693%62, 194–199, 247
choice_learn/basket_models/utils
   __init__.py00100% 
   permutation.py22195%37
choice_learn/data
   __init__.py30100% 
   choice_dataset.py6493395%198, 250, 283, 421, 463–464, 589, 724, 738, 840, 842, 937, 957–961, 1140, 1159–1161, 1179–1181, 1209, 1214, 1223, 1240, 1281, 1293, 1307, 1346, 1361, 1366, 1395, 1408, 1443–1444
   indexer.py2412390%20, 31, 45, 60–67, 202–204, 219–230, 265, 291, 582
   storage.py161696%22, 33, 51, 56, 61, 71
   store.py72720%3–275
choice_learn/datasets
   __init__.py40100% 
   base.py400599%42–43, 153–154, 714
   expedia.py1028319%37–301
   tafeng.py490100% 
choice_learn/datasets/data
   __init__.py00100% 
choice_learn/models
   __init__.py14286%15–16
   base_model.py3092592%144, 186, 283, 291, 297, 306, 346, 350–351, 356, 385, 389–390, 407, 438–439, 448–449, 550, 552, 568, 572, 574, 697–698
   baseline_models.py490100% 
   conditional_logit.py2692690%49, 52, 54, 85, 88, 91–95, 98–102, 136, 206, 212–216, 351, 388, 445, 520–526, 651, 685, 822, 826
   halo_mnl.py124298%186, 374
   latent_class_base_model.py2863986%55–61, 273–279, 288, 325–330, 497–500, 605, 624, 665–701, 715, 720, 751–752, 774–775, 869–870, 974
   latent_class_mnl.py62690%257–261, 296
   learning_mnl.py67396%157, 182, 188
   nested_logit.py2911296%55, 77, 160, 269, 351, 484, 530, 600, 679, 848, 900, 904
   reslogit.py132695%285, 360, 369, 374, 382, 432
   rumnet.py236399%748–751, 982
   simple_mnl.py139696%167, 275, 347, 355, 357, 359
   tastenet.py94397%142, 180, 188
choice_learn/toolbox
   __init__.py00100% 
   assortment_optimizer.py27678%28–30, 93–95, 160–162
   gurobi_opt.py2382380%3–675
   or_tools_opt.py2301195%103, 107, 296–305, 315, 319, 607, 611
choice_learn/utils
   metrics.py854349%74, 126–130, 147–166, 176, 190–199, 211–232, 242
TOTAL562084385% 

Tests Skipped Failures Errors Time
221 0 💤 0 ❌ 0 🔥 9m 19s ⏱️

@VincentAuriau
VincentAuriau merged commit 0f0ce10 into main Dec 13, 2025
8 checks passed
@VincentAuriau
VincentAuriau deleted the enh-early-stop branch December 13, 2025 14:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant