Skip to content

Add log_train_loss_on_step toggle to EasySyntax #886

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sevmag wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Add log_train_loss_on_step toggle to EasySyntax #886

Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
470d077
Add log_train_loss_on_step toggle to EasySyntax
sevmag May 4, 2026
d4e3508
Apply suggestions from code review
christianlocatelli May 26, 2026
1152aa7
Fix botched code-review suggestion apply
sevmag May 27, 2026
4dd59b7
using canonical on_step and on_epoch kwargs
sevmag May 29, 2026
07b2c72
setting the correct attributes
sevmag May 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 14 additions & 1 deletion src/graphnet/models/easy_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ def __init__(
scheduler_class: Optional[type] = None,
scheduler_kwargs: Optional[Dict] = None,
scheduler_config: Optional[Dict] = None,
also_log_train_loss_per_step: bool = False,
) -> None:
"""Construct `StandardModel`."""
# Base class constructor
Expand All @@ -52,6 +53,7 @@ def __init__(
self._scheduler_class = scheduler_class
self._scheduler_kwargs = scheduler_kwargs or dict()
self._scheduler_config = scheduler_config or dict()
self._also_log_train_loss_per_step = also_log_train_loss_per_step

self.validate_tasks()

Expand Down Expand Up @@ -243,15 +245,26 @@ def training_step(
if isinstance(train_batch, Data):
train_batch = [train_batch]
loss = self.shared_step(train_batch, batch_idx)
batch_size = self._get_batch_size(train_batch)
self.log(
"train_loss",
loss,
batch_size=self._get_batch_size(train_batch),
batch_size=batch_size,
prog_bar=True,
on_epoch=True,
on_step=False,
Copy link
Copy Markdown
Collaborator

@Aske-Rosted Aske-Rosted May 29, 2026

Choose a reason for hiding this comment

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

And then expose the logging setting to the class. This removes duplicate code and then we don't have to define the batch_size.

Suggested change
on_epoch=True,
on_step=False,
on_epoch=log_on_epoch,
on_step=log_on_step,

Copy link
Copy Markdown
Collaborator Author

@sevmag sevmag May 29, 2026

Choose a reason for hiding this comment

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

makes sense, but how do we want to handle the logging of the val loss? log_on_epoch and log_on_step for me sounds like you log both val and train on epoch and or step, which I think could also be valid (personally I log train on log and step and val only on epoch). As long as we agree on something together I think either way is fine

Copy link
Copy Markdown
Collaborator

@Aske-Rosted Aske-Rosted May 29, 2026 •
edited
Loading

Choose a reason for hiding this comment

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

I think it is fine to have the logging of the validation and train loss behave in the same way. In principle we could separate the arguments for validation and training, but I think that is a little too many arguments, and moving towards instances where people should just create their own torch-lightning callbacks.

sync_dist=True,
)
if self._also_log_train_loss_on_step:
self.log(
"train_loss_step",
loss,
batch_size=batch_size,
prog_bar=False,
on_epoch=False,
on_step=True,
sync_dist=False,
)

current_lr = self.trainer.optimizers[0].param_groups[0]["lr"]
self.log("lr", current_lr, prog_bar=True, on_step=True)
Expand Down
Loading