Skip to content

Follow-up changes to #10575 #10957

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

Merged
merged 10 commits into from
Dec 7, 2021
Merged

Follow-up changes to #10575 #10957

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
d37e8da
Follow-up changes to #10575
carmocca Dec 6, 2021
ab7a76e
Fix types
carmocca Dec 6, 2021
ac07666
Fix types
carmocca Dec 6, 2021
b34c77a
Progress
carmocca Dec 6, 2021
4f91270
Any return
carmocca Dec 6, 2021
c3513a0
Missing hooks
carmocca Dec 6, 2021
d8a4158
Fix one tests, profilers failing
carmocca Dec 6, 2021
9559167
Remove dupe profiler call
carmocca Dec 6, 2021
96417c0
Merge branch 'master' into refactor/follow-up-call-hook
awaelchli Dec 6, 2021
30244f4
update _call_accelerator_hook to _call_ttp_hook
awaelchli Dec 6, 2021
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
5 changes: 2 additions & 3 deletions pytorch_lightning/loops/epoch/training_epoch_loop.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,11 +180,10 @@ def advance(self, data_fetcher: AbstractDataFetcher) -> None: # type: ignore[ov

# hook
self.trainer._call_callback_hooks("on_train_batch_start", batch, batch_idx, **extra_kwargs)
model_response = self.trainer._call_lightning_module_hook(
response = self.trainer._call_lightning_module_hook(
"on_train_batch_start", batch, batch_idx, **extra_kwargs
)
ttp_response = self.trainer._call_ttp_hook("on_train_batch_start", batch, batch_idx, **extra_kwargs)
response = ttp_response if model_response is None else model_response
self.trainer._call_ttp_hook("on_train_batch_start", batch, batch_idx, **extra_kwargs)
if response == -1:
self.batch_progress.increment_processed()
raise StopIteration
Expand Down
8 changes: 7 additions & 1 deletion pytorch_lightning/trainer/trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1537,14 +1537,20 @@ def _call_accelerator_hook(
*args: Any,
**kwargs: Any,
) -> Optional[Any]:
self.lightning_module._current_fx_name = hook_name
pl_module = self.lightning_module
prev_fx_name = pl_module._current_fx_name
pl_module._current_fx_name = hook_name

fn = getattr(self.accelerator, hook_name)
if not callable(fn):
return None

with self.profiler.profile(hook_name):
output = fn(*args, **kwargs)

# restore current_fx when nested context
pl_module._current_fx_name = prev_fx_name

return output

@staticmethod
Expand Down