Refactor setup_training and remove test_mode (#5388)

* ref and fix call for on_pretrained_routine

* avoid failing tests

* unnecessary_call

* unnecessary call in accelerators

* tmpdir

* rm test_mode

* pep

* updates

* more ref

* Revert "more ref"

This reverts commit 5d9e95f.

* more refac

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: rohitgr7

pytorch_lightning/accelerators/accelerator.py

@@ -52,6 +52,10 @@ def __init__(self,
     def setup(self, model):
         pass
 
+    def train(self):
+        self.trainer.setup_trainer(self.trainer.model)
+        return self.train_or_test()
+
     def teardown(self):
         # Ensure if necessary all processes are finished
         self.barrier()
@@ -66,6 +70,7 @@ def train_or_test(self):
         if self.trainer.testing:
             results = self.trainer.run_test()
         else:
+            self.trainer.train_loop.setup_training()
             results = self.trainer.train()
         return results
 

pytorch_lightning/accelerators/cpu_accelerator.py

@@ -50,16 +50,6 @@ def setup(self, model):
 
         self.trainer.model = model
 
-    def train(self):
-        model = self.trainer.model
-
-        # set up training routine
-        self.trainer.train_loop.setup_training(model)
-
-        # train or test
-        results = self.train_or_test()
-        return results
-
     def _step(self, model_step: Callable, args):
         if self.trainer.amp_backend == AMPType.NATIVE:
             with torch.cuda.amp.autocast():

pytorch_lightning/accelerators/ddp2_accelerator.py

@@ -186,9 +186,6 @@ def ddp_train(self, process_idx, mp_queue, model):
 
         self.ddp_plugin.on_after_setup_optimizers(self.trainer)
 
-        # set model properties before going into wrapper
-        self.trainer.model_connector.copy_trainer_model_properties(model)
-
         # 16-bit
         model = self.trainer.precision_connector.connect(model)
 
@@ -198,8 +195,7 @@ def ddp_train(self, process_idx, mp_queue, model):
         # allow user to configure ddp
         model = self.configure_ddp(model, device_ids)
 
-        # set up training routine
-        self.trainer.train_loop.setup_training(model)
+        self.trainer.setup_trainer(model)
 
         # train or test
         results = self.train_or_test()

pytorch_lightning/accelerators/ddp_accelerator.py

@@ -285,9 +285,6 @@ def ddp_train(self, process_idx, model):
         # allow for lr schedulers as well
         self.setup_optimizers(model)
 
-        # set model properties before going into wrapper
-        self.trainer.model_connector.copy_trainer_model_properties(model)
-
         # 16-bit
         model = self.trainer.precision_connector.connect(model)
 
@@ -297,9 +294,8 @@ def ddp_train(self, process_idx, model):
         # allow user to configure ddp
         model = self.configure_ddp(model, device_ids)
 
-        # set up training routine
         self.barrier('ddp_setup')
-        self.trainer.train_loop.setup_training(model)
+        self.trainer.setup_trainer(model)
 
         # train or test
         results = self.train_or_test()

pytorch_lightning/accelerators/ddp_cpu_spawn_accelerator.py

@@ -146,9 +146,6 @@ def ddp_train(self, process_idx, mp_queue, model):
 
         self.ddp_plugin.on_after_setup_optimizers(self.trainer)
 
-        # set model properties before going into wrapper
-        self.trainer.model_connector.copy_trainer_model_properties(model)
-
         # 16-bit
         model = self.trainer.precision_connector.connect(model)
 
@@ -158,8 +155,7 @@ def ddp_train(self, process_idx, mp_queue, model):
         # allow user to configure ddp
         model = self.configure_ddp(model, device_ids)
 
-        # set up training routine
-        self.trainer.train_loop.setup_training(model)
+        self.trainer.setup_trainer(model)
 
         # train or test
         results = self.train_or_test()

pytorch_lightning/accelerators/ddp_hpc_accelerator.py

@@ -177,9 +177,6 @@ def ddp_train(self, process_idx, model):
 
         self.ddp_plugin.on_after_setup_optimizers(self.trainer)
 
-        # set model properties before going into wrapper
-        self.trainer.model_connector.copy_trainer_model_properties(model)
-
         # 16-bit
         model = self.trainer.precision_connector.connect(model)
 
@@ -189,8 +186,7 @@ def ddp_train(self, process_idx, model):
         # allow user to configure ddp
         model = self.configure_ddp(model, device_ids)
 
-        # set up training routine
-        self.trainer.train_loop.setup_training(model)
+        self.trainer.setup_trainer(model)
 
         # train or test
         results = self.train_or_test()

pytorch_lightning/accelerators/ddp_spawn_accelerator.py

@@ -161,9 +161,6 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
 
         self.ddp_plugin.on_after_setup_optimizers(self.trainer)
 
-        # set model properties before going into wrapper
-        self.trainer.model_connector.copy_trainer_model_properties(model)
-
         # 16-bit
         model = self.trainer.precision_connector.connect(model)
 
@@ -173,8 +170,7 @@ def ddp_train(self, process_idx, mp_queue, model, is_master=False, proc_offset=0
         # allow user to configure ddp
         model = self.configure_ddp(model, device_ids)
 
-        # set up training routine
-        self.trainer.train_loop.setup_training(model)
+        self.trainer.setup_trainer(model)
 
         # train or test
         results = self.train_or_test()

pytorch_lightning/accelerators/dp_accelerator.py

@@ -101,16 +101,6 @@ def __init_nvidia_apex(self, model):
 
         return model
 
-    def train(self):
-        model = self.trainer.model
-        # set up training routine
-        self.trainer.train_loop.setup_training(model)
-
-        # train or test
-        results = self.train_or_test()
-
-        return results
-
     def teardown(self):
         # replace the original fwd function
         self.trainer.model.forward = self.model_autocast_original_forward

pytorch_lightning/accelerators/gpu_accelerator.py

@@ -56,16 +56,6 @@ def setup(self, model):
 
         self.trainer.model = model
 
-    def train(self):
-        model = self.trainer.model
-
-        # set up training routine
-        self.trainer.train_loop.setup_training(model)
-
-        # train or test
-        results = self.train_or_test()
-        return results
-
     def _step(self, model_step: Callable, args):
         args[0] = self.to_device(args[0])
 

pytorch_lightning/accelerators/horovod_accelerator.py

@@ -104,8 +104,7 @@ def train(self):
                 # Synchronization will be performed explicitly following backward()
                 stack.enter_context(optimizer.skip_synchronize())
 
-            # set up training routine
-            self.trainer.train_loop.setup_training(self.trainer.model)
+            self.trainer.setup_trainer(self.trainer.model)
 
             # train or test
             results = self.train_or_test()