Merge pull request #16 from williamFalcon/tests

Tests

pytorch_lightning/examples/__init__.py

@@ -0,0 +1 @@
+from .new_project_templates.lightning_module_template import LightningTemplateModel

pytorch_lightning/examples/new_project_templates/lightning_module_template.py

@@ -28,6 +28,9 @@ class LightningTemplateModel(LightningModule):
 
         self.batch_size = hparams.batch_size
 
+        # if you specify an example input, the summary will show input/output for each layer
+        self.example_input_array = torch.rand(5, 28 * 28)
+
         # build model
         self.__build_model()
 
@@ -78,11 +81,16 @@ class LightningTemplateModel(LightningModule):
         # forward pass
         x, y = data_batch
         x = x.view(x.size(0), -1)
+
         y_hat = self.forward(x)
 
         # calculate loss
         loss_val = self.loss(y, y_hat)
 
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+
         output = OrderedDict({
             'loss': loss_val
         })
@@ -105,10 +113,19 @@ class LightningTemplateModel(LightningModule):
         # acc
         labels_hat = torch.argmax(y_hat, dim=1)
         val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        val_acc = torch.tensor(val_acc)
+
+        if self.on_gpu:
+            val_acc = val_acc.cuda(loss_val.device.index)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+            val_acc = val_acc.unsqueeze(0)
 
         output = OrderedDict({
             'val_loss': loss_val,
-            'val_acc': torch.tensor(val_acc).cuda(loss_val.device.index),
+            'val_acc': val_acc,
         })
 
         # can also return just a scalar instead of a dict (return loss_val)
@@ -135,9 +152,6 @@ class LightningTemplateModel(LightningModule):
         tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
         return tqdm_dic
 
-    def update_tng_log_metrics(self, logs):
-        return logs
-
     # ---------------------
     # MODEL SAVING
     # ---------------------
@@ -217,7 +231,7 @@ class LightningTemplateModel(LightningModule):
         return self._test_dataloader
 
     @staticmethod
-    def add_model_specific_args(parent_parser, root_dir):
+    def add_model_specific_args(parent_parser, root_dir): # pragma: no cover
         """
         Parameters you define here will be available to your model through self.hparams
         :param parent_parser:

pytorch_lightning/models/trainer.py

@@ -19,12 +19,12 @@ import tqdm
 from pytorch_lightning.root_module.memory import get_gpu_memory_map
 from pytorch_lightning.root_module.model_saving import TrainerIO
 from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
-
+from pytorch_lightning.utils.debugging import MisconfigurationException
 
 try:
     from apex import amp
     APEX_AVAILABLE = True
-except ModuleNotFoundError:
+except ModuleNotFoundError: # pragma: no cover
     APEX_AVAILABLE = False
 
 
@@ -114,6 +114,7 @@ class Trainer(TrainerIO):
         """
 
         # Transfer params
+
         self.nb_gpu_nodes = nb_gpu_nodes
         self.gradient_clip = gradient_clip
         self.check_val_every_n_epoch = check_val_every_n_epoch
@@ -150,6 +151,16 @@ class Trainer(TrainerIO):
         self.use_ddp = False
         self.use_dp = False
 
+        # training bookeeping
+        self.total_batch_nb = 0
+        self.running_loss = []
+        self.avg_loss = 0
+        self.batch_nb = 0
+        self.tqdm_metrics = {}
+        self.nb_val_batches = None
+        self.nb_tng_batches = None
+        self.nb_test_batches = None
+
 
         # gpus come in as a string.
         # if gpus = -1 then use all available devices
@@ -178,7 +189,7 @@ class Trainer(TrainerIO):
             self.use_ddp = distributed_backend == 'ddp'
 
             # use ddp automatically if nb_gpu_nodes > 1
-            if nb_gpu_nodes > 1 and self.use_dp:
+            if nb_gpu_nodes > 1 and self.use_dp:  # pragma: no cover
                 self.use_ddp = True
                 self.use_dp = False
                 w = 'DataParallel does not support nb_gpu_nodes > 1. ' \
@@ -186,6 +197,19 @@ class Trainer(TrainerIO):
                     'To silence this warning set distributed_backend=ddp'
                 warnings.warn(w)
 
+        # extract SLURM flag vars
+        # whenever we have the correct number of tasks, we let slurm manage processes
+        # otherwise we launch the required number of processes
+        if self.use_ddp:
+            self.nb_requested_gpus = len(self.data_parallel_device_ids) * self.nb_gpu_nodes
+            self.nb_slurm_tasks = 0
+            try:
+                self.nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
+                self.is_slurm_managing_tasks = self.nb_slurm_tasks == self.nb_requested_gpus
+            except Exception as e:
+                # likely not on slurm, so set the slurm managed flag to false
+                self.is_slurm_managing_tasks = False
+
         # process info
         self.proc_rank = 0
 
@@ -215,7 +239,7 @@ class Trainer(TrainerIO):
         if self.use_amp:
             print('using 16bit precision')
 
-        if use_amp and not APEX_AVAILABLE:
+        if use_amp and not APEX_AVAILABLE:  # pragma: no cover
             msg = '''
             You set use_amp=True but do not have apex installed.
             Install apex first using this guide and rerun with use_amp=True: 
@@ -223,7 +247,7 @@ class Trainer(TrainerIO):
             
             this run will NOT use 16 bit precision
             '''
-            warnings.warn(msg)
+            raise ModuleNotFoundError(msg)
 
     @property
     def data_parallel(self):
@@ -251,6 +275,7 @@ class Trainer(TrainerIO):
 
     @property
     def __tng_tqdm_dic(self):
+        # ForkedPdb().set_trace()
         tqdm_dic = {
             'tng_loss': '{0:.3f}'.format(self.avg_loss),
             'v_nb': '{}'.format(self.experiment.version),
@@ -264,13 +289,15 @@ class Trainer(TrainerIO):
 
         return tqdm_dic
 
+    @property
+    def tng_tqdm_dic(self):
+        """
+        Read-only for tqdm metrics
+        :return:
+        """
+        return self.__tng_tqdm_dic
+
     def __layout_bookeeping(self):
-        # training bookeeping
-        self.total_batch_nb = 0
-        self.running_loss = []
-        self.avg_loss = 0
-        self.batch_nb = 0
-        self.tqdm_metrics = {}
 
         # determine number of training batches
         self.nb_tng_batches = len(self.tng_dataloader)
@@ -317,7 +344,7 @@ class Trainer(TrainerIO):
         # run training
         for batch_i, data_batch in enumerate(dataloader):
 
-            if data_batch is None:
+            if data_batch is None:  # pragma: no cover
                 continue
 
             # stop short when on fast dev run
@@ -356,7 +383,7 @@ class Trainer(TrainerIO):
 
         return val_results
 
-    def __get_dataloaders(self, model):
+    def get_dataloaders(self, model):
         """
         Dataloaders are provided by the model
         :param model:
@@ -379,7 +406,7 @@ class Trainer(TrainerIO):
             dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
             dataloader = Dataloader(dataset, sampler=dist_sampler)
             '''
-            raise Exception(msg)
+            raise MisconfigurationException(msg)
 
     # -----------------------------
     # MODEL TRAINING
@@ -391,25 +418,14 @@ class Trainer(TrainerIO):
             # must copy only the meta of the exp so it survives pickle/unpickle when going to new process
             self.experiment = self.experiment.get_meta_copy()
 
-            # whenever we have the correct number of tasks, we let slurm manage processes
-            # otherwise we launch the required number of processes
-            nb_requested_gpus = len(self.data_parallel_device_ids) * self.nb_gpu_nodes
-            nb_slurm_tasks = 0
-            try:
-                nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
-                is_slurm_managing_tasks = nb_slurm_tasks == nb_requested_gpus
-            except Exception as e:
-                # likely not on slurm, so set the slurm managed flag to false
-                is_slurm_managing_tasks = False
-
-            if is_slurm_managing_tasks:
+            if self.is_slurm_managing_tasks:
                 task = int(os.environ['SLURM_LOCALID'])
                 self.ddp_train(task, model)
             else:
                 msg = f"""
-                You requested {nb_requested_gpus} GPUs but launched {nb_slurm_tasks} slurm tasks. 
-                We will launch {nb_requested_gpus} processes for you. 
-                We recommend you let slurm manage the processes by setting: --ntasks-per-node={nb_requested_gpus}
+                You requested {self.nb_requested_gpus} GPUs but launched {self.nb_slurm_tasks} slurm tasks. 
+                We will launch {self.nb_requested_gpus} processes for you. 
+                We recommend you let slurm manage the processes by setting: --ntasks-per-node={self.nb_requested_gpus}
                 If you're not using SLURM, ignore this message!
                 """
                 warnings.warn(msg)
@@ -418,7 +434,7 @@ class Trainer(TrainerIO):
         # 1 gpu or dp option triggers training using DP module
         # easier to avoid NCCL issues
         elif self.use_dp:
-            self.dp_train(model)
+            self.__dp_train(model)
 
         # ON CPU
         else:
@@ -428,15 +444,15 @@ class Trainer(TrainerIO):
 
             # 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
+                raise MisconfigurationException('amp + cpu is not supported. Please use a GPU option')
 
             self.__run_pretrain_routine(model)
 
-    def dp_train(self, model):
+        # return 1 when finished
+        # used for testing or when we need to know that training succeeded
+        return 1
+
+    def __dp_train(self, model):
 
         # CHOOSE OPTIMIZER
         # filter out the weights that were done on gpu so we can load on good old cpus
@@ -444,13 +460,13 @@ class Trainer(TrainerIO):
 
         model.cuda(self.data_parallel_device_ids[0])
 
-        # 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
+        # check for this bug (amp + dp + !01 doesn't work)
+        # https://github.com/NVIDIA/apex/issues/227
+        if self.use_dp and self.use_amp:
+            m = f'amp level {self.amp_level} with DataParallel is not supported. ' \
+                f'See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227. ' \
+                f'We recommend you switch to ddp if you want to use amp'
+            raise MisconfigurationException(m)
 
         model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
 
@@ -520,22 +536,22 @@ class Trainer(TrainerIO):
         :param tries:
         :return:
         """
+        # sets the appropriate port
         try:
             port = os.environ['MASTER_PORT']
         except Exception as e:
             port = 12910
             os.environ['MASTER_PORT'] = f'{port}'
 
-        root_node = self.__resolve_root_node_address()
+        # figure out the root node addr
+        root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
+        root_node = self.resolve_root_node_address(root_node)
         os.environ['MASTER_ADDR'] = root_node
 
         dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
 
-
-    def __resolve_root_node_address(self):
+    def resolve_root_node_address(self, root_node):
         try:
-            root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
-
             if '[' in root_node:
                 name = root_node.split('[')[0]
                 number = root_node.split(',')[0]
@@ -566,7 +582,7 @@ class Trainer(TrainerIO):
         ref_model.on_gpu = self.on_gpu
 
         # transfer data loaders from model
-        self.__get_dataloaders(ref_model)
+        self.get_dataloaders(ref_model)
 
         # init training constants
         self.__layout_bookeeping()
@@ -593,7 +609,7 @@ class Trainer(TrainerIO):
             self.experiment.save()
 
         # enable cluster checkpointing
-        if self.cluster is not None:
+        if self.cluster is not None:  # pragma: no cover
             self.enable_auto_hpc_walltime_manager()
 
         # ---------------------------
@@ -662,7 +678,7 @@ class Trainer(TrainerIO):
                     # nb_params, nb_tensors = count_mem_items()
 
                     model = self.__get_model()
-                    metrics = model.update_tng_log_metrics(self.__tng_tqdm_dic)
+                    metrics = self.__tng_tqdm_dic
 
                     # add gpu memory
                     if self.on_gpu:

pytorch_lightning/pt_overrides/override_data_parallel.py

@@ -9,7 +9,7 @@ from torch.cuda._utils import _get_device_index
 import pdb
 
 
-def _find_tensors(obj):
+def _find_tensors(obj):  # pragma: no cover
     r"""
     Recursively find all tensors contained in the specified object.
     """
@@ -22,8 +22,7 @@ def _find_tensors(obj):
     return []
 
 
-
-def get_a_var(obj):
+def get_a_var(obj):  # pragma: no cover
     if isinstance(obj, torch.Tensor):
         return obj
 
@@ -78,7 +77,7 @@ class LightningDistributedDataParallel(DistributedDataParallel):
     def parallel_apply(self, replicas, inputs, kwargs):
         return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
 
-    def forward(self, *inputs, **kwargs):
+    def forward(self, *inputs, **kwargs):  # pragma: no cover
         self._sync_params()
         if self.device_ids:
             inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
@@ -113,7 +112,7 @@ class LightningDistributedDataParallel(DistributedDataParallel):
         return output
 
 
-def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
+def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):  # pragma: no cover
     r"""Applies each `module` in :attr:`modules` in parallel on arguments
     contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword)
     on each of :attr:`devices`.

pytorch_lightning/root_module/grads.py

@@ -27,14 +27,3 @@ class GradInformation(nn.Module):
         results['grad_{}_norm_total'.format(norm_type)] = round(total_norm.data.cpu().numpy().flatten()[0], 3)
         return results
 
-
-    def describe_grads(self):
-        for p in self.parameters():
-            g = p.grad.data.numpy().flatten()
-            print(np.max(g), np.min(g), np.mean(g))
-
-
-    def describe_params(self):
-        for p in self.parameters():
-            g = p.data.numpy().flatten()
-            print(np.max(g), np.min(g), np.mean(g))

pytorch_lightning/root_module/memory.py

@@ -33,33 +33,42 @@ class ModelSummary(object):
         mods = list(self.model.modules())
         in_sizes = []
         out_sizes = []
-        input_ = self.example_input_array
-        for i in range(1, len(mods)):
-            m = mods[i]
-            if type(input_) is list or type(input_) is tuple:
-                out = m(*input_)
-            else:
-                out = m(input_)
+        input_ = self.model.example_input_array
 
-            if type(input_) is tuple or type(input_) is list:
-                in_size = []
-                for x in input_:
-                    if type(x) is list:
-                        in_size.append(len(x))
-                    else:
-                        in_size.append(x.size())
-            else:
-                in_size = np.array(input_.size())
+        if self.model.on_gpu:
+            input_ = input_.cuda(0)
 
-            in_sizes.append(in_size)
+        if self.model.trainer.use_amp:
+            input_ = input_.half()
 
-            if type(out) is tuple or type(out) is list:
-                out_size = np.asarray([x.size() for x in out])
-            else:
-                out_size = np.array(out.size())
+        with torch.no_grad():
 
-            out_sizes.append(out_size)
-            input_ = out
+            for i in range(1, len(mods)):
+                m = mods[i]
+                if type(input_) is list or type(input_) is tuple:  # pragma: no cover
+                    out = m(*input_)
+                else:
+                    out = m(input_)
+
+                if type(input_) is tuple or type(input_) is list:  # pragma: no cover
+                    in_size = []
+                    for x in input_:
+                        if type(x) is list:
+                            in_size.append(len(x))
+                        else:
+                            in_size.append(x.size())
+                else:
+                    in_size = np.array(input_.size())
+
+                in_sizes.append(in_size)
+
+                if type(out) is tuple or type(out) is list:  # pragma: no cover
+                    out_size = np.asarray([x.size() for x in out])
+                else:
+                    out_size = np.array(out.size())
+
+                out_sizes.append(out_size)
+                input_ = out
 
         self.in_sizes = in_sizes
         self.out_sizes = out_sizes
@@ -114,13 +123,22 @@ class ModelSummary(object):
         Layer Name, Layer Type, Input Size, Output Size, Number of Parameters
         '''
 
-        df = pd.DataFrame( np.zeros( (len(self.layer_names), 3) ) )
-        df.columns = ['Name', 'Type', 'Params']
+        cols = ['Name', 'Type', 'Params']
+        if self.model.example_input_array is not None:
+            cols.extend(['In_sizes', 'Out_sizes'])
+
+        df = pd.DataFrame(np.zeros( (len(self.layer_names), len(cols))))
+        df.columns = cols
 
         df['Name'] = self.layer_names
         df['Type'] = self.layer_types
         df['Params'] = self.param_nums
 
+        if self.model.example_input_array is not None:
+
+            df['In_sizes'] = self.in_sizes
+            df['Out_sizes'] = self.out_sizes
+
         self.summary = df
         return
 
@@ -128,10 +146,13 @@ class ModelSummary(object):
         self.get_layer_names()
         self.get_parameter_sizes()
         self.get_parameter_nums()
+
+        if self.model.example_input_array is not None:
+            self.get_variable_sizes()
         self.make_summary()
 
 
-def print_mem_stack():
+def print_mem_stack(): # pragma: no cover
     for obj in gc.get_objects():
         try:
             if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
@@ -140,7 +161,7 @@ def print_mem_stack():
             pass
 
 
-def count_mem_items():
+def count_mem_items(): # pragma: no cover
     nb_params = 0
     nb_tensors = 0
     for obj in gc.get_objects():

pytorch_lightning/root_module/model_saving.py

@@ -41,6 +41,12 @@ class ModelIO(object):
 
 class TrainerIO(object):
 
+    def __get_model(self):
+        print(type(self.model))
+        is_dp_module = type(self.model) is LightningDistributedDataParallel or type(self.model) is LightningDataParallel
+        model = self.model.module if is_dp_module else self.model
+        return model
+
     # --------------------
     # MODEL SAVE CHECKPOINT
     # --------------------
@@ -51,14 +57,19 @@ class TrainerIO(object):
         torch.save(checkpoint, filepath)
 
     def dump_checkpoint(self):
+
         checkpoint = {
             'epoch': self.current_epoch,
-            'checkpoint_callback_best': self.checkpoint_callback.best,
-            'early_stop_callback_wait': self.early_stop_callback.wait,
-            'early_stop_callback_patience': self.early_stop_callback.patience,
             'global_step': self.global_step
         }
 
+        if self.checkpoint_callback is not None:
+            checkpoint['checkpoint_callback_best'] = self.checkpoint_callback.best
+
+        if self.early_stop_callback is not None:
+            checkpoint['early_stop_callback_wait'] = self.early_stop_callback.wait
+            checkpoint['early_stop_callback_patience'] = self.early_stop_callback.patience
+
         optimizer_states = []
         for i, optimizer in enumerate(self.optimizers):
             optimizer_states.append(optimizer.state_dict())
@@ -66,8 +77,7 @@ class TrainerIO(object):
         checkpoint['optimizer_states'] = optimizer_states
 
         # request what to save from the model
-        is_dp_module = type(self.model) is LightningDistributedDataParallel or type(self.model) is LightningDataParallel
-        model = self.model.module if is_dp_module else self.model
+        model = self.__get_model()
         checkpoint_dict = model.get_save_dict()
 
         # merge trainer and model saving items
@@ -77,7 +87,7 @@ class TrainerIO(object):
     # --------------------
     # HPC IO
     # --------------------
-    def enable_auto_hpc_walltime_manager(self):
+    def enable_auto_hpc_walltime_manager(self):  # pragma: no cover
         if self.cluster is None:
             return
 
@@ -104,9 +114,13 @@ class TrainerIO(object):
         :param checkpoint:
         :return:
         """
-        self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
-        self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
-        self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
+        if self.checkpoint_callback is not None:
+            self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
+
+        if self.early_stop_callback is not None:
+            self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
+            self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
+
         self.global_step = checkpoint['global_step']
         self.current_epoch = checkpoint['epoch']
 
@@ -135,7 +149,8 @@ class TrainerIO(object):
         filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number)
 
         # give model a chance to do something on hpc_save
-        self.on_hpc_save()
+        model = self.__get_model()
+        model.on_hpc_save()
 
         # request what to save from the model
         checkpoint_dict = self.dump_checkpoint()
@@ -155,11 +170,11 @@ class TrainerIO(object):
         self.restore_training_state(checkpoint)
 
         # load model state
-        model = self.model.module if type(self.model) is LightningDataParallel else self.model
+        model = self.__get_model()
         model.load_model_specific(checkpoint)
 
         # call model hook
-        self.on_hpc_load()
+        model.on_hpc_load()
 
     def max_ckpt_in_folder(self, path):
         files = os.listdir(path)

pytorch_lightning/root_module/optimization.py

@@ -1,22 +0,0 @@
-from torch import nn
-from torch import optim
-
-
-class OptimizerConfig(nn.Module):
-
-    def choose_optimizer(self, optimizer, params, optimizer_params, opt_name_key):
-        if optimizer == 'adam':
-            optimizer = optim.Adam(params, **optimizer_params)
-        if optimizer == 'sparse_adam':
-            optimizer = optim.SparseAdam(params, **optimizer_params)
-        if optimizer == 'sgd':
-            optimizer = optim.SGD(params, **optimizer_params)
-        if optimizer == 'adadelta':
-            optimizer = optim.Adadelta(params, **optimizer_params)
-
-        # transfer opt state if loaded
-        if opt_name_key in self.loaded_optimizer_states_dict:
-            state = self.loaded_optimizer_states_dict[opt_name_key]
-            optimizer.load_state_dict(state)
-
-        return optimizer

pytorch_lightning/root_module/root_module.py

@@ -5,11 +5,10 @@ import math
 from pytorch_lightning.root_module.memory import ModelSummary
 from pytorch_lightning.root_module.grads import GradInformation
 from pytorch_lightning.root_module.model_saving import ModelIO, load_hparams_from_tags_csv
-from pytorch_lightning.root_module.optimization import OptimizerConfig
 from pytorch_lightning.root_module.hooks import ModelHooks
 
 
-class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
+class LightningModule(GradInformation, ModelIO, ModelHooks):
 
     def __init__(self, hparams):
         super(LightningModule, self).__init__()
@@ -22,6 +21,7 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         self.loaded_optimizer_states_dict = {}
         self.trainer = None
         self.experiment = None
+        self.example_input_array = None
 
         # track if gpu was requested for checkpointing
         self.on_gpu = False
@@ -71,15 +71,6 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         """
         raise NotImplementedError
 
-    def update_tng_log_metrics(self, logs):
-        """
-        Chance to update metrics to be logged for training step.
-        For example, add music, images, etc... to log
-        :param logs:
-        :return:
-        """
-        return logs
-
     def loss(self, *args, **kwargs):
         """
         Expand model_out into your components
@@ -92,7 +83,6 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         model_summary = ModelSummary(self)
         print(model_summary)
 
-
     def freeze(self):
         for param in self.parameters():
             param.requires_grad = False
@@ -125,16 +115,6 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         """
         raise NotImplementedError
 
-    @staticmethod
-    def get_process_position(gpus):
-        try:
-            current_gpu = os.environ["CUDA_VISIBLE_DEVICES"]
-            gpu_ids = gpus.split(',')
-            process_position = gpu_ids.index(current_gpu)
-            return process_position, current_gpu
-        except Exception as e:
-            return 0, 0
-
     @classmethod
     def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
         """

pytorch_lightning/testing_models/lm_test_module.py

@@ -0,0 +1,280 @@
+import os
+from collections import OrderedDict
+import torch.nn as nn
+from torchvision.datasets import MNIST
+import torchvision.transforms as transforms
+import torch
+import torch.nn.functional as F
+from test_tube import HyperOptArgumentParser
+from torch import optim
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from pytorch_lightning.root_module.root_module import LightningModule
+
+
+class LightningTestModel(LightningModule):
+    """
+    Sample model to show how to define a template
+    """
+
+    def __init__(self, hparams, force_remove_distributed_sampler=False):
+        """
+        Pass in parsed HyperOptArgumentParser to the model
+        :param hparams:
+        """
+        # init superclass
+        super(LightningTestModel, self).__init__(hparams)
+
+        self.batch_size = hparams.batch_size
+
+        # if you specify an example input, the summary will show input/output for each layer
+        self.example_input_array = torch.rand(5, 28 * 28)
+
+        # remove to test warning for dist sampler
+        self.force_remove_distributed_sampler = force_remove_distributed_sampler
+
+        # build model
+        self.__build_model()
+
+    # ---------------------
+    # MODEL SETUP
+    # ---------------------
+    def __build_model(self):
+        """
+        Layout model
+        :return:
+        """
+        self.c_d1 = nn.Linear(in_features=self.hparams.in_features, out_features=self.hparams.hidden_dim)
+        self.c_d1_bn = nn.BatchNorm1d(self.hparams.hidden_dim)
+        self.c_d1_drop = nn.Dropout(self.hparams.drop_prob)
+
+        self.c_d2 = nn.Linear(in_features=self.hparams.hidden_dim, out_features=self.hparams.out_features)
+
+    # ---------------------
+    # TRAINING
+    # ---------------------
+    def forward(self, x):
+        """
+        No special modification required for lightning, define as you normally would
+        :param x:
+        :return:
+        """
+
+        x = self.c_d1(x)
+        x = torch.tanh(x)
+        x = self.c_d1_bn(x)
+        x = self.c_d1_drop(x)
+
+        x = self.c_d2(x)
+        logits = F.log_softmax(x, dim=1)
+
+        return logits
+
+    def loss(self, labels, logits):
+        nll = F.nll_loss(logits, labels)
+        return nll
+
+    def training_step(self, data_batch, batch_i):
+        """
+        Lightning calls this inside the training loop
+        :param data_batch:
+        :return:
+        """
+        # forward pass
+        x, y = data_batch
+        x = x.view(x.size(0), -1)
+
+        y_hat = self.forward(x)
+
+        # calculate loss
+        loss_val = self.loss(y, y_hat)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+
+        output = OrderedDict({
+            'loss': loss_val
+        })
+
+        # can also return just a scalar instead of a dict (return loss_val)
+        return output
+
+    def validation_step(self, data_batch, batch_i):
+        """
+        Lightning calls this inside the validation loop
+        :param data_batch:
+        :return:
+        """
+        x, y = data_batch
+        x = x.view(x.size(0), -1)
+        y_hat = self.forward(x)
+
+        loss_val = self.loss(y, y_hat)
+
+        # acc
+        labels_hat = torch.argmax(y_hat, dim=1)
+        val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+        val_acc = torch.tensor(val_acc)
+
+        if self.on_gpu:
+            val_acc = val_acc.cuda(loss_val.device.index)
+
+        # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+        if self.trainer.use_dp:
+            loss_val = loss_val.unsqueeze(0)
+            val_acc = val_acc.unsqueeze(0)
+
+        # alternate possible outputs to test
+        if self.trainer.batch_nb % 1 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+            })
+            return output
+        if self.trainer.batch_nb % 2 == 0:
+            return val_acc
+
+        if self.trainer.batch_nb % 3 == 0:
+            output = OrderedDict({
+                'val_loss': loss_val,
+                'val_acc': val_acc,
+                'test_dic': {'val_loss_a': loss_val}
+            })
+            return output
+
+    def validation_end(self, outputs):
+        """
+        Called at the end of validation to aggregate outputs
+        :param outputs: list of individual outputs of each validation step
+        :return:
+        """
+        # if returned a scalar from validation_step, outputs is a list of tensor scalars
+        # we return just the average in this case (if we want)
+        # return torch.stack(outputs).mean()
+
+        val_loss_mean = 0
+        val_acc_mean = 0
+        for output in outputs:
+            val_loss_mean += output['val_loss']
+            val_acc_mean += output['val_acc']
+
+        val_loss_mean /= len(outputs)
+        val_acc_mean /= len(outputs)
+
+        tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+        return tqdm_dic
+
+    def on_tng_metrics(self, logs):
+        logs['some_tensor_to_test'] = torch.rand(1)
+
+    # ---------------------
+    # MODEL SAVING
+    # ---------------------
+    def get_save_dict(self):
+        checkpoint = {'state_dict': self.state_dict()}
+        return checkpoint
+
+    def load_model_specific(self, checkpoint):
+        self.load_state_dict(checkpoint['state_dict'])
+        pass
+
+    # ---------------------
+    # TRAINING SETUP
+    # ---------------------
+    def configure_optimizers(self):
+        """
+        return whatever optimizers we want here
+        :return: list of optimizers
+        """
+        optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
+        return [optimizer]
+
+    def __dataloader(self, train):
+        # init data generators
+        transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
+        dataset = MNIST(root=self.hparams.data_root, train=train, transform=transform, download=True)
+
+        # when using multi-node we need to add the datasampler
+        train_sampler = None
+        batch_size = self.hparams.batch_size
+
+        try:
+            if self.on_gpu and not self.force_remove_distributed_sampler:
+                train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
+                batch_size = batch_size // self.trainer.world_size  # scale batch size
+        except Exception as e:
+            pass
+
+        should_shuffle = train_sampler is None
+        loader = DataLoader(
+            dataset=dataset,
+            batch_size=batch_size,
+            shuffle=should_shuffle,
+            sampler=train_sampler
+        )
+
+        return loader
+
+    @property
+    def tng_dataloader(self):
+        if self._tng_dataloader is None:
+            try:
+                self._tng_dataloader = self.__dataloader(train=True)
+            except Exception as e:
+                print(e)
+                raise e
+        return self._tng_dataloader
+
+    @property
+    def val_dataloader(self):
+        if self._val_dataloader is None:
+            try:
+                self._val_dataloader = self.__dataloader(train=False)
+            except Exception as e:
+                print(e)
+                raise e
+        return self._val_dataloader
+
+    @property
+    def test_dataloader(self):
+        if self._test_dataloader is None:
+            try:
+                self._test_dataloader = self.__dataloader(train=False)
+            except Exception as e:
+                print(e)
+                raise e
+        return self._test_dataloader
+
+    @staticmethod
+    def add_model_specific_args(parent_parser, root_dir):
+        """
+        Parameters you define here will be available to your model through self.hparams
+        :param parent_parser:
+        :param root_dir:
+        :return:
+        """
+        parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
+
+        # param overwrites
+        # parser.set_defaults(gradient_clip=5.0)
+
+        # network params
+        parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
+        parser.add_argument('--in_features', default=28*28, type=int)
+        parser.add_argument('--out_features', default=10, type=int)
+        parser.add_argument('--hidden_dim', default=50000, type=int) # use 500 for CPU, 50000 for GPU to see speed difference
+
+        # data
+        parser.add_argument('--data_root', default=os.path.join(root_dir, 'mnist'), type=str)
+
+        # training params (opt)
+        parser.opt_list('--learning_rate', default=0.001*8, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
+                        tunable=False)
+        parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False)
+
+        # if using 2 nodes with 4 gpus each the batch size here (256) will be 256 / (2*8) = 16 per gpu
+        parser.opt_list('--batch_size', default=256*8, type=int, options=[32, 64, 128, 256], tunable=False,
+                        help='batch size will be divided over all the gpus being used across all nodes')
+        return parser

pytorch_lightning/trainer_main.py

@@ -52,10 +52,6 @@ def main(hparams, cluster, results_dict):
     hparams.__setattr__('nb_gpus', torch.cuda.device_count())
     hparams.__setattr__('inference_mode', hparams.model_load_weights_path is not None)
 
-    # delay each training start to not overwrite logs
-    process_position, current_gpu = TRAINING_MODEL.get_process_position(hparams.gpus)
-    sleep(process_position + 1)
-
     # init experiment
     exp = Experiment(
         name=hparams.tt_name,

pytorch_lightning/utils/debugging.py

@@ -0,0 +1,5 @@
+import pdb
+import sys
+
+class MisconfigurationException(Exception):
+    pass

pytorch_lightning/utils/embeddings.py

@@ -1,104 +0,0 @@
-import torch
-import numpy as np
-from copy import deepcopy
-
-
-class PretrainedEmbedding(torch.nn.Embedding):
-
-    def __init__(self, embedding_path, embedding_dim, task_vocab, freeze=True, *args, **kwargs):
-        """
-        Loads a prebuilt pytorch embedding from any embedding formated file.
-        Padding=0 by default.
-
-        >>> emb = PretrainedEmbedding(embedding_path='glove.840B.300d.txt',embedding_dim=300, task_vocab={'hello': 1, 'world': 2})
-        >>> data = torch.Tensor([[0, 1], [0, 2]]).long()
-        >>> embedded = emb(data)
-        
-
-
-        :param embedding_path:
-        :param emb_dim:
-        :param task_vocab:
-        :param freeze:
-        :return:
-        """
-        # count the vocab
-        self.vocab_size = max(task_vocab.values()) + 1
-        super(PretrainedEmbedding, self).__init__(self.vocab_size, embedding_dim, padding_idx=0, *args, **kwargs)
-
-        # load pretrained embeddings
-        new_emb = self.__load_task_specific_embeddings(deepcopy(task_vocab), embedding_path, embedding_dim, freeze)
-
-        # transfer weights
-        self.weight = new_emb.weight
-
-        # apply freeze
-        should_freeze = not freeze
-        self.weight.requires_grad = should_freeze
-
-    def __load_task_specific_embeddings(self, vocab_words, embedding_path, emb_dim, freeze):
-        """
-        Iterates embedding file to only pull out task specific embeddings
-        :param vocab_words:
-        :param embedding_path:
-        :param emb_dim:
-        :param freeze:
-        :return:
-        """
-
-        # holds final embeddings for relevant words
-        embeddings = np.zeros(shape=(self.vocab_size, emb_dim))
-
-        # load embedding line by line and extract relevant embeddings
-        with open(embedding_path, encoding='utf-8') as f:
-            for line in f:
-                tokens = line.split(' ')
-                word = tokens[0]
-                embedding = tokens[1:]
-                embedding[-1] = embedding[-1][:-1]  # remove last new line
-
-                if word in vocab_words:
-                    vocab_word_i = vocab_words[word]
-
-                    # skip words that try to overwrite pad idx
-                    if vocab_word_i == 0:
-                        del vocab_words[word]
-                        continue
-
-                    emb_vals = np.asarray([float(x) for x in embedding])
-                    embeddings[vocab_word_i] = emb_vals
-
-                    # remove vocab word to early terminate
-                    del vocab_words[word]
-
-                # early break
-                if len(vocab_words) == 0:
-                    break
-
-        # add random vectors for the non-pretrained words
-        # these are vocab words NOT found in the pretrained embeddings
-        for w, i in vocab_words.items():
-            # skip words that try to overwrite pad idx
-            if i == 0:
-                continue
-
-            embedding = np.random.normal(size=emb_dim)
-            embeddings[i] = embedding
-
-        # turn into pt embedding
-        embeddings = torch.FloatTensor(embeddings)
-        embeddings = torch.nn.Embedding.from_pretrained(embeddings, freeze=freeze)
-
-        return embeddings
-
-
-if __name__ == '__main__':
-    emb = PretrainedEmbedding(
-        embedding_path='/Users/waf/Developer',
-        embedding_dim=300,
-        task_vocab={'hello': 1, 'world': 2}
-    )
-
-    data = torch.Tensor([[0, 1], [0, 2]]).long()
-    embedded = emb(data)
-    print(embedded)

pytorch_lightning/utils/plotting.py

@@ -1,28 +0,0 @@
-import numpy as np
-np.seterr(divide='ignore', invalid='ignore')
-
-
-def plot_confusion_matrix(cm,
-                          save_path,
-                          normalize=False,
-                          title='Confusion matrix',
-                          ylabel='y',
-                          xlabel='x'):
-    """
-    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")
-    else:
-        print('Confusion matrix, without normalization')
-
-    fig = plt.figure()
-    plt.matshow(cm)
-    plt.title(title)
-    plt.colorbar()
-    plt.ylabel(ylabel)
-    plt.xlabel(xlabel)
-    plt.savefig(save_path)

requirements.txt

@@ -1,10 +1,11 @@
-
 atomicwrites==1.2.1
 attrs==18.2.0
 certifi==2018.11.29
 cffi==1.11.5
+coverage==4.5.3
 imageio==2.4.1
 mkl-fft==1.0.6
+mkdocs==1.0.4
 mkl-random==1.0.2
 more-itertools==5.0.0
 numpy==1.15.4
@@ -14,7 +15,7 @@ Pillow==5.3.0
 pluggy==0.8.0
 py==1.7.0
 pycparser==2.19
-pytest==4.0.2
+pytest==5.0.1
 python-dateutil==2.7.5
 pytz==2018.7
 scikit-learn==0.20.2
@@ -25,7 +26,7 @@ tensorboard==1.14.0
 tensorboardX==1.7
 tensorflow==1.14.0
 test-tube==0.643
-torch==1.0.0
+torch==1.1.0
 torchvision==0.2.1
 tqdm==4.32.1
 twine==1.13.0

setup.cfg

@@ -16,6 +16,33 @@ markers =
 ignore = E731,W504
 max-line-length = 120
 
+[coverage:report]
+exclude_lines =
+    pragma: no cover
+    def __repr__
+    if self.debug:
+    if settings.DEBUG
+    raise AssertionError
+    raise NotImplementedError
+    if 0:
+    if __name__ == .__main__.:
+    except Exception as e
+    print(e)
+    print(traceback.print_exc())
+    return *
+    raise Exception
+    warnings
+    print
+    raise RuntimeError
+    break
+    pass
+    os.makedirs
+
+omit =
+    pytorch_lightning/callbacks/pt_callbacks.py
+    tests/test_models.py
+    pytorch_lightning/testing_models/lm_test_module.py
+
 [flake8]
 ignore = E731,W504,F401,F841
 max-line-length = 120

tests/README.md

@@ -0,0 +1,46 @@
+# Pytorch-Lightning Tests
+
+## Running tests
+
+To run all tests do the following:
+```bash
+git clone https://github.com/williamFalcon/pytorch-lightning
+cd pytorch-lightning
+
+# install module locally
+pip install -e .
+
+# install dev deps
+pip install -r requirements.txt
+
+# run tests
+py.test
+
+# or to generate coverage 
+pip install coverage
+coverage run tests/test_models.py   
+```
+
+To test models that require GPU make sure to run the above command on a GPU machine.
+The GPU machine must have:
+1. At least 2 GPUs.
+2. [NVIDIA-apex](https://github.com/NVIDIA/apex#linux) installed.
+
+
+### test_models.py
+This file fits a tiny model on MNIST using these different set-ups.
+1. CPU only.
+2. Single GPU with DP.
+3. Multiple (2) GPUs using DP.
+3. Multiple (2) GPUs using DDP.
+3. Multiple (2) GPUs using DP + apex (for 16-bit precision).
+3. Multiple (2) GPUs using DDP + apex (for 16-bit precision).   
+
+For each set up it also tests:
+1. model saving.
+2. model loading.
+3. predicting with a loaded model.
+4. simulated save from HPC signal.
+5. simulated load from HPC signal.
+
+

tests/debug.py

@@ -0,0 +1,142 @@
+import pytest
+from pytorch_lightning import Trainer
+from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel
+from argparse import Namespace
+from test_tube import Experiment
+from pytorch_lightning.callbacks import ModelCheckpoint
+import numpy as np
+import warnings
+import torch
+import os
+import shutil
+import pdb
+
+
+def get_model():
+    # set up model with these hyperparams
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    hparams = Namespace(**{'drop_prob': 0.2,
+                           'batch_size': 32,
+                           'in_features': 28*28,
+                           'learning_rate': 0.001*8,
+                           'optimizer_name': 'adam',
+                           'data_root': os.path.join(root_dir, 'mnist'),
+                           'out_features': 10,
+                           'hidden_dim': 1000})
+    model = LightningTemplateModel(hparams)
+
+    return model, hparams
+
+
+def get_exp(debug=True):
+    # set up exp object without actually saving logs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
+    return exp
+
+
+def init_save_dir():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    save_dir = os.path.join(root_dir, 'save_dir')
+
+    if os.path.exists(save_dir):
+        shutil.rmtree(save_dir)
+
+    os.makedirs(save_dir, exist_ok=True)
+
+    return save_dir
+
+
+def clear_save_dir():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    save_dir = os.path.join(root_dir, 'save_dir')
+    if os.path.exists(save_dir):
+        shutil.rmtree(save_dir)
+
+
+def load_model(exp, save_dir):
+
+    # load trained model
+    tags_path = exp.get_data_path(exp.name, exp.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+
+    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
+    weights_dir = os.path.join(save_dir, checkpoints[0])
+
+    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=True)
+
+    assert trained_model is not None, 'loading model failed'
+
+    return trained_model
+
+
+def run_prediction(dataloader, trained_model):
+    # run prediction on 1 batch
+    for batch in dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    y_hat = trained_model(x)
+
+    # acc
+    labels_hat = torch.argmax(y_hat, dim=1)
+    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+    val_acc = torch.tensor(val_acc)
+    val_acc = val_acc.item()
+
+    print(val_acc)
+
+    assert val_acc > 0.70, f'this model is expected to get > 0.7 in test set (it got {val_acc})'
+
+
+def mainasdf():
+
+    save_dir = init_save_dir()
+    model, hparams = get_model()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.save()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    trainer = Trainer(
+        experiment=exp,
+        checkpoint_callback=checkpoint,
+        progress_bar=True,
+        max_nb_epochs=1,
+        gpus=[0, 1],
+        distributed_backend='dp',
+        use_amp=True
+    )
+
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # test model loading
+    pretrained_model = load_model(exp, save_dir)
+
+    # test model preds
+    run_prediction(model.test_dataloader, pretrained_model)
+
+    clear_save_dir()
+
+
+
+
+if __name__ == '__main__':
+    import subprocess
+    import re
+
+    print('getting pid')
+    command = "lsof -i :%s | awk '{print $2}'" % 12910
+    pids = subprocess.check_output(command, shell=True)
+    pids = pids.strip()
+
+    print(len(pids))

tests/test_models.py

@@ -0,0 +1,506 @@
+import pytest
+from pytorch_lightning import Trainer
+from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel
+from pytorch_lightning.testing_models.lm_test_module import LightningTestModel
+from argparse import Namespace
+from test_tube import Experiment
+from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
+from pytorch_lightning.utils.debugging import MisconfigurationException
+from pytorch_lightning.root_module import memory
+from pytorch_lightning.models.trainer import reduce_distributed_output
+import numpy as np
+import warnings
+import torch
+import os
+import shutil
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+
+# ------------------------------------------------------------------------
+# TESTS
+# ------------------------------------------------------------------------
+def test_dp_output_reduce():
+
+    # test identity when we have a single gpu
+    out = torch.rand(3, 1)
+    assert reduce_distributed_output(out, nb_gpus=1) is out
+
+    # average when we have multiples
+    assert reduce_distributed_output(out, nb_gpus=2) == out.mean()
+
+    # when we have a dict of vals
+    out = {
+        'a': out,
+        'b': {
+            'c': out
+        }
+    }
+    reduced = reduce_distributed_output(out, nb_gpus=3)
+    assert reduced['a'] == out['a']
+    assert reduced['b']['c'] == out['b']['c']
+
+
+def test_amp_gpu_ddp_slurm_managed():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        progress_bar=True,
+        max_nb_epochs=1,
+        gpus=[0],
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.save()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    # add these to the trainer options
+    trainer_options['checkpoint_callback'] = checkpoint
+    trainer_options['experiment'] = exp
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    trainer.is_slurm_managing_tasks = True
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # test root model address
+    assert trainer.resolve_root_node_address('abc') == 'abc'
+    assert trainer.resolve_root_node_address('abc[23]') == 'abc23'
+    assert trainer.resolve_root_node_address('abc[23-24]') == 'abc23'
+    assert trainer.resolve_root_node_address('abc[23-24, 45-40, 40]') == 'abc23'
+
+    # test model loading with a map_location
+    map_location = 'cuda:1'
+    pretrained_model = load_model(exp, save_dir, True, map_location)
+
+    # test model preds
+    run_prediction(model.test_dataloader, pretrained_model)
+
+    if trainer.use_ddp:
+        # on hpc this would work fine... but need to hack it for the purpose of the test
+        trainer.model = pretrained_model
+        trainer.optimizers = pretrained_model.configure_optimizers()
+
+    # test HPC loading / saving
+    trainer.hpc_save(save_dir, exp)
+    trainer.hpc_load(save_dir, on_gpu=True)
+
+    # test freeze on gpu
+    model.freeze()
+    model.unfreeze()
+
+    clear_save_dir()
+
+
+def test_early_stopping_cpu_model():
+    """
+    Test each of the trainer options
+    :return:
+    """
+
+    stopping = EarlyStopping()
+    trainer_options = dict(
+        early_stop_callback=stopping,
+        gradient_clip=1.0,
+        overfit_pct=0.20,
+        track_grad_norm=2,
+        print_nan_grads=True,
+        progress_bar=False,
+        experiment=get_exp(),
+        train_percent_check=0.1,
+        val_percent_check=0.1
+    )
+
+    model, hparams = get_model()
+    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+    # test freeze on cpu
+    model.freeze()
+    model.unfreeze()
+
+
+def test_cpu_model_with_amp():
+    """
+    Make sure model trains on CPU
+    :return:
+    """
+
+    trainer_options = dict(
+        progress_bar=False,
+        experiment=get_exp(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4,
+        use_amp=True
+    )
+
+    model, hparams = get_model()
+
+    with pytest.raises(MisconfigurationException):
+        run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_cpu_model():
+    """
+    Make sure model trains on CPU
+    :return:
+    """
+
+    trainer_options = dict(
+        progress_bar=False,
+        experiment=get_exp(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4
+    )
+
+    model, hparams = get_model()
+
+    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_all_features_cpu_model():
+    """
+    Test each of the trainer options
+    :return:
+    """
+
+    trainer_options = dict(
+        gradient_clip=1.0,
+        overfit_pct=0.20,
+        track_grad_norm=2,
+        print_nan_grads=True,
+        progress_bar=False,
+        experiment=get_exp(),
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.4
+    )
+
+    model, hparams = get_model()
+    run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_single_gpu_model():
+    """
+    Make sure single GPU works (DP mode)
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_single_gpu_model cannot run. Rerun on a GPU node to run this test')
+        return
+    model, hparams = get_model()
+
+    trainer_options = dict(
+        progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.1,
+        val_percent_check=0.1,
+        gpus=[0]
+    )
+
+    run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_multi_gpu_model_dp():
+    """
+    Make sure DP works
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+    model, hparams = get_model()
+    trainer_options = dict(
+        progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.1,
+        val_percent_check=0.1,
+        gpus='-1'
+    )
+
+    run_gpu_model_test(trainer_options, model, hparams)
+
+    # test memory helper functions
+    memory.get_gpu_memory_map()
+
+
+def test_amp_gpu_dp():
+    """
+    Make sure DP + AMP work
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+    model, hparams = get_model()
+    trainer_options = dict(
+        max_nb_epochs=1,
+        gpus='0, 1',  # test init with gpu string
+        distributed_backend='dp',
+        use_amp=True
+    )
+    with pytest.raises(MisconfigurationException):
+        run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_multi_gpu_model_ddp():
+    """
+    Make sure DDP works
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+    model, hparams = get_model()
+    trainer_options = dict(
+        progress_bar=False,
+        max_nb_epochs=1,
+        train_percent_check=0.4,
+        val_percent_check=0.2,
+        gpus=[0, 1],
+        distributed_backend='ddp'
+    )
+
+    run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_amp_gpu_ddp():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    trainer_options = dict(
+        progress_bar=True,
+        max_nb_epochs=1,
+        gpus=[0, 1],
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_ddp_sampler_error():
+    """
+    Make sure DDP + AMP work
+    :return:
+    """
+    if not torch.cuda.is_available():
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
+        return
+    if not torch.cuda.device_count() > 1:
+        warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+        return
+
+    os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+
+    hparams = get_hparams()
+    model = LightningTestModel(hparams, force_remove_distributed_sampler=True)
+
+    exp = get_exp(True)
+    exp.save()
+
+    trainer = Trainer(
+        experiment=exp,
+        progress_bar=False,
+        max_nb_epochs=1,
+        gpus=[0, 1],
+        distributed_backend='ddp',
+        use_amp=True
+    )
+
+    with pytest.raises(MisconfigurationException):
+        trainer.get_dataloaders(model)
+
+    clear_save_dir()
+
+
+# ------------------------------------------------------------------------
+# UTILS
+# ------------------------------------------------------------------------
+def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.save()
+
+    # exp file to get weights
+    checkpoint = ModelCheckpoint(save_dir)
+
+    # add these to the trainer options
+    trainer_options['checkpoint_callback'] = checkpoint
+    trainer_options['experiment'] = exp
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # correct result and ok accuracy
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # test model loading
+    pretrained_model = load_model(exp, save_dir, on_gpu)
+
+    # test model preds
+    run_prediction(model.test_dataloader, pretrained_model)
+
+    if trainer.use_ddp:
+        # on hpc this would work fine... but need to hack it for the purpose of the test
+        trainer.model = pretrained_model
+        trainer.optimizers = pretrained_model.configure_optimizers()
+
+    # test HPC loading / saving
+    trainer.hpc_save(save_dir, exp)
+    trainer.hpc_load(save_dir, on_gpu=on_gpu)
+
+    clear_save_dir()
+
+
+def get_hparams():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    hparams = Namespace(**{'drop_prob': 0.2,
+                           'batch_size': 32,
+                           'in_features': 28*28,
+                           'learning_rate': 0.001*8,
+                           'optimizer_name': 'adam',
+                           'data_root': os.path.join(root_dir, 'mnist'),
+                           'out_features': 10,
+                           'hidden_dim': 1000})
+    return hparams
+
+
+def get_model():
+    # set up model with these hyperparams
+    hparams = get_hparams()
+    model = LightningTemplateModel(hparams)
+
+    return model, hparams
+
+
+def get_exp(debug=True):
+    # set up exp object without actually saving logs
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
+    return exp
+
+
+def init_save_dir():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    save_dir = os.path.join(root_dir, 'save_dir')
+
+    if os.path.exists(save_dir):
+        shutil.rmtree(save_dir)
+
+    os.makedirs(save_dir, exist_ok=True)
+
+    return save_dir
+
+
+def clear_save_dir():
+    root_dir = os.path.dirname(os.path.realpath(__file__))
+    save_dir = os.path.join(root_dir, 'save_dir')
+    if os.path.exists(save_dir):
+        shutil.rmtree(save_dir)
+
+
+def load_model(exp, save_dir, on_gpu, map_location=None):
+
+    # load trained model
+    tags_path = exp.get_data_path(exp.name, exp.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+
+    checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
+    weights_dir = os.path.join(save_dir, checkpoints[0])
+
+    trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir,
+                                                             tags_csv=tags_path,
+                                                             on_gpu=on_gpu,
+                                                             map_location=map_location)
+
+    assert trained_model is not None, 'loading model failed'
+
+    return trained_model
+
+
+def run_prediction(dataloader, trained_model):
+    # run prediction on 1 batch
+    for batch in dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    y_hat = trained_model(x)
+
+    # acc
+    labels_hat = torch.argmax(y_hat, dim=1)
+    val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+    val_acc = torch.tensor(val_acc)
+    val_acc = val_acc.item()
+
+    print(val_acc)
+
+    assert val_acc > 0.50, f'this model is expected to get > 0.50 in test set (it got {val_acc})'
+
+
+def assert_ok_acc(trainer):
+    # this model should get 0.80+ acc
+    acc = trainer.tng_tqdm_dic['val_acc']
+    assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])