release v0.3.6.4

Loading

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)
 ```

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
@@ -611,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):

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:
@@ -75,12 +77,13 @@ 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
 
     # --------------------
@@ -149,13 +152,12 @@ 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
 
@@ -167,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.1',
+    version='0.3.6.4',
     description="The Keras for ML researchers using PyTorch",
     author="William Falcon",
     author_email="waf2107@columbia.edu",

tests/test_models.py

@@ -24,6 +24,88 @@ 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()
 
@@ -43,6 +125,7 @@ def test_loading_meta_tags():
 
     clear_save_dir()
 
+
 def test_dp_output_reduce():
 
     # test identity when we have a single gpu
@@ -64,9 +147,9 @@ def test_dp_output_reduce():
     assert reduced['b']['c'] == out['b']['c']
 
 
-def test_cpu_slurm_saving_loading():
+def test_model_saving_loading():
     """
-    Verify model save/load/checkpoint on CPU
+    Tests use case where trainer saves the model, and user loads it from tags independently
     :return:
     """
     hparams = get_hparams()
@@ -89,43 +172,49 @@ def test_cpu_slurm_saving_loading():
     # 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'
 
-    # test saving checkpoint
-    ckpt_test = os.path.join(save_dir, 'test.ckpt')
-    trainer.save_checkpoint(ckpt_test)
+    # make a prediction
+    for batch in model.test_dataloader:
+        break
 
-    # test registering a save function
-    trainer.enable_auto_hpc_walltime_manager()
+    x, y = batch
+    x = x.view(x.size(0), -1)
 
-    # test model loading with a map_location
-    pretrained_model = load_model(exp, save_dir, True)
+    # generate preds before saving model
+    model.eval()
+    pred_before_saving = model(x)
 
-    # test model preds
-    run_prediction(model.test_dataloader, pretrained_model)
+    # save model
+    new_weights_path = os.path.join(save_dir, 'save_test.ckpt')
+    trainer.save_checkpoint(new_weights_path)
 
-    trainer.model = pretrained_model
-    trainer.optimizers = pretrained_model.configure_optimizers()
+    # 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()
 
-    # test HPC saving
-    saved_filepath = trainer.hpc_save(save_dir, exp)
-    assert os.path.exists(saved_filepath)
-
-    # test HPC loading
-    trainer.global_step = 20000000
-    trainer.hpc_load(save_dir, on_gpu=False)
-    assert trainer.global_step == real_global_step and trainer.global_step != 20000000
-
-    # test freeze on gpu
-    model.freeze()
-    model.unfreeze()
+    # 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
@@ -494,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