pytorch-lightning: A Python repository from cinjon

Pytorch Lightning

The Keras for ML researchers using PyTorch. More control. Less boilerplate.

pip install pytorch-lightning

Docs

View the docs here

What is it?

Keras and fast.ai are too abstract for researchers. Lightning abstracts the full training loop but gives you control in the critical points.

Why do I want to use lightning?

Because you don't want to define a training loop, validation loop, gradient clipping, checkpointing, loading, gpu training, etc... every time you start a project. Let lightning handle all of that for you! Just define your data and what happens in the training, testing and validation loop and lightning will do the rest.

To use lightning do 2 things:

What does lightning control for me?

Everything!
Except for these 6 core functions which you define:

# what to do in the training loop
def training_step(self, data_batch, batch_nb):

# what to do in the validation loop
def validation_step(self, data_batch, batch_nb):

# how to aggregate validation_step outputs
def validation_end(self, outputs):

# and your dataloaders
def tng_dataloader():
def val_dataloader():
def test_dataloader():

Could be as complex as seq-2-seq + attention

# define what happens for training here
def training_step(self, data_batch, batch_nb):
    x, y = data_batch
    
    # define your own forward and loss calculation
    hidden_states = self.encoder(x)
     
    # even as complex as a seq-2seq + attn model
    # (this is just a toy, non-working example to illustrate)
    start_token = '<SOS>'
    last_hidden = torch.zeros(...)
    loss = 0
    for step in range(max_seq_len):
        attn_context = self.attention_nn(hidden_states, start_token)
        pred = self.decoder(start_token, attn_context, last_hidden) 
        last_hidden = pred
        pred = self.predict_nn(pred)
        loss += self.loss(last_hidden, y[step])
        
    #toy example as well
    loss = loss / max_seq_len
    return {'loss': loss}

Or as basic as CNN image classification

# define what happens for validation here
def validation_step(self, data_batch, batch_nb):    
    x, y = data_batch
    
    # or as basic as a CNN classification
    out = self.forward(x)
    loss = my_loss(out, y)
    return {'loss': loss}

And you also decide how to collate the output of all validation steps

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
    val_acc_mean = 0
    for output in outputs:
        val_loss_mean += output['val_loss']
        val_acc_mean += output['val_acc']

    val_loss_mean /= len(outputs)
    val_acc_mean /= len(outputs)
    tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
    return tqdm_dic

TensorboardX

Lightning is fully integrated with tensorboardX.

Lightning also adds a text column with all the hyperparameters for this experiment.

Simply note the path you set for the Experiment

from test_tube import Experiment
from pytorch-lightning import  Trainer

exp = Experiment(save_dir='/some/path')
trainer = Trainer(experiment=exp)
...

And run tensorboard from that dir

tensorboard --logdir /some/path

Lightning automatically automates all of the following (each is also configurable):

Checkpointing

Computing cluster (SLURM)

Debugging

Distributed training

Experiment Logging

Training loop

Validation loop

Demo

# install lightning
pip install pytorch-lightning

# clone lightning for the demo
git clone https://github.com/williamFalcon/pytorch-lightning.git
cd examples/new_project_templates/

# run demo (on cpu)
python trainer_gpu_cluster_template.py

Without changing the model AT ALL, you can run the model on a single gpu, over multiple gpus, or over multiple nodes.

# 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

cinjon/pytorch-lightning