Added max/min number of steps in Trainer

pytorch_lightning/trainer/trainer.py

@@ -71,6 +71,7 @@ def __init__(
  min_nb_epochs=None, # backward compatible, todo: remove in v0.8.0
  max_epochs=1000,
  min_epochs=1,
+ max_steps=None,
  train_percent_check=1.0,
  val_percent_check=1.0,
  test_percent_check=1.0,
@@ -289,6 +290,15 @@ def __init__(
  .. deprecated:: 0.5.0
  Use `min_nb_epochs` instead. Will remove 0.8.0.
 
+ max_steps (int): Stop training after this number of steps.
+ Example::
+
+ # default used by the Trainer (disabled)
+ trainer = Trainer(max_steps=None)
+
+ # Stop after 100 steps (batches)
+ trainer = Trainer(max_steps=100)
+
  train_percent_check (int): How much of training dataset to check.
  Useful when debugging or testing something that happens at the end of an epoch.
  Example::
@@ -505,6 +515,8 @@ def __init__(
  min_epochs = min_nb_epochs
  self.min_epochs = min_epochs
 
+ self.max_steps = max_steps
+
  # Backward compatibility
  if nb_sanity_val_steps is not None:
  warnings.warn("`nb_sanity_val_steps` has renamed to `num_sanity_val_steps` since v0.5.0"

pytorch_lightning/trainer/training_loop.py

@@ -345,6 +345,11 @@ def train(self):
  raise MisconfigurationException(m)
  self.reduce_lr_on_plateau_scheduler.step(val_loss, epoch=self.current_epoch)
 
+ if self.max_steps is not None and self.max_steps == self.global_step:
+ self.main_progress_bar.close()
+ model.on_train_end()
+ return
+
  # early stopping
  met_min_epochs = epoch >= self.min_epochs - 1
  if (self.enable_early_stop and not self.disable_validation and is_val_epoch and
@@ -421,6 +426,10 @@ def run_training_epoch(self):
  self.global_step += 1
  self.total_batch_idx += 1
 
+ # max steps reached, end training
+ if self.max_steps is not None and self.max_steps == self.global_step:
+ break
+
  # end epoch early
  # stop when the flag is changed or we've gone past the amount
  # requested in the batches