Skip to content
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

Jump to bottom

[Fix] Ensure we set the eval/train flag correctly on accelerator model #6877

Merged
merged 10 commits into from
Apr 8, 2021
Merged

[Fix] Ensure we set the eval/train flag correctly on accelerator model #6877

Changes from 2 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
f9fdcea
Ensure we move the model to eval mode before running evaluation
Apr 7, 2021
5555ca8
Ensure we set the flag appropriately across all stages
Apr 7, 2021
1ecdd23
Add test, move hooks logic
Apr 8, 2021
4b10fdb
Apply same fix to the validate loop
Apr 8, 2021
d642932
Update pytorch_lightning/trainer/trainer.py
Apr 8, 2021
c3abba5
Fix function name
Apr 8, 2021
42b99e3
Fix order, add predict
Apr 8, 2021
778e026
Shorten the name
Apr 8, 2021
d6b8607
Fix input dm, drop duplicate on predict start hook call, as it's call…
Apr 8, 2021
b8d473f
Use hook, remove double call
Apr 8, 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
7 changes: 5 additions & 2 deletions pytorch_lightning/trainer/trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -582,11 +582,11 @@ def run_train(self) -> None:
self.checkpoint_connector.has_trained = False

# enable train mode
model = self.lightning_module
model.train()
self.model.train()
torch.set_grad_enabled(True)

# reload data when needed
model = self.lightning_module
self.train_loop.reset_train_val_dataloaders(model)

# hook
Expand Down Expand Up @@ -665,6 +665,8 @@ def run_evaluation(self, on_epoch=False):
# ref model
model = self.lightning_module
model.zero_grad()

self.model.eval()
SeanNaren marked this conversation as resolved.
Show resolved Hide resolved
torch.set_grad_enabled(False)

# hook
Expand Down Expand Up @@ -786,6 +788,7 @@ def run_predict(self):

# enable eval mode + no grads
self.predict_loop.on_predict_model_eval()
SeanNaren marked this conversation as resolved.
Show resolved Hide resolved
self.model.eval()
model.zero_grad()
torch.set_grad_enabled(False)

Expand Down