wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-09 11:32:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

support for native amp (#1561)

Browse Source 
* adding native amp suppport

* adding native amp suppport

* adding native amp suppport

* adding native amp suppport

* autocast

* autocast

* autocast

* autocast

* autocast

* autocast

* removed comments

* removed comments

* added state saving

* added state saving

* try install amp again

* added state saving

* drop Apex reinstall

Co-authored-by: J. Borovec <jirka.borovec@seznam.cz>
Co-authored-by: Jirka Borovec <Borda@users.noreply.github.com>
This commit is contained in:
William Falcon
2020-04-23 14:47:08 -04:00
committed by GitHub
co-authored by J. Borovec Jirka Borovec
parent 41b6cbb3ca
commit 29ebe92208
11 changed files with 100 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/trainer/evaluation_loop.py
+5 -1
View File
@@ -268,7 +268,11 @@ class TrainerEvaluationLoopMixin(ABC):
# -----------------
# RUN EVALUATION STEP
# -----------------
output = self.evaluation_forward(model, batch, batch_idx, dataloader_idx, test_mode)
if self.use_amp and self.use_native_amp:
with torch.cuda.amp.autocast():
output = self.evaluation_forward(model, batch, batch_idx, dataloader_idx, test_mode)
else:
output = self.evaluation_forward(model, batch, batch_idx, dataloader_idx, test_mode)
# on dp / ddp2 might still want to do something with the batch parts
if test_mode: