release v

Fix some link bugs in the README.md

.gitignore

@@ -7,6 +7,7 @@ test_tube_data/
 datasets/
 model_weights/
 app/models/
+pip-wheel-metadata/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

MANIFEST.in

@@ -0,0 +1,9 @@
+graft docs
+
+include COPYING
+include AUTHORS
+
+recursive-include src/einsteinpy/tests *.py *.html
+
+prune docs/source/examples/.ipynb_checkpoints
+global-exclude *.py[cod] __pycache__ *.so *.dylib

README.md

@@ -1,18 +1,18 @@
 <p align="center">
   <a href="https://williamfalcon.github.io/pytorch-lightning/">
-    <img alt="" src="https://github.com/williamFalcon/pytorch-lightning/blob/master/imgs/lightning_logo.png" width="50">
+    <img alt="" src="https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/_static/lightning_logo.png" width="50">
   </a>
 </p>
 <h3 align="center">
   Pytorch Lightning
 </h3>
 <p align="center">
-  The Keras for ML-researchers in PyTorch. More control. Less boilerplate.    
+  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"></a>
-  <a href="https://travis-ci.org/williamFalcon/test-tube"><img src="https://travis-ci.org/williamFalcon/pytorch-lightning.svg?branch=master"></a>
-  <a href="https://github.com/williamFalcon/pytorch-lightning/blob/master/LICENSE"><img src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
+  <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://travis-ci.org/williamFalcon/test-tube"><img src="https://travis-ci.org/williamFalcon/pytorch-lightning.svg?branch=master"></a> -->
+  <a href="https://github.com/williamFalcon/pytorch-lightning/blob/master/COPYING"><img src="https://img.shields.io/badge/License-MIT-yellow.svg"></a>
 </p>   
 
 ```bash
@@ -22,37 +22,43 @@ pip install pytorch-lightning
 ## Docs   
 In progress. Documenting now!  
 
-## What is it?  
-All you do is define the forward passes, your data and **lightning runs everything else for you:**
+## Disclaimer
+This is a research tool I built for myself internally while doing my PhD. The API is not 100% production quality, but my hope is that by open-sourcing, we can all get it there (I don't have too much time nowadays to write production-level code).
 
-1. Running the training loop.   
-2. Running the validation loop.   
-3. Running the testing loop.   
+## What is it?  
+Keras is too abstract for researchers. Lightning makes it so you only have to define your model but still control all details of training if you need to.   
+
+Pytorch   
+<-- Lightning   
+Your model.   
+
+**Lightning will do the following for you:**   
+
+1. Run the training loop.   
+2. Run the validation loop.   
+3. Run the testing loop.   
 4. Early stopping.   
-5. Learning rate annealing. 
+5. Learning rate annealing.
 6. Can train complex models like GANs or anything with multiple optimizers.
 7. Weight checkpointing.   
 8. Model saving.   
 9. Model loading.   
-10. Logging training details (through test-tube).  
-11. Running training on multiple GPUs (through test-tube).   
-12. Running training on a GPU cluster managed by SLURM (through test-tube).   
-13. Distributing memory-bound models on multiple GPUs.
-14. Gives your model hyperparameters parsed from the command line OR a JSON file.   
-15. Runs your model in a dev environment where nothing logs.      
-  
+10. Log training details (through test-tube).  
+11. Run training on multiple GPUs (through test-tube).   
+12. Run training on a GPU cluster managed by SLURM (through test-tube).   
+13. Distribute memory-bound models on multiple GPUs.
+14. Give your model hyperparameters parsed from the command line OR a JSON file.   
+15. Run your model in a dev environment where nothing logs.      
+
 ## Usage
 To use lightning do 2 things:  
-1. [Define a trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/trainer_main.py) (which will run ALL your models).   
-2. [Define a model](https://github.com/williamFalcon/pytorch-lightning/blob/master/pytorch_lightning/models/sample_model_template/model_template.py).     
+1. [Define a trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/examples/basic_trainer.py) (which will run ALL your models).   
+2. [Define a model](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/examples/example_model.py).     
 
-### Example:   
-
-#### Define the trainer   
+#### Basic trainer example  
+See [this demo](https://github.com/williamFalcon/pytorch-lightning/blob/master/docs/source/examples/fully_featured_trainer.py) for a more robust trainer example.   
 
 ```python
-# trainer.py
-
 import os
 import sys
 
@@ -82,35 +88,17 @@ def main(hparams):
     exp.argparse(hparams)
     exp.save()
 
+    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+
     # build model
-    print('loading model...')
     model = ExampleModel(hparams)
-    print('model built')
 
     # callbacks
-    early_stop = EarlyStopping(
-        monitor=hparams.early_stop_metric,
-        patience=hparams.early_stop_patience,
-        verbose=True,
-        mode=hparams.early_stop_mode
-    )
-
-    model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
-    checkpoint = ModelCheckpoint(
-        filepath=model_save_path,
-        save_function=None,
-        save_best_only=True,
-        verbose=True,
-        monitor=hparams.model_save_monitor_value,
-        mode=hparams.model_save_monitor_mode
-    )
+    early_stop = EarlyStopping(monitor='val_acc', patience=3, mode='min', verbose=True)
+    checkpoint = ModelCheckpoint(filepath=model_save_path, save_function=None, save_best_only=True, verbose=True, monitor='val_acc', mode='min')
 
     # configure trainer
-    trainer = Trainer(
-        experiment=exp,
-        checkpoint_callback=checkpoint,
-        early_stop_callback=early_stop,
-    )
+    trainer = Trainer(experiment=exp, checkpoint_callback=checkpoint, early_stop_callback=early_stop)
 
     # train model
     trainer.fit(model)
@@ -129,9 +117,11 @@ if __name__ == '__main__':
 
     # train model
     main(hyperparams)
+
 ```
 
-#### Define the model   
+#### Basic model example  
+Here we only show the method signatures. It's up to you to define the content.   
 
 ```python
 from torch import nn
@@ -140,32 +130,32 @@ class My_Model(RootModule):
     def __init__(self):
         # define model
         self.l1 = nn.Linear(200, 10)
-    
+
     # ---------------
     # TRAINING
     def training_step(self, data_batch):
         x, y = data_batch
         y_hat = self.l1(x)
         loss = some_loss(y_hat)
-        
+
         return loss_val, {'train_loss': loss}
-    
+
     def validation_step(self, data_batch):
         x, y = data_batch
         y_hat = self.l1(x)
         loss = some_loss(y_hat)
-        
+
         return loss_val, {'val_loss': loss}
- 
+
      def validation_end(self, outputs):
         total_accs = []
-        
+
         for output in outputs:
             total_accs.append(output['val_acc'].item())
-        
+
         # return a dict
         return {'total_acc': np.mean(total_accs)}
-     
+
      # ---------------
      # SAVING
      def get_save_dict(self):
@@ -177,7 +167,7 @@ class My_Model(RootModule):
     def load_model_specific(self, checkpoint):
         # lightning loads for you. Here's your chance to say what you want to load
         self.load_state_dict(checkpoint['state_dict'])
-    
+
     # ---------------
     # TRAINING CONFIG
     def configure_optimizers(self):
@@ -185,7 +175,7 @@ class My_Model(RootModule):
         # lightning will call automatically
         optimizer = self.choose_optimizer('adam', self.parameters(), {'lr': self.hparams.learning_rate}, 'optimizer')
         return [optimizer]
-    
+
     @property
     def tng_dataloader(self):
         return pytorch_dataloader('train')
@@ -197,7 +187,7 @@ class My_Model(RootModule):
     @property
     def test_dataloader(self):
         return pytorch_dataloader('test')
-        
+
     # ---------------
     # MODIFY YOUR COMMAND LINE ARGS
     @staticmethod
@@ -206,6 +196,8 @@ class My_Model(RootModule):
         parser.add_argument('--out_features', default=20)
         return parser
 ```
+
+
 ### Details    
 
 #### Model definition   
@@ -214,7 +206,7 @@ class My_Model(RootModule):
 | training_step  |  Called with a batch of data during training | data from your dataloaders  | tuple: scalar, dict  |
 | validation_step  |  Called with a batch of data during validation  | data from your dataloaders  | tuple: scalar, dict |
 | validation_end  |  Collate metrics from all validation steps |  outputs: array where each item is the output of a validation step | dict: for logging |
-| get_save_dict  | called when your model needs to be saved (checkpoints, hpc save, etc...)  | None  |  dict to be saved | 
+| get_save_dict  | called when your model needs to be saved (checkpoints, hpc save, etc...)  | None  |  dict to be saved |
 
 #### Model training   
 | Name  | Description  |  Input |  Return |
@@ -230,7 +222,7 @@ class My_Model(RootModule):
 |---|---|---|---|
 | get_save_dict  | called when your model needs to be saved (checkpoints, hpc save, etc...)  | None  |  dict to be saved |
 |  load_model_specific |  called when loading a model | checkpoint: dict you created in get_save_dict  | dict: modified in whatever way you want  |
-    
+
 ## Optional model hooks.   
 Add these to the model whenever you want to configure training behavior.    
 

docs/source/examples/basic_trainer.py

@@ -2,9 +2,9 @@ import os
 import sys
 
 from test_tube import HyperOptArgumentParser, Experiment
-from pytorch_lightning.models.trainer import Trainer
-from pytorch_lightning.utils.arg_parse import add_default_args
-from pytorch_lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
+from pytorch-lightning.models.trainer import Trainer
+from pytorch-lightning.utils.arg_parse import add_default_args
+from pytorch-lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
 from demo.example_model import ExampleModel
 
 
@@ -28,16 +28,14 @@ def main(hparams):
     exp.save()
 
     # build model
-    print('loading model...')
     model = ExampleModel(hparams)
-    print('model built')
 
     # callbacks
     early_stop = EarlyStopping(
-        monitor=hparams.early_stop_metric,
-        patience=hparams.early_stop_patience,
+        monitor='val_acc',
+        patience=3,
+        mode='min',
         verbose=True,
-        mode=hparams.early_stop_mode
     )
 
     model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
@@ -46,8 +44,8 @@ def main(hparams):
         save_function=None,
         save_best_only=True,
         verbose=True,
-        monitor=hparams.model_save_monitor_value,
-        mode=hparams.model_save_monitor_mode
+        monitor='val_acc',
+        mode='min'
     )
 
     # configure trainer

docs/source/examples/example_model.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 import numpy as np
-from pytorch_lightning.root_module.root_module import RootModule
+from pytorch-lightning.root_module.root_module import RootModule
 from test_tube import HyperOptArgumentParser
 from torchvision.datasets import MNIST
 import torchvision.transforms as transforms

docs/source/examples/fully_featured_trainer.py

@@ -5,10 +5,10 @@ from time import sleep
 import torch
 
 from test_tube import HyperOptArgumentParser, Experiment, SlurmCluster
-from pytorch_lightning.models.trainer import Trainer
-from pytorch_lightning.utils.arg_parse import add_default_args
+from pytorch-lightning.models.trainer import Trainer
+from pytorch-lightning.utils.arg_parse import add_default_args
 
-from pytorch_lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
+from pytorch-lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
 
 SEED = 2334
 torch.manual_seed(SEED)
@@ -17,7 +17,7 @@ np.random.seed(SEED)
 # ---------------------
 # DEFINE MODEL HERE
 # ---------------------
-from demo.example_model import ExampleModel
+from example_model import ExampleModel
 # ---------------------
 
 AVAILABLE_MODELS = {
@@ -27,7 +27,7 @@ AVAILABLE_MODELS = {
 
 """
 Allows training by using command line arguments
-Run by: 
+Run by:
 # TYPE YOUR RUN COMMAND HERE
 """
 
@@ -57,13 +57,14 @@ def main(hparams, cluster, results_dict):
     sleep(process_position + 1)
 
     # init experiment
+    log_dir = os.path.dirname(os.path.realpath(__file__))
     exp = Experiment(
-        name=hparams.tt_name,
-        debug=hparams.debug,
-        save_dir=hparams.tt_save_path,
-        version=hparams.hpc_exp_number,
+        name='test_tube_exp',
+        debug=True,
+        save_dir=log_dir,
+        version=0,
         autosave=False,
-        description=hparams.tt_description
+        description='test demo'
     )
 
     exp.argparse(hparams)
@@ -108,7 +109,7 @@ def get_default_parser(strategy, root_dir):
 
     possible_model_names = list(AVAILABLE_MODELS.keys())
     parser = HyperOptArgumentParser(strategy=strategy, add_help=False)
-    add_default_args(parser, root_dir, possible_model_names, SEED)
+    add_default_args(parser, root_dir, possible_model_names=possible_model_names, rand_seed=SEED)
     return parser
 
 
@@ -198,4 +199,3 @@ if __name__ == '__main__':
             nb_trials=hyperparams.nb_hopt_trials,
             nb_workers=len(gpu_ids)
         )
-

pyproject.toml

@@ -0,0 +1,5 @@
+[build-system]
+requires = [
+    "setuptools",
+    "wheel",
+]

pytorch_lightning.egg-info/PKG-INFO

@@ -1,10 +0,0 @@
-Metadata-Version: 1.0
-Name: pytorch-lightning
-Version: 0.0.1
-Summary: Rapid research framework
-Home-page: https://github.com/williamFalcon/pytorch-lightning
-Author: UNKNOWN
-Author-email: UNKNOWN
-License: UNKNOWN
-Description: UNKNOWN
-Platform: UNKNOWN

pytorch_lightning.egg-info/SOURCES.txt

@@ -1,32 +0,0 @@
-README.md
-setup.py
-notebooks/__init__.py
-pytorch_lightning/__init__.py
-pytorch_lightning/trainer_main.py
-pytorch_lightning.egg-info/PKG-INFO
-pytorch_lightning.egg-info/SOURCES.txt
-pytorch_lightning.egg-info/dependency_links.txt
-pytorch_lightning.egg-info/requires.txt
-pytorch_lightning.egg-info/top_level.txt
-pytorch_lightning/models/__init__.py
-pytorch_lightning/models/trainer.py
-pytorch_lightning/models/model_examples/__init__.py
-pytorch_lightning/models/model_examples/bilstm.py
-pytorch_lightning/models/sample_model_template/__init__.py
-pytorch_lightning/models/sample_model_template/model_template.py
-pytorch_lightning/root_module/__init__.py
-pytorch_lightning/root_module/grads.py
-pytorch_lightning/root_module/hooks.py
-pytorch_lightning/root_module/memory.py
-pytorch_lightning/root_module/model_saving.py
-pytorch_lightning/root_module/optimization.py
-pytorch_lightning/root_module/root_module.py
-pytorch_lightning/utils/__init__.py
-pytorch_lightning/utils/arg_parse.py
-pytorch_lightning/utils/embeddings.py
-pytorch_lightning/utils/plotting.py
-pytorch_lightning/utils/pt_callbacks.py
-tests/__init__.py
-tests/research_proj/__init__.py
-tests/research_proj/sample_model_template/__init__.py
-tests/research_proj/sample_model_template/model_template_test.py

pytorch_lightning.egg-info/dependency_links.txt

@@ -1 +0,0 @@
-

pytorch_lightning.egg-info/top_level.txt

@@ -1,3 +0,0 @@
-notebooks
-pytorch_lightning
-tests

pytorch_lightning/models/model_examples/bilstm.py

@@ -1,167 +0,0 @@
-import torch.nn as nn
-import numpy as np
-
-from test_tube import HyperOptArgumentParser
-import torch
-from torch.autograd import Variable
-from sklearn.metrics import confusion_matrix, f1_score
-from torch.nn import functional as F
-
-
-class BiLSTMPack(nn.Module):
-    """
-    Sample model to show how to define a template
-    """
-    def __init__(self, hparams):
-        # init superclass
-        super(BiLSTMPack, self).__init__(hparams)
-
-        self.hidden = None
-
-        # trigger tag building
-        self.ner_tagset = {'O': 0, 'I-Bio': 1}
-        self.nb_tags = len(self.ner_tagset)
-
-        # build model
-        print('building model...')
-        if hparams.model_load_weights_path is None:
-            self.__build_model()
-            print('model built')
-        else:
-            self = BiLSTMPack.load(hparams.model_load_weights_path, hparams.on_gpu, hparams)
-            print('model loaded from: {}'.format(hparams.model_load_weights_path))
-
-    def __build_model(self):
-        """
-        Layout model
-        :return:
-        """
-        # design the number of final units
-        self.output_dim = self.hparams.nb_lstm_units
-
-        # when it's bidirectional our weights double
-        if self.hparams.bidirectional:
-            self.output_dim *= 2
-
-        # total number of words
-        total_words = len(self.tng_dataloader.dataset.words_token_to_idx)
-
-        # word embeddings
-        self.word_embedding = nn.Embedding(
-            num_embeddings=total_words + 1,
-            embedding_dim=self.hparams.embedding_dim,
-            padding_idx=0
-        )
-
-        # design the LSTM
-        self.lstm = nn.LSTM(
-            self.hparams.embedding_dim,
-            self.hparams.nb_lstm_units,
-            num_layers=self.hparams.nb_lstm_layers,
-            bidirectional=self.hparams.bidirectional,
-            dropout=self.hparams.drop_prob,
-            batch_first=True,
-        )
-
-        # map to tag space
-        self.fc_out = nn.Linear(self.output_dim, self.out_dim)
-        self.hidden_to_tag = nn.Linear(self.output_dim, self.nb_tags)
-
-
-    def init_hidden(self, batch_size):
-
-        # the weights are of the form (nb_layers * 2 if bidirectional, batch_size, nb_lstm_units)
-        mult = 2 if self.hparams.bidirectional else 1
-        hidden_a = torch.randn(self.hparams.nb_layers * mult, batch_size, self.nb_rnn_units)
-        hidden_b = torch.randn(self.hparams.nb_layers * mult, batch_size, self.nb_rnn_units)
-
-        if self.hparams.on_gpu:
-            hidden_a = hidden_a.cuda()
-            hidden_b = hidden_b.cuda()
-
-        hidden_a = Variable(hidden_a)
-        hidden_b = Variable(hidden_b)
-
-        return (hidden_a, hidden_b)
-
-    def forward(self, model_in):
-        # layout data (expand it, etc...)
-        # x = sequences
-        x, seq_lengths = model_in
-        batch_size, seq_len = x.size()
-
-        # reset RNN hidden state
-        self.hidden = self.init_hidden(batch_size)
-
-        # embed
-        x = self.word_embedding(x)
-
-        # run through rnn using packed sequences
-        x = torch.nn.utils.rnn.pack_padded_sequence(x, seq_lengths, batch_first=True)
-        x, self.hidden = self.lstm(x, self.hidden)
-        x, _ = torch.nn.utils.rnn.pad_packed_sequence(x, batch_first=True)
-
-        # if asked for only last state, use the h_n which is the same as out(t=n)
-        if not self.return_sequence:
-            # pull out hidden states
-            # h_n = (nb_directions * nb_layers, batch_size, emb_size)
-            nb_directions = 2 if self.bidirectional else 1
-            (h_n, _) = self.hidden
-
-            # reshape to make indexing easier
-            # forward = 0, backward = 1 (of nb_directions)
-            h_n = h_n.view(self.nb_layers, nb_directions, batch_size, self.nb_rnn_units)
-
-            # pull out last forward
-            forward_h_n = h_n[-1, 0, :, :]
-            x = forward_h_n
-
-            # if bidirectional, also pull out the last hidden of backward network
-            if self.bidirectional:
-                backward_h_n = h_n[-1, 1, :, :]
-                x = torch.cat([forward_h_n, backward_h_n], dim=1)
-
-        # project to tag space
-        x = x.contiguous()
-        x = x.view(-1, self.output_dim)
-        x = self.hidden_to_tag(x)
-
-        return x
-
-    def loss(self, model_out):
-        # cross entropy loss
-        logits, y = model_out
-        y, y_lens = y
-
-        # flatten y and logits
-        y = y.view(-1)
-        logits = logits.view(-1, self.nb_tags)
-
-        # calculate a mask to remove padding tokens
-        mask = (y >= 0).float()
-
-        # count how many tokens we have
-        num_tokens = int(torch.sum(mask).data[0])
-
-        # pick the correct values and mask out
-        logits = logits[range(logits.shape[0]), y] * mask
-
-        # compute the ce loss
-        ce_loss = -torch.sum(logits)/num_tokens
-
-        return ce_loss
-
-    def pull_out_last_embedding(self, x, seq_lengths, batch_size, on_gpu):
-        # grab only the last activations from the non-padded ouput
-        x_last = torch.zeros([batch_size, 1, x.size(-1)])
-        for i, seq_len in enumerate(seq_lengths):
-            x_last[i, :, :] = x[i, seq_len-1, :]
-
-        # put on gpu when requested
-        if on_gpu:
-            x_last = x_last.cuda()
-
-        # turn into torch var
-        x_last = Variable(x_last)
-
-        return x_last

pytorch_lightning/models/trainer.py

@@ -5,7 +5,13 @@ from pytorch_lightning.root_module.memory import get_gpu_memory_map
 import traceback
 from pytorch_lightning.root_module.model_saving import TrainerIO
 from torch.optim.lr_scheduler import MultiStepLR
+import pdb
 
+try:
+    from apex import amp
+    APEX_AVAILABLE = True
+except ModuleNotFoundError:
+    APEX_AVAILABLE = False
 
 class Trainer(TrainerIO):
 
@@ -26,6 +32,9 @@ class Trainer(TrainerIO):
                  train_percent_check=1.0, val_percent_check=1.0, test_percent_check=1.0, val_check_interval=0.95,
                  log_save_interval=1, add_log_row_interval=1,
                  lr_scheduler_milestones=None,
+                 use_amp=False,
+                 check_grad_nans=False,
+                 amp_level='O2',
                  nb_sanity_val_steps=5):
 
         # Transfer params
@@ -51,6 +60,8 @@ class Trainer(TrainerIO):
         self.nb_sanity_val_steps = nb_sanity_val_steps
         self.lr_scheduler_milestones = [] if lr_scheduler_milestones is None else [int(x.strip()) for x in lr_scheduler_milestones.split(',')]
         self.lr_schedulers = []
+        self.amp_level = amp_level
+        self.check_grad_nans = check_grad_nans
 
         # training state
         self.optimizers = None
@@ -73,6 +84,11 @@ class Trainer(TrainerIO):
         self.__determine_data_use_amount(train_percent_check, val_percent_check, test_percent_check, overfit_pct)
         print('gpu available: {}, used: {}'.format(torch.cuda.is_available(), self.on_gpu))
 
+        # apex test
+        self.use_amp = use_amp and APEX_AVAILABLE
+        if self.use_amp:
+            print('using 16bit precision')
+
     def __determine_data_use_amount(self, train_percent_check, val_percent_check, test_percent_check, overfit_pct):
         """
         Use less data for debugging purposes
@@ -86,7 +102,7 @@ class Trainer(TrainerIO):
             self.test_percent_check = overfit_pct
 
     def __is_function_implemented(self, f_name):
-        f_op = getattr(self, f_name, None)
+        f_op = getattr(self.model, f_name, None)
         return callable(f_op)
 
     @property
@@ -110,21 +126,21 @@ class Trainer(TrainerIO):
         self.tqdm_metrics = {}
 
         # determine number of training batches
-        nb_tng_batches = self.model.nb_batches(self.tng_dataloader)
-        self.nb_tng_batches = int(nb_tng_batches * self.train_percent_check)
+        self.nb_tng_batches = self.model.nb_batches(self.tng_dataloader)
+        self.nb_tng_batches = int(self.nb_tng_batches * self.train_percent_check)
 
         # determine number of validation batches
-        nb_val_batches = self.model.nb_batches(self.val_dataloader)
-        nb_val_batches = int(nb_val_batches * self.val_percent_check)
-        nb_val_batches = max(1, nb_val_batches)
-        self.nb_val_batches = nb_val_batches
+        self.nb_val_batches = self.model.nb_batches(self.val_dataloader)
+        self.nb_val_batches = int(self.nb_val_batches * self.val_percent_check)
+        self.nb_val_batches = max(1, self.nb_val_batches)
+        self.nb_val_batches = self.nb_val_batches
 
         # determine number of test batches
-        nb_test_batches = self.model.nb_batches(self.test_dataloader)
-        self.nb_test_batches = int(nb_test_batches * self.test_percent_check)
+        self.nb_test_batches = self.model.nb_batches(self.test_dataloader)
+        self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check)
 
         # determine when to check validation
-        self.val_check_batch = int(nb_tng_batches * self.val_check_interval)
+        self.val_check_batch = int(self.nb_tng_batches * self.val_check_interval)
 
     def __add_tqdm_metrics(self, metrics):
         for k, v in metrics.items():
@@ -151,19 +167,19 @@ class Trainer(TrainerIO):
         outputs = []
 
         # run training
-        for i, data_batch in enumerate(dataloader):
+        for batch_i, data_batch in enumerate(dataloader):
 
             if data_batch is None:
                 continue
 
             # stop short when on fast dev run
-            if max_batches is not None and i >= max_batches:
+            if max_batches is not None and batch_i >= max_batches:
                 break
 
             # -----------------
             # RUN VALIDATION STEP
             # -----------------
-            output = model.validation_step(data_batch)
+            output = model.validation_step(data_batch, batch_i)
             outputs.append(output)
 
             # batch done
@@ -195,6 +211,7 @@ class Trainer(TrainerIO):
     # -----------------------------
     def fit(self, model):
         self.model = model
+        model.trainer = self
 
         # transfer data loaders from model
         self.__get_dataloaders(model)
@@ -206,6 +223,14 @@ class Trainer(TrainerIO):
         # filter out the weights that were done on gpu so we can load on good old cpus
         self.optimizers = model.configure_optimizers()
 
+        if self.use_amp:
+            # An example
+            self.model, optimizer = amp.initialize(
+                self.model, self.optimizers[0], opt_level=self.amp_level,
+            )
+            self.optimizers[0] = optimizer
+            model.trainer = self
+
         # add lr schedulers
         if self.lr_scheduler_milestones is not None:
             for optimizer in self.optimizers:
@@ -269,21 +294,22 @@ class Trainer(TrainerIO):
                 # ---------------
                 # RUN TRAIN STEP
                 # ---------------
-                self.__run_tng_batch(data_batch)
+                batch_result = self.__run_tng_batch(data_batch, batch_nb)
+                early_stop_epoch = batch_result == -1
 
                 # ---------------
                 # RUN VAL STEP
                 # ---------------
                 is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0
-                if self.fast_dev_run or is_val_check_batch:
+                if self.fast_dev_run or is_val_check_batch or early_stop_epoch:
                     self.__run_validation()
 
                 # when batch should be saved
-                if (batch_nb + 1) % self.log_save_interval == 0:
+                if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
                     self.experiment.save()
 
                 # when metrics should be logged
-                if batch_nb % self.add_log_row_interval == 0:
+                if batch_nb % self.add_log_row_interval == 0 or early_stop_epoch:
                     # count items in memory
                     # nb_params, nb_tensors = count_mem_items()
 
@@ -307,6 +333,10 @@ class Trainer(TrainerIO):
                 if self.__is_function_implemented('on_batch_end'):
                     self.model.on_batch_end()
 
+                # end epoch early
+                if early_stop_epoch:
+                    break
+
             # hook
             if self.__is_function_implemented('on_epoch_end'):
                 self.model.on_epoch_end()
@@ -321,24 +351,37 @@ class Trainer(TrainerIO):
                 if stop:
                     return
 
-    def __run_tng_batch(self, data_batch):
+
+    def __run_tng_batch(self, data_batch, batch_nb):
         if data_batch is None:
-            return
+            return 0
 
         # hook
         if self.__is_function_implemented('on_batch_start'):
-            self.model.on_batch_start()
+            response = self.model.on_batch_start(data_batch)
+            if response == -1:
+                return -1
 
         if self.enable_tqdm:
             self.prog_bar.update(1)
 
         # forward pass
         # return a scalar value and a dic with tqdm metrics
-        loss, model_specific_tqdm_metrics_dic = self.model.training_step(data_batch)
+        loss, model_specific_tqdm_metrics_dic = self.model.training_step(data_batch, batch_nb)
         self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
 
         # backward pass
-        loss.backward()
+        if self.use_amp:
+            for optimizer in self.optimizers:
+                with amp.scale_loss(loss, optimizer) as scaled_loss:
+                    scaled_loss.backward()
+        else:
+            loss.backward()
+
+        if self.check_grad_nans:
+            for param in self.model.parameters():
+                print(param.grad.float().sum())
+
         self.batch_loss_value += loss.item()
 
         # gradient update with accumulated gradients
@@ -369,6 +412,8 @@ class Trainer(TrainerIO):
         if self.__is_function_implemented('on_batch_end'):
             self.model.on_batch_end()
 
+        return 0
+
     def __run_validation(self):
         # decide if can check epochs
         can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0

pytorch_lightning/root_module/hooks.py

@@ -1,7 +1,7 @@
 import torch
 
 class ModelHooks(torch.nn.Module):
-    def on_batch_start(self):
+    def on_batch_start(self, data_batch):
         pass
 
     def on_batch_end(self):
@@ -18,3 +18,4 @@ class ModelHooks(torch.nn.Module):
 
     def on_post_performance_check(self):
         pass
+

pytorch_lightning/root_module/model_saving.py

@@ -88,6 +88,7 @@ class TrainerIO(object):
         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']
 
         # restore the optimizers
         optimizer_states = checkpoint['optimizer_states']

pytorch_lightning/root_module/root_module.py

@@ -24,6 +24,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         self.overfit = hparams.overfit
         self.gradient_clip = hparams.gradient_clip
         self.num = 2
+        self.trainer = None
 
         # track if gpu was requested for checkpointing
         self.on_gpu = False
@@ -39,8 +40,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
 
         if self.on_gpu:
             print('running on gpu...')
-            self.dtype = torch.cuda.FloatTensor
-            torch.set_default_tensor_type('torch.cuda.FloatTensor')
+            torch.set_default_tensor_type(hparams.default_tensor_type)
 
     def forward(self, *args, **kwargs):
         """
@@ -51,7 +51,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         """
         raise NotImplementedError
 
-    def validation_step(self, data_batch):
+    def validation_step(self, data_batch, batch_nb):
         """
         return whatever outputs will need to be aggregated in validation_end
         :param data_batch:
@@ -67,7 +67,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
         """
         raise NotImplementedError
 
-    def training_step(self, data_batch):
+    def training_step(self, data_batch, batch_nb):
         """
         return loss, dict with metrics for tqdm
         :param data_batch:
@@ -150,19 +150,23 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
             return 0, 0
 
     @classmethod
-    def load_from_metrics(cls, weights_path, tags_csv, on_gpu):
+    def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
         """
         Primary way of loading model from csv weights path
         :param weights_path:
         :param tags_csv:
         :param on_gpu:
+        :param map_location: dic for mapping storage {'cuda:1':'cuda:0'}
         :return:
         """
         hparams = load_hparams_from_tags_csv(tags_csv)
         hparams.__setattr__('on_gpu', on_gpu)
 
         if on_gpu:
-            checkpoint = torch.load(weights_path)
+            if map_location is not None:
+                checkpoint = torch.load(weights_path, map_location=map_location)
+            else:
+                checkpoint = torch.load(weights_path)
         else:
             checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)
 

pytorch_lightning/trainer_main.py

@@ -9,7 +9,6 @@ from pytorch_lightning.utils.arg_parse import add_default_args
 from time import sleep
 
 from pytorch_lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint
-
 SEED = 2334
 torch.manual_seed(SEED)
 np.random.seed(SEED)

pytorch_lightning/utils/arg_parse.py

@@ -49,6 +49,11 @@ def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None
     parser.add_argument('--gpus', default='0', type=str)
     parser.add_argument('--single_run_gpu', dest='single_run_gpu', action='store_true')
     parser.add_argument('--disable_cuda', dest='disable_cuda', action='store_true')
+    parser.add_argument('--default_tensor_type', default='torch.cuda.FloatTensor', type=str)
+    parser.add_argument('--use_amp', dest='use_amp', action='store_true')
+    parser.add_argument('--check_grad_nans', dest='check_grad_nans', action='store_true')
+    parser.add_argument('--amp_level', default='O2',type=str)
+
 
     # run on hpc
     parser.add_argument('--on_cluster', dest='on_cluster', action='store_true')

pytorch_lightning/utils/embeddings.py

@@ -13,11 +13,7 @@ class PretrainedEmbedding(torch.nn.Embedding):
         >>> 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)
-        tensor([[[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
-         [ 0.2523,  0.1018, -0.6748,  ...,  0.1787, -0.5192,  0.3359]],
-
-        [[ 0.0000,  0.0000,  0.0000,  ...,  0.0000,  0.0000,  0.0000],
-         [-0.0067,  0.2224,  0.2771,  ...,  0.0594,  0.0014,  0.0987]]])
+        
 
 
         :param embedding_path:
@@ -37,7 +33,8 @@ class PretrainedEmbedding(torch.nn.Embedding):
         self.weight = new_emb.weight
 
         # apply freeze
-        self.weight.requires_grad = not freeze
+        should_freeze = not freeze
+        self.weight.requires_grad = should_freeze
 
     def __load_task_specific_embeddings(self, vocab_words, embedding_path, emb_dim, freeze):
         """
@@ -97,11 +94,11 @@ class PretrainedEmbedding(torch.nn.Embedding):
 
 if __name__ == '__main__':
     emb = PretrainedEmbedding(
-        embedding_path='/Users/waf/Developer/NGV/research-fermat/fermat/.vector_cache/glove.840B.300d.txt',
+        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)
+    print(embedded)

setup.cfg

@@ -0,0 +1,21 @@
+[tool:pytest]
+norecursedirs =
+    .git
+    dist
+    build
+python_files =
+    test_*.py
+doctest_plus = disabled
+addopts = --strict
+markers =
+    slow
+    remote_data
+    filterwarnings
+
+[pycodestyle]
+ignore = E731,W504
+max-line-length = 120
+
+[flake8]
+ignore = E731,W504,F401,F841
+max-line-length = 120

setup.py

@@ -2,12 +2,28 @@
 
 from setuptools import setup, find_packages
 
-setup(name='pytorch-lightning',
-      version='0.0.2',
-      description='Rapid research framework',
-      author='',
-      author_email='',
-      url='https://github.com/williamFalcon/pytorch-lightning',
-      install_requires=['test-tube', 'torch', 'tqdm'],
-      packages=find_packages()
-      )
+# https://packaging.python.org/guides/single-sourcing-package-version/
+
+# http://blog.ionelmc.ro/2014/05/25/python-packaging/
+setup(
+    name="pytorch-lightning",
+    version='0.1.dev1832',
+    description="The Keras for ML researchers using PyTorch",
+    author="William Falcon",
+    author_email="waf2107@columbia.edu",
+    url="https://github.com/williamFalcon/pytorch-lightning",
+    download_url="https://github.com/williamFalcon/pytorch-lightning",
+    license="MIT",
+    keywords=["deep learning", "pytorch", "AI"],
+    python_requires=">=3.5",
+    install_requires=[
+        "torch",
+        "tqdm",
+        "test-tube",
+    ],
+    packages=find_packages(),
+    long_description=open("README.md", encoding="utf-8").read(),
+    long_description_content_type='text/markdown',
+    include_package_data=True,
+    zip_safe=False,
+)

tests/README.md

@@ -1,65 +0,0 @@
-# Testing setup
-
-## A. Enable CircleCI for your project
-1. Integrate CircleCI by clicking "Set up Project" at [this link](https://circleci.com/add-projects/gh/NextGenVest).
-
-## B. Add your own tests
-1. In the /tests, emulate exactly the folder structure for your module found under /bot_seed
-2. To create a test for file ```/bot_seed/folder/example.py```:
-    - create the file ```/tests/folder/example_test.py```
-    - notice the **_test**
-    - notice the mirror path under **/tests**
-
-3. Your ```example_test.py``` file should have these main components
-
-```python
-# example.py
-
-def function_i_want_to_test(x):
-    return x*2
-
-def square(x):
-    return x*x
-
-```
-
-```python
-# example_test.py
-
-import pytest
-
-# do whatever imports you need
-from app.bot_seed.folder.example import function_i_want_to_test, square
-
-def test_function_i_want_to_test():
-    answer = function_i_want_to_test(4)
-    assert answer == 8
-
-# -----------------------------------
-# Your function must start with test_
-# -----------------------------------
-def test_square():
-    answer = square(3)
-    assert answer == 9
-
-# -----------------------------------
-# boilerplate (link this file to pytest)
-# -----------------------------------
-if __name__ == '__main__':
-    pytest.main([__file__])
-```
-
-## C. Add build passing badge
-1. Create a CircleCI status token:
-    - Go here: https://circleci.com/gh/NextGenVest/your-project-name/edit#api
-    - Click create token
-    - Select status
-    - Type "badge status"
-
-2. Get a copy of the markdown code:
-    - Go here: https://circleci.com/gh/NextGenVest/your-project-name/edit#badges
-    - Select master
-    - Select "badge status" token
-    - Select image URL
-    - Copy the image url link and change the html at the top of the root README.md file for your project
-    

tests/research_proj/sample_model_template/model_template_test.py

@@ -1,13 +0,0 @@
-import pytest
-
-"""
-Example test to show how to add a test for anything in the project.  
-Look at the README for more instructions   
-"""
-
-
-def test_cube():
-    assert 27 == 27
-
-if __name__ == '__main__':
-    pytest.main([__file__])