release v0.2.2

pytorch_lightning/models/trainer.py

@@ -288,17 +288,6 @@ class Trainer(TrainerIO):
     # MODEL TRAINING
     # -----------------------------
     def fit(self, model):
-        # CHOOSE OPTIMIZER
-        # filter out the weights that were done on gpu so we can load on good old cpus
-        self.optimizers = model.configure_optimizers()
-
-        # run through amp wrapper
-        if self.use_amp:
-            # An example
-            model, optimizers = amp.initialize(
-                model, self.optimizers, opt_level=self.amp_level,
-            )
-            self.optimizers = optimizers
 
         # when using gpus, first thing we do is spawn a new process between each worker
         # applies to single gpu, multi-gpu and multi-nodes
@@ -306,6 +295,18 @@ class Trainer(TrainerIO):
             self.experiment = self.experiment.get_meta_copy()
             mp.spawn(self.dp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
         else:
+            # CHOOSE OPTIMIZER
+            # filter out the weights that were done on gpu so we can load on good old cpus
+            self.optimizers = model.configure_optimizers()
+
+            # run through amp wrapper
+            if self.use_amp:
+                # An example
+                model, optimizers = amp.initialize(
+                    model, self.optimizers, opt_level=self.amp_level,
+                )
+                self.optimizers = optimizers
+
             self.__run_pretrain_routine(model)
 
     def dp_train(self, gpu_nb, model):
@@ -334,17 +335,44 @@ class Trainer(TrainerIO):
         self.world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids)
 
         # set up server using proc 0's ip address
+        # try to init for 20 times at max in case ports are taken
         ip = self.__get_root_node_ip(self.proc_rank, self.nb_gpu_nodes)
-        dist.init_process_group("nccl", init_method=f'tcp://{ip}:12001', rank=self.proc_rank, world_size=self.world_size)
+        self.__init_tcp_connection(ip)
 
+        # CHOOSE OPTIMIZER
+        # filter out the weights that were done on gpu so we can load on good old cpus
+        self.optimizers = model.configure_optimizers()
+
+        # MODEL
         # copy model to each gpu
         torch.cuda.set_device(gpu_nb)
         model.cuda(gpu_nb)
+
+        # AMP
+        # run through amp wrapper before going to distributed DP
+        if self.use_amp:
+            # An example
+            model, optimizers = amp.initialize(
+                model, self.optimizers, opt_level=self.amp_level,
+            )
+            self.optimizers = optimizers
+
         model = LightningDistributedDataParallel(model, device_ids=[gpu_nb])
 
         # continue training routine
         self.__run_pretrain_routine(model)
 
+    def __init_tcp_connection(self, ip, port=12000, tries=0):
+        if tries > 20:
+            raise RuntimeError('Failed to connect using 20 different ip addresses')
+
+        try:
+            dist.init_process_group("nccl", init_method=f'tcp://{ip}:{port}', rank=self.proc_rank, world_size=self.world_size)
+        except RuntimeError as e:
+            # port taken
+            warnings.warn(f'port {port} taken, trying port {port}...')
+            self.__init_tcp_connection(ip, port + 1, tries + 1)
+
     def __get_root_node_ip(self, world_gpu_nb, nb_gpu_nodes):
         """
         Resolves the ip address of proc 0.
@@ -588,8 +616,18 @@ class Trainer(TrainerIO):
         else:
             output = self.model.training_step(data_batch, batch_nb)
 
-        model_specific_tqdm_metrics_dic = output['tqdm_metrics']
-        loss = output['loss']
+        try:
+            model_specific_tqdm_metrics_dic = output['tqdm_metrics']
+        except Exception as e:
+            model_specific_tqdm_metrics_dic = {}
+
+        # if output dict doesn't have the keyword loss
+        # then assume the output=loss if scalar
+        try:
+            loss = output['loss']
+        except Exception as e:
+            if type(loss) is torch.Tensor:
+                loss = output
 
         self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
 
@@ -607,6 +645,7 @@ class Trainer(TrainerIO):
             for param in model.parameters():
                 print(param.grad.float().sum())
 
+        # avoid memory leaks
         self.batch_loss_value += loss.item()
 
         # gradient update with accumulated gradients

setup.py

@@ -7,7 +7,7 @@ from setuptools import setup, find_packages
 # http://blog.ionelmc.ro/2014/05/25/python-packaging/
 setup(
     name="pytorch-lightning",
-    version='0.2',
+    version='0.2.2',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",