From d286206e86fca39a7f44c78691430bd25a203034 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Sun, 31 Mar 2019 16:29:50 -0400 Subject: [PATCH] added example and verified --- demo/basic_trainer.py | 76 +++++++ demo/example_model.py | 200 ++++++++++++++++++ .../{trainer.py => fully_featured_trainer.py} | 39 ++-- pytorch_lightning/models/trainer.py | 13 +- pytorch_lightning/root_module/root_module.py | 9 +- pytorch_lightning/utils/arg_parse.py | 10 +- 6 files changed, 311 insertions(+), 36 deletions(-) create mode 100644 demo/basic_trainer.py rename demo/{trainer.py => fully_featured_trainer.py} (82%) diff --git a/demo/basic_trainer.py b/demo/basic_trainer.py new file mode 100644 index 00000000..0ce8a28b --- /dev/null +++ b/demo/basic_trainer.py @@ -0,0 +1,76 @@ +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.utils.pt_callbacks import EarlyStopping, ModelCheckpoint +from demo.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 + print('loading model...') + model = ExampleModel(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 + ) + + # 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) diff --git a/demo/example_model.py b/demo/example_model.py index e69de29b..f7197a48 100644 --- a/demo/example_model.py +++ b/demo/example_model.py @@ -0,0 +1,200 @@ +import torch.nn as nn +import numpy as np +from pytorch_lightning.root_module.root_module import RootModule +from test_tube import HyperOptArgumentParser +from torchvision.datasets import MNIST +import torchvision.transforms as transforms +import torch +import torch.nn.functional as F + + +class ExampleModel(RootModule): + """ + Sample model to show how to define a template + """ + + def __init__(self, hparams): + # init superclass + super(ExampleModel, self).__init__(hparams) + + self.batch_size = hparams.batch_size + + # 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): + x = self.c_d1(x) + x = F.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): + """ + Called 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) + + tqdm_dic = {'tng_loss': loss_val.item()} + return loss_val, tqdm_dic + + def validation_step(self, data_batch): + """ + Called 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) + + output = {'y_hat': y_hat, 'val_loss': loss_val.item(), 'val_acc': val_acc} + 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: + """ + val_loss_mean = 0 + accs = [] + for output in outputs: + val_loss_mean += output['val_loss'] + accs.append(output['val_acc']) + + val_loss_mean /= len(outputs) + tqdm_dic = {'val_loss': val_loss_mean, 'val_acc': np.mean(accs)} + 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 + # --------------------- + def configure_optimizers(self): + """ + 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 + + 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) + + loader = torch.utils.data.DataLoader( + dataset=dataset, + batch_size=self.hparams.batch_size, + shuffle=True + ) + + return loader + + @property + 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 + + @property + 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 + + @property + 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 + + @staticmethod + def add_model_specific_args(parent_parser): + 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) + parser.add_argument('--hidden_dim', default=500) + parser.add_argument('--out_features', default=10) + + # data + parser.add_argument('--data_root', default='/Users/williamfalcon/Developer/personal/research_lib/research_proj/datasets/mnist', type=str) + + # training params (opt) + parser.opt_list('--learning_rate', default=0.001, type=float, options=[0.0001, 0.0005, 0.001, 0.005], + tunable=False) + 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 diff --git a/demo/trainer.py b/demo/fully_featured_trainer.py similarity index 82% rename from demo/trainer.py rename to demo/fully_featured_trainer.py index b6742352..b5484c5e 100644 --- a/demo/trainer.py +++ b/demo/fully_featured_trainer.py @@ -1,12 +1,12 @@ import os import sys - -import torch 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 time import sleep from pytorch_lightning.utils.pt_callbacks import EarlyStopping, ModelCheckpoint @@ -17,11 +17,11 @@ np.random.seed(SEED) # --------------------- # DEFINE MODEL HERE # --------------------- -from pytorch_lightning.models.sample_model_template.model_template import ExampleModel1 +from demo.example_model import ExampleModel # --------------------- AVAILABLE_MODELS = { - 'model_1': ExampleModel1 + 'model_template': ExampleModel } @@ -95,28 +95,9 @@ def main(hparams, cluster, results_dict): # configure trainer trainer = Trainer( experiment=exp, - on_gpu=on_gpu, cluster=cluster, - enable_tqdm=hparams.enable_tqdm, - overfit_pct=hparams.overfit, - track_grad_norm=hparams.track_grad_norm, - fast_dev_run=hparams.fast_dev_run, - check_val_every_n_epoch=hparams.check_val_every_n_epoch, - accumulate_grad_batches=hparams.accumulate_grad_batches, - process_position=process_position, - current_gpu_name=current_gpu, checkpoint_callback=checkpoint, early_stop_callback=early_stop, - enable_early_stop=hparams.enable_early_stop, - max_nb_epochs=hparams.max_nb_epochs, - min_nb_epochs=hparams.min_nb_epochs, - train_percent_check=hparams.train_percent_check, - val_percent_check=hparams.val_percent_check, - test_percent_check=hparams.test_percent_check, - val_check_interval=hparams.val_check_interval, - log_save_interval=hparams.log_save_interval, - add_log_row_interval=hparams.add_log_row_interval, - lr_scheduler_milestones=hparams.lr_scheduler_milestones ) # train model @@ -173,6 +154,8 @@ def optimize_on_cluster(hyperparams): 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.split(os.path.dirname(sys.modules['__main__'].__file__))[0] @@ -181,7 +164,6 @@ if __name__ == '__main__': # allow model to overwrite or extend args TRAINING_MODEL = AVAILABLE_MODELS[model_name] parser = TRAINING_MODEL.add_model_specific_args(parent_parser) - parser.json_config('-c', '--config', default=root_dir + '/run_configs/local.json') hyperparams = parser.parse_args() # format GPU layout @@ -190,25 +172,30 @@ if __name__ == '__main__': # RUN TRAINING if hyperparams.on_cluster: + # Gets called when running via HPC cluster print('RUNNING ON SLURM CLUSTER') os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(gpu_ids) optimize_on_cluster(hyperparams) elif hyperparams.single_run_gpu: + # run on 1 gpu print(f'RUNNING 1 TRIAL ON GPU. gpu: {gpu_ids[0]}') os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids[0] main(hyperparams, None, None) elif hyperparams.local or hyperparams.single_run: + # run 1 trial but on CPU os.environ["CUDA_VISIBLE_DEVICES"] = '0' print('RUNNING LOCALLY') 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) - ) \ No newline at end of file + ) + diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index fb7a1dee..34c9e27b 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -1,9 +1,9 @@ import torch import tqdm import numpy as np -from research_lib.root_module.memory import get_gpu_memory_map +from pytorch_lightning.root_module.memory import get_gpu_memory_map import traceback -from research_lib.root_module.model_saving import TrainerIO +from pytorch_lightning.root_module.model_saving import TrainerIO from torch.optim.lr_scheduler import MultiStepLR @@ -11,17 +11,17 @@ class Trainer(TrainerIO): def __init__(self, experiment, - cluster, checkpoint_callback, early_stop_callback, + cluster=None, process_position=0, current_gpu_name=0, on_gpu=False, enable_tqdm=True, - overfit_pct=None, + overfit_pct=0.0, track_grad_norm=-1, check_val_every_n_epoch=1, fast_dev_run=False, - accumulate_grad_batches=False, + accumulate_grad_batches=1, enable_early_stop=True, max_nb_epochs=5, 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=1, add_log_row_interval=1, @@ -226,7 +226,8 @@ class Trainer(TrainerIO): self.experiment.save() # enable cluster checkpointing - self.enable_auto_hpc_walltime_manager() + if self.cluster is not None: + self.enable_auto_hpc_walltime_manager() # --------------------------- # CORE TRAINING LOOP diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py index 3833f1f4..ab49a9a1 100644 --- a/pytorch_lightning/root_module/root_module.py +++ b/pytorch_lightning/root_module/root_module.py @@ -14,7 +14,7 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks): def __init__(self, hparams): super(RootModule, self).__init__() self.hparams = hparams - self.on_gpu = hparams.on_gpu + self.dtype = torch.FloatTensor self.exp_save_path = None self.current_epoch = 0 @@ -25,6 +25,13 @@ class RootModule(GradInformation, ModelIO, OptimizerConfig, ModelHooks): self.gradient_clip = hparams.gradient_clip self.num = 2 + # track if gpu was requested for checkpointing + self.on_gpu = False + try: + self.on_gpu = hparams.on_gpu + except Exception as e: + pass + # computed vars for the dataloaders self._tng_dataloader = None self._val_dataloader = None diff --git a/pytorch_lightning/utils/arg_parse.py b/pytorch_lightning/utils/arg_parse.py index a3302b00..5e416565 100644 --- a/pytorch_lightning/utils/arg_parse.py +++ b/pytorch_lightning/utils/arg_parse.py @@ -1,4 +1,4 @@ -def add_default_args(parser, root_dir, possible_model_names, rand_seed): +def add_default_args(parser, root_dir, rand_seed=None, possible_model_names=None): # tng, test, val check intervals parser.add_argument('--eval_test_set', dest='eval_test_set', action='store_true', help='true = run test set also') @@ -32,7 +32,9 @@ def add_default_args(parser, root_dir, possible_model_names, rand_seed): # model paths parser.add_argument('--model_load_weights_path', default=None, type=str) - parser.add_argument('--model_name', default='', help=','.join(possible_model_names)) + + if possible_model_names is not None: + parser.add_argument('--model_name', default='', help=','.join(possible_model_names)) # test_tube settings parser.add_argument('-en', '--tt_name', default='r_lib_') @@ -58,7 +60,9 @@ def add_default_args(parser, root_dir, possible_model_names, rand_seed): parser.add_argument('--overfit', default=-1, type=float, help='% of dataset to use with this option. float, or -1 for none') # debug args - parser.add_argument('--random_seed', default=rand_seed, type=int) + if rand_seed is not None: + parser.add_argument('--random_seed', default=rand_seed, type=int) + parser.add_argument('--live', dest='live', action='store_true', help='runs on gpu without cluster') parser.add_argument('--enable_debug', dest='debug', action='store_true', help='enables/disables test tube') parser.add_argument('--enable_local', dest='local', action='store_true', help='enables local tng')