diff --git a/.gitignore b/.gitignore index 54ead9a4..654e70b1 100644 --- a/.gitignore +++ b/.gitignore @@ -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__/ diff --git a/examples/basic_examples/lightning_module_template.py b/examples/basic_examples/lightning_module_template.py index c977fc14..517bf9a0 100644 --- a/examples/basic_examples/lightning_module_template.py +++ b/examples/basic_examples/lightning_module_template.py @@ -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 # --------------------- diff --git a/examples/domain_templates/gan.py b/examples/domain_templates/gan.py index bb150162..5a94fcd4 100644 --- a/examples/domain_templates/gan.py +++ b/examples/domain_templates/gan.py @@ -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) diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index b8b7f4db..1b89ca4b 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -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) diff --git a/pytorch_lightning/trainer/trainer_io.py b/pytorch_lightning/trainer/trainer_io.py index f6305308..8852e7d8 100644 --- a/pytorch_lightning/trainer/trainer_io.py +++ b/pytorch_lightning/trainer/trainer_io.py @@ -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'] diff --git a/tests/test_logging.py b/tests/test_logging.py index b8436a9d..2458639b 100644 --- a/tests/test_logging.py +++ b/tests/test_logging.py @@ -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(): diff --git a/tests/test_models.py b/tests/test_models.py index d95f926e..20b33595 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -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)