diff --git a/.gitignore b/.gitignore
index 085acbe9..cbe7d5a1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,6 +9,7 @@ model_weights/
app/models/
pip-wheel-metadata/
test_tube_exp/
+tests/tests_tt_dir/
# Byte-compiled / optimized / DLL files
__pycache__/
@@ -119,4 +120,4 @@ ENV/
.mypy_cache/
# data
-mnist/
\ No newline at end of file
+mnist/
diff --git a/.readthedocs.yml b/.readthedocs.yml
new file mode 100644
index 00000000..b85b4bc2
--- /dev/null
+++ b/.readthedocs.yml
@@ -0,0 +1,19 @@
+# .readthedocs.yml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+# Required
+version: 2
+
+# Build documentation with MkDocs
+mkdocs:
+ configuration: mkdocs.yml
+
+# Optionally build your docs in additional formats such as PDF and ePub
+formats: all
+
+# Optionally set the version of Python and requirements required to build your docs
+python:
+ version: 3.7
+ install:
+ - requirements: docs/doc_requirements.txt
\ No newline at end of file
diff --git a/.travis.yml b/.travis.yml
new file mode 100644
index 00000000..f4e08df6
--- /dev/null
+++ b/.travis.yml
@@ -0,0 +1,16 @@
+language: python
+python:
+ - "3.7"
+# command to install dependencies
+cache: pip
+install:
+ - pip install -e .
+ - pip install -r requirements.txt
+ - pip install -U numpy
+
+# keep build from timing out
+dist: xenial
+
+# command to run tests
+script:
+ - py.test # or py.test for Python versions 3.5 and below
\ No newline at end of file
diff --git a/README.md b/README.md
index 6e11c297..85b27c20 100644
--- a/README.md
+++ b/README.md
@@ -9,31 +9,103 @@
The Keras for ML researchers using PyTorch. More control. Less boilerplate.
+
```bash
-pip install pytorch-lightning
+pip install pytorch-lightning
```
## Docs
**[View the docs here](https://williamfalcon.github.io/pytorch-lightning/)**
## 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.
+Lightning defers training and validation loop logic to you. It guarantees correct, modern best practices for the core training logic.
## 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.
+When starting a new project the last thing you want to do is recode a training loop, model loading/saving, distributed training, when to validate, etc... You're likely to spend a long time ironing out all the bugs without even getting to the core of your research.
+
+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 Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/trainer_cpu_template.py).
-2. [Define a LightningModel](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/lightning_module_template.py).
+1. [Define a LightningModel](https://williamfalcon.github.io/pytorch-lightning/LightningModule/RequiredTrainerInterface/)
+```python
+import pytorch_lightning as ptl
+import torch
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torchvision.datasets import MNIST
+
+class CoolModel(ptl.LightningModule):
+
+ def __init(self):
+ super(CoolModel, self).__init__()
+ # not the best model...
+ self.l1 = torch.nn.Linear(28 * 28, 10)
+
+ def forward(self, x):
+ return torch.relu(self.l1(x))
+
+ def my_loss(self, y_hat, y):
+ return F.cross_entropy(y_hat, y)
+
+ def training_step(self, batch, batch_nb):
+ x, y = batch
+ y_hat = self.forward(x)
+ return {'tng_loss': self.my_loss(y_hat, y)}
+
+ def validation_step(self, batch, batch_nb):
+ x, y = batch
+ y_hat = self.forward(x)
+ return {'val_loss': self.my_loss(y_hat, y)}
+
+ def validation_end(self, outputs):
+ avg_loss = torch.stack([x for x in outputs['val_loss']]).mean()
+ return avg_loss
+
+ def configure_optimizers(self):
+ return [torch.optim.Adam(self.parameters(), lr=0.02)]
+
+ @ptl.data_loader
+ def tng_dataloader(self):
+ return DataLoader(MNIST('path/to/save', train=True), batch_size=32)
+
+ @ptl.data_loader
+ def val_dataloader(self):
+ return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+
+ @ptl.data_loader
+ def test_dataloader(self):
+ return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+```
+
+2. Fit with a [trainer](https://williamfalcon.github.io/pytorch-lightning/Trainer/)
+```python
+from pytorch_lightning import Trainer
+from test_tube import Experiment
+
+model = CoolModel()
+
+# fit on 32 gpus across 4 nodes
+exp = Experiment(save_dir='some/dir')
+trainer = Trainer(experiment=exp, nb_gpu_nodes=4, gpus=[0,1,2,3,4,5,6,7])
+
+trainer.fit(model)
+
+# see all experiment metrics here
+# tensorboard --log_dir some/dir
+```
+
## What does lightning control for me?
Everything!
@@ -116,8 +188,8 @@ def validation_end(self, outputs):
return tqdm_dic
```
-## TensorboardX
-Lightning is fully integrated with tensorboardX.
+## Tensorboard
+Lightning is fully integrated with tensorboard.
@@ -148,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
@@ -215,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
+```
\ No newline at end of file
diff --git a/coverage.svg b/coverage.svg
new file mode 100644
index 00000000..6bfc8faf
--- /dev/null
+++ b/coverage.svg
@@ -0,0 +1,21 @@
+
+
+
+
+
+
+
+
+
+ coverage
+ coverage
+ 99%
+ 99%
+
+
diff --git a/docs/LightningModule/RequiredTrainerInterface.md b/docs/LightningModule/RequiredTrainerInterface.md
index b709c3a8..b9bb1d36 100644
--- a/docs/LightningModule/RequiredTrainerInterface.md
+++ b/docs/LightningModule/RequiredTrainerInterface.md
@@ -3,7 +3,7 @@
A lightning module is a strict superclass of nn.Module, it provides a standard interface for the trainer to interact with the model.
-The easiest thing to do is copy [this template](../../examples/new_project_templates/lightning_module_template.py) and modify accordingly.
+The easiest thing to do is copy [this template](../../pytorch_lightning/examples/new_project_templates/lightning_module_template.py) and modify accordingly.
Otherwise, to Define a Lightning Module, implement the following methods:
@@ -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,9 +21,63 @@ 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)
+---
+**Minimal example**
+```python
+import pytorch_lightning as ptl
+import torch
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torchvision.datasets import MNIST
+
+class CoolModel(ptl.LightningModule):
+
+ def __init(self):
+ super(CoolModel, self).__init__()
+ # not the best model...
+ self.l1 = torch.nn.Linear(28 * 28, 10)
+
+ def forward(self, x):
+ return torch.relu(self.l1(x))
+
+ def my_loss(self, y_hat, y):
+ return F.cross_entropy(y_hat, y)
+
+ def training_step(self, batch, batch_nb):
+ x, y = batch
+ y_hat = self.forward(x)
+ return {'tng_loss': self.my_loss(y_hat, y)}
+
+ def validation_step(self, batch, batch_nb):
+ x, y = batch
+ y_hat = self.forward(x)
+ return {'val_loss': self.my_loss(y_hat, y)}
+
+ def validation_end(self, outputs):
+ avg_loss = torch.stack([x for x in outputs['val_loss']]).mean()
+ return avg_loss
+
+ def configure_optimizers(self):
+ return [torch.optim.Adam(self.parameters(), lr=0.02)]
+
+ @ptl.data_loader
+ def tng_dataloader(self):
+ return DataLoader(MNIST('path/to/save', train=True), batch_size=32)
+
+ @ptl.data_loader
+ def val_dataloader(self):
+ return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+
+ @ptl.data_loader
+ def test_dataloader(self):
+ return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+```
+
---
### training_step
@@ -193,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
@@ -228,19 +281,19 @@ 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']
```
---
### tng_dataloader
``` {.python}
-@property
+@ptl.data_loader
def tng_dataloader(self)
```
-Called by lightning during training loop. Define it as a property.
+Called by lightning during training loop. Make sure to use the @ptl.data_loader decorator, this ensures not calling this function until the data are needed.
##### Return
Pytorch DataLoader
@@ -248,32 +301,26 @@ Pytorch DataLoader
**Example**
``` {.python}
-@property
+@ptl.data_loader
def tng_dataloader(self):
- if self._tng_dataloader is None:
- try:
- transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
- dataset = MNIST(root='/path/to/mnist/', train=True, transform=transform, download=True)
- loader = torch.utils.data.DataLoader(
- dataset=dataset,
- batch_size=self.hparams.batch_size,
- shuffle=True
- )
- self._tng_dataloader = loader
- except Exception as e:
- raise e
-
- return self._tng_dataloader
+ transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
+ dataset = MNIST(root='/path/to/mnist/', train=True, transform=transform, download=True)
+ loader = torch.utils.data.DataLoader(
+ dataset=dataset,
+ batch_size=self.hparams.batch_size,
+ shuffle=True
+ )
+ return loader
```
---
### val_dataloader
``` {.python}
-@property
+@ptl.data_loader
def tng_dataloader(self)
```
-Called by lightning during validation loop. Define it as a property.
+Called by lightning during validation loop. Make sure to use the @ptl.data_loader decorator, this ensures not calling this function until the data are needed.
##### Return
Pytorch DataLoader
@@ -281,32 +328,27 @@ Pytorch DataLoader
**Example**
``` {.python}
-@property
+@ptl.data_loader
def val_dataloader(self):
- if self._val_dataloader is None:
- try:
- transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
- dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform, download=True)
- loader = torch.utils.data.DataLoader(
- dataset=dataset,
- batch_size=self.hparams.batch_size,
- shuffle=True
- )
- self._val_dataloader = loader
- except Exception as e:
- raise e
-
- return self._val_dataloader
+ transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
+ dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform, download=True)
+ loader = torch.utils.data.DataLoader(
+ dataset=dataset,
+ batch_size=self.hparams.batch_size,
+ shuffle=True
+ )
+
+ return loader
```
---
### test_dataloader
``` {.python}
-@property
+@ptl.data_loader
def test_dataloader(self)
```
-Called by lightning during test loop. Define it as a property.
+Called by lightning during test loop. Make sure to use the @ptl.data_loader decorator, this ensures not calling this function until the data are needed.
##### Return
Pytorch DataLoader
@@ -314,22 +356,17 @@ Pytorch DataLoader
**Example**
``` {.python}
-@property
+@ptl.data_loader
def test_dataloader(self):
- if self._test_dataloader is None:
- try:
- transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
- dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform, download=True)
- loader = torch.utils.data.DataLoader(
- dataset=dataset,
- batch_size=self.hparams.batch_size,
- shuffle=True
- )
- self._test_dataloader = loader
- except Exception as e:
- raise e
-
- return self._test_dataloader
+ transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
+ dataset = MNIST(root='/path/to/mnist/', train=False, transform=transform, download=True)
+ loader = torch.utils.data.DataLoader(
+ dataset=dataset,
+ batch_size=self.hparams.batch_size,
+ shuffle=True
+ )
+
+ return loader
```
---
diff --git a/docs/LightningModule/methods.md b/docs/LightningModule/methods.md
index 9163326e..d57c6950 100644
--- a/docs/LightningModule/methods.md
+++ b/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)
```
diff --git a/docs/Trainer/Distributed training.md b/docs/Trainer/Distributed training.md
index dcd8a422..a7d487b5 100644
--- a/docs/Trainer/Distributed training.md
+++ b/docs/Trainer/Distributed training.md
@@ -3,6 +3,26 @@ Lightning makes multi-gpu training and 16 bit training trivial.
*Note:*
None of the flags below require changing anything about your lightningModel definition.
+---
+#### Choosing a backend
+Lightning supports two backends. DataParallel and DistributedDataParallel. Both can be used for single-node multi-GPU training.
+For multi-node training you must use DistributedDataParallel.
+
+You can toggle between each mode by setting this flag.
+``` {.python}
+# DEFAULT uses DataParallel
+trainer = Trainer(distributed_backend='dp')
+
+# change to distributed data parallel
+trainer = Trainer(distributed_backend='ddp')
+```
+
+If you request multiple nodes, the back-end will auto-switch to ddp.
+We recommend you use DistributedDataparallel even for single-node multi-GPU training. It is MUCH faster than DP but *may*
+have configuration issues depending on your cluster.
+
+For a deeper understanding of what lightning is doing, feel free to read [this guide](https://medium.com/@_willfalcon/9-tips-for-training-lightning-fast-neural-networks-in-pytorch-8e63a502f565).
+
---
#### 16-bit mixed precision
16 bit precision can cut your memory footprint by half. If using volta architecture GPUs it can give a dramatic training speed-up as well.
@@ -37,16 +57,69 @@ Make sure you're on a GPU machine. You can set as many GPUs as you want.
In this setting, the model will run on all 8 GPUs at once using DataParallel under the hood.
```python
# set these flags
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
+# lightning sets these flags for you automatically
+# no need to set yourself
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,5,6,7"
-# DEFAULT
-trainer = Trainer(gpus=[0,1,2,3,4,5,6,7])
+
+# to use DataParallel (default)
+trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='dp')
+
+# RECOMMENDED use DistributedDataParallel
+trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], distributed_backend='ddp')
```
---
#### Multi-node
-COMING SOON.
+Multi-node training is easily done by specifying these flags.
+```python
+# train on 12*8 GPUs
+trainer = Trainer(gpus=[0,1,2,3,4,5,6,7], nb_gpu_nodes=12)
+```
+
+In addition, make sure to set up your SLURM job correctly via the [SlurmClusterObject](https://williamfalcon.github.io/test-tube/hpc/SlurmCluster/). In particular, specify the number of tasks per node correctly.
+
+```python
+cluster = SlurmCluster(
+ hyperparam_optimizer=test_tube.HyperOptArgumentParser(),
+ log_path='/some/path/to/save',
+)
+
+# OPTIONAL FLAGS WHICH MAY BE CLUSTER DEPENDENT
+# which interface your nodes use for communication
+cluster.add_command('export NCCL_SOCKET_IFNAME=^docker0,lo')
+
+# see output of the NCCL connection process
+# NCCL is how the nodes talk to each other
+cluster.add_command('export NCCL_DEBUG=INFO')
+
+# setting a master port here is a good idea.
+cluster.add_command(f'export MASTER_PORT={PORT}')
+
+# good to load the latest NCCL version
+cluster.load_modules(['NCCL/2.4.7-1-cuda.10.0'])
+
+# configure cluster
+cluster.per_experiment_nb_nodes = 12
+cluster.per_experiment_nb_gpus = 8
+
+cluster.add_slurm_cmd(cmd='ntasks-per-node', value=8, comment='1 task per gpu')
+```
+
+Finally, make sure to add a distributed sampler to your dataset. The distributed sampler copies a
+portion of your dataset onto each GPU. (World_size = gpus_per_node * nb_nodes).
+
+```python
+# ie: this:
+dataset = myDataset()
+dataloader = Dataloader(dataset)
+
+# becomes:
+dataset = myDataset()
+dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+dataloader = Dataloader(dataset, sampler=dist_sampler)
+```
---
#### Self-balancing architecture
diff --git a/docs/Trainer/Training Loop.md b/docs/Trainer/Training Loop.md
index 9b05fd58..2be8da9e 100644
--- a/docs/Trainer/Training Loop.md
+++ b/docs/Trainer/Training Loop.md
@@ -19,7 +19,7 @@ Cut the learning rate by 10 at every epoch listed in this list.
trainer = Trainer(lr_scheduler_milestones=None)
# cut LR by 10 at 100, 200, and 300 epochs
-trainer = Trainer(lr_scheduler_milestones=[100, 200, 300])
+trainer = Trainer(lr_scheduler_milestones='100, 200, 300')
```
---
diff --git a/docs/doc_requirements.txt b/docs/doc_requirements.txt
new file mode 100644
index 00000000..d52ea0ac
--- /dev/null
+++ b/docs/doc_requirements.txt
@@ -0,0 +1 @@
+mkdocs-material==4.4.0
diff --git a/docs/index.md b/docs/index.md
index 01e3b292..0e25fa79 100644
--- a/docs/index.md
+++ b/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/#template-model-definition)
-2. Pick a trainer
- - [Basic CPU Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/trainer_cpu_template.py)
- - [GPU cluster Trainer](https://github.com/williamFalcon/pytorch-lightning/blob/master/examples/new_project_templates/trainer_gpu_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/)
diff --git a/examples/new_project_templates/__init__.py b/examples/new_project_templates/__init__.py
deleted file mode 100644
index bc8ec6d7..00000000
--- a/examples/new_project_templates/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from .lightning_module_template import LightningTemplateModel
\ No newline at end of file
diff --git a/pytorch_lightning/__init__.py b/pytorch_lightning/__init__.py
index 5893eb4a..e7a03ba8 100644
--- a/pytorch_lightning/__init__.py
+++ b/pytorch_lightning/__init__.py
@@ -1,2 +1,3 @@
from .models import Trainer
-from .root_module.root_module import LightningModule
\ No newline at end of file
+from .root_module.root_module import LightningModule
+from .root_module.decorators import data_loader
\ No newline at end of file
diff --git a/pytorch_lightning/examples/__init__.py b/pytorch_lightning/examples/__init__.py
new file mode 100644
index 00000000..6743d7f9
--- /dev/null
+++ b/pytorch_lightning/examples/__init__.py
@@ -0,0 +1 @@
+from .new_project_templates.lightning_module_template import LightningTemplateModel
\ No newline at end of file
diff --git a/pytorch_lightning/examples/new_project_templates/__init__.py b/pytorch_lightning/examples/new_project_templates/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/examples/new_project_templates/lightning_module_template.py b/pytorch_lightning/examples/new_project_templates/lightning_module_template.py
similarity index 79%
rename from examples/new_project_templates/lightning_module_template.py
rename to pytorch_lightning/examples/new_project_templates/lightning_module_template.py
index 5245e2cd..608e534e 100644
--- a/examples/new_project_templates/lightning_module_template.py
+++ b/pytorch_lightning/examples/new_project_templates/lightning_module_template.py
@@ -10,6 +10,7 @@ from torch import optim
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
+import pytorch_lightning as ptl
from pytorch_lightning.root_module.root_module import LightningModule
@@ -24,10 +25,14 @@ class LightningTemplateModel(LightningModule):
:param hparams:
"""
# init superclass
- super(LightningTemplateModel, self).__init__(hparams)
+ super(LightningTemplateModel, self).__init__()
+ self.hparams = hparams
self.batch_size = hparams.batch_size
+ # if you specify an example input, the summary will show input/output for each layer
+ self.example_input_array = torch.rand(5, 28 * 28)
+
# build model
self.__build_model()
@@ -78,15 +83,21 @@ class LightningTemplateModel(LightningModule):
# forward pass
x, y = data_batch
x = x.view(x.size(0), -1)
+
y_hat = self.forward(x)
# calculate loss
loss_val = self.loss(y, y_hat)
+ # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+ if self.trainer.use_dp:
+ loss_val = loss_val.unsqueeze(0)
+
output = OrderedDict({
- 'loss': loss_val,
- 'tqdm_metrics': {}
+ 'loss': loss_val
})
+
+ # can also return just a scalar instead of a dict (return loss_val)
return output
def validation_step(self, data_batch, batch_i):
@@ -104,11 +115,22 @@ class LightningTemplateModel(LightningModule):
# acc
labels_hat = torch.argmax(y_hat, dim=1)
val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+ val_acc = torch.tensor(val_acc)
+
+ if self.on_gpu:
+ val_acc = val_acc.cuda(loss_val.device.index)
+
+ # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+ if self.trainer.use_dp:
+ loss_val = loss_val.unsqueeze(0)
+ val_acc = val_acc.unsqueeze(0)
output = OrderedDict({
'val_loss': loss_val,
- 'val_acc': torch.tensor(val_acc),
+ 'val_acc': val_acc,
})
+
+ # can also return just a scalar instead of a dict (return loss_val)
return output
def validation_end(self, outputs):
@@ -117,6 +139,10 @@ class LightningTemplateModel(LightningModule):
:param outputs: list of individual outputs of each validation step
:return:
"""
+ # if returned a scalar from validation_step, outputs is a list of tensor scalars
+ # we return just the average in this case (if we want)
+ # return torch.stack(outputs).mean()
+
val_loss_mean = 0
val_acc_mean = 0
for output in outputs:
@@ -128,20 +154,6 @@ class LightningTemplateModel(LightningModule):
tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
return tqdm_dic
- def update_tng_log_metrics(self, logs):
- return logs
-
- # ---------------------
- # MODEL SAVING
- # ---------------------
- 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
# ---------------------
@@ -179,38 +191,23 @@ class LightningTemplateModel(LightningModule):
return loader
- @property
+ @ptl.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
+ print('tng data loader called')
+ return self.__dataloader(train=True)
- @property
+ @ptl.data_loader
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
+ print('val data loader called')
+ return self.__dataloader(train=False)
- @property
+ @ptl.data_loader
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
+ print('test data loader called')
+ return self.__dataloader(train=False)
@staticmethod
- def add_model_specific_args(parent_parser, root_dir):
+ def add_model_specific_args(parent_parser, root_dir): # pragma: no cover
"""
Parameters you define here will be available to your model through self.hparams
:param parent_parser:
diff --git a/examples/new_project_templates/multi_node_cluster_template.py b/pytorch_lightning/examples/new_project_templates/multi_node_cluster_template.py
similarity index 100%
rename from examples/new_project_templates/multi_node_cluster_template.py
rename to pytorch_lightning/examples/new_project_templates/multi_node_cluster_template.py
diff --git a/examples/new_project_templates/single_cpu_template.py b/pytorch_lightning/examples/new_project_templates/single_cpu_template.py
similarity index 100%
rename from examples/new_project_templates/single_cpu_template.py
rename to pytorch_lightning/examples/new_project_templates/single_cpu_template.py
diff --git a/examples/new_project_templates/single_gpu_node_16bit_template.py b/pytorch_lightning/examples/new_project_templates/single_gpu_node_16bit_template.py
similarity index 100%
rename from examples/new_project_templates/single_gpu_node_16bit_template.py
rename to pytorch_lightning/examples/new_project_templates/single_gpu_node_16bit_template.py
diff --git a/examples/new_project_templates/single_gpu_node_template.py b/pytorch_lightning/examples/new_project_templates/single_gpu_node_dp_template.py
similarity index 100%
rename from examples/new_project_templates/single_gpu_node_template.py
rename to pytorch_lightning/examples/new_project_templates/single_gpu_node_dp_template.py
diff --git a/pytorch_lightning/examples/new_project_templates/single_gpu_node_template.py b/pytorch_lightning/examples/new_project_templates/single_gpu_node_template.py
new file mode 100644
index 00000000..34dc4441
--- /dev/null
+++ b/pytorch_lightning/examples/new_project_templates/single_gpu_node_template.py
@@ -0,0 +1,112 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+import os
+import sys
+import numpy as np
+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.callbacks import EarlyStopping, ModelCheckpoint
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+from lightning_module_template import LightningTemplateModel
+
+
+def main(hparams):
+ """
+ Main training routine specific for this project
+ :param hparams:
+ :return:
+ """
+ # ------------------------
+ # 1 INIT LIGHTNING MODEL
+ # ------------------------
+ print('loading model...')
+ model = LightningTemplateModel(hparams)
+ print('model built')
+
+ # ------------------------
+ # 2 INIT TEST TUBE EXP
+ # ------------------------
+
+ # init experiment
+ exp = Experiment(
+ name=hyperparams.experiment_name,
+ save_dir=hyperparams.test_tube_save_path,
+ autosave=False,
+ description='test demo'
+ )
+
+ exp.argparse(hparams)
+ exp.save()
+
+ # ------------------------
+ # 3 DEFINE CALLBACKS
+ # ------------------------
+ model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+ early_stop = EarlyStopping(
+ monitor='val_acc',
+ patience=3,
+ verbose=True,
+ mode='max'
+ )
+
+ checkpoint = ModelCheckpoint(
+ filepath=model_save_path,
+ save_best_only=True,
+ verbose=True,
+ monitor='val_loss',
+ mode='min'
+ )
+
+ # ------------------------
+ # 4 INIT TRAINER
+ # ------------------------
+ trainer = Trainer(
+ experiment=exp,
+ checkpoint_callback=checkpoint,
+ early_stop_callback=early_stop,
+ gpus=hparams.gpus,
+ )
+
+ # ------------------------
+ # 5 START TRAINING
+ # ------------------------
+ trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+ # dirs
+ root_dir = os.path.dirname(os.path.realpath(__file__))
+ demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+ checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+ test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+
+ # although we user hyperOptParser, we are using it only as argparse right now
+ parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+ # gpu args
+ parent_parser.add_argument('--gpus', type=str, default='-1', help='how many gpus to use in the node. -1 uses all the gpus on the node')
+ parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs')
+ parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir, help='where to save model')
+ parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+
+ # allow model to overwrite or extend args
+ parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+ hyperparams = parser.parse_args()
+
+ # ---------------------
+ # RUN TRAINING
+ # ---------------------
+ # run on HPC cluster
+ print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {hyperparams.gpus}')
+ main(hyperparams)
diff --git a/pytorch_lightning/examples/new_project_templates/single_gpu_template.py b/pytorch_lightning/examples/new_project_templates/single_gpu_template.py
new file mode 100644
index 00000000..66714230
--- /dev/null
+++ b/pytorch_lightning/examples/new_project_templates/single_gpu_template.py
@@ -0,0 +1,112 @@
+"""
+Runs a model on a single node across N-gpus.
+"""
+import os
+import sys
+import numpy as np
+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.callbacks import EarlyStopping, ModelCheckpoint
+
+SEED = 2334
+torch.manual_seed(SEED)
+np.random.seed(SEED)
+
+from lightning_module_template import LightningTemplateModel
+
+
+def main(hparams):
+ """
+ Main training routine specific for this project
+ :param hparams:
+ :return:
+ """
+ # ------------------------
+ # 1 INIT LIGHTNING MODEL
+ # ------------------------
+ print('loading model...')
+ model = LightningTemplateModel(hparams)
+ print('model built')
+
+ # ------------------------
+ # 2 INIT TEST TUBE EXP
+ # ------------------------
+
+ # init experiment
+ exp = Experiment(
+ name=hyperparams.experiment_name,
+ save_dir=hyperparams.test_tube_save_path,
+ autosave=False,
+ description='test demo'
+ )
+
+ exp.argparse(hparams)
+ exp.save()
+
+ # ------------------------
+ # 3 DEFINE CALLBACKS
+ # ------------------------
+ model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
+ early_stop = EarlyStopping(
+ monitor='val_acc',
+ patience=3,
+ verbose=True,
+ mode='max'
+ )
+
+ checkpoint = ModelCheckpoint(
+ filepath=model_save_path,
+ save_best_only=True,
+ verbose=True,
+ monitor='val_loss',
+ mode='min'
+ )
+
+ # ------------------------
+ # 4 INIT TRAINER
+ # ------------------------
+ trainer = Trainer(
+ experiment=exp,
+ checkpoint_callback=checkpoint,
+ early_stop_callback=early_stop,
+ gpus=hparams.gpus,
+ )
+
+ # ------------------------
+ # 5 START TRAINING
+ # ------------------------
+ trainer.fit(model)
+
+
+if __name__ == '__main__':
+
+ # dirs
+ root_dir = os.path.dirname(os.path.realpath(__file__))
+ demo_log_dir = os.path.join(root_dir, 'pt_lightning_demo_logs')
+ checkpoint_dir = os.path.join(demo_log_dir, 'model_weights')
+ test_tube_dir = os.path.join(demo_log_dir, 'test_tube_data')
+
+ # although we user hyperOptParser, we are using it only as argparse right now
+ parent_parser = HyperOptArgumentParser(strategy='grid_search', add_help=False)
+
+ # gpu args
+ parent_parser.add_argument('--gpus', type=str, default='0', help='how many gpus to use in the node. -1 uses all the gpus on the node')
+ parent_parser.add_argument('--test_tube_save_path', type=str, default=test_tube_dir, help='where to save logs')
+ parent_parser.add_argument('--model_save_path', type=str, default=checkpoint_dir, help='where to save model')
+ parent_parser.add_argument('--experiment_name', type=str, default='pt_lightning_exp_a', help='test tube exp name')
+
+ # allow model to overwrite or extend args
+ parser = LightningTemplateModel.add_model_specific_args(parent_parser, root_dir)
+ hyperparams = parser.parse_args()
+
+ # ---------------------
+ # RUN TRAINING
+ # ---------------------
+ # run on HPC cluster
+ print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {hyperparams.gpus}')
+ main(hyperparams)
diff --git a/examples/new_project_templates/trainer_cpu_template.py b/pytorch_lightning/examples/new_project_templates/trainer_cpu_template.py
similarity index 100%
rename from examples/new_project_templates/trainer_cpu_template.py
rename to pytorch_lightning/examples/new_project_templates/trainer_cpu_template.py
diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py
new file mode 100644
index 00000000..d686e39f
--- /dev/null
+++ b/pytorch_lightning/models/trainer.py
@@ -0,0 +1,891 @@
+"""
+The trainer handles all the logic for running a val loop, training loop, distributing, etc...
+"""
+import subprocess
+import traceback
+import warnings
+import os
+import pdb
+import re
+
+import torch
+from torch.utils.data.distributed import DistributedSampler
+from torch.optim.lr_scheduler import MultiStepLR
+import torch.multiprocessing as mp
+import torch.distributed as dist
+import numpy as np
+import tqdm
+
+from pytorch_lightning.root_module.memory import get_gpu_memory_map
+from pytorch_lightning.root_module.model_saving import TrainerIO
+from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel, LightningDataParallel
+from pytorch_lightning.utils.debugging import MisconfigurationException
+
+try:
+ from apex import amp
+ APEX_AVAILABLE = True
+except Exception:
+ APEX_AVAILABLE = False
+
+
+def reduce_distributed_output(output, nb_gpus):
+ if nb_gpus <= 1:
+ return output
+
+ # when using DP, we get one output per gpu
+ # average outputs and return
+ if type(output) is torch.Tensor:
+ return output.mean()
+
+ for k, v in output.items():
+ # recurse on nested dics
+ if isinstance(output[k], dict):
+ output[k] = reduce_distributed_output(output[k], nb_gpus)
+
+ # reduce only metrics that have the same nb of gpus
+ elif output[k].size(0) == nb_gpus:
+ reduced = torch.mean(output[k])
+ output[k] = reduced
+ return output
+
+
+class Trainer(TrainerIO):
+
+ def __init__(self,
+ experiment,
+ early_stop_callback=None,
+ checkpoint_callback=None,
+ gradient_clip=0,
+ cluster=None,
+ process_position=0,
+ current_gpu_name=0,
+ nb_gpu_nodes=1,
+ gpus=None,
+ progress_bar=True,
+ overfit_pct=0.0,
+ track_grad_norm=-1,
+ check_val_every_n_epoch=1,
+ fast_dev_run=False,
+ accumulate_grad_batches=1,
+ max_nb_epochs=1000, min_nb_epochs=1,
+ train_percent_check=1.0, val_percent_check=1.0, test_percent_check=1.0,
+ val_check_interval=0.95,
+ log_save_interval=100, add_log_row_interval=10,
+ lr_scheduler_milestones=None,
+ distributed_backend='dp',
+ use_amp=False,
+ print_nan_grads=False,
+ print_weights_summary=True,
+ amp_level='O2',
+ nb_sanity_val_steps=5):
+
+ """
+
+ :param experiment: Test-tube experiment
+ :param early_stop_callback: from pytorch_lightning import EarlyStopping
+ :param checkpoint_callback: from pytorch_lightning import Checkpoint
+ :param gradient_clip:
+ :param cluster:
+ :param process_position:
+ :param current_gpu_name:
+ :param nb_gpu_nodes:
+ :param gpus:
+ :param progress_bar:
+ :param overfit_pct:
+ :param track_grad_norm:
+ :param check_val_every_n_epoch:
+ :param fast_dev_run:
+ :param accumulate_grad_batches:
+ :param max_nb_epochs:
+ :param min_nb_epochs:
+ :param train_percent_check:
+ :param val_percent_check:
+ :param test_percent_check:
+ :param val_check_interval:
+ :param log_save_interval:
+ :param add_log_row_interval:
+ :param lr_scheduler_milestones:
+ :param distributed_backend: 'np' to use DistributedParallel, 'ddp' to use DistributedDataParallel
+ :param use_amp:
+ :param print_nan_grads:
+ :param print_weights_summary:
+ :param amp_level:
+ :param nb_sanity_val_steps:
+ """
+
+ # Transfer params
+
+ self.nb_gpu_nodes = nb_gpu_nodes
+ self.gradient_clip = gradient_clip
+ self.check_val_every_n_epoch = check_val_every_n_epoch
+ self.enable_early_stop = early_stop_callback is not None
+ self.track_grad_norm = track_grad_norm
+ self.fast_dev_run = fast_dev_run
+ self.on_gpu = gpus is not None and torch.cuda.is_available()
+ self.progress_bar = progress_bar
+ self.experiment = experiment
+ self.exp_save_path = experiment.get_data_path(experiment.name, experiment.version)
+ self.cluster = cluster
+ self.process_position = process_position
+ self.current_gpu_name = current_gpu_name
+ self.print_weights_summary = print_weights_summary
+ self.checkpoint_callback = checkpoint_callback
+
+ if self.checkpoint_callback is not None:
+ self.checkpoint_callback.save_function = self.save_checkpoint
+
+ self.early_stop = early_stop_callback
+ self.model = None
+ self.max_nb_epochs = max_nb_epochs
+ self.accumulate_grad_batches = accumulate_grad_batches
+ self.early_stop_callback = early_stop_callback
+ self.min_nb_epochs = min_nb_epochs
+ 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.print_nan_grads = print_nan_grads
+ self.data_parallel_device_ids = None
+ self.world_size = 1
+ self.node_rank = 0
+ self.use_ddp = False
+ self.use_dp = False
+
+ # training bookeeping
+ self.total_batch_nb = 0
+ self.running_loss = []
+ self.avg_loss = 0
+ self.batch_nb = 0
+ self.tqdm_metrics = {}
+ self.nb_val_batches = None
+ self.nb_tng_batches = None
+ self.nb_test_batches = None
+
+ # gpus come in as a string.
+ # if gpus = -1 then use all available devices
+ # otherwise, split the string using commas
+ if gpus is not None:
+ if type(gpus) is list:
+ self.data_parallel_device_ids = gpus
+ elif type(gpus) is str:
+ if gpus == '-1':
+ self.data_parallel_device_ids = list(range(0, torch.cuda.device_count()))
+ else:
+ self.data_parallel_device_ids = [int(x.strip()) for x in gpus.split(',')]
+ else:
+ raise Exception('gpus has to be a string or list of ids')
+
+ # set the correct cuda visible devices (using pci order)
+ os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+ os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(x) for x in self.data_parallel_device_ids])
+ print(f'VISIBLE GPUS: {os.environ["CUDA_VISIBLE_DEVICES"]}')
+
+ # make DP and DDP mutually exclusive
+ # single GPU will also use DP with devices=[0]
+ have_gpus = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 0
+ if have_gpus:
+ self.use_dp = distributed_backend == 'dp'
+ self.use_ddp = distributed_backend == 'ddp'
+
+ # use ddp automatically if nb_gpu_nodes > 1
+ if nb_gpu_nodes > 1 and self.use_dp: # pragma: no cover
+ self.use_ddp = True
+ self.use_dp = False
+ w = 'DataParallel does not support nb_gpu_nodes > 1. ' \
+ 'Switching to DistributedDataParallel for you. ' \
+ 'To silence this warning set distributed_backend=ddp'
+ warnings.warn(w)
+
+ # extract SLURM flag vars
+ # whenever we have the correct number of tasks, we let slurm manage processes
+ # otherwise we launch the required number of processes
+ if self.use_ddp:
+ self.nb_requested_gpus = len(self.data_parallel_device_ids) * self.nb_gpu_nodes
+ self.nb_slurm_tasks = 0
+ try:
+ self.nb_slurm_tasks = int(os.environ['SLURM_NTASKS'])
+ self.is_slurm_managing_tasks = self.nb_slurm_tasks == self.nb_requested_gpus
+ except Exception as e:
+ # likely not on slurm, so set the slurm managed flag to false
+ self.is_slurm_managing_tasks = False
+
+ # process info
+ self.proc_rank = 0
+
+ # training state
+ self.optimizers = None
+ self.prog_bar = None
+ self.global_step = 0
+ self.current_epoch = 0
+ self.total_batches = 0
+
+ # logging
+ self.log_save_interval = log_save_interval
+ self.val_check_interval = val_check_interval
+ self.add_log_row_interval = add_log_row_interval
+
+ # dataloaders
+ self.tng_dataloader = None
+ self.test_dataloader = None
+ self.val_dataloader = None
+
+ # how much of the data to use
+ 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))
+
+ # 16 bit mixed precision training using apex
+ self.use_amp = use_amp and APEX_AVAILABLE
+ if self.use_amp:
+ print('using 16bit precision')
+
+ if use_amp and not APEX_AVAILABLE: # pragma: no cover
+ msg = '''
+ You set use_amp=True but do not have apex installed.
+ Install apex first using this guide and rerun with use_amp=True:
+ https://github.com/NVIDIA/apex#linux
+
+ this run will NOT use 16 bit precision
+ '''
+ raise ModuleNotFoundError(msg)
+
+ @property
+ def data_parallel(self):
+ return self.use_dp or self.use_ddp
+
+ def __determine_data_use_amount(self, train_percent_check, val_percent_check, test_percent_check, overfit_pct):
+ """
+ Use less data for debugging purposes
+ """
+ self.train_percent_check = train_percent_check
+ self.val_percent_check = val_percent_check
+ self.test_percent_check = test_percent_check
+ if overfit_pct > 0:
+ self.train_percent_check = overfit_pct
+ self.val_percent_check = overfit_pct
+ self.test_percent_check = overfit_pct
+
+ def __get_model(self):
+ return self.model.module if self.data_parallel else self.model
+
+ def __is_function_implemented(self, f_name):
+ model = self.__get_model()
+ f_op = getattr(model, f_name, None)
+ return callable(f_op)
+
+ @property
+ def __tng_tqdm_dic(self):
+ # ForkedPdb().set_trace()
+ tqdm_dic = {
+ 'tng_loss': '{0:.3f}'.format(self.avg_loss),
+ 'v_nb': '{}'.format(self.experiment.version),
+ 'epoch': '{}'.format(self.current_epoch),
+ 'batch_nb':'{}'.format(self.batch_nb),
+ }
+ tqdm_dic.update(self.tqdm_metrics)
+
+ if self.on_gpu:
+ tqdm_dic['gpu'] = '{}'.format(self.current_gpu_name)
+
+ return tqdm_dic
+
+ @property
+ def tng_tqdm_dic(self):
+ """
+ Read-only for tqdm metrics
+ :return:
+ """
+ return self.__tng_tqdm_dic
+
+ def __layout_bookeeping(self):
+
+ # determine number of training batches
+ self.nb_tng_batches = len(self.tng_dataloader)
+ self.nb_tng_batches = int(self.nb_tng_batches * self.train_percent_check)
+
+ # determine number of validation batches
+ self.nb_val_batches = len(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
+ self.nb_test_batches = len(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(self.nb_tng_batches * self.val_check_interval)
+
+ def __add_tqdm_metrics(self, metrics):
+ for k, v in metrics.items():
+ if type(v) is torch.Tensor:
+ v = v.item()
+
+ self.tqdm_metrics[k] = v
+
+ def validate(self, model, dataloader, max_batches):
+ """
+ Run validation code
+ :param model: PT model
+ :param dataloader: PT dataloader
+ :param max_batches: Scalar
+ :return:
+ """
+ # enable eval mode
+ model.zero_grad()
+ model.eval()
+
+ # disable gradients to save memory
+ torch.set_grad_enabled(False)
+
+ # bookkeeping
+ outputs = []
+
+ # run training
+ for batch_i, data_batch in enumerate(dataloader):
+
+ if data_batch is None: # pragma: no cover
+ continue
+
+ # stop short when on fast dev run
+ if max_batches is not None and batch_i >= max_batches:
+ break
+
+ # -----------------
+ # RUN VALIDATION STEP
+ # -----------------
+ if self.use_ddp:
+ output = model(data_batch, batch_i)
+ elif self.use_dp:
+ output = model(data_batch, batch_i)
+ output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
+
+ else:
+ output = model.validation_step(data_batch, batch_i)
+
+ outputs.append(output)
+
+ # batch done
+ if self.progress_bar and self.prog_bar is not None:
+ self.prog_bar.update(1)
+
+ # give model a chance to do something with the outputs
+ if self.data_parallel:
+ val_results = model.module.validation_end(outputs)
+ else:
+ val_results = model.validation_end(outputs)
+
+ # enable train mode again
+ model.train()
+
+ # enable gradients to save memory
+ torch.set_grad_enabled(True)
+
+ return val_results
+
+ def get_dataloaders(self, model):
+ """
+ Dataloaders are provided by the model
+ :param model:
+ :return:
+ """
+ self.tng_dataloader = model.tng_dataloader
+ self.test_dataloader = model.test_dataloader
+ self.val_dataloader = model.val_dataloader
+
+ if self.use_ddp and not isinstance(self.tng_dataloader.sampler, DistributedSampler):
+ msg = '''
+ when using multiple gpus and multiple nodes you must pass a DistributedSampler to DataLoader(sampler).
+
+ ie: this:
+ dataset = myDataset()
+ dataloader = Dataloader(dataset)
+
+ becomes:
+ dataset = myDataset()
+ dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+ dataloader = Dataloader(dataset, sampler=dist_sampler)
+ '''
+ raise MisconfigurationException(msg)
+
+ # -----------------------------
+ # MODEL TRAINING
+ # -----------------------------
+ def fit(self, model):
+
+ # when using multi-node or DDP within a node start each module in a separate process
+ if self.use_ddp:
+ # must copy only the meta of the exp so it survives pickle/unpickle when going to new process
+ self.experiment = self.experiment.get_meta_copy()
+
+ if self.is_slurm_managing_tasks:
+ task = int(os.environ['SLURM_LOCALID'])
+ self.ddp_train(task, model)
+ else:
+ msg = f"""
+ You requested {self.nb_requested_gpus} GPUs but launched {self.nb_slurm_tasks} slurm tasks.
+ We will launch {self.nb_requested_gpus} processes for you.
+ We recommend you let slurm manage the processes by setting: --ntasks-per-node={self.nb_requested_gpus}
+ If you're not using SLURM, ignore this message!
+ """
+ warnings.warn(msg)
+ mp.spawn(self.ddp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
+
+ # 1 gpu or dp option triggers training using DP module
+ # easier to avoid NCCL issues
+ elif self.use_dp:
+ self.__dp_train(model)
+
+ # ON CPU
+ else:
+ # run through amp wrapper
+ if self.use_amp:
+ raise MisconfigurationException('amp + cpu is not supported. Please use a GPU option')
+
+ # CHOOSE OPTIMIZER
+ # filter out the weights that were done on gpu so we can load on good old cpus
+ self.optimizers = model.configure_optimizers()
+
+ self.__run_pretrain_routine(model)
+
+ # return 1 when finished
+ # used for testing or when we need to know that training succeeded
+ return 1
+
+ def __dp_train(self, model):
+
+ # CHOOSE OPTIMIZER
+ # filter out the weights that were done on gpu so we can load on good old cpus
+ self.optimizers = model.configure_optimizers()
+
+ model.cuda(self.data_parallel_device_ids[0])
+
+ # check for this bug (amp + dp + !01 doesn't work)
+ # https://github.com/NVIDIA/apex/issues/227
+ if self.use_dp and self.use_amp:
+ m = f'amp level {self.amp_level} with DataParallel is not supported. ' \
+ f'See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227. ' \
+ f'We recommend you switch to ddp if you want to use amp'
+ raise MisconfigurationException(m)
+
+ model = LightningDataParallel(model, device_ids=self.data_parallel_device_ids)
+
+ self.__run_pretrain_routine(model)
+
+ def ddp_train(self, gpu_nb, model):
+ """
+ Entry point into a DP thread
+ :param gpu_nb:
+ :param model:
+ :param cluster_obj:
+ :return:
+ """
+ # node rank using relative slurm id
+ # otherwise default to node rank 0
+ try:
+ node_id = os.environ['SLURM_NODEID']
+ self.node_rank = int(node_id)
+ except Exception as e:
+ self.node_rank = 0
+
+ # recover original exp before went into process
+ # init in write mode only on proc 0
+ self.experiment.debug = self.proc_rank > 0
+ self.experiment = self.experiment.get_non_ddp_exp()
+
+ # show progbar only on prog_rank 0
+ self.prog_bar = self.prog_bar and self.node_rank == 0 and gpu_nb == 0
+
+ # determine which process we are and world size
+ 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
+ self.__init_tcp_connection()
+
+ # CHOOSE OPTIMIZER
+ # filter out the weights that were done on gpu so we can load on good old cpus
+ self.optimizers = model.configure_optimizers()
+
+ # MODEL
+ # copy model to each gpu
+ torch.cuda.set_device(gpu_nb)
+ model.cuda(gpu_nb)
+
+ # AMP
+ # run through amp wrapper before going to distributed DP
+ if self.use_amp:
+ # An example
+ model, optimizers = amp.initialize(
+ model, self.optimizers, opt_level=self.amp_level,
+ )
+ self.optimizers = optimizers
+
+ model = LightningDistributedDataParallel(model, device_ids=[gpu_nb], find_unused_parameters=True)
+
+ # continue training routine
+ self.__run_pretrain_routine(model)
+
+ def __init_tcp_connection(self):
+ """
+ Connect all procs in the world using the env:// init
+ Use the first node as the root address
+ :param port:
+ :param tries:
+ :return:
+ """
+ # sets the appropriate port
+ try:
+ port = os.environ['MASTER_PORT']
+ except Exception as e:
+ port = 12910
+ os.environ['MASTER_PORT'] = f'{port}'
+
+ # figure out the root node addr
+ try:
+ root_node = os.environ['SLURM_NODELIST'].split(' ')[0]
+ except Exception as e:
+ root_node = '127.0.0.2'
+
+ root_node = self.resolve_root_node_address(root_node)
+ os.environ['MASTER_ADDR'] = root_node
+
+ dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size)
+
+ def resolve_root_node_address(self, root_node):
+ if '[' in root_node:
+ name = root_node.split('[')[0]
+ number = root_node.split(',')[0]
+ if '-' in number:
+ number = number.split('-')[0]
+
+ number = re.sub('[^0-9]', '', number)
+ root_node = name + number
+
+ return root_node
+
+ def __run_pretrain_routine(self, model):
+ """
+ Sanity check a few things before starting actual training
+ :param model:
+ :return:
+ """
+ ref_model = model
+ if self.data_parallel:
+ ref_model = model.module
+
+ ref_model.trainer = self
+
+ # set local properties on the model
+ ref_model.on_gpu = self.on_gpu
+
+ # transfer data loaders from model
+ self.get_dataloaders(ref_model)
+
+ # init training constants
+ self.__layout_bookeeping()
+
+ # add lr schedulers
+ if self.lr_scheduler_milestones is not None:
+ for optimizer in self.optimizers:
+ scheduler = MultiStepLR(optimizer, self.lr_scheduler_milestones)
+ self.lr_schedulers.append(scheduler)
+
+ # print model summary
+ if self.proc_rank == 0 and self.print_weights_summary:
+ ref_model.summarize()
+
+ # give model convenience properties
+ ref_model.trainer = self
+ ref_model.experiment = self.experiment
+
+ # run tiny validation to make sure program won't crash during val
+ _ = self.validate(model, self.val_dataloader, max_batches=self.nb_sanity_val_steps)
+
+ # save exp to get started
+ 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.__train()
+
+ def __train(self):
+ # run all epochs
+ for epoch_nb in range(self.current_epoch, self.max_nb_epochs):
+ # update the lr scheduler
+ for lr_scheduler in self.lr_schedulers:
+ lr_scheduler.step()
+
+ model = self.__get_model()
+ model.current_epoch = epoch_nb
+
+ # hook
+ if self.__is_function_implemented('on_epoch_start'):
+ model = self.__get_model()
+ model.on_epoch_start()
+
+ self.current_epoch = epoch_nb
+ self.total_batches = self.nb_tng_batches + self.nb_val_batches
+ self.batch_loss_value = 0 # accumulated grads
+
+ # init progbar when requested
+ if self.progress_bar:
+ self.prog_bar = tqdm.tqdm(range(self.total_batches), position=self.process_position)
+
+ for batch_nb, data_batch in enumerate(self.tng_dataloader):
+ self.batch_nb = batch_nb
+ self.global_step += 1
+
+ model = self.__get_model()
+ model.global_step = self.global_step
+
+ # stop when the flag is changed or we've gone past the amount requested in the batches
+ self.total_batch_nb += 1
+ met_batch_limit = batch_nb > self.nb_tng_batches
+ if met_batch_limit:
+ break
+
+ # ---------------
+ # RUN TRAIN STEP
+ # ---------------
+ 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 or early_stop_epoch:
+ self.__run_validation()
+
+ # when batch should be saved
+ if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch:
+ if self.proc_rank == 0:
+ self.experiment.save()
+
+ # when metrics should be logged
+ if batch_nb % self.add_log_row_interval == 0 or early_stop_epoch:
+ # count items in memory
+ # nb_params, nb_tensors = count_mem_items()
+
+ model = self.__get_model()
+ metrics = self.__tng_tqdm_dic
+
+ # add gpu memory
+ if self.on_gpu:
+ mem_map = get_gpu_memory_map()
+ metrics.update(mem_map)
+
+ # add norms
+ if self.track_grad_norm > 0:
+ model = self.__get_model()
+ grad_norm_dic = model.grad_norm(self.track_grad_norm)
+ metrics.update(grad_norm_dic)
+
+ if self.__is_function_implemented('on_tng_metrics'):
+ model.on_tng_metrics(metrics)
+
+ # log metrics
+ scalar_metrics = self.__metrics_to_scalars(metrics, blacklist=self.__log_vals_blacklist())
+ if self.proc_rank == 0:
+ self.experiment.log(scalar_metrics, global_step=self.global_step)
+ self.experiment.save()
+
+ # hook
+ if self.__is_function_implemented('on_batch_end'):
+ model = self.__get_model()
+ model.on_batch_end()
+
+ # end epoch early
+ if early_stop_epoch:
+ break
+
+ # hook
+ if self.__is_function_implemented('on_epoch_end'):
+ model = self.__get_model()
+ model.on_epoch_end()
+
+ # early stopping
+ met_min_epochs = epoch_nb > self.min_nb_epochs
+ if self.enable_early_stop and met_min_epochs:
+ should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb, logs=self.__tng_tqdm_dic)
+
+ # stop training
+ stop = should_stop and met_min_epochs
+ if stop:
+ return
+
+ def __metrics_to_scalars(self, metrics, blacklist=[]):
+ new_metrics = {}
+ for k, v in metrics.items():
+ if type(v) is torch.Tensor:
+ v = v.item()
+
+ if type(v) is dict:
+ v = self.__metrics_to_scalars(v)
+
+ if k not in blacklist:
+ new_metrics[k] = float(v)
+
+ return new_metrics
+
+ def __log_vals_blacklist(self):
+ """avoid logging some vals lightning uses to maintain state"""
+ blacklist = {'batch_nb', 'v_nb', 'gpu'}
+ return blacklist
+
+ def __run_tng_batch(self, data_batch, batch_nb):
+ if data_batch is None:
+ return 0
+
+ # hook
+ if self.__is_function_implemented('on_batch_start'):
+ model_ref = self.__get_model()
+ response = model_ref.on_batch_start(data_batch)
+
+ if response == -1:
+ return -1
+
+ if self.progress_bar:
+ self.prog_bar.update(1)
+
+ # forward pass
+ # return a scalar value and a dic with tqdm metrics
+ if self.use_ddp:
+ output = self.model(data_batch, batch_nb)
+ elif self.use_dp:
+ output = self.model(data_batch, batch_nb)
+ output = reduce_distributed_output(output, len(self.data_parallel_device_ids))
+ else:
+ output = self.model.training_step(data_batch, batch_nb)
+
+ try:
+ model_specific_tqdm_metrics_dic = output['tqdm_metrics']
+ except Exception as e:
+ model_specific_tqdm_metrics_dic = {}
+
+ # if output dict doesn't have the keyword loss
+ # then assume the output=loss if scalar
+ try:
+ loss = output['loss']
+ except Exception as e:
+ if type(output) is torch.Tensor:
+ loss = output
+
+ self.__add_tqdm_metrics(model_specific_tqdm_metrics_dic)
+
+ # backward pass
+ if self.use_amp:
+ # scale loss when using amp
+ for optimizer in self.optimizers:
+ with amp.scale_loss(loss, optimizer) as scaled_loss:
+ scaled_loss.backward()
+ else:
+ loss.backward()
+
+ # insert after step hook
+ if self.__is_function_implemented('on_after_backward'):
+ model_ref = self.__get_model()
+ response = model_ref.on_after_backward()
+
+ if self.print_nan_grads:
+ model = self.__get_model()
+ for param in model.parameters():
+ print(param.grad.float().sum())
+
+ # avoid memory leaks
+ self.batch_loss_value += loss.item()
+
+ # gradient update with accumulated gradients
+ if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
+
+ # clip gradients
+ if self.gradient_clip > 0:
+ model = self.__get_model()
+ torch.nn.utils.clip_grad_norm(model.parameters(), self.gradient_clip)
+
+ # update gradients across all optimizers
+ for optimizer in self.optimizers:
+ optimizer.step()
+
+ # insert after step hook
+ if self.__is_function_implemented('on_before_zero_grad'):
+ model_ref = self.__get_model()
+ response = model_ref.on_before_zero_grad(optimizer)
+
+ # clear gradients
+ optimizer.zero_grad()
+
+ # queuing loss across batches blows it up proportionally... divide out the number accumulated
+ self.batch_loss_value = self.batch_loss_value / self.accumulate_grad_batches
+
+ # track loss
+ self.running_loss.append(self.batch_loss_value)
+ self.batch_loss_value = 0
+ self.avg_loss = np.mean(self.running_loss[-100:])
+
+ # update progbar
+ if self.progress_bar:
+ # add model specific metrics
+ tqdm_metrics = self.__tng_tqdm_dic
+ self.prog_bar.set_postfix(**tqdm_metrics)
+
+ # activate batch end hook
+ if self.__is_function_implemented('on_batch_end'):
+ model = self.__get_model()
+ 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
+ if self.fast_dev_run:
+ print('skipping to check performance bc of --fast_dev_run')
+ elif not can_check_epoch:
+ return
+
+ # 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)
+
+ # 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
+ tqdm_metrics = self.__tng_tqdm_dic
+ self.prog_bar.set_postfix(**tqdm_metrics)
+
+ # model checkpointing
+ 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)
\ No newline at end of file
diff --git a/pytorch_lightning/pt_overrides/override_data_parallel.py b/pytorch_lightning/pt_overrides/override_data_parallel.py
index ff4fce8c..89b550fd 100644
--- a/pytorch_lightning/pt_overrides/override_data_parallel.py
+++ b/pytorch_lightning/pt_overrides/override_data_parallel.py
@@ -1,6 +1,7 @@
from torch.nn import DataParallel
from torch.nn.parallel import DistributedDataParallel
import itertools
+from itertools import chain
import threading
import torch
@@ -8,7 +9,7 @@ from torch.cuda._utils import _get_device_index
import pdb
-def _find_tensors(obj):
+def _find_tensors(obj): # pragma: no cover
r"""
Recursively find all tensors contained in the specified object.
"""
@@ -21,8 +22,7 @@ def _find_tensors(obj):
return []
-
-def get_a_var(obj):
+def get_a_var(obj): # pragma: no cover
if isinstance(obj, torch.Tensor):
return obj
@@ -42,6 +42,29 @@ class LightningDataParallel(DataParallel):
Override the forward call in lightning so it goes to training and validation step respectively
"""
+ def forward(self, *inputs, **kwargs):
+ if not self.device_ids:
+ return self.module(*inputs, **kwargs)
+
+ for t in chain(self.module.parameters(), self.module.buffers()):
+ if t.device != self.src_device_obj:
+ raise RuntimeError("module must have its parameters and buffers "
+ "on device {} (device_ids[0]) but found one of "
+ "them on device: {}".format(self.src_device_obj, t.device))
+
+ inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
+ if len(self.device_ids) == 1:
+ # lightning
+ if self.module.training:
+ return self.module.training_step(*inputs[0], **kwargs[0])
+ else:
+ return self.module.validation_step(*inputs[0], **kwargs[0])
+
+ replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
+ outputs = self.parallel_apply(replicas, inputs, kwargs)
+ return self.gather(outputs, self.output_device)
+
+
def parallel_apply(self, replicas, inputs, kwargs):
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
@@ -54,7 +77,7 @@ class LightningDistributedDataParallel(DistributedDataParallel):
def parallel_apply(self, replicas, inputs, kwargs):
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
- def forward(self, *inputs, **kwargs):
+ def forward(self, *inputs, **kwargs): # pragma: no cover
self._sync_params()
if self.device_ids:
inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
@@ -89,7 +112,7 @@ class LightningDistributedDataParallel(DistributedDataParallel):
return output
-def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
+def parallel_apply(modules, inputs, kwargs_tup=None, devices=None): # pragma: no cover
r"""Applies each `module` in :attr:`modules` in parallel on arguments
contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword)
on each of :attr:`devices`.
diff --git a/pytorch_lightning/root_module/decorators.py b/pytorch_lightning/root_module/decorators.py
new file mode 100644
index 00000000..ef7bd502
--- /dev/null
+++ b/pytorch_lightning/root_module/decorators.py
@@ -0,0 +1,17 @@
+
+def data_loader(fn):
+ """
+ Decorator to make any fx with this use the lazy property
+ :param fn:
+ :return:
+ """
+
+ attr_name = '_lazy_' + fn.__name__
+
+ @property
+ def _data_loader(self):
+ if not hasattr(self, attr_name):
+ setattr(self, attr_name, fn(self))
+ return getattr(self, attr_name)
+
+ return _data_loader
diff --git a/pytorch_lightning/root_module/grads.py b/pytorch_lightning/root_module/grads.py
index e4d1701b..8ed17a3e 100644
--- a/pytorch_lightning/root_module/grads.py
+++ b/pytorch_lightning/root_module/grads.py
@@ -27,14 +27,3 @@ class GradInformation(nn.Module):
results['grad_{}_norm_total'.format(norm_type)] = round(total_norm.data.cpu().numpy().flatten()[0], 3)
return results
-
- def describe_grads(self):
- for p in self.parameters():
- g = p.grad.data.numpy().flatten()
- print(np.max(g), np.min(g), np.mean(g))
-
-
- def describe_params(self):
- for p in self.parameters():
- g = p.data.numpy().flatten()
- print(np.max(g), np.min(g), np.mean(g))
\ No newline at end of file
diff --git a/pytorch_lightning/root_module/hooks.py b/pytorch_lightning/root_module/hooks.py
index 99155ab9..88abe80d 100644
--- a/pytorch_lightning/root_module/hooks.py
+++ b/pytorch_lightning/root_module/hooks.py
@@ -19,3 +19,27 @@ class ModelHooks(torch.nn.Module):
def on_post_performance_check(self):
pass
+ def on_tng_metrics(self, metrics):
+ pass
+
+ def on_before_zero_grad(self, optimizer):
+ """
+ Called after optimizer.step() and before optimizer.zero_grad()
+
+ for optimizer in optimizers:
+ optimizer.step()
+ model.on_before_zero_grad(optimizer) # < ---- called here
+ optimizer.zero_grad
+
+ :param optimizer:
+ :return:
+ """
+ pass
+
+ def on_after_backward(self):
+ """
+ Called after loss.backward() and before optimizers do anything
+ :return:
+ """
+ pass
+
diff --git a/pytorch_lightning/root_module/memory.py b/pytorch_lightning/root_module/memory.py
index 17f20efe..389c8680 100644
--- a/pytorch_lightning/root_module/memory.py
+++ b/pytorch_lightning/root_module/memory.py
@@ -33,33 +33,42 @@ class ModelSummary(object):
mods = list(self.model.modules())
in_sizes = []
out_sizes = []
- input_ = self.example_input_array
- for i in range(1, len(mods)):
- m = mods[i]
- if type(input_) is list or type(input_) is tuple:
- out = m(*input_)
- else:
- out = m(input_)
+ input_ = self.model.example_input_array
- if type(input_) is tuple or type(input_) is list:
- in_size = []
- for x in input_:
- if type(x) is list:
- in_size.append(len(x))
- else:
- in_size.append(x.size())
- else:
- in_size = np.array(input_.size())
+ if self.model.on_gpu:
+ input_ = input_.cuda(0)
- in_sizes.append(in_size)
+ if self.model.trainer.use_amp:
+ input_ = input_.half()
- if type(out) is tuple or type(out) is list:
- out_size = np.asarray([x.size() for x in out])
- else:
- out_size = np.array(out.size())
+ with torch.no_grad():
- out_sizes.append(out_size)
- input_ = out
+ for i in range(1, len(mods)):
+ m = mods[i]
+ if type(input_) is list or type(input_) is tuple: # pragma: no cover
+ out = m(*input_)
+ else:
+ out = m(input_)
+
+ if type(input_) is tuple or type(input_) is list: # pragma: no cover
+ in_size = []
+ for x in input_:
+ if type(x) is list:
+ in_size.append(len(x))
+ else:
+ in_size.append(x.size())
+ else:
+ in_size = np.array(input_.size())
+
+ in_sizes.append(in_size)
+
+ if type(out) is tuple or type(out) is list: # pragma: no cover
+ out_size = np.asarray([x.size() for x in out])
+ else:
+ out_size = np.array(out.size())
+
+ out_sizes.append(out_size)
+ input_ = out
self.in_sizes = in_sizes
self.out_sizes = out_sizes
@@ -114,13 +123,22 @@ class ModelSummary(object):
Layer Name, Layer Type, Input Size, Output Size, Number of Parameters
'''
- df = pd.DataFrame( np.zeros( (len(self.layer_names), 3) ) )
- df.columns = ['Name', 'Type', 'Params']
+ cols = ['Name', 'Type', 'Params']
+ if self.model.example_input_array is not None:
+ cols.extend(['In_sizes', 'Out_sizes'])
+
+ df = pd.DataFrame(np.zeros( (len(self.layer_names), len(cols))))
+ df.columns = cols
df['Name'] = self.layer_names
df['Type'] = self.layer_types
df['Params'] = self.param_nums
+ if self.model.example_input_array is not None:
+
+ df['In_sizes'] = self.in_sizes
+ df['Out_sizes'] = self.out_sizes
+
self.summary = df
return
@@ -128,10 +146,13 @@ class ModelSummary(object):
self.get_layer_names()
self.get_parameter_sizes()
self.get_parameter_nums()
+
+ if self.model.example_input_array is not None:
+ self.get_variable_sizes()
self.make_summary()
-def print_mem_stack():
+def print_mem_stack(): # pragma: no cover
for obj in gc.get_objects():
try:
if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
@@ -140,7 +161,7 @@ def print_mem_stack():
pass
-def count_mem_items():
+def count_mem_items(): # pragma: no cover
nb_params = 0
nb_tensors = 0
for obj in gc.get_objects():
diff --git a/pytorch_lightning/root_module/model_saving.py b/pytorch_lightning/root_module/model_saving.py
index 0fca9161..c5831317 100644
--- a/pytorch_lightning/root_module/model_saving.py
+++ b/pytorch_lightning/root_module/model_saving.py
@@ -2,36 +2,38 @@ import torch
import os
import re
import pdb
-from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel
+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:
@@ -41,6 +43,11 @@ class ModelIO(object):
class TrainerIO(object):
+ def __get_model(self):
+ is_dp_module = type(self.model) is LightningDistributedDataParallel or type(self.model) is LightningDataParallel
+ model = self.model.module if is_dp_module else self.model
+ return model
+
# --------------------
# MODEL SAVE CHECKPOINT
# --------------------
@@ -51,26 +58,32 @@ class TrainerIO(object):
torch.save(checkpoint, filepath)
def dump_checkpoint(self):
+
checkpoint = {
'epoch': self.current_epoch,
- 'checkpoint_callback_best': self.checkpoint_callback.best,
- 'early_stop_callback_wait': self.early_stop_callback.wait,
- 'early_stop_callback_patience': self.early_stop_callback.patience,
'global_step': self.global_step
}
+ if self.checkpoint_callback is not None:
+ checkpoint['checkpoint_callback_best'] = self.checkpoint_callback.best
+
+ if self.early_stop_callback is not None:
+ checkpoint['early_stop_callback_wait'] = self.early_stop_callback.wait
+ checkpoint['early_stop_callback_patience'] = self.early_stop_callback.patience
+
optimizer_states = []
for i, optimizer in enumerate(self.optimizers):
optimizer_states.append(optimizer.state_dict())
checkpoint['optimizer_states'] = optimizer_states
- # request what to save from the model
- model = self.model.module if type(self.model) is LightningDistributedDataParallel else self.model
- checkpoint_dict = model.get_save_dict()
+ # add the state_dict from the model
+ model = self.__get_model()
+ 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
# --------------------
@@ -103,9 +116,13 @@ class TrainerIO(object):
:param checkpoint:
:return:
"""
- self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
- self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
- self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
+ if self.checkpoint_callback is not None:
+ self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
+
+ if self.early_stop_callback is not None:
+ self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
+ self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
+
self.global_step = checkpoint['global_step']
self.current_epoch = checkpoint['epoch']
@@ -134,13 +151,15 @@ class TrainerIO(object):
filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number)
# give model a chance to do something on hpc_save
- self.on_hpc_save()
+ model = self.__get_model()
+ 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))
@@ -150,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.model.module if type(self.model) is LightningDataParallel else self.model
- model.load_model_specific(checkpoint)
+ model = self.__get_model()
+
+ # load the state_dict on the model automatically
+ model.load_state_dict(checkpoint['state_dict'])
# call model hook
- self.on_hpc_load()
+ model.on_hpc_load(checkpoint)
def max_ckpt_in_folder(self, path):
files = os.listdir(path)
diff --git a/pytorch_lightning/root_module/optimization.py b/pytorch_lightning/root_module/optimization.py
deleted file mode 100644
index 3172e1a1..00000000
--- a/pytorch_lightning/root_module/optimization.py
+++ /dev/null
@@ -1,22 +0,0 @@
-from torch import nn
-from torch import optim
-
-
-class OptimizerConfig(nn.Module):
-
- def choose_optimizer(self, optimizer, params, optimizer_params, opt_name_key):
- if optimizer == 'adam':
- optimizer = optim.Adam(params, **optimizer_params)
- if optimizer == 'sparse_adam':
- optimizer = optim.SparseAdam(params, **optimizer_params)
- if optimizer == 'sgd':
- optimizer = optim.SGD(params, **optimizer_params)
- if optimizer == 'adadelta':
- optimizer = optim.Adadelta(params, **optimizer_params)
-
- # transfer opt state if loaded
- if opt_name_key in self.loaded_optimizer_states_dict:
- state = self.loaded_optimizer_states_dict[opt_name_key]
- optimizer.load_state_dict(state)
-
- return optimizer
diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py
index ef43e09a..b49dd3af 100644
--- a/pytorch_lightning/root_module/root_module.py
+++ b/pytorch_lightning/root_module/root_module.py
@@ -1,19 +1,15 @@
-import os
import torch
-import math
-
from pytorch_lightning.root_module.memory import ModelSummary
from pytorch_lightning.root_module.grads import GradInformation
from pytorch_lightning.root_module.model_saving import ModelIO, load_hparams_from_tags_csv
-from pytorch_lightning.root_module.optimization import OptimizerConfig
from pytorch_lightning.root_module.hooks import ModelHooks
+from pytorch_lightning.root_module.decorators import data_loader
-class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
+class LightningModule(GradInformation, ModelIO, ModelHooks):
- def __init__(self, hparams):
- super(LightningModule, self).__init__()
- self.hparams = hparams
+ def __init__(self, *args, **kwargs):
+ super(LightningModule, self).__init__(*args, **kwargs)
self.dtype = torch.FloatTensor
self.exp_save_path = None
@@ -22,15 +18,11 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
self.loaded_optimizer_states_dict = {}
self.trainer = None
self.experiment = None
+ self.example_input_array = None
# track if gpu was requested for checkpointing
self.on_gpu = False
- # computed vars for the dataloaders
- self._tng_dataloader = None
- self._val_dataloader = None
- self._test_dataloader = None
-
def forward(self, *args, **kwargs):
"""
Expand model in into whatever you need.
@@ -71,37 +63,7 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
"""
raise NotImplementedError
- def update_tng_log_metrics(self, logs):
- """
- Chance to update metrics to be logged for training step.
- For example, add music, images, etc... to log
- :param logs:
- :return:
- """
- raise NotImplementedError
-
- def loss(self, *args, **kwargs):
- """
- Expand model_out into your components
- :param model_out:
- :return:
- """
- raise NotImplementedError
-
- def summarize(self):
- model_summary = ModelSummary(self)
- print(model_summary)
-
-
- def freeze(self):
- for param in self.parameters():
- param.requires_grad = False
-
- def unfreeze(self):
- for param in self.parameters():
- param.requires_grad = True
-
- @property
+ @data_loader
def tng_dataloader(self):
"""
Implement a function to load an h5py of this data
@@ -109,7 +71,7 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
"""
raise NotImplementedError
- @property
+ @data_loader
def test_dataloader(self):
"""
Implement a function to load an h5py of this data
@@ -117,7 +79,7 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
"""
raise NotImplementedError
- @property
+ @data_loader
def val_dataloader(self):
"""
Implement a function to load an h5py of this data
@@ -125,16 +87,6 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks):
"""
raise NotImplementedError
- @staticmethod
- def get_process_position(gpus):
- try:
- current_gpu = os.environ["CUDA_VISIBLE_DEVICES"]
- gpu_ids = gpus.split(';')
- process_position = gpu_ids.index(current_gpu)
- return process_position, current_gpu
- except Exception as e:
- return 0, 0
-
@classmethod
def load_from_metrics(cls, weights_path, tags_csv, on_gpu, map_location=None):
"""
@@ -156,9 +108,26 @@ class LightningModule(GradInformation, ModelIO, OptimizerConfig, 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):
+ model_summary = ModelSummary(self)
+ print(model_summary)
+
+ def freeze(self):
+ for param in self.parameters():
+ param.requires_grad = False
+
+ def unfreeze(self):
+ for param in self.parameters():
+ param.requires_grad = True
+
+
+
diff --git a/pytorch_lightning/testing_models/__init__.py b/pytorch_lightning/testing_models/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/pytorch_lightning/testing_models/lm_test_module.py b/pytorch_lightning/testing_models/lm_test_module.py
new file mode 100644
index 00000000..39fcbb4c
--- /dev/null
+++ b/pytorch_lightning/testing_models/lm_test_module.py
@@ -0,0 +1,253 @@
+import os
+from collections import OrderedDict
+import torch.nn as nn
+from torchvision.datasets import MNIST
+import torchvision.transforms as transforms
+import torch
+import torch.nn.functional as F
+from test_tube import HyperOptArgumentParser
+from torch import optim
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from pytorch_lightning.root_module.root_module import LightningModule
+import pytorch_lightning as ptl
+
+
+class LightningTestModel(LightningModule):
+ """
+ Sample model to show how to define a template
+ """
+
+ def __init__(self, hparams, force_remove_distributed_sampler=False):
+ """
+ Pass in parsed HyperOptArgumentParser to the model
+ :param hparams:
+ """
+ # init superclass
+ super(LightningTestModel, self).__init__()
+ self.hparams = hparams
+
+ self.batch_size = hparams.batch_size
+
+ # if you specify an example input, the summary will show input/output for each layer
+ self.example_input_array = torch.rand(5, 28 * 28)
+
+ # remove to test warning for dist sampler
+ self.force_remove_distributed_sampler = force_remove_distributed_sampler
+
+ # build model
+ self.__build_model()
+
+ # ---------------------
+ # MODEL SETUP
+ # ---------------------
+ def __build_model(self):
+ """
+ Layout model
+ :return:
+ """
+ self.c_d1 = nn.Linear(in_features=self.hparams.in_features, out_features=self.hparams.hidden_dim)
+ self.c_d1_bn = nn.BatchNorm1d(self.hparams.hidden_dim)
+ self.c_d1_drop = nn.Dropout(self.hparams.drop_prob)
+
+ self.c_d2 = nn.Linear(in_features=self.hparams.hidden_dim, out_features=self.hparams.out_features)
+
+ # ---------------------
+ # TRAINING
+ # ---------------------
+ def forward(self, x):
+ """
+ No special modification required for lightning, define as you normally would
+ :param x:
+ :return:
+ """
+
+ x = self.c_d1(x)
+ x = torch.tanh(x)
+ x = self.c_d1_bn(x)
+ x = self.c_d1_drop(x)
+
+ x = self.c_d2(x)
+ logits = F.log_softmax(x, dim=1)
+
+ return logits
+
+ def loss(self, labels, logits):
+ nll = F.nll_loss(logits, labels)
+ return nll
+
+ def training_step(self, data_batch, batch_i):
+ """
+ Lightning calls this inside the training loop
+ :param data_batch:
+ :return:
+ """
+ # forward pass
+ x, y = data_batch
+ x = x.view(x.size(0), -1)
+
+ y_hat = self.forward(x)
+
+ # calculate loss
+ loss_val = self.loss(y, y_hat)
+
+ # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+ if self.trainer.use_dp:
+ loss_val = loss_val.unsqueeze(0)
+
+ output = OrderedDict({
+ 'loss': loss_val
+ })
+
+ # can also return just a scalar instead of a dict (return loss_val)
+ return output
+
+ def validation_step(self, data_batch, batch_i):
+ """
+ Lightning calls this inside the validation loop
+ :param data_batch:
+ :return:
+ """
+ x, y = data_batch
+ x = x.view(x.size(0), -1)
+ y_hat = self.forward(x)
+
+ loss_val = self.loss(y, y_hat)
+
+ # acc
+ labels_hat = torch.argmax(y_hat, dim=1)
+ val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+ val_acc = torch.tensor(val_acc)
+
+ if self.on_gpu:
+ val_acc = val_acc.cuda(loss_val.device.index)
+
+ # in DP mode (default) make sure if result is scalar, there's another dim in the beginning
+ if self.trainer.use_dp:
+ loss_val = loss_val.unsqueeze(0)
+ val_acc = val_acc.unsqueeze(0)
+
+ # alternate possible outputs to test
+ if self.trainer.batch_nb % 1 == 0:
+ output = OrderedDict({
+ 'val_loss': loss_val,
+ 'val_acc': val_acc,
+ })
+ return output
+ if self.trainer.batch_nb % 2 == 0:
+ return val_acc
+
+ if self.trainer.batch_nb % 3 == 0:
+ output = OrderedDict({
+ 'val_loss': loss_val,
+ 'val_acc': val_acc,
+ 'test_dic': {'val_loss_a': loss_val}
+ })
+ return output
+
+ def validation_end(self, outputs):
+ """
+ Called at the end of validation to aggregate outputs
+ :param outputs: list of individual outputs of each validation step
+ :return:
+ """
+ # if returned a scalar from validation_step, outputs is a list of tensor scalars
+ # we return just the average in this case (if we want)
+ # return torch.stack(outputs).mean()
+
+ val_loss_mean = 0
+ val_acc_mean = 0
+ for output in outputs:
+ val_loss_mean += output['val_loss']
+ val_acc_mean += output['val_acc']
+
+ val_loss_mean /= len(outputs)
+ val_acc_mean /= len(outputs)
+
+ tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()}
+ return tqdm_dic
+
+ def on_tng_metrics(self, logs):
+ logs['some_tensor_to_test'] = torch.rand(1)
+
+ # ---------------------
+ # TRAINING SETUP
+ # ---------------------
+ def configure_optimizers(self):
+ """
+ return whatever optimizers we want here
+ :return: list of optimizers
+ """
+ optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
+ return [optimizer]
+
+ def __dataloader(self, train):
+ # init data generators
+ transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (1.0,))])
+ dataset = MNIST(root=self.hparams.data_root, train=train, transform=transform, download=True)
+
+ # when using multi-node we need to add the datasampler
+ train_sampler = None
+ batch_size = self.hparams.batch_size
+
+ try:
+ if self.on_gpu and not self.force_remove_distributed_sampler:
+ train_sampler = DistributedSampler(dataset, rank=self.trainer.proc_rank)
+ batch_size = batch_size // self.trainer.world_size # scale batch size
+ except Exception as e:
+ pass
+
+ should_shuffle = train_sampler is None
+ loader = DataLoader(
+ dataset=dataset,
+ batch_size=batch_size,
+ shuffle=should_shuffle,
+ sampler=train_sampler
+ )
+
+ return loader
+
+ @ptl.data_loader
+ def tng_dataloader(self):
+ return self.__dataloader(train=True)
+
+ @ptl.data_loader
+ def val_dataloader(self):
+ return self.__dataloader(train=False)
+
+ @ptl.data_loader
+ def test_dataloader(self):
+ return self.__dataloader(train=False)
+
+ @staticmethod
+ def add_model_specific_args(parent_parser, root_dir):
+ """
+ Parameters you define here will be available to your model through self.hparams
+ :param parent_parser:
+ :param root_dir:
+ :return:
+ """
+ parser = HyperOptArgumentParser(strategy=parent_parser.strategy, parents=[parent_parser])
+
+ # param overwrites
+ # parser.set_defaults(gradient_clip=5.0)
+
+ # network params
+ parser.opt_list('--drop_prob', default=0.2, options=[0.2, 0.5], type=float, tunable=False)
+ parser.add_argument('--in_features', default=28*28, type=int)
+ parser.add_argument('--out_features', default=10, type=int)
+ parser.add_argument('--hidden_dim', default=50000, type=int) # use 500 for CPU, 50000 for GPU to see speed difference
+
+ # data
+ parser.add_argument('--data_root', default=os.path.join(root_dir, 'mnist'), type=str)
+
+ # training params (opt)
+ parser.opt_list('--learning_rate', default=0.001*8, type=float, options=[0.0001, 0.0005, 0.001, 0.005],
+ tunable=False)
+ parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False)
+
+ # if using 2 nodes with 4 gpus each the batch size here (256) will be 256 / (2*8) = 16 per gpu
+ parser.opt_list('--batch_size', default=256*8, type=int, options=[32, 64, 128, 256], tunable=False,
+ help='batch size will be divided over all the gpus being used across all nodes')
+ return parser
diff --git a/pytorch_lightning/trainer_main.py b/pytorch_lightning/trainer_main.py
index 8389b5ec..0c30a419 100644
--- a/pytorch_lightning/trainer_main.py
+++ b/pytorch_lightning/trainer_main.py
@@ -52,10 +52,6 @@ def main(hparams, cluster, results_dict):
hparams.__setattr__('nb_gpus', torch.cuda.device_count())
hparams.__setattr__('inference_mode', hparams.model_load_weights_path is not None)
- # delay each training start to not overwrite logs
- process_position, current_gpu = TRAINING_MODEL.get_process_position(hparams.gpus)
- sleep(process_position + 1)
-
# init experiment
exp = Experiment(
name=hparams.tt_name,
diff --git a/pytorch_lightning/utils/debugging.py b/pytorch_lightning/utils/debugging.py
new file mode 100644
index 00000000..3091ff3b
--- /dev/null
+++ b/pytorch_lightning/utils/debugging.py
@@ -0,0 +1,5 @@
+import pdb
+import sys
+
+class MisconfigurationException(Exception):
+ pass
\ No newline at end of file
diff --git a/pytorch_lightning/utils/embeddings.py b/pytorch_lightning/utils/embeddings.py
deleted file mode 100644
index 4e96c61b..00000000
--- a/pytorch_lightning/utils/embeddings.py
+++ /dev/null
@@ -1,104 +0,0 @@
-import torch
-import numpy as np
-from copy import deepcopy
-
-
-class PretrainedEmbedding(torch.nn.Embedding):
-
- def __init__(self, embedding_path, embedding_dim, task_vocab, freeze=True, *args, **kwargs):
- """
- Loads a prebuilt pytorch embedding from any embedding formated file.
- Padding=0 by default.
-
- >>> emb = PretrainedEmbedding(embedding_path='glove.840B.300d.txt',embedding_dim=300, task_vocab={'hello': 1, 'world': 2})
- >>> data = torch.Tensor([[0, 1], [0, 2]]).long()
- >>> embedded = emb(data)
-
-
-
- :param embedding_path:
- :param emb_dim:
- :param task_vocab:
- :param freeze:
- :return:
- """
- # count the vocab
- self.vocab_size = max(task_vocab.values()) + 1
- super(PretrainedEmbedding, self).__init__(self.vocab_size, embedding_dim, padding_idx=0, *args, **kwargs)
-
- # load pretrained embeddings
- new_emb = self.__load_task_specific_embeddings(deepcopy(task_vocab), embedding_path, embedding_dim, freeze)
-
- # transfer weights
- self.weight = new_emb.weight
-
- # apply freeze
- should_freeze = not freeze
- self.weight.requires_grad = should_freeze
-
- def __load_task_specific_embeddings(self, vocab_words, embedding_path, emb_dim, freeze):
- """
- Iterates embedding file to only pull out task specific embeddings
- :param vocab_words:
- :param embedding_path:
- :param emb_dim:
- :param freeze:
- :return:
- """
-
- # holds final embeddings for relevant words
- embeddings = np.zeros(shape=(self.vocab_size, emb_dim))
-
- # load embedding line by line and extract relevant embeddings
- with open(embedding_path, encoding='utf-8') as f:
- for line in f:
- tokens = line.split(' ')
- word = tokens[0]
- embedding = tokens[1:]
- embedding[-1] = embedding[-1][:-1] # remove last new line
-
- if word in vocab_words:
- vocab_word_i = vocab_words[word]
-
- # skip words that try to overwrite pad idx
- if vocab_word_i == 0:
- del vocab_words[word]
- continue
-
- emb_vals = np.asarray([float(x) for x in embedding])
- embeddings[vocab_word_i] = emb_vals
-
- # remove vocab word to early terminate
- del vocab_words[word]
-
- # early break
- if len(vocab_words) == 0:
- break
-
- # add random vectors for the non-pretrained words
- # these are vocab words NOT found in the pretrained embeddings
- for w, i in vocab_words.items():
- # skip words that try to overwrite pad idx
- if i == 0:
- continue
-
- embedding = np.random.normal(size=emb_dim)
- embeddings[i] = embedding
-
- # turn into pt embedding
- embeddings = torch.FloatTensor(embeddings)
- embeddings = torch.nn.Embedding.from_pretrained(embeddings, freeze=freeze)
-
- return embeddings
-
-
-if __name__ == '__main__':
- emb = PretrainedEmbedding(
- embedding_path='/Users/waf/Developer',
- embedding_dim=300,
- task_vocab={'hello': 1, 'world': 2}
- )
-
- data = torch.Tensor([[0, 1], [0, 2]]).long()
- embedded = emb(data)
- print(embedded)
diff --git a/pytorch_lightning/utils/plotting.py b/pytorch_lightning/utils/plotting.py
deleted file mode 100644
index 95fe64cc..00000000
--- a/pytorch_lightning/utils/plotting.py
+++ /dev/null
@@ -1,28 +0,0 @@
-from matplotlib import pyplot as plt
-import numpy as np
-np.seterr(divide='ignore', invalid='ignore')
-
-
-def plot_confusion_matrix(cm,
- save_path,
- normalize=False,
- title='Confusion matrix',
- ylabel='y',
- xlabel='x'):
- """
- This function prints and plots the confusion matrix.
- Normalization can be applied by setting `normalize=True`.
- """
- if normalize:
- cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
- print("Normalized confusion matrix")
- else:
- print('Confusion matrix, without normalization')
-
- fig = plt.figure()
- plt.matshow(cm)
- plt.title(title)
- plt.colorbar()
- plt.ylabel(ylabel)
- plt.xlabel(xlabel)
- plt.savefig(save_path)
diff --git a/requirements.txt b/requirements.txt
index 730be99c..3e863f05 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,35 +1,9 @@
-
-atomicwrites==1.2.1
-attrs==18.2.0
-certifi==2018.11.29
-cffi==1.11.5
-imageio==2.4.1
-mkl-fft==1.0.6
-mkl-random==1.0.2
-more-itertools==5.0.0
-numpy==1.15.4
-olefile==0.46
-pandas==0.23.4
-Pillow==5.3.0
-pluggy==0.8.0
-py==1.7.0
-pycparser==2.19
-pytest==4.0.2
-python-dateutil==2.7.5
-pytz==2018.7
+coverage==4.5.3
+mkdocs==1.0.4
+pytest==5.0.1
scikit-learn==0.20.2
-scipy==1.2.0
-six==1.12.0
-sklearn==0.0
-tensorboard==1.14.0
-tensorboardX==1.7
-tensorflow==1.14.0
-test-tube==0.643
-torch==1.0.0
-torchvision==0.2.1
tqdm==4.32.1
twine==1.13.0
-urllib3==1.25.3
-webencodings==0.5.1
-Werkzeug==0.15.4
-wrapt==1.11.2
+numpy==1.16.4
+torch>=1.1.0
+torchvision==0.3.0
diff --git a/setup.cfg b/setup.cfg
index c7616dab..5b21c683 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -16,6 +16,33 @@ markers =
ignore = E731,W504
max-line-length = 120
+[coverage:report]
+exclude_lines =
+ pragma: no cover
+ def __repr__
+ if self.debug:
+ if settings.DEBUG
+ raise AssertionError
+ raise NotImplementedError
+ if 0:
+ if __name__ == .__main__.:
+ except Exception as e
+ print(e)
+ print(traceback.print_exc())
+ return *
+ raise Exception
+ warnings
+ print
+ raise RuntimeError
+ break
+ pass
+ os.makedirs
+
+omit =
+ pytorch_lightning/callbacks/pt_callbacks.py
+ tests/test_models.py
+ pytorch_lightning/testing_models/lm_test_module.py
+
[flake8]
ignore = E731,W504,F401,F841
max-line-length = 120
diff --git a/setup.py b/setup.py
index 5fe9ac8b..50a05057 100755
--- a/setup.py
+++ b/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.21',
+ version='0.3.6.4',
description="The Keras for ML researchers using PyTorch",
author="William Falcon",
author_email="waf2107@columbia.edu",
@@ -19,8 +19,7 @@ setup(
install_requires=[
"torch>=1.1.0",
"tqdm",
- "test-tube>=0.653",
- "tensorflow>=1.14.0"
+ "test-tube>=0.6.7.4",
],
packages=find_packages(),
long_description=open("README.md", encoding="utf-8").read(),
diff --git a/tests/README.md b/tests/README.md
new file mode 100644
index 00000000..20f783bf
--- /dev/null
+++ b/tests/README.md
@@ -0,0 +1,58 @@
+# Pytorch-Lightning Tests
+
+## Running tests
+The automatic travis tests ONLY run CPU-based tests. Although these cover most of the use cases,
+run on a 2-GPU machine to validate the full test-suite.
+
+
+To run all tests do the following:
+```bash
+git clone https://github.com/williamFalcon/pytorch-lightning
+cd pytorch-lightning
+
+# install module locally
+pip install -e .
+
+# install dev deps
+pip install -r requirements.txt
+
+# run tests
+py.test
+```
+
+To test models that require GPU make sure to run the above command on a GPU machine.
+The GPU machine must have:
+1. At least 2 GPUs.
+2. [NVIDIA-apex](https://github.com/NVIDIA/apex#linux) installed.
+
+
+### test_models.py
+This file fits a tiny model on MNIST using these different set-ups.
+1. CPU only.
+2. Single GPU with DP.
+3. Multiple (2) GPUs using DP.
+3. Multiple (2) GPUs using DDP.
+3. Multiple (2) GPUs using DP + apex (for 16-bit precision).
+3. Multiple (2) GPUs using DDP + apex (for 16-bit precision).
+
+For each set up it also tests:
+1. model saving.
+2. model loading.
+3. predicting with a loaded model.
+4. simulated save from HPC signal.
+5. simulated load from HPC signal.
+
+## Running Coverage
+
+```bash
+cd pytorch-lightning
+
+# generate coverage
+pip install coverage
+coverage run tests/test_models.py
+
+# print coverage stats
+coverage report -m
+```
+
+
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/debug.py b/tests/debug.py
new file mode 100644
index 00000000..c4c3ffd1
--- /dev/null
+++ b/tests/debug.py
@@ -0,0 +1,180 @@
+import pytest
+from pytorch_lightning import Trainer
+from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel
+from argparse import Namespace
+from test_tube import Experiment
+from pytorch_lightning.callbacks import ModelCheckpoint
+import numpy as np
+import warnings
+import torch
+import os
+import shutil
+import pdb
+
+import pytorch_lightning as ptl
+import torch
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torchvision.datasets import MNIST
+
+
+class CoolModel(ptl.LightningModule):
+
+ def __init(self):
+ super(CoolModel, self).__init__()
+ # not the best model...
+ self.l1 = torch.nn.Linear(28 * 28, 10)
+
+ def forward(self, x):
+ return torch.relu(self.l1(x))
+
+ def my_loss(self, y_hat, y):
+ return F.cross_entropy(y_hat, y)
+
+ def training_step(self, batch, batch_nb):
+ x, y = batch
+ y_hat = self.forward(x)
+ return {'tng_loss': self.my_loss(y_hat, y)}
+
+ def validation_step(self, batch, batch_nb):
+ x, y = batch
+ y_hat = self.forward(x)
+ return {'val_loss': self.my_loss(y_hat, y)}
+
+ def validation_end(self, outputs):
+ avg_loss = torch.stack([x for x in outputs['val_loss']]).mean()
+ return avg_loss
+
+ def configure_optimizers(self):
+ return [torch.optim.Adam(self.parameters(), lr=0.02)]
+
+ @ptl.data_loader
+ def tng_dataloader(self):
+ return DataLoader(MNIST('path/to/save', train=True), batch_size=32)
+
+ @ptl.data_loader
+ def val_dataloader(self):
+ return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+
+ @ptl.data_loader
+ def test_dataloader(self):
+ return DataLoader(MNIST('path/to/save', train=False), batch_size=32)
+
+
+def get_model():
+ # set up model with these hyperparams
+ root_dir = os.path.dirname(os.path.realpath(__file__))
+ hparams = Namespace(**{'drop_prob': 0.2,
+ 'batch_size': 32,
+ 'in_features': 28*28,
+ 'learning_rate': 0.001*8,
+ 'optimizer_name': 'adam',
+ 'data_root': os.path.join(root_dir, 'mnist'),
+ 'out_features': 10,
+ 'hidden_dim': 1000})
+ model = LightningTemplateModel(hparams)
+
+ return model, hparams
+
+
+def get_exp(debug=True):
+ # set up exp object without actually saving logs
+ root_dir = os.path.dirname(os.path.realpath(__file__))
+ exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
+ return exp
+
+
+def init_save_dir():
+ root_dir = os.path.dirname(os.path.realpath(__file__))
+ save_dir = os.path.join(root_dir, 'save_dir')
+
+ if os.path.exists(save_dir):
+ shutil.rmtree(save_dir)
+
+ os.makedirs(save_dir, exist_ok=True)
+
+ return save_dir
+
+
+def clear_save_dir():
+ root_dir = os.path.dirname(os.path.realpath(__file__))
+ save_dir = os.path.join(root_dir, 'save_dir')
+ if os.path.exists(save_dir):
+ shutil.rmtree(save_dir)
+
+
+def load_model(exp, save_dir):
+
+ # load trained model
+ tags_path = exp.get_data_path(exp.name, exp.version)
+ tags_path = os.path.join(tags_path, 'meta_tags.csv')
+
+ checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
+ weights_dir = os.path.join(save_dir, checkpoints[0])
+
+ trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir, tags_csv=tags_path, on_gpu=True)
+
+ assert trained_model is not None, 'loading model failed'
+
+ return trained_model
+
+
+def run_prediction(dataloader, trained_model):
+ # run prediction on 1 batch
+ for batch in dataloader:
+ break
+
+ x, y = batch
+ x = x.view(x.size(0), -1)
+
+ y_hat = trained_model(x)
+
+ # acc
+ labels_hat = torch.argmax(y_hat, dim=1)
+ val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+ val_acc = torch.tensor(val_acc)
+ val_acc = val_acc.item()
+
+ print(val_acc)
+
+ assert val_acc > 0.70, f'this model is expected to get > 0.7 in test set (it got {val_acc})'
+
+
+def main():
+
+ save_dir = init_save_dir()
+
+ # exp file to get meta
+ exp = get_exp(False)
+ exp.save()
+
+ # exp file to get weights
+ checkpoint = ModelCheckpoint(save_dir)
+
+ trainer = Trainer(
+ experiment=exp,
+ checkpoint_callback=checkpoint,
+ progress_bar=True,
+ max_nb_epochs=1,
+ gpus=[0, 1],
+ distributed_backend='dp',
+ )
+
+ model = CoolModel()
+
+ result = trainer.fit(model)
+
+ # correct result and ok accuracy
+ assert result == 1, 'amp + ddp model failed to complete'
+
+ # test model loading
+ pretrained_model = load_model(exp, save_dir)
+
+ # test model preds
+ run_prediction(model.test_dataloader, pretrained_model)
+
+ clear_save_dir()
+
+
+if __name__ == '__main__':
+ main()
diff --git a/tests/test_models.py b/tests/test_models.py
new file mode 100644
index 00000000..e8fa339b
--- /dev/null
+++ b/tests/test_models.py
@@ -0,0 +1,688 @@
+import pytest
+from pytorch_lightning import Trainer
+from pytorch_lightning.examples.new_project_templates.lightning_module_template import LightningTemplateModel
+from pytorch_lightning.testing_models.lm_test_module import LightningTestModel
+from argparse import Namespace
+from test_tube import Experiment, 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)
+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
+ out = torch.rand(3, 1)
+ assert reduce_distributed_output(out, nb_gpus=1) is out
+
+ # average when we have multiples
+ assert reduce_distributed_output(out, nb_gpus=2) == out.mean()
+
+ # when we have a dict of vals
+ out = {
+ 'a': out,
+ 'b': {
+ 'c': out
+ }
+ }
+ reduced = reduce_distributed_output(out, nb_gpus=3)
+ assert reduced['a'] == out['a']
+ assert reduced['b']['c'] == out['b']['c']
+
+
+def test_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
+ :return:
+ """
+ if not torch.cuda.is_available():
+ warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
+ return
+ if not torch.cuda.device_count() > 1:
+ warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+ return
+
+ # simulate setting slurm flags
+ os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+ os.environ['SLURM_LOCALID'] = str(0)
+
+ hparams = get_hparams()
+ model = LightningTestModel(hparams)
+
+ trainer_options = dict(
+ progress_bar=True,
+ max_nb_epochs=1,
+ gpus=[0],
+ distributed_backend='ddp',
+ use_amp=True
+ )
+
+ save_dir = init_save_dir()
+
+ # exp file to get meta
+ exp = get_exp(False)
+ exp.argparse(hparams)
+ exp.save()
+
+ # exp file to get weights
+ checkpoint = ModelCheckpoint(save_dir)
+
+ # add these to the trainer options
+ trainer_options['checkpoint_callback'] = checkpoint
+ trainer_options['experiment'] = exp
+
+ # fit model
+ trainer = Trainer(**trainer_options)
+ trainer.is_slurm_managing_tasks = True
+ result = trainer.fit(model)
+
+ # correct result and ok accuracy
+ assert result == 1, 'amp + ddp model failed to complete'
+
+ # test root model address
+ assert trainer.resolve_root_node_address('abc') == 'abc'
+ assert trainer.resolve_root_node_address('abc[23]') == 'abc23'
+ assert trainer.resolve_root_node_address('abc[23-24]') == 'abc23'
+ assert trainer.resolve_root_node_address('abc[23-24, 45-40, 40]') == 'abc23'
+
+ # test model loading with a map_location
+ map_location = 'cuda:1'
+ pretrained_model = load_model(exp, save_dir, True, map_location)
+
+ # test model preds
+ run_prediction(model.test_dataloader, pretrained_model)
+
+ if trainer.use_ddp:
+ # on hpc this would work fine... but need to hack it for the purpose of the test
+ trainer.model = pretrained_model
+ trainer.optimizers = pretrained_model.configure_optimizers()
+
+ # test HPC loading / saving
+ trainer.hpc_save(save_dir, exp)
+ trainer.hpc_load(save_dir, on_gpu=True)
+
+ # test freeze on gpu
+ model.freeze()
+ model.unfreeze()
+
+ clear_save_dir()
+
+
+def test_early_stopping_cpu_model():
+ """
+ Test each of the trainer options
+ :return:
+ """
+
+ stopping = EarlyStopping()
+ trainer_options = dict(
+ early_stop_callback=stopping,
+ gradient_clip=1.0,
+ overfit_pct=0.20,
+ track_grad_norm=2,
+ print_nan_grads=True,
+ progress_bar=False,
+ experiment=get_exp(),
+ train_percent_check=0.1,
+ val_percent_check=0.1
+ )
+
+ model, hparams = get_model()
+ run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+ # test freeze on cpu
+ model.freeze()
+ model.unfreeze()
+
+
+def test_cpu_model_with_amp():
+ """
+ Make sure model trains on CPU
+ :return:
+ """
+
+ trainer_options = dict(
+ progress_bar=False,
+ experiment=get_exp(),
+ max_nb_epochs=1,
+ train_percent_check=0.4,
+ val_percent_check=0.4,
+ use_amp=True
+ )
+
+ model, hparams = get_model()
+
+ with pytest.raises((MisconfigurationException, ModuleNotFoundError)):
+ run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_cpu_model():
+ """
+ Make sure model trains on CPU
+ :return:
+ """
+
+ trainer_options = dict(
+ progress_bar=False,
+ experiment=get_exp(),
+ max_nb_epochs=1,
+ train_percent_check=0.4,
+ val_percent_check=0.4
+ )
+
+ model, hparams = get_model()
+
+ run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_all_features_cpu_model():
+ """
+ Test each of the trainer options
+ :return:
+ """
+
+ trainer_options = dict(
+ gradient_clip=1.0,
+ overfit_pct=0.20,
+ track_grad_norm=2,
+ print_nan_grads=True,
+ progress_bar=False,
+ experiment=get_exp(),
+ max_nb_epochs=1,
+ train_percent_check=0.4,
+ val_percent_check=0.4
+ )
+
+ model, hparams = get_model()
+ run_gpu_model_test(trainer_options, model, hparams, on_gpu=False)
+
+
+def test_single_gpu_model():
+ """
+ Make sure single GPU works (DP mode)
+ :return:
+ """
+ if not torch.cuda.is_available():
+ warnings.warn('test_single_gpu_model cannot run. Rerun on a GPU node to run this test')
+ return
+ model, hparams = get_model()
+
+ trainer_options = dict(
+ progress_bar=False,
+ max_nb_epochs=1,
+ train_percent_check=0.1,
+ val_percent_check=0.1,
+ gpus=[0]
+ )
+
+ run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_multi_gpu_model_dp():
+ """
+ Make sure DP works
+ :return:
+ """
+ if not torch.cuda.is_available():
+ warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a GPU node to run this test')
+ return
+ if not torch.cuda.device_count() > 1:
+ warnings.warn('test_multi_gpu_model_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
+ return
+ model, hparams = get_model()
+ trainer_options = dict(
+ progress_bar=False,
+ max_nb_epochs=1,
+ train_percent_check=0.1,
+ val_percent_check=0.1,
+ gpus='-1'
+ )
+
+ run_gpu_model_test(trainer_options, model, hparams)
+
+ # test memory helper functions
+ memory.get_gpu_memory_map()
+
+
+def test_amp_gpu_dp():
+ """
+ Make sure DP + AMP work
+ :return:
+ """
+ if not torch.cuda.is_available():
+ warnings.warn('test_amp_gpu_dp cannot run. Rerun on a GPU node to run this test')
+ return
+ if not torch.cuda.device_count() > 1:
+ warnings.warn('test_amp_gpu_dp cannot run. Rerun on a node with 2+ GPUs to run this test')
+ return
+ model, hparams = get_model()
+ trainer_options = dict(
+ max_nb_epochs=1,
+ gpus='0, 1', # test init with gpu string
+ distributed_backend='dp',
+ use_amp=True
+ )
+ with pytest.raises(MisconfigurationException):
+ run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_multi_gpu_model_ddp():
+ """
+ Make sure DDP works
+ :return:
+ """
+ if not torch.cuda.is_available():
+ warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a GPU node to run this test')
+ return
+ if not torch.cuda.device_count() > 1:
+ warnings.warn('test_multi_gpu_model_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+ return
+
+ os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+ model, hparams = get_model()
+ trainer_options = dict(
+ progress_bar=False,
+ max_nb_epochs=1,
+ train_percent_check=0.4,
+ val_percent_check=0.2,
+ gpus=[0, 1],
+ distributed_backend='ddp'
+ )
+
+ run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_amp_gpu_ddp():
+ """
+ Make sure DDP + AMP work
+ :return:
+ """
+ if not torch.cuda.is_available():
+ warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
+ return
+ if not torch.cuda.device_count() > 1:
+ warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+ return
+
+ os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+
+ hparams = get_hparams()
+ model = LightningTestModel(hparams)
+
+ trainer_options = dict(
+ progress_bar=True,
+ max_nb_epochs=1,
+ gpus=[0, 1],
+ distributed_backend='ddp',
+ use_amp=True
+ )
+
+ run_gpu_model_test(trainer_options, model, hparams)
+
+
+def test_ddp_sampler_error():
+ """
+ Make sure DDP + AMP work
+ :return:
+ """
+ if not torch.cuda.is_available():
+ warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a GPU node to run this test')
+ return
+ if not torch.cuda.device_count() > 1:
+ warnings.warn('test_amp_gpu_ddp cannot run. Rerun on a node with 2+ GPUs to run this test')
+ return
+
+ os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0])
+
+ hparams = get_hparams()
+ model = LightningTestModel(hparams, force_remove_distributed_sampler=True)
+
+ exp = get_exp(True)
+ exp.save()
+
+ trainer = Trainer(
+ experiment=exp,
+ progress_bar=False,
+ max_nb_epochs=1,
+ gpus=[0, 1],
+ distributed_backend='ddp',
+ use_amp=True
+ )
+
+ with pytest.raises(MisconfigurationException):
+ trainer.get_dataloaders(model)
+
+ clear_save_dir()
+
+
+# ------------------------------------------------------------------------
+# UTILS
+# ------------------------------------------------------------------------
+def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
+ save_dir = init_save_dir()
+
+ # exp file to get meta
+ exp = get_exp(False)
+ exp.argparse(hparams)
+ exp.save()
+
+ # exp file to get weights
+ checkpoint = ModelCheckpoint(save_dir)
+
+ # add these to the trainer options
+ trainer_options['checkpoint_callback'] = checkpoint
+ trainer_options['experiment'] = exp
+
+ # fit model
+ trainer = Trainer(**trainer_options)
+ result = trainer.fit(model)
+
+ # correct result and ok accuracy
+ assert result == 1, 'amp + ddp model failed to complete'
+
+ # test model loading
+ pretrained_model = load_model(exp, save_dir, on_gpu)
+
+ # test model preds
+ run_prediction(model.test_dataloader, pretrained_model)
+
+ if trainer.use_ddp:
+ # on hpc this would work fine... but need to hack it for the purpose of the test
+ trainer.model = pretrained_model
+ trainer.optimizers = pretrained_model.configure_optimizers()
+
+ # test HPC loading / saving
+ trainer.hpc_save(save_dir, exp)
+ trainer.hpc_load(save_dir, on_gpu=on_gpu)
+
+ clear_save_dir()
+
+
+def get_hparams(continue_training=False, hpc_exp_number=0):
+ root_dir = os.path.dirname(os.path.realpath(__file__))
+
+ 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
+
+
+def get_model():
+ # set up model with these hyperparams
+ hparams = get_hparams()
+ model = LightningTemplateModel(hparams)
+
+ return model, hparams
+
+
+def get_exp(debug=True):
+ # set up exp object without actually saving logs
+ root_dir = os.path.dirname(os.path.realpath(__file__))
+ exp = Experiment(debug=debug, save_dir=root_dir, name='tests_tt_dir')
+ return exp
+
+
+def init_save_dir():
+ root_dir = os.path.dirname(os.path.realpath(__file__))
+ save_dir = os.path.join(root_dir, 'save_dir')
+
+ if os.path.exists(save_dir):
+ shutil.rmtree(save_dir)
+
+ os.makedirs(save_dir, exist_ok=True)
+
+ return save_dir
+
+
+def clear_save_dir():
+ root_dir = os.path.dirname(os.path.realpath(__file__))
+ save_dir = os.path.join(root_dir, 'save_dir')
+ if os.path.exists(save_dir):
+ shutil.rmtree(save_dir)
+
+
+def load_model(exp, save_dir, on_gpu, map_location=None):
+
+ # load trained model
+ tags_path = exp.get_data_path(exp.name, exp.version)
+ tags_path = os.path.join(tags_path, 'meta_tags.csv')
+
+ checkpoints = [x for x in os.listdir(save_dir) if '.ckpt' in x]
+ weights_dir = os.path.join(save_dir, checkpoints[0])
+
+ trained_model = LightningTemplateModel.load_from_metrics(weights_path=weights_dir,
+ tags_csv=tags_path,
+ on_gpu=on_gpu,
+ map_location=map_location)
+
+ assert trained_model is not None, 'loading model failed'
+
+ return trained_model
+
+
+def run_prediction(dataloader, trained_model):
+ # run prediction on 1 batch
+ for batch in dataloader:
+ break
+
+ x, y = batch
+ x = x.view(x.size(0), -1)
+
+ y_hat = trained_model(x)
+
+ # acc
+ labels_hat = torch.argmax(y_hat, dim=1)
+ val_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
+ val_acc = torch.tensor(val_acc)
+ val_acc = val_acc.item()
+
+ print(val_acc)
+
+ assert val_acc > 0.50, f'this model is expected to get > 0.50 in test set (it got {val_acc})'
+
+
+def assert_ok_acc(trainer):
+ # this model should get 0.80+ acc
+ acc = trainer.tng_tqdm_dic['val_acc']
+ assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}'
+
+
+if __name__ == '__main__':
+ pytest.main([__file__])
-
+ 
+ 
+ 