release v0.3.6.4

Loading

README.md

@@ -9,12 +9,10 @@
 <p align="center">
   The Keras for ML researchers using PyTorch. More control. Less boilerplate.    
 </p>
+
 <p align="center">
   <a href="https://badge.fury.io/py/pytorch-lightning"><img src="https://badge.fury.io/py/pytorch-lightning.svg" alt="PyPI version" height="18"></a>
   <a href="https://pepy.tech/project/pytorch-lightning"><img src="https://pepy.tech/badge/pytorch-lightning" alt="PyPI version" height="18"></a>
-</p>   
-
-<p align="center">
   <a href="https://github.com/williamFalcon/pytorch-lightning/tree/master/tests"><img src="https://github.com/williamFalcon/pytorch-lightning/blob/master/coverage.svg"></a>
   <a href="https://travis-ci.org/williamFalcon/pytorch-lightning"><img src="https://travis-ci.org/williamFalcon/pytorch-lightning.svg?branch=master"></a>
   <a href="https://williamfalcon.github.io/pytorch-lightning/"><img src="https://readthedocs.org/projects/pytorch-lightning/badge/?version=latest"></a>
@@ -37,6 +35,8 @@ When starting a new project the last thing you want to do is recode a training l
 
 With lightning, you guarantee those parts of your code work so you can focus on what the meat of the research: Data and training, validation loop logic. Don't worry about multiple gpus or speeding up your code, lightning will do that for you!
 
+## How do I do use it?   
+
 To use lightning do 2 things:  
 1. [Define a LightningModel](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/)         
 ```python
@@ -220,7 +220,7 @@ And run tensorboard from that dir
 tensorboard --logdir /some/path     
 ```    
 
-## Lightning automatically automates all of the following ([each is also configurable](https://williamfalcon.github.io/pytorch-lightning/Trainer/)):
+## Lightning automates all of the following ([each is also configurable](https://williamfalcon.github.io/pytorch-lightning/Trainer/)):
 
 ###### Checkpointing    
 
@@ -287,19 +287,22 @@ pip install pytorch-lightning
 
 # clone lightning for the demo
 git clone https://github.com/williamFalcon/pytorch-lightning.git
-cd examples/new_project_templates/
+cd pytorch_lightning/examples/new_project_templates/
 
-# run demo (on cpu)
-python trainer_gpu_cluster_template.py
+# all of the following demos use the SAME model to show no modification needs to be made to your code
+
+# train on cpu 
+python single_cpu_template.py
+
+# train on multiple-gpus 
+python single_gpu_node_template.py --gpus "0,1"
+
+# train on 32 gpus on a cluster (run on a SLURM managed cluster)
+python multi_node_cluster_template.py --nb_gpu_nodes 4 --gpus '0,1,2,3,4,5,6,7'
 ```
 
-Without changing the model AT ALL, you can run the model on a single gpu, over multiple gpus, or over multiple nodes.
+## Bleeding edge
+If you can't wait for the next release, install the most up to date code with:  
 ```bash
-# run a grid search on two gpus
-python fully_featured_trainer.py --gpus "0;1"
-
-# run single model on multiple gpus
-python fully_featured_trainer.py --gpus "0;1" --interactive
-```
-
-
+pip install git+https://github.com/williamFalcon/pytorch-lightning.git@master --upgrade
+```

docs/LightningModule/RequiredTrainerInterface.md

@@ -14,8 +14,6 @@ Otherwise, to Define a Lightning Module, implement the following methods:
 - [validation_end](RequiredTrainerInterface.md#validation_end)
     
 - [configure_optimizers](RequiredTrainerInterface.md#configure_optimizers)
-- [get_save_dict](RequiredTrainerInterface.md#get_save_dict)
-- [load_model_specific](RequiredTrainerInterface.md#load_model_specific)
 
 - [tng_dataloader](RequiredTrainerInterface.md#tng_dataloader)
 - [tng_dataloader](RequiredTrainerInterface.md#tng_dataloader)
@@ -23,6 +21,8 @@ Otherwise, to Define a Lightning Module, implement the following methods:
 
 **Optional**:   
 
+- [on_save_checkpoint](RequiredTrainerInterface.md#on_save_checkpoint)
+- [on_load_checkpoint](RequiredTrainerInterface.md#on_load_checkpoint)
 - [update_tng_log_metrics](RequiredTrainerInterface.md#update_tng_log_metrics)
 - [add_model_specific_args](RequiredTrainerInterface.md#add_model_specific_args)
 
@@ -245,34 +245,35 @@ def configure_optimizers(self):
 ```
 
 --- 
-### get_save_dict 
+### on_save_checkpoint 
 
 ``` {.python}
-def get_save_dict(self)
+def on_save_checkpoint(self, checkpoint)
 ```
-Called by lightning to checkpoint your model. Lightning saves current epoch, current batch nb, etc...
-All you have to return is what specifically about your lightning model you want to checkpoint.
+Called by lightning to checkpoint your model. Lightning saves the training state (current epoch, global_step, etc)
+and also saves the model state_dict. If you want to save anything else, use this method to add your own
+key-value pair.
 
 ##### Return
-Dictionary - No required keys. Most of the time as described in this example.   
+Nothing
 
 **Example**
 
 ``` {.python}
-def get_save_dict(self):
-    # 99% of use cases this is all you need to return
-    checkpoint = {'state_dict': self.state_dict()}
-    return checkpoint
+def on_save_checkpoint(self, checkpoint):
+    # 99% of use cases you don't need to implement this method 
+    checkpoint['something_cool_i_want_to_save'] = my_cool_pickable_object
 ```
 
 --- 
-### load_model_specific 
+### on_load_checkpoint 
 
 ``` {.python}
-def load_model_specific(self, checkpoint)
+def on_load_checkpoint(self, checkpoint)
 ```
-Called by lightning to restore your model. This is your chance to restore your model using the keys you added in get_save_dict.
-Lightning will automatically restore current epoch, batch nb, etc. 
+Called by lightning to restore your model. Lighting auto-restores global step, epoch, etc...
+It also restores the model state_dict.
+If you saved something with **on_save_checkpoint** this is your chance to restore this.
 
 ##### Return
 Nothing  
@@ -280,9 +281,9 @@ Nothing
 **Example**
 
 ``` {.python}
-def load_model_specific(self, checkpoint):
-    # you defined 'state_dict' in get_save_dict()
-    self.load_state_dict(checkpoint['state_dict'])
+def on_load_checkpoint(self, checkpoint):
+    # 99% of the time you don't need to implement this method
+    self.something_cool_i_want_to_save = checkpoint['something_cool_i_want_to_save']
 ```
 
 --- 

docs/LightningModule/methods.md

@@ -21,7 +21,8 @@ pretrained_model = MyLightningModule.load_from_metrics(
     map_location=None
 )
     
-# predict 
+# predict
+pretrained_model.eval()
 pretrained_model.freeze()
 y_hat = pretrained_model(x)
 ```

docs/index.md

@@ -1,10 +1,11 @@
 ###### New project Quick Start    
 To start a new project define these two files.   
 
-1. [Define a LightningModule](/LightningModule/RequiredTrainerInterface/) 
-2. Pick a trainer      
-    - [Basic CPU Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/examples/new_project_templates/single_cpu_template.py) 
-    - [GPU cluster Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/examples/new_project_templates/multi_node_cluster_template.py)
+1. [Define a LightningModule](/pytorch-lightning/LightningModule/RequiredTrainerInterface/) 
+2. [Define a trainer](https://williamfalcon.github.io/pytorch-lightning/Trainer/)
+    - [Basic CPU Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/examples/new_project_templates/single_cpu_template.py) 
+    - [Multi-GPU Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/examples/new_project_templates/single_gpu_node_template.py)
+    - [GPU cluster Trainer Template](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/examples/new_project_templates/multi_node_cluster_template.py)
 
 ###### Docs shortcuts
 - [LightningModule](LightningModule/RequiredTrainerInterface/)  

pytorch_lightning/examples/new_project_templates/lightning_module_template.py

@@ -154,17 +154,6 @@ class LightningTemplateModel(LightningModule):
         tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
         return tqdm_dic
 
-    # ---------------------
-    # 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
     # ---------------------

pytorch_lightning/models/sample_model_template/model_template.py

@@ -1,203 +0,0 @@
-import torch.nn as nn
-import numpy as np
-from pytorch_lightning import LightningModule
-from test_tube import HyperOptArgumentParser
-from torchvision.datasets import MNIST
-import torchvision.transforms as transforms
-import torch
-import torch.nn.functional as F
-
-
-class ExampleModel1(LightningModule):
-    """
-    Sample model to show how to define a template
-    """
-
-    def __init__(self, hparams):
-        # init superclass
-        super(ExampleModel1, self).__init__(hparams)
-
-        self.batch_size = hparams.batch_size
-
-        # 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):
-        x = self.c_d1(x)
-        x = F.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):
-        """
-        Called 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)
-
-        tqdm_dic = {'jefe': 1}
-        return loss_val, tqdm_dic
-
-    def validation_step(self, data_batch):
-        """
-        Called 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)
-
-        output = {'y_hat': y_hat, 'val_loss': loss_val.item(), 'val_acc': val_acc}
-        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:
-        """
-        val_loss_mean = 0
-        accs = []
-        for output in outputs:
-            val_loss_mean += output['val_loss']
-            accs.append(output['val_acc'])
-
-        val_loss_mean /= len(outputs)
-        tqdm_dic = {'val_loss': val_loss_mean, 'val_acc': np.mean(accs)}
-        return tqdm_dic
-
-    def update_tng_log_metrics(self, logs):
-        return logs
-
-    # ---------------------
-    # 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 = self.choose_optimizer(self.hparams.optimizer_name, self.parameters(), {'lr': self.hparams.learning_rate}, 'optimizer')
-        self.optimizers = [optimizer]
-        return self.optimizers
-
-    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)
-
-        loader = torch.utils.data.DataLoader(
-            dataset=dataset,
-            batch_size=self.hparams.batch_size,
-            shuffle=True
-        )
-
-        return loader
-
-    @data_loader
-    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):
-        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)
-        parser.add_argument('--hidden_dim', default=500)
-        parser.add_argument('--out_features', default=10)
-
-        # data
-        parser.add_argument('--data_root', default='/Users/williamfalcon/Developer/personal/research_lib/research_proj/datasets/mnist', type=str)
-
-        # training params (opt)
-        parser.opt_list('--learning_rate', default=0.001, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
-                        tunable=False)
-        parser.opt_list('--batch_size', default=256, type=int, options=[32, 64, 128, 256], tunable=False)
-        parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False)
-        return parser

pytorch_lightning/models/trainer.py

@@ -24,7 +24,7 @@ from pytorch_lightning.utils.debugging import MisconfigurationException
 try:
     from apex import amp
     APEX_AVAILABLE = True
-except ModuleNotFoundError: # pragma: no cover
+except Exception:
     APEX_AVAILABLE = False
 
 
@@ -161,7 +161,6 @@ class Trainer(TrainerIO):
         self.nb_tng_batches = None
         self.nb_test_batches = None
 
-
         # gpus come in as a string.
         # if gpus = -1 then use all available devices
         # otherwise, split the string using commas
@@ -500,6 +499,9 @@ class Trainer(TrainerIO):
         self.proc_rank = self.node_rank * len(self.data_parallel_device_ids) + gpu_nb
         self.world_size = self.nb_gpu_nodes * len(self.data_parallel_device_ids)
 
+        # let the exp know the rank to avoid overwriting logs
+        self.experiment.rank = self.proc_rank
+
         # set up server using proc 0's ip address
         # try to init for 20 times at max in case ports are taken
         # where to store ip_table
@@ -608,14 +610,18 @@ class Trainer(TrainerIO):
         if self.proc_rank == 0:
             self.experiment.save()
 
+        # track model now.
+        # if cluster resets state, the model will update with the saved weights
+        self.model = model
+
         # enable cluster checkpointing
+        # also restores training state
         if self.cluster is not None:  # pragma: no cover
             self.enable_auto_hpc_walltime_manager()
 
         # ---------------------------
         # CORE TRAINING LOOP
         # ---------------------------
-        self.model = model
         self.__train()
 
     def __train(self):
@@ -854,30 +860,25 @@ class Trainer(TrainerIO):
         elif not can_check_epoch:
             return
 
-        try:
-            # hook
-            if self.__is_function_implemented('on_pre_performance_check'):
-                model = self.__get_model()
-                model.on_pre_performance_check()
+        # hook
+        if self.__is_function_implemented('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
-            max_batches = None if not self.fast_dev_run else 1
-            model_specific_tqdm_metrics_dic = self.validate(
-                self.model,
-                self.val_dataloader,
-                max_batches
-            )
-            self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
+        # use full val set on end of epoch
+        # use a small portion otherwise
+        max_batches = None if not self.fast_dev_run else 1
+        model_specific_tqdm_metrics_dic = self.validate(
+            self.model,
+            self.val_dataloader,
+            max_batches
+        )
+        self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
 
-            # hook
-            if self.__is_function_implemented('on_post_performance_check'):
-                model = self.__get_model()
-                model.on_post_performance_check()
-
-        except Exception as e:
-            print(e)
-            print(traceback.print_exc())
+        # hook
+        if self.__is_function_implemented('on_post_performance_check'):
+            model = self.__get_model()
+            model.on_post_performance_check()
 
         if self.progress_bar:
             # add model specific metrics
@@ -885,6 +886,6 @@ class Trainer(TrainerIO):
             self.prog_bar.set_postfix(**tqdm_metrics)
 
         # model checkpointing
-        if self.proc_rank == 0 and self.checkpoint_callback:
+        if self.proc_rank == 0 and self.checkpoint_callback is not None:
             print('save callback...')
             self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic)

pytorch_lightning/root_module/model_saving.py

@@ -4,34 +4,36 @@ import re
 import pdb
 from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
 
+
 class ModelIO(object):
 
-    def load_model_specific(self, checkpoint):
+    def on_load_checkpoint(self, checkpoint):
         """
         Do something with the checkpoint
+        Gives model a chance to load something before state_dict is restored
         :param checkpoint:
         :return:
         """
-        raise NotImplementedError
+        pass
 
-    def get_save_dict(self):
+    def on_save_checkpoint(self, checkpoint):
         """
-        Return specific things for the model
-        :return:
+        Give the model a chance to add something to the checkpoint.
+        state_dict is already there
         """
-        raise NotImplementedError
+        pass
 
     # -------------------------
     # OPTIONAL HOOKS
     # -------------------------
-    def on_hpc_save(self):
+    def on_hpc_save(self, checkpoint):
         """
         Hook to do whatever you need right before Slurm manager saves the model
         :return:
         """
         pass
 
-    def on_hpc_load(self):
+    def on_hpc_load(self, checkpoint):
         """
         Hook to do whatever you need right before Slurm manager loads the model
         :return:
@@ -42,7 +44,6 @@ 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
@@ -76,18 +77,19 @@ class TrainerIO(object):
 
         checkpoint['optimizer_states'] = optimizer_states
 
-        # request what to save from the model
+        # add the state_dict from the model
         model = self.__get_model()
-        checkpoint_dict = model.get_save_dict()
+        checkpoint['state_dict'] = model.state_dict()
+
+        # give the model a chance to add a few things
+        model.on_save_checkpoint(checkpoint)
 
-        # merge trainer and model saving items
-        checkpoint.update(checkpoint_dict)
         return checkpoint
 
     # --------------------
     # HPC IO
     # --------------------
-    def enable_auto_hpc_walltime_manager(self):  # pragma: no cover
+    def enable_auto_hpc_walltime_manager(self):
         if self.cluster is None:
             return
 
@@ -150,13 +152,14 @@ class TrainerIO(object):
 
         # give model a chance to do something on hpc_save
         model = self.__get_model()
-        model.on_hpc_save()
+        checkpoint = self.dump_checkpoint()
 
-        # request what to save from the model
-        checkpoint_dict = self.dump_checkpoint()
+        model.on_hpc_save(checkpoint)
 
         # do the actual save
-        torch.save(checkpoint_dict, filepath)
+        torch.save(checkpoint, filepath)
+
+        return filepath
 
     def hpc_load(self, folderpath, on_gpu):
         filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, self.max_ckpt_in_folder(folderpath))
@@ -166,15 +169,17 @@ class TrainerIO(object):
         else:
             checkpoint = torch.load(filepath, map_location=lambda storage, loc: storage)
 
-        # load training state
+        # load training state (affects trainer only)
         self.restore_training_state(checkpoint)
 
         # load model state
         model = self.__get_model()
-        model.load_model_specific(checkpoint)
+
+        # load the state_dict on the model automatically
+        model.load_state_dict(checkpoint['state_dict'])
 
         # call model hook
-        model.on_hpc_load()
+        model.on_hpc_load(checkpoint)
 
     def max_ckpt_in_folder(self, path):
         files = os.listdir(path)

pytorch_lightning/root_module/root_module.py

@@ -108,11 +108,13 @@ class LightningModule(GradInformation, ModelIO, ModelHooks):
         else:
             checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
 
+        # load the state_dict on the model automatically
         model = cls(hparams)
+        model.load_state_dict(checkpoint['state_dict'])
+
+        # give model a chance to load something
+        model.on_load_checkpoint(checkpoint)
 
-        # allow model to load
-        model.load_model_specific(checkpoint)
-        model.load_state_dict(checkpoint['state_dict'], strict=False)
         return model
 
     def summarize(self):

pytorch_lightning/testing_models/lm_test_module.py

@@ -171,17 +171,6 @@ class LightningTestModel(LightningModule):
     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
     # ---------------------

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.3.6',
+    version='0.3.6.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.6.7.1",
+        "test-tube>=0.6.7.4",
     ],
     packages=find_packages(),
     long_description=open("README.md", encoding="utf-8").read(),

tests/test_models.py

@@ -3,16 +3,18 @@ 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 test_tube import Experiment, SlurmCluster
 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
+from pytorch_lightning.root_module import model_saving
 import numpy as np
 import warnings
 import torch
 import os
 import shutil
+import pdb
 
 SEED = 2334
 torch.manual_seed(SEED)
@@ -22,6 +24,108 @@ np.random.seed(SEED)
 # ------------------------------------------------------------------------
 # TESTS
 # ------------------------------------------------------------------------
+
+def test_cpu_slurm_save_load():
+    """
+    Verify model save/load/checkpoint on CPU
+    :return:
+    """
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.save()
+
+    cluster_a = SlurmCluster()
+    trainer_options = dict(
+        max_nb_epochs=1,
+        cluster=cluster_a,
+        experiment=exp,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+    real_global_step = trainer.global_step
+
+    # traning complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # predict with trained model before saving
+    # make a prediction
+    for batch in model.test_dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    model.eval()
+    pred_before_saving = model(x)
+
+    # test registering a save function
+    trainer.enable_auto_hpc_walltime_manager()
+
+    # test HPC saving
+    # simulate snapshot on slurm
+    saved_filepath = trainer.hpc_save(save_dir, exp)
+    assert os.path.exists(saved_filepath)
+
+    # wipe-out trainer and model
+    # retrain with not much data... this simulates picking training back up after slurm
+    # we want to see if the weights come back correctly
+    continue_tng_hparams = get_hparams(continue_training=True, hpc_exp_number=cluster_a.hpc_exp_number)
+    trainer_options = dict(
+        max_nb_epochs=1,
+        cluster=SlurmCluster(continue_tng_hparams),
+        experiment=exp,
+        checkpoint_callback=ModelCheckpoint(save_dir),
+    )
+    trainer = Trainer(**trainer_options)
+    model = LightningTestModel(hparams)
+
+    # set the epoch start hook so we can predict before the model does the full training
+    def assert_pred_same():
+        assert trainer.global_step == real_global_step and trainer.global_step > 0
+
+        # predict with loaded model to make sure answers are the same
+        trainer.model.eval()
+        new_pred = trainer.model(x)
+        assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
+
+    model.on_epoch_start = assert_pred_same
+
+    # by calling fit again, we trigger training, loading weights from the cluster
+    # and our hook to predict using current model before any more weight updates
+    trainer.fit(model)
+
+    clear_save_dir()
+
+
+def test_loading_meta_tags():
+    hparams = get_hparams()
+
+    save_dir = init_save_dir()
+
+    # save tags
+    exp = get_exp(False)
+    exp.tag({'some_str':'a_str', 'an_int': 1, 'a_float': 2.0})
+    exp.argparse(hparams)
+    exp.save()
+
+    # load tags
+    tags_path = exp.get_data_path(exp.name, exp.version) + '/meta_tags.csv'
+    tags = model_saving.load_hparams_from_tags_csv(tags_path)
+
+    assert tags.batch_size == 32 and tags.hidden_dim == 1000
+
+    clear_save_dir()
+
+
 def test_dp_output_reduce():
 
     # test identity when we have a single gpu
@@ -43,6 +147,74 @@ def test_dp_output_reduce():
     assert reduced['b']['c'] == out['b']['c']
 
 
+def test_model_saving_loading():
+    """
+    Tests use case where trainer saves the model, and user loads it from tags independently
+    :return:
+    """
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    save_dir = init_save_dir()
+
+    # exp file to get meta
+    exp = get_exp(False)
+    exp.argparse(hparams)
+    exp.save()
+
+    trainer_options = dict(
+        max_nb_epochs=1,
+        cluster=SlurmCluster(),
+        experiment=exp,
+        checkpoint_callback=ModelCheckpoint(save_dir)
+    )
+
+    # fit model
+    trainer = Trainer(**trainer_options)
+    result = trainer.fit(model)
+
+    # traning complete
+    assert result == 1, 'amp + ddp model failed to complete'
+
+    # make a prediction
+    for batch in model.test_dataloader:
+        break
+
+    x, y = batch
+    x = x.view(x.size(0), -1)
+
+    # generate preds before saving model
+    model.eval()
+    pred_before_saving = model(x)
+
+    # save model
+    new_weights_path = os.path.join(save_dir, 'save_test.ckpt')
+    trainer.save_checkpoint(new_weights_path)
+
+    # load new model
+    tags_path = exp.get_data_path(exp.name, exp.version)
+    tags_path = os.path.join(tags_path, 'meta_tags.csv')
+    model_2 = LightningTestModel.load_from_metrics(weights_path=new_weights_path, tags_csv=tags_path, on_gpu=False)
+    model_2.eval()
+
+    # make prediction
+    # assert that both predictions are the same
+    new_pred = model_2(x)
+    assert torch.all(torch.eq(pred_before_saving, new_pred)).item() == 1
+
+    clear_save_dir()
+
+
+
+
+def test_model_freeze_unfreeze():
+    hparams = get_hparams()
+    model = LightningTestModel(hparams)
+
+    model.freeze()
+    model.unfreeze()
+
+
 def test_amp_gpu_ddp_slurm_managed():
     """
     Make sure DDP + AMP work
@@ -411,16 +583,24 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
     clear_save_dir()
 
 
-def get_hparams():
+def get_hparams(continue_training=False, hpc_exp_number=0):
     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})
+
+    args = {
+        '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}
+
+    if continue_training:
+        args['test_tube_do_checkpoint_load'] = True
+        args['hpc_exp_number'] = hpc_exp_number
+
+    hparams = Namespace(**args)
     return hparams