ref: moving train loop to own object 2/n (intermediate steps) (Lightning-AI#3314)

* ref: moving train loop to own object 2/n (intermediate steps)

* ref: moving train loop to own object 2/n (intermediate steps)

Author: williamFalcon

pytorch_lightning/trainer/training_loop.py

@@ -341,33 +341,16 @@ def run_sanity_check(self, *args):
     def train(self):
         self.run_sanity_check(self.get_model())
 
-        # TODO: shrink
-        # clear cache before training
-        if self.on_gpu and self.root_gpu is not None:
-            # use context because of:
-            # https://discuss.pytorch.org/t/out-of-memory-when-i-use-torch-cuda-empty-cache/57898
-            with torch.cuda.device(f'cuda:{self.root_gpu}'):
-                torch.cuda.empty_cache()
-
-        # get model
-        model = self.get_model()
-
         # enable train mode
+        model = self.get_model()
         model.train()
-
-        # enable gradients
         torch.set_grad_enabled(True)
 
-        # load data
-        # if reload_dataloaders_every_epoch, this is moved to the epoch loop
-        if not self.reload_dataloaders_every_epoch:
-            self.reset_train_dataloader(model)
-
-        if self.val_dataloaders is None and not self.reload_dataloaders_every_epoch:
-            self.reset_val_dataloader(model)
+        # reload data when needed
+        self.train_loop.reset_train_val_dataloaders(model)
 
         # hook
-        self.call_hook('on_train_start')
+        self.train_loop.on_train_start()
 
         try:
             # run all epochs
@@ -399,7 +382,9 @@ def train(self):
                 self.run_training_epoch()
 
                 if self.max_steps and self.max_steps <= self.global_step:
-                    self.run_training_teardown()
+
+                    # hook
+                    self.train_loop.on_train_end()
                     return
 
                 # update LR schedulers
@@ -411,14 +396,15 @@ def train(self):
 
                 if self.should_stop:
                     if (met_min_epochs and met_min_steps):
-                        self.run_training_teardown()
+                        self.train_loop.on_train_end()
                         return
                     else:
                         log.info('Trainer was signaled to stop but required minimum epochs'
                                  f' ({self.min_epochs}) or minimum steps ({self.min_steps}) has'
                                  ' not been met. Training will continue...')
 
-            self.run_training_teardown()
+            # hook
+            self.train_loop.on_train_end()
 
         except KeyboardInterrupt:
             rank_zero_warn('Detected KeyboardInterrupt, attempting graceful shutdown...')
@@ -429,7 +415,8 @@ def train(self):
                 self._state = TrainerState.INTERRUPTED
                 self.on_keyboard_interrupt()
 
-                self.run_training_teardown()
+                # hook
+                self.train_loop.on_train_end()
 
     def run_training_epoch(self):
 
@@ -1053,47 +1040,6 @@ def optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens)
         )
         return result
 
-
-    # @atexit.register
-    def run_training_teardown(self):
-        if hasattr(self, '_teardown_already_run') and self._teardown_already_run:
-            return
-
-        self._teardown_already_run = True
-
-        # Save latest checkpoint
-        log.info('Saving latest checkpoint..')
-        self.check_checkpoint_callback(should_check_val=False)
-
-        # Train end events
-        with self.profiler.profile('on_train_end'):
-            # callbacks
-            self.on_train_end()
-            # model hooks
-            if self.is_function_implemented('on_train_end'):
-                self.get_model().on_train_end()
-
-        if self.logger is not None:
-            self.logger.finalize("success")
-
-        # summarize profile results
-        if self.global_rank == 0:
-            self.profiler.describe()
-
-        if self.global_rank == 0:
-            for proc in self.interactive_ddp_procs:
-                subprocess.Popen.kill(proc)
-
-        # clean up dist group
-        if self.use_ddp or self.use_ddp2:
-            torch_distrib.destroy_process_group()
-
-        # clear mem
-        if self.on_gpu:
-            model = self.get_model()
-            model.cpu()
-            torch.cuda.empty_cache()
-
     def build_train_args(self, batch, batch_idx, opt_idx, hiddens):
         # enable not needing to add opt_idx to training_step
         args = [batch, batch_idx]

pytorch_lightning/trainer/training_loop_temp.py

@@ -1,6 +1,12 @@
-from pytorch_lightning.trainer.supporters import Accumulator
+import subprocess
 import numpy as np
+import torch
+import torch.distributed as torch_distrib
+from pytorch_lightning.utilities.model_utils import is_overridden
+from pytorch_lightning.trainer.supporters import Accumulator
+from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.core.step_result import Result
+from pytorch_lightning import _logger as log
 
 
 class TrainLoop:
@@ -10,12 +16,69 @@ def __init__(self, trainer):
         self.should_check_val = False
         self.early_stopping_accumulator = None
         self.checkpoint_accumulator = None
+        self._teardown_already_run = False
 
     @property
     def num_optimizers(self):
         num_optimizers = len(self.get_optimizers_iterable())
         return num_optimizers
 
+    def on_train_start(self):
+        # clear cache before training
+        if self.trainer.on_gpu and self.trainer.root_gpu is not None:
+            # use context because of:
+            # https://discuss.pytorch.org/t/out-of-memory-when-i-use-torch-cuda-empty-cache/57898
+            with torch.cuda.device(f'cuda:{self.trainer.root_gpu}'):
+                torch.cuda.empty_cache()
+
+        # hook
+        self.trainer.call_hook('on_train_start')
+
+    def on_train_end(self):
+        if self._teardown_already_run:
+            return
+
+        self._teardown_already_run = True
+
+        # Save latest checkpoint
+        log.info('Saving latest checkpoint..')
+        self.check_checkpoint_callback(should_check_val=False)
+
+        # hook
+        self.trainer.call_hook('on_train_end')
+
+        # kill loggers
+        if self.trainer.logger is not None:
+            self.trainer.logger.finalize("success")
+
+        # summarize profile results
+        if self.trainer.global_rank == 0:
+            self.trainer.profiler.describe()
+
+        if self.trainer.global_rank == 0:
+            for proc in self.trainer.interactive_ddp_procs:
+                subprocess.Popen.kill(proc)
+
+        # clean up dist group
+        if self.trainer.use_ddp or self.trainer.use_ddp2:
+            torch_distrib.destroy_process_group()
+
+        # clear mem
+        if self.trainer.on_gpu:
+            model = self.trainer.get_model()
+            model.cpu()
+            torch.cuda.empty_cache()
+
+    def check_checkpoint_callback(self, should_check_val):
+        model = self.trainer.get_model()
+
+        # when no val loop is present or fast-dev-run still need to call checkpoints
+        # TODO bake this logic into the checkpoint callback
+        should_activate = not is_overridden('validation_step', model) and not should_check_val
+        if should_activate:
+            checkpoint_callbacks = [c for c in self.trainer.callbacks if isinstance(c, ModelCheckpoint)]
+            [c.on_validation_end(self.trainer, model) for c in checkpoint_callbacks]
+
     def on_train_epoch_start(self):
         # hook
         self.trainer.call_hook('on_epoch_start')
@@ -28,6 +91,7 @@ def on_train_epoch_start(self):
         self.early_stopping_accumulator = Accumulator()
         self.checkpoint_accumulator = Accumulator()
 
+
     def on_train_batch_end(self, epoch_output, epoch_end_outputs, batch, batch_idx, dataloader_idx):
         # figure out what to track for epoch end
         self.track_epoch_end_reduce_metrics(epoch_output, epoch_end_outputs)
@@ -36,6 +100,13 @@ def on_train_batch_end(self, epoch_output, epoch_end_outputs, batch, batch_idx,
         self.trainer.call_hook('on_batch_end')
         self.trainer.call_hook('on_train_batch_end', batch, batch_idx, dataloader_idx)
 
+    def reset_train_val_dataloaders(self, model):
+        if not self.trainer.reload_dataloaders_every_epoch:
+            self.trainer.reset_train_dataloader(model)
+
+        if self.trainer.val_dataloaders is None and not self.trainer.reload_dataloaders_every_epoch:
+            self.trainer.reset_val_dataloader(model)
+
     def track_epoch_end_reduce_metrics(self, epoch_output, epoch_end_outputs):
         # track the outputs to reduce at the end of the epoch
         for opt_idx, opt_outputs in enumerate(epoch_end_outputs):