wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-09 11:32:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

added gradient clipping

Browse Source
This commit is contained in:
William Falcon
2019-06-28 18:00:57 -04:00
parent d9e7174a7b
commit a86ce398a9
7 changed files with 51 additions and 298 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/examples/example_model.py → docs/Examples.md
+33 -13
View File
@@ -1,23 +1,32 @@
#### Template model definition
In 99% of cases you want to just copy this template to start a new lightningModule and change the core of what your model is actually trying to do.
``` {.python}
import os
from collections import OrderedDict
import torch.nn as nn
import numpy as np
from pytorch_lightning.root_module.root_module import LightningModule
from test_tube import HyperOptArgumentParser
from torchvision.datasets import MNIST
import torchvision.transforms as transforms
import torch
import torch.nn.functional as F
import os, pdb
from collections import OrderedDict
from test_tube import HyperOptArgumentParser
from torch import optim
from pytorch_lightning.root_module.root_module import LightningModule
class ExampleModel(LightningModule):
class LightningTemplateModel(LightningModule):
"""
Sample model to show how to define a template
"""
def __init__(self, hparams):
"""
Pass in parsed HyperOptArgumentParser to the model
:param hparams:
"""
# init superclass
super(ExampleModel, self).__init__(hparams)
super(LightningTemplateModel, self).__init__(hparams)
self.batch_size = hparams.batch_size
@@ -42,6 +51,11 @@ class ExampleModel(LightningModule):
# 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)
@@ -59,7 +73,7 @@ class ExampleModel(LightningModule):
def training_step(self, data_batch, batch_i):
"""
Called inside the training loop
Lightning calls this inside the training loop
:param data_batch:
:return:
"""
@@ -79,7 +93,7 @@ class ExampleModel(LightningModule):
def validation_step(self, data_batch, batch_i):
"""
Called inside the validation loop
Lightning calls this inside the validation loop
:param data_batch:
:return:
"""
@@ -99,7 +113,6 @@ class ExampleModel(LightningModule):
})
return output
def validation_end(self, outputs):
"""
Called at the end of validation to aggregate outputs
@@ -139,9 +152,8 @@ class ExampleModel(LightningModule):
return whatever optimizers we want here
:return: list of optimizers
"""
optimizer = self.choose_optimizer(self.hparams.optimizer_name, self.parameters(), {'lr': self.hparams.learning_rate}, 'optimizer')
self.optimizers = [optimizer]
return self.optimizers
optimizer = optim.Adam(self.parameters(), lr=self.hparams.learning_rate)
return [optimizer]
def __dataloader(self, train):
# init data generators
@@ -189,6 +201,12 @@ class ExampleModel(LightningModule):
@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
@@ -209,3 +227,5 @@ class ExampleModel(LightningModule):
parser.opt_list('--batch_size', default=256, type=int, options=[32, 64, 128, 256], tunable=False)
parser.opt_list('--optimizer_name', default='adam', type=str, options=['adam'], tunable=False)
return parser
```
docs/Trainer/Training Loop.md
+10
View File
@@ -38,6 +38,16 @@ Use this to turn off early stopping and run training to the [max_epoch](#force-t
trainer = Trainer(enable_early_stop=True)
```
---
#### Gradient Clipping
Use this to turn off early stopping and run training to the [max_epoch](#force-training-for-min-or-max-epochs)
``` {.python}
# DEFAULT (ie: don't clip)
trainer = Trainer(gradient_clip=0)
```
---
#### Inspect gradient norms
Looking at grad norms can help you figure out where training might be going wrong.
docs/index.md
+1
View File
@@ -62,6 +62,7 @@
- [Anneal Learning rate](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#anneal-learning-rate)
- [Force training for min or max epochs](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-training-for-min-or-max-epochs)
- [Force disable early stop](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#force-disable-early-stop)
- [Gradient Clipping: DOC TODO](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#gradient-clipping)
- [Use multiple optimizers (like GANs)](https://williamfalcon.github.io/pytorch-lightning/Pytorch-Lightning/LightningModule/#configure_optimizers)
- [Set how much of the training set to check (1-100%)](https://williamfalcon.github.io/pytorch-lightning/Trainer/Training%20Loop/#set-how-much-of-the-training-set-to-check)
docs/source/examples/__init__.py
-1
View File
@@ -1 +0,0 @@
from .example_model import ExampleModel
docs/source/examples/basic_trainer.py
-74
View File
@@ -1,74 +0,0 @@
import os
import sys
from test_tube import HyperOptArgumentParser, Experiment
from pytorch_lightning.models.trainer import Trainer
from pytorch_lightning.utils.arg_parse import add_default_args
from pytorch_lightning.callbacks.pt_callbacks import EarlyStopping, ModelCheckpoint
from docs.source.examples.example_model import ExampleModel
def main(hparams):
"""
Main training routine specific for this project
:param hparams:
:return:
"""
# init experiment
exp = Experiment(
name=hparams.tt_name,
debug=hparams.debug,
save_dir=hparams.tt_save_path,
version=hparams.hpc_exp_number,
autosave=False,
description=hparams.tt_description
)
exp.argparse(hparams)
exp.save()
# build model
model = ExampleModel(hparams)
# callbacks
early_stop = EarlyStopping(
monitor='val_acc',
patience=3,
mode='min',
verbose=True,
)
model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
checkpoint = ModelCheckpoint(
filepath=model_save_path,
save_function=None,
save_best_only=True,
verbose=True,
monitor='val_acc',
mode='min'
)
# configure trainer
trainer = Trainer(
experiment=exp,
checkpoint_callback=checkpoint,
early_stop_callback=early_stop,
)
# train model
trainer.fit(model)
if __name__ == '__main__':
# use default args given by lightning
root_dir = os.path.split(os.path.dirname(sys.modules['__main__'].__file__))[0]
parent_parser = HyperOptArgumentParser(strategy='random_search', add_help=False)
add_default_args(parent_parser, root_dir)
# allow model to overwrite or extend args
parser = ExampleModel.add_model_specific_args(parent_parser)
hyperparams = parser.parse_args()
# train model
main(hyperparams)
docs/source/examples/fully_featured_trainer.py
-210
View File
@@ -1,210 +0,0 @@
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)
# ---------------------
# DEFINE MODEL HERE
# ---------------------
from docs.source.examples.example_model import ExampleModel
# ---------------------
AVAILABLE_MODELS = {
'model_template': ExampleModel
}
"""
Allows training by using command line arguments
Run by:
# TYPE YOUR RUN COMMAND HERE
"""
def main_local(hparams):
main(hparams, None, None)
def main(hparams, cluster, results_dict):
"""
Main training routine specific for this project
:param hparams:
:return:
"""
on_gpu = hparams.gpus is not None and torch.cuda.is_available()
device = 'cuda' if on_gpu else 'cpu'
hparams.__setattr__('device', device)
hparams.__setattr__('on_gpu', on_gpu)
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
log_dir = os.path.dirname(os.path.realpath(__file__))
exp = Experiment(
name='test_tube_exp',
debug=True,
save_dir=log_dir,
version=0,
autosave=False,
description='test demo'
)
exp.argparse(hparams)
exp.save()
# build model
print('loading model...')
model = TRAINING_MODEL(hparams)
print('model built')
# callbacks
early_stop = EarlyStopping(
monitor=hparams.early_stop_metric,
patience=hparams.early_stop_patience,
verbose=True,
mode=hparams.early_stop_mode
)
model_save_path = '{}/{}/{}'.format(hparams.model_save_path, exp.name, exp.version)
checkpoint = ModelCheckpoint(
filepath=model_save_path,
save_function=None,
save_best_only=True,
verbose=True,
monitor=hparams.model_save_monitor_value,
mode=hparams.model_save_monitor_mode
)
# gpus are ; separated for inside a node and , within nodes
gpu_list = None
if hparams.gpus is not None:
gpu_list = [int(x) for x in hparams.gpus.split(';')]
# configure trainer
trainer = Trainer(
experiment=exp,
cluster=cluster,
checkpoint_callback=checkpoint,
early_stop_callback=early_stop,
gpus=gpu_list
)
# train model
trainer.fit(model)
def get_default_parser(strategy, root_dir):
possible_model_names = list(AVAILABLE_MODELS.keys())
parser = HyperOptArgumentParser(strategy=strategy, add_help=False)
add_default_args(parser, root_dir, possible_model_names=possible_model_names, rand_seed=SEED)
return parser
def get_model_name(args):
for i, arg in enumerate(args):
if 'model_name' in arg:
return args[i+1]
def optimize_on_cluster(hyperparams):
# enable cluster training
cluster = SlurmCluster(
hyperparam_optimizer=hyperparams,
log_path=hyperparams.tt_save_path,
test_tube_exp_name=hyperparams.tt_name
)
# email for cluster coms
cluster.notify_job_status(email='add_email_here', on_done=True, on_fail=True)
# configure cluster
cluster.per_experiment_nb_gpus = hyperparams.per_experiment_nb_gpus
cluster.job_time = '48:00:00'
cluster.gpu_type = '1080ti'
cluster.memory_mb_per_node = 48000
# any modules for code to run in env
cluster.add_command('source activate pytorch_lightning')
# name of exp
job_display_name = hyperparams.tt_name.split('_')[0]
job_display_name = job_display_name[0:3]
# run hopt
print('submitting jobs...')
cluster.optimize_parallel_cluster_gpu(
main,
nb_trials=hyperparams.nb_hopt_trials,
job_name=job_display_name
)
if __name__ == '__main__':
model_name = get_model_name(sys.argv)
if model_name is None:
model_name = 'model_template'
# use default args
root_dir = os.path.dirname(os.path.realpath(__file__))
parent_parser = get_default_parser(strategy='random_search', root_dir=root_dir)
# allow model to overwrite or extend args
TRAINING_MODEL = AVAILABLE_MODELS[model_name]
parser = TRAINING_MODEL.add_model_specific_args(parent_parser, root_dir)
hyperparams = parser.parse_args()
# format GPU layout
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# ---------------------
# RUN TRAINING
# ---------------------
# cluster and CPU
if hyperparams.on_cluster:
# run on HPC cluster
print('RUNNING ON SLURM CLUSTER')
gpu_ids = hyperparams.gpus.split(';')
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids)
optimize_on_cluster(hyperparams)
elif hyperparams.gpus is None:
# run on cpu
print('RUNNING ON CPU')
main(hyperparams, None, None)
# single or multiple GPUs on same machine
gpu_ids = hyperparams.gpus.split(';')
if hyperparams.interactive:
# run on 1 gpu
print(f'RUNNING INTERACTIVE MODE ON GPUS. gpu ids: {gpu_ids}')
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids)
main(hyperparams, None, None)
else:
# multiple GPUs on same machine
print(f'RUNNING MULTI GPU. GPU ids: {gpu_ids}')
hyperparams.optimize_parallel_gpu(
main_local,
gpu_ids=gpu_ids,
nb_trials=hyperparams.nb_hopt_trials,
nb_workers=len(gpu_ids)
)
pytorch_lightning/models/trainer.py
+7
View File
@@ -33,6 +33,7 @@ class Trainer(TrainerIO):
def __init__(self,
experiment,
checkpoint_callback, early_stop_callback,
gradient_clip=0,
cluster=None,
process_position=0,
current_gpu_name=0,
@@ -53,6 +54,7 @@ class Trainer(TrainerIO):
nb_sanity_val_steps=5):
# Transfer params
self.gradient_clip = gradient_clip
self.check_val_every_n_epoch = check_val_every_n_epoch
self.enable_early_stop = enable_early_stop
self.track_grad_norm = track_grad_norm
@@ -441,6 +443,11 @@ class Trainer(TrainerIO):
# gradient update with accumulated gradients
if (self.batch_nb + 1) % self.accumulate_grad_batches == 0:
# clip gradients
if self.gradient_clip > 0:
model = self.model.module if self.data_parallel else self.model
torch.nn.utils.clip_grad_norm(model.parameters(), self.gradient_clip)
# update gradients across all optimizers
for optimizer in self.optimizers:
optimizer.step()