release v0.2.4

trainer: module fix.

README.md

@@ -116,8 +116,8 @@ def validation_end(self, outputs):
     return tqdm_dic
 ```
    
-## TensorboardX    
-Lightning is fully integrated with tensorboardX.   
+## Tensorboard    
+Lightning is fully integrated with tensorboard.   
 
 <p align="center">
   <a href="https://williamfalcon.github.io/pytorch-lightning/">

pytorch_lightning/models/trainer.py

@@ -31,7 +31,8 @@ class Trainer(TrainerIO):
 
     def __init__(self,
                  experiment,
-                 checkpoint_callback, early_stop_callback,
+                 early_stop_callback=None,
+                 checkpoint_callback=None,
                  gradient_clip=0,
                  cluster=None,
                  process_position=0,
@@ -44,12 +45,14 @@ class Trainer(TrainerIO):
                  check_val_every_n_epoch=1,
                  fast_dev_run=False,
                  accumulate_grad_batches=1,
-                 enable_early_stop=True, max_nb_epochs=1000, min_nb_epochs=1,
-                 train_percent_check=1.0, val_percent_check=1.0, test_percent_check=1.0, val_check_interval=0.95,
+                 max_nb_epochs=1000, min_nb_epochs=1,
+                 train_percent_check=1.0, val_percent_check=1.0, test_percent_check=1.0,
+                 val_check_interval=0.95,
                  log_save_interval=100, add_log_row_interval=10,
                  lr_scheduler_milestones=None,
                  use_amp=False,
                  print_nan_grads=False,
+                 print_weights_summary=True,
                  amp_level='O2',
                  nb_sanity_val_steps=5):
 
@@ -57,7 +60,7 @@ class Trainer(TrainerIO):
         self.nb_gpu_nodes = nb_gpu_nodes
         self.gradient_clip = gradient_clip
         self.check_val_every_n_epoch = check_val_every_n_epoch
-        self.enable_early_stop = enable_early_stop
+        self.enable_early_stop = early_stop_callback is not None
         self.track_grad_norm = track_grad_norm
         self.fast_dev_run = fast_dev_run
         self.on_gpu = gpus is not None and torch.cuda.is_available()
@@ -67,8 +70,12 @@ class Trainer(TrainerIO):
         self.cluster = cluster
         self.process_position = process_position
         self.current_gpu_name = current_gpu_name
+        self.print_weights_summary = print_weights_summary
         self.checkpoint_callback = checkpoint_callback
-        self.checkpoint_callback.save_function = self.save_checkpoint
+
+        if self.checkpoint_callback is not None:
+            self.checkpoint_callback.save_function = self.save_checkpoint
+
         self.early_stop = early_stop_callback
         self.model = None
         self.max_nb_epochs = max_nb_epochs
@@ -96,7 +103,7 @@ class Trainer(TrainerIO):
             os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
             os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(x) for x in self.data_parallel_device_ids])
 
-        self.data_parallel = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 0
+        self.data_parallel = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 1
 
         # process info
         self.proc_rank = 0
@@ -149,8 +156,12 @@ class Trainer(TrainerIO):
             self.val_percent_check = overfit_pct
             self.test_percent_check = overfit_pct
 
+    def __get_model(self):
+        return self.model.module if self.data_parallel else self.model
+
     def __is_function_implemented(self, f_name):
-        f_op = getattr(self.model, f_name, None)
+        model = self.__get_model()
+        f_op = getattr(model, f_name, None)
         return callable(f_op)
 
     @property
@@ -208,9 +219,6 @@ class Trainer(TrainerIO):
         :param max_batches: Scalar
         :return:
         """
-        if self.proc_rank == 0:
-            print('validating...')
-
         # enable eval mode
         model.zero_grad()
         model.eval()
@@ -269,7 +277,7 @@ class Trainer(TrainerIO):
         self.test_dataloader = model.test_dataloader
         self.val_dataloader = model.val_dataloader
 
-        if self.on_gpu and type(self.tng_dataloader.sampler) is not DistributedSampler:
+        if self.data_parallel and type(self.tng_dataloader.sampler) is not DistributedSampler:
             msg = '''
             when using multiple gpus and multiple nodes you must pass a DistributedSampler to DataLoader(sampler).
             
@@ -288,6 +296,36 @@ class Trainer(TrainerIO):
     # MODEL TRAINING
     # -----------------------------
     def fit(self, model):
+
+        # when using gpus, first thing we do is spawn a new process between each worker
+        #  multi-gpu and multi-nodes
+        if self.data_parallel:
+            self.experiment = self.experiment.get_meta_copy()
+            mp.spawn(self.dp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
+
+        # treat 1 gpu as a different case to avoid nccl bugs
+        elif self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) == 1:
+            self.single_gpu_train(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 single_gpu_train(self, model):
+        # torch.cuda.set_device(0)
+        model.cuda(0)
+
         # CHOOSE OPTIMIZER
         # filter out the weights that were done on gpu so we can load on good old cpus
         self.optimizers = model.configure_optimizers()
@@ -300,13 +338,7 @@ class Trainer(TrainerIO):
             )
             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
-        if self.on_gpu:
-            self.experiment = self.experiment.get_meta_copy()
-            mp.spawn(self.dp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
-        else:
-            self.__run_pretrain_routine(model)
+        self.__run_pretrain_routine(model)
 
     def dp_train(self, gpu_nb, model):
         """
@@ -324,6 +356,8 @@ class Trainer(TrainerIO):
             node_rank = 0
 
         # recover original exp before went into process
+        # init in write mode only on proc 0
+        self.experiment.debug = self.proc_rank > 0
         self.experiment = self.experiment.get_non_ddp_exp()
 
         # show progbar only on prog_rank 0
@@ -334,60 +368,56 @@ 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
-        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)
+        # try to init for 20 times at max in case ports are taken
+        # where to store ip_table
+        self.__init_tcp_connection()
 
+        # 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 __get_root_node_ip(self, world_gpu_nb, nb_gpu_nodes):
+    def __init_tcp_connection(self):
         """
-        Resolves the ip address of proc 0.
-        Proc 0 writes address to a file. Every other process waits until the ip is available before it starts
-
-        :param world_gpu_nb: gpu number amongst all the world gpus
-        :param nb_gpu_nodes:
-        :param ip_file_dir:
+        Connect all procs in the world using the env:// init
+        Use the first node as the root address
+        :param port:
+        :param tries:
         :return:
         """
-        # on one node we use localhost
-        if nb_gpu_nodes == 1:
-            return '127.0.0.1'
+        try:
+            port = os.environ['MASTER_PORT']
+        except Exception as e:
+            port = 12910
+            os.environ['MASTER_PORT'] = f'{port}'
 
-        # where to store ip_table
-        ip_file_dir = os.path.join(self.cluster.log_path, 'ip_tables')
+        try:
+            root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
+        except Exception as e:
+            root_node = '127.0.0.2'
 
-        # the first gpu in the world becomes the host
-        # this is based on its global rank
-        # it communicates its ip by saving an ip_table to the slurm cluster logging dir
-        # every other process waits for this ip to appear before continuing
-        ip_table_name = f'.ip_meta_' + os.environ['SLURM_JOB_ID']
-        ip_file = os.path.join(ip_file_dir, ip_table_name)
-        os.makedirs(ip_file_dir, exist_ok=True)
+        os.environ['MASTER_ADDR'] = root_node
 
-        if world_gpu_nb == 0:
-            # get the proc 0 IP
-            root_ip = subprocess.run(['hostname', '-I'], stdout=subprocess.PIPE).stdout.decode('utf-8')
-            root_ip = root_ip.split(' ')[0]
-
-            # save the ip to the file
-            with open(file=ip_file, mode='w') as f:
-                f.write(root_ip)
-
-            return root_ip
-        else:
-            # wait up to 120 seconds until proc 0 writes
-            # once written, read proc 0's address and use it to configure server
-            for i in range(0, 120):
-                sleep(1.0)
-                if os.path.exists(ip_file):
-                    ip = list(open(file=ip_file, mode='r'))[0]
-                    return ip
+        sleep(self.proc_rank*0.5)
+        dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
 
     def __run_pretrain_routine(self, model):
         """
@@ -396,7 +426,7 @@ class Trainer(TrainerIO):
         :return:
         """
         ref_model = model
-        if self.on_gpu:
+        if self.data_parallel:
             ref_model = model.module
 
         ref_model.trainer = self
@@ -417,7 +447,7 @@ class Trainer(TrainerIO):
                 self.lr_schedulers.append(scheduler)
 
         # print model summary
-        if self.proc_rank == 0:
+        if self.proc_rank == 0 and self.print_weights_summary:
             ref_model.summarize()
 
         # give model convenience properties
@@ -448,12 +478,12 @@ class Trainer(TrainerIO):
             for lr_scheduler in self.lr_schedulers:
                 lr_scheduler.step()
 
-            model = self.model.module if self.data_parallel else self.model
+            model = self.__get_model()
             model.current_epoch = epoch_nb
 
             # hook
             if self.__is_function_implemented('on_epoch_start'):
-                model = self.model.module if self.data_parallel else self.model
+                model = self.__get_model()
                 model.on_epoch_start()
 
             self.current_epoch = epoch_nb
@@ -468,7 +498,7 @@ class Trainer(TrainerIO):
                 self.batch_nb = batch_nb
                 self.global_step += 1
 
-                model = self.model.module if self.data_parallel else self.model
+                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
@@ -500,10 +530,8 @@ class Trainer(TrainerIO):
                     # count items in memory
                     # nb_params, nb_tensors = count_mem_items()
 
-                    if self.data_parallel:
-                        metrics = self.model.module.update_tng_log_metrics(self.__tng_tqdm_dic)
-                    else:
-                        metrics = self.model.update_tng_log_metrics(self.__tng_tqdm_dic)
+                    model = self.__get_model()
+                    metrics = model.update_tng_log_metrics(self.__tng_tqdm_dic)
 
                     # add gpu memory
                     if self.on_gpu:
@@ -512,7 +540,7 @@ class Trainer(TrainerIO):
 
                     # add norms
                     if self.track_grad_norm > 0:
-                        model = self.model.module if self.data_parallel else self.model
+                        model = self.__get_model()
                         grad_norm_dic = model.grad_norm(self.track_grad_norm)
 
                         metrics.update(grad_norm_dic)
@@ -525,7 +553,7 @@ class Trainer(TrainerIO):
 
                 # hook
                 if self.__is_function_implemented('on_batch_end'):
-                    model = self.model.module if self.data_parallel else self.model
+                    model = self.__get_model()
                     model.on_batch_end()
 
                 # end epoch early
@@ -534,13 +562,13 @@ class Trainer(TrainerIO):
 
             # hook
             if self.__is_function_implemented('on_epoch_end'):
-                model = self.model.module if self.data_parallel else self.model
+                model = self.__get_model()
                 model.on_epoch_end()
 
             # early stopping
-            if self.enable_early_stop:
+            met_min_epochs = epoch_nb > self.min_nb_epochs
+            if self.enable_early_stop and met_min_epochs:
                 should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb, logs=self.__tng_tqdm_dic)
-                met_min_epochs = epoch_nb > self.min_nb_epochs
 
                 # stop training
                 stop = should_stop and met_min_epochs
@@ -572,7 +600,7 @@ class Trainer(TrainerIO):
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
-            model = self.model.module if self.data_parallel else self.model
+            model = self.__get_model()
             response = model.on_batch_start(data_batch)
 
             if response == -1:
@@ -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(output) is torch.Tensor:
+                loss = output
 
         self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
 
@@ -603,10 +641,11 @@ class Trainer(TrainerIO):
             loss.backward()
 
         if self.print_nan_grads:
-            model = self.model.module if self.data_parallel else self.model
+            model = self.__get_model()
             for param in model.parameters():
                 print(param.grad.float().sum())
 
+        # avoid memory leaks
         self.batch_loss_value += loss.item()
 
         # gradient update with accumulated gradients
@@ -614,7 +653,7 @@ class Trainer(TrainerIO):
 
             # clip gradients
             if self.gradient_clip > 0:
-                model = self.model.module if self.data_parallel else self.model
+                model = self.__get_model()
                 torch.nn.utils.clip_grad_norm(model.parameters(), self.gradient_clip)
 
             # update gradients across all optimizers
@@ -640,7 +679,8 @@ class Trainer(TrainerIO):
 
         # activate batch end hook
         if self.__is_function_implemented('on_batch_end'):
-            self.model.on_batch_end()
+            model = self.__get_model()
+            model.on_batch_end()
 
         return 0
 
@@ -655,7 +695,8 @@ class Trainer(TrainerIO):
         try:
             # hook
             if self.__is_function_implemented('on_pre_performance_check'):
-                self.model.on_pre_performance_check()
+                model = self.__get_model()
+                model.on_pre_performance_check()
 
             # use full val set on end of epoch
             # use a small portion otherwise
@@ -669,7 +710,8 @@ class Trainer(TrainerIO):
 
             # hook
             if self.__is_function_implemented('on_post_performance_check'):
-                self.model.on_post_performance_check()
+                model = self.__get_model()
+                model.on_post_performance_check()
 
         except Exception as e:
             print(e)
@@ -682,5 +724,6 @@ class Trainer(TrainerIO):
 
         # model checkpointing
         if self.proc_rank == 0:
-            print('save callback...')
-            self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)
+            if self.checkpoint_callback:
+                print('save callback...')
+                self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)

pytorch_lightning/root_module/root_module.py

@@ -78,7 +78,7 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         :param logs:
         :return:
         """
-        raise NotImplementedError
+        return logs
 
     def loss(self, *args, **kwargs):
         """
@@ -129,7 +129,7 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
     def get_process_position(gpus):
         try:
             current_gpu = os.environ["CUDA_VISIBLE_DEVICES"]
-            gpu_ids = gpus.split(';')
+            gpu_ids = gpus.split(',')
             process_position = gpu_ids.index(current_gpu)
             return process_position, current_gpu
         except Exception as e:

pytorch_lightning/utils/plotting.py

@@ -1,4 +1,3 @@
-from matplotlib import pyplot as plt
 import numpy as np
 np.seterr(divide='ignore', invalid='ignore')
 
@@ -13,6 +12,7 @@ def plot_confusion_matrix(cm,
     This function prints and plots the confusion matrix.
     Normalization can be applied by setting `normalize=True`.
     """
+    from matplotlib import pyplot as plt
     if normalize:
         cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
         print("Normalized confusion matrix")

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.4',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",
@@ -19,7 +19,7 @@ setup(
     install_requires=[
         "torch>=1.1.0",
         "tqdm",
-        "test-tube>=0.653",
+        "test-tube>=0.6.6",
         "tensorflow>=1.14.0"
     ],
     packages=find_packages(),