Skip to content

Commit

Permalink
Use PyTorch API logging for Lightning Trainer (#6771)
Browse files Browse the repository at this point in the history
* Update trainer.py

* Update trainer.py

* Update trainer.py
  • Loading branch information
ananthsub authored Apr 15, 2021
1 parent f29ecbf commit 4c07ab5
Showing 1 changed file with 10 additions and 1 deletion.
11 changes: 10 additions & 1 deletion pytorch_lightning/trainer/trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -295,7 +295,7 @@ def __init__(
"""
super().__init__()

Trainer._log_api_event("init")
distributed_backend = distributed_backend or accelerator

# init connectors
Expand Down Expand Up @@ -416,6 +416,7 @@ def fit(
If the model has a predefined val_dataloaders method this will be skipped
"""
Trainer._log_api_event("fit")
# we reuse fit for other functions. When already set, it shouldn't be modified.
if not self.state.running:
self.state = TrainerState.FITTING
Expand Down Expand Up @@ -881,6 +882,7 @@ def validate(
# --------------------
# SETUP HOOK
# --------------------
Trainer._log_api_event("validate")
self.verbose_evaluate = verbose

self.state = TrainerState.VALIDATING
Expand Down Expand Up @@ -943,6 +945,7 @@ def test(
# --------------------
# SETUP HOOK
# --------------------
Trainer._log_api_event("test")
self.verbose_evaluate = verbose

self.state = TrainerState.TESTING
Expand Down Expand Up @@ -1039,6 +1042,7 @@ def predict(
# SETUP HOOK
# --------------------
# If you supply a datamodule you can't supply dataloaders
Trainer._log_api_event("predict")

model = model or self.lightning_module

Expand Down Expand Up @@ -1084,6 +1088,7 @@ def tune(
If the model has a predefined val_dataloaders method this will be skipped
"""
Trainer._log_api_event("tune")
self.state = TrainerState.TUNING
self.tuning = True

Expand Down Expand Up @@ -1174,3 +1179,7 @@ def call_hook(self, hook_name, *args, **kwargs):
if not skip:
self._cache_logged_metrics()
return output

@staticmethod
def _log_api_event(event: str) -> None:
torch._C._log_api_usage_once("lightning.trainer." + event)

0 comments on commit 4c07ab5

Please sign in to comment.