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

Merge branch 'master' of https://github.com/williamFalcon/pytorch-lightning

Browse Source
This commit is contained in:
William Falcon
2019-10-17 23:32:30 +02:00
parent 93222cf86d 188b598f29
commit ec87c0efe1
7 changed files with 153 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+2
View File
@@ -12,6 +12,8 @@ test_tube_exp/
tests/tests_tt_dir/
tests/save_dir
default/
lightning_logs/
tests/tests/
# Byte-compiled / optimized / DLL files
__pycache__/
examples/basic_examples/lightning_module_template.py
+1 -1
View File
@@ -171,7 +171,7 @@ class LightningTemplateModel(LightningModule):
val_loss_mean /= len(outputs)
val_acc_mean /= len(outputs)
tqdm_dict = {'val_loss': val_loss_mean, 'val_acc': val_acc_mean}
result = {'progress_bar': tqdm_dict, 'log': tqdm_dict}
result = {'progress_bar': tqdm_dict, 'log': tqdm_dict, 'val_loss': val_loss_mean}
return result
# ---------------------
examples/domain_templates/gan.py
+43 -12
View File
@@ -9,6 +9,7 @@ tensorboard --logdir default
from argparse import ArgumentParser
import os
import numpy as np
from collections import OrderedDict
import torchvision
import torchvision.transforms as transforms
@@ -21,7 +22,6 @@ import torch.nn.functional as F
import torch
import pytorch_lightning as pl
from test_tube import Experiment
class Generator(nn.Module):
@@ -84,6 +84,7 @@ class GAN(pl.LightningModule):
# cache for generated images
self.generated_imgs = None
self.last_imgs = None
def forward(self, z):
return self.generator(z)
@@ -93,6 +94,7 @@ class GAN(pl.LightningModule):
def training_step(self, batch, batch_nb, optimizer_i):
imgs, _ = batch
self.last_imgs = imgs
# train generator
if optimizer_i == 0:
@@ -107,17 +109,22 @@ class GAN(pl.LightningModule):
self.generated_imgs = self.forward(z)
# log sampled images
sample_imgs = self.generated_imgs[:6]
grid = torchvision.utils.make_grid(sample_imgs)
self.logger.experiment.add_image('generated_images', grid, 0)
# sample_imgs = self.generated_imgs[:6]
# grid = torchvision.utils.make_grid(sample_imgs)
# self.logger.experiment.add_image('generated_images', grid, 0)
# ground truth result (ie: all fake)
valid = torch.ones(imgs.size(0), 1)
# adversarial loss is binary cross-entropy
g_loss = self.adversarial_loss(self.discriminator(self.generated_imgs), valid)
return g_loss
tqdm_dict = {'g_loss': g_loss}
output = OrderedDict({
'loss': g_loss,
'progress_bar': tqdm_dict,
'log': tqdm_dict
})
return output
# train discriminator
if optimizer_i == 1:
@@ -133,8 +140,15 @@ class GAN(pl.LightningModule):
# discriminator loss is the average of these
d_loss = (real_loss + fake_loss) / 2
tqdm_dict = {'d_loss': d_loss}
output = OrderedDict({
'loss': d_loss,
'progress_bar': tqdm_dict,
'log': tqdm_dict
})
return d_loss
return output
def configure_optimizers(self):
lr = self.hparams.lr
@@ -152,16 +166,33 @@ class GAN(pl.LightningModule):
dataset = MNIST(os.getcwd(), train=True, download=True, transform=transform)
return DataLoader(dataset, batch_size=self.hparams.batch_size)
def on_epoch_end(self):
z = torch.randn(8, self.hparams.latent_dim)
# match gpu device (or keep as cpu)
if self.on_gpu:
z = z.cuda(self.last_imgs.device.index)
# log sampled images
sample_imgs = self.forward(z)
grid = torchvision.utils.make_grid(sample_imgs)
self.logger.experiment.add_image(f'generated_images', grid, self.current_epoch)
def main(hparams):
# save tensorboard logs
exp = Experiment(save_dir=os.getcwd())
# init model
# ------------------------
# 1 INIT LIGHTNING MODEL
# ------------------------
model = GAN(hparams)
# fit trainer on CPU
trainer = pl.Trainer(experiment=exp, max_nb_epochs=200)
# ------------------------
# 2 INIT TRAINER
# ------------------------
trainer = pl.Trainer()
# ------------------------
# 3 START TRAINING
# ------------------------
trainer.fit(model)
pytorch_lightning/trainer/trainer.py
+32 -18
View File
@@ -17,7 +17,7 @@ from torch.optim.optimizer import Optimizer
from pytorch_lightning.root_module.root_module import LightningModule
from pytorch_lightning.root_module import memory
from pytorch_lightning.logging import TestTubeLogger
from pytorch_lightning.trainer.trainer_io import TrainerIO
from pytorch_lightning.trainer.trainer_io import TrainerIOMixin
from pytorch_lightning.pt_overrides.override_data_parallel import (
LightningDistributedDataParallel, LightningDataParallel)
from pytorch_lightning.callbacks import GradientAccumulationScheduler, \
@@ -55,12 +55,12 @@ def reduce_distributed_output(output, nb_gpus):
return output
class Trainer(TrainerIO):
class Trainer(TrainerIOMixin):
def __init__(self,
logger=None,
checkpoint_callback=None,
early_stop_callback=None,
logger=True,
checkpoint_callback=True,
early_stop_callback=True,
default_save_path=None,
gradient_clip_val=0,
process_position=0,
@@ -126,7 +126,6 @@ class Trainer(TrainerIO):
self.log_gpu_memory = log_gpu_memory
self.gradient_clip_val = gradient_clip_val
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.on_gpu = gpus is not None and torch.cuda.is_available()
self.process_position = process_position
@@ -176,24 +175,36 @@ class Trainer(TrainerIO):
# configure early stop callback
# creates a default one if none passed in
self.early_stop_callback = early_stop_callback
if self.early_stop_callback is None:
self.early_stop = EarlyStopping(
self.early_stop_callback = None
if early_stop_callback is True:
self.early_stop_callback = EarlyStopping(
monitor='val_loss',
patience=3,
verbose=True,
mode='min'
)
self.enable_early_stop = True
elif not early_stop_callback:
self.early_stop_callback = None
self.enable_early_stop = False
else:
self.early_stop_callback = early_stop_callback
self.enable_early_stop = True
# configure logger
self.logger = logger
if self.logger is None:
if logger is True:
# default logger
self.logger = TestTubeLogger(
save_dir=self.default_save_path,
version=self.slurm_job_id,
name='lightning_logs'
)
self.logger.rank = 0
self.logger.rank = 0
elif logger is False:
self.logger = None
else:
self.logger = logger
self.logger.rank = 0
# configure checkpoint callback
self.checkpoint_callback = checkpoint_callback
@@ -257,7 +268,7 @@ class Trainer(TrainerIO):
User provided weights_saved_path
Otherwise use os.getcwd()
"""
if self.checkpoint_callback is None:
if self.checkpoint_callback is True:
# init a default one
if isinstance(self.logger, TestTubeLogger):
ckpt_path = '{}/{}/version_{}/{}'.format(
@@ -271,12 +282,15 @@ class Trainer(TrainerIO):
self.checkpoint_callback = ModelCheckpoint(
filepath=ckpt_path
)
elif self.checkpoint_callback is False:
self.checkpoint_callback = None
# set the path for the callbacks
self.checkpoint_callback.save_function = self.save_checkpoint
if self.checkpoint_callback:
# set the path for the callbacks
self.checkpoint_callback.save_function = self.save_checkpoint
# if checkpoint callback used, then override the weights path
self.weights_save_path = self.checkpoint_callback.filepath
# if checkpoint callback used, then override the weights path
self.weights_save_path = self.checkpoint_callback.filepath
# if weights_save_path is still none here, set to current working dir
if self.weights_save_path is None:
@@ -1321,7 +1335,7 @@ class Trainer(TrainerIO):
log_output = output['log']
# reduce progress metrics for tqdm when using dp
if train and self.use_dp or self.use_ddp2:
if train and(self.use_dp or self.use_ddp2):
nb_gpus = self.num_gpus
log_output = reduce_distributed_output(log_output, nb_gpus)
pytorch_lightning/trainer/trainer_io.py
+6 -6
View File
@@ -10,7 +10,7 @@ from pytorch_lightning.pt_overrides.override_data_parallel import (
LightningDistributedDataParallel, LightningDataParallel)
class TrainerIO(object):
class TrainerIOMixin(object):
def __get_model(self):
is_dp_module = isinstance(self.model, (LightningDistributedDataParallel,
@@ -42,7 +42,7 @@ class TrainerIO(object):
def restore_state_if_checkpoint_exists(self, model):
# do nothing if there's not dir or callback
no_ckpt_callback = self.checkpoint_callback is None
no_ckpt_callback = (self.checkpoint_callback is None) or (not self.checkpoint_callback)
if no_ckpt_callback or not os.path.exists(self.checkpoint_callback.filepath):
return
@@ -151,10 +151,10 @@ class TrainerIO(object):
'global_step': self.global_step
}
if self.checkpoint_callback is not None:
if self.checkpoint_callback is not None or self.checkpoint_callback is not False:
checkpoint['checkpoint_callback_best'] = self.checkpoint_callback.best
if self.early_stop_callback is not None:
if self.early_stop_callback is not None or self.checkpoint_callback is not False:
checkpoint['early_stop_callback_wait'] = self.early_stop_callback.wait
checkpoint['early_stop_callback_patience'] = self.early_stop_callback.patience
@@ -207,10 +207,10 @@ class TrainerIO(object):
:param checkpoint:
:return:
"""
if self.checkpoint_callback is not None:
if self.checkpoint_callback is not None or self.checkpoint_callback is not False:
self.checkpoint_callback.best = checkpoint['checkpoint_callback_best']
if self.early_stop_callback is not None:
if self.early_stop_callback is not None or self.early_stop_callback is not False:
self.early_stop_callback.wait = checkpoint['early_stop_callback_wait']
self.early_stop_callback.patience = checkpoint['early_stop_callback_patience']
tests/test_logging.py
+5
View File
@@ -66,6 +66,8 @@ def test_testtube_pickle():
trainer2 = pickle.loads(pkl_bytes)
trainer2.logger.log_metrics({"acc": 1.0})
clear_save_dir()
def test_mlflow_logger():
"""
@@ -134,6 +136,9 @@ def test_mlflow_pickle():
trainer2 = pickle.loads(pkl_bytes)
trainer2.logger.log_metrics({"acc": 1.0})
n = np.random.randint(0, 10000000, 1)[0]
shutil.move(mlflow_dir, mlflow_dir + f'_{n}')
def test_custom_logger():
tests/test_models.py
+64 -65
View File
@@ -32,6 +32,7 @@ from pytorch_lightning.logging import TestTubeLogger
from examples import LightningTemplateModel
# generate a list of random seeds for each test
RANDOM_FILE_PATHS = list(np.random.randint(12000, 19000, 1000))
RANDOM_PORTS = list(np.random.randint(12000, 19000, 1000))
ROOT_SEED = 1234
torch.manual_seed(ROOT_SEED)
@@ -42,6 +43,34 @@ RANDOM_SEEDS = list(np.random.randint(0, 10000, 1000))
# ------------------------------------------------------------------------
# TESTS
# ------------------------------------------------------------------------
def test_early_stopping_cpu_model():
"""
Test each of the trainer options
:return:
"""
reset_seed()
stopping = EarlyStopping(monitor='val_loss')
trainer_options = dict(
early_stop_callback=stopping,
gradient_clip_val=1.0,
overfit_pct=0.20,
track_grad_norm=2,
print_nan_grads=True,
show_progress_bar=True,
logger=get_test_tube_logger(),
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_running_test_pretrained_model_ddp():
"""Verify test() on pretrained model"""
if not can_run_gpu_test():
@@ -90,7 +119,29 @@ def test_running_test_pretrained_model_ddp():
[run_prediction(dataloader, pretrained_model) for dataloader in model.test_dataloader()]
# test we have good test accuracy
clear_save_dir()
def test_lbfgs_cpu_model():
"""
Test each of the trainer options
:return:
"""
reset_seed()
trainer_options = dict(
max_nb_epochs=1,
gradient_clip_val=1.0,
print_nan_grads=True,
show_progress_bar=False,
weights_summary='top',
train_percent_check=1.0,
val_percent_check=0.2
)
model, hparams = get_model(use_test_model=True, lbfgs=True)
run_model_test_no_loggers(trainer_options, model, hparams, on_gpu=False)
clear_save_dir()
@@ -118,30 +169,7 @@ def test_default_logger_callbacks_cpu_model():
model.freeze()
model.unfreeze()
def test_lbfgs_cpu_model():
"""
Test each of the trainer options
:return:
"""
reset_seed()
trainer_options = dict(
max_nb_epochs=1,
gradient_clip_val=1.0,
print_nan_grads=True,
show_progress_bar=False,
weights_summary='top',
train_percent_check=1.0,
val_percent_check=0.2
)
model, hparams = get_model(use_test_model=True, lbfgs=True)
run_model_test_no_loggers(trainer_options, model, hparams, on_gpu=False)
# test freeze on cpu
model.freeze()
model.unfreeze()
clear_save_dir()
def test_multi_gpu_model_ddp2():
@@ -365,7 +393,7 @@ def test_running_test_pretrained_model():
# correct result and ok accuracy
assert result == 1, 'training failed to complete'
pretrained_model = load_model(
logger.experiment, save_dir, module_class=LightningTestModel
logger.experiment, trainer.checkpoint_callback.filepath, module_class=LightningTestModel
)
new_trainer = Trainer(**trainer_options)
@@ -595,34 +623,6 @@ def test_single_gpu_batch_parse():
assert batch[1][0]['b'].type() == 'torch.cuda.FloatTensor'
def test_early_stopping_cpu_model():
"""
Test each of the trainer options
:return:
"""
reset_seed()
stopping = EarlyStopping(monitor='val_loss')
trainer_options = dict(
early_stop_callback=stopping,
gradient_clip_val=1.0,
overfit_pct=0.20,
track_grad_norm=2,
print_nan_grads=True,
show_progress_bar=True,
logger=get_test_tube_logger(),
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_no_val_module():
"""
Tests use case where trainer saves the model, and user loads it from tags independently
@@ -1431,7 +1431,6 @@ def test_multiple_test_dataloader():
# ------------------------------------------------------------------------
def run_model_test_no_loggers(trainer_options, model, hparams, on_gpu=True):
save_dir = init_save_dir()
trainer_options['default_save_path'] = save_dir
# fit model
@@ -1442,7 +1441,8 @@ def run_model_test_no_loggers(trainer_options, model, hparams, on_gpu=True):
assert result == 1, 'amp + ddp model failed to complete'
# test model loading
pretrained_model = load_model(trainer.logger.experiment, save_dir)
pretrained_model = load_model(trainer.logger.experiment,
trainer.checkpoint_callback.filepath)
# test new model accuracy
[run_prediction(dataloader, pretrained_model) for dataloader in model.test_dataloader()]
@@ -1476,7 +1476,7 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True):
assert result == 1, 'amp + ddp model failed to complete'
# test model loading
pretrained_model = load_model(logger.experiment, save_dir)
pretrained_model = load_model(logger.experiment, trainer.checkpoint_callback.filepath)
# test new model accuracy
[run_prediction(dataloader, pretrained_model) for dataloader in model.test_dataloader()]
@@ -1519,6 +1519,7 @@ def get_model(use_test_model=False, lbfgs=False):
hparams = get_hparams()
if lbfgs:
setattr(hparams, 'optimizer_name', 'lbfgs')
setattr(hparams, 'learning_rate', 0.001)
if use_test_model:
model = LightningTestModel(hparams)
@@ -1538,10 +1539,10 @@ def get_test_tube_logger(debug=True, version=None):
def init_save_dir():
root_dir = os.path.dirname(os.path.realpath(__file__))
save_dir = os.path.join(root_dir, 'save_dir')
save_dir = os.path.join(root_dir, 'tests', 'save_dir')
if os.path.exists(save_dir):
n = np.random.randint(0, 10000000, 1)[0]
n = RANDOM_FILE_PATHS.pop()
shutil.move(save_dir, save_dir + f'_{n}')
os.makedirs(save_dir, exist_ok=True)
@@ -1553,19 +1554,17 @@ 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):
n = np.random.randint(0, 10000000, 1)[0]
n = RANDOM_FILE_PATHS.pop()
shutil.move(save_dir, save_dir + f'_{n}')
def load_model(exp, save_dir, module_class=LightningTemplateModel):
def load_model(exp, root_weights_dir, module_class=LightningTemplateModel):
# load trained model
tags_path = exp.get_data_path(exp.name, exp.version)
checkpoint_folder = os.path.join(tags_path, 'checkpoints')
tags_path = os.path.join(tags_path, 'meta_tags.csv')
checkpoints = [x for x in os.listdir(checkpoint_folder) if '.ckpt' in x]
weights_dir = os.path.join(checkpoint_folder, checkpoints[0])
checkpoints = [x for x in os.listdir(root_weights_dir) if '.ckpt' in x]
weights_dir = os.path.join(root_weights_dir, checkpoints[0])
trained_model = module_class.load_from_metrics(weights_path=weights_dir,
tags_csv=tags_path)