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

Move global_step incrementing (#412)

Browse Source 
* Move global_step incrementing to the end of a batch loop, per https://github.com/williamFalcon/pytorch-lightning/issues/411

* Move met_batch_limit condition to the end

* cleanup whitespace

* Update train_loop_mixin.py
This commit is contained in:
David Kossnick
2019-10-23 06:11:18 -04:00
committed by William Falcon
parent 5afae59715
commit 56fa2075a5
1 changed files with 12 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/trainer/train_loop_mixin.py
+12 -8
View File
@@ -74,23 +74,17 @@ class TrainerTrainLoopMixin(object):
# run epoch
for batch_nb, batch in enumerate(self.get_train_dataloader()):
self.batch_nb = batch_nb
self.global_step += 1
model = self.get_model()
model.global_step = self.global_step
# stop when the flag is changed or we've gone past the amount
# requested in the batches
self.total_batch_nb += 1
met_batch_limit = batch_nb >= self.nb_training_batches
if met_batch_limit:
break
# ---------------
# RUN TRAIN STEP
# ---------------
output = self.run_training_batch(batch, batch_nb)
batch_result, grad_norm_dic, batch_step_metrics = output
# when returning -1 from train_step, we end epoch early
early_stop_epoch = batch_result == -1
# ---------------
@@ -116,10 +110,20 @@ class TrainerTrainLoopMixin(object):
# logs user requested information to logger
self.log_metrics(batch_step_metrics, grad_norm_dic)
self.global_step += 1
self.total_batch_nb += 1
# end epoch early
# stop when the flag is changed or we've gone past the amount
# requested in the batches
if early_stop_epoch or self.fast_dev_run:
break
# stop epoch if we limited nb batches
met_batch_limit = batch_nb >= self.nb_training_batches
if met_batch_limit:
break
# epoch end hook
if self.is_function_implemented('on_epoch_end'):
model = self.get_model()