From f51b45933b8fd90302f5e5e7469d25902d2cd5df Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 30 Aug 2019 18:56:09 -0400 Subject: [PATCH] Expectopatronum implement #89 (#182) * rename validate -> evaluate; implement test logic; allow multiple test_loaders * add test_step and test_end to LightningModule * add in_test_mode to pretraining to implement case 2 (test pretrained model) * fix code style issues * LightningTestModel: add optional second test set, implement test_step and test_end * implemented test for multiple test_dataloaders; fixed typo * add two test cases for #89 * add documentation for test_step, test_end; fix computation of loss in validation_step example * Update trainer.py * Update trainer.py * Update trainer.py * Update trainer.py * Update trainer.py * Update trainer.py * Added proper dp ddp routing calls for test mode * Update trainer.py * Update test_models.py * Update trainer.py * Update trainer.py * Update override_data_parallel.py * Update test_models.py * Update test_models.py * Update trainer.py * Update trainer.py * Update trainer.py * Update test_models.py * Update test_models.py * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * debug * Update trainer.py * Update override_data_parallel.py * Update debug.py * Update lm_test_module.py * Update test_models.py --- .../RequiredTrainerInterface.md | 138 +++++++- pytorch_lightning/models/trainer.py | 181 ++++++---- .../pt_overrides/override_data_parallel.py | 8 + pytorch_lightning/root_module/hooks.py | 2 +- pytorch_lightning/root_module/root_module.py | 19 ++ pytorch_lightning/testing/lm_test_module.py | 102 +++++- tests/debug.py | 87 ++++- tests/test_models.py | 317 ++++++++++++++---- 8 files changed, 707 insertions(+), 147 deletions(-) diff --git a/docs/LightningModule/RequiredTrainerInterface.md b/docs/LightningModule/RequiredTrainerInterface.md index a3698410..e53e4d05 100644 --- a/docs/LightningModule/RequiredTrainerInterface.md +++ b/docs/LightningModule/RequiredTrainerInterface.md @@ -16,7 +16,9 @@ Otherwise, to Define a Lightning Module, implement the following methods: **Optional**: - [validation_step](RequiredTrainerInterface.md#validation_step) -- [validation_end](RequiredTrainerInterface.md#validation_end) +- [validation_end](RequiredTrainerInterface.md#validation_end) +- [test_step](RequiredTrainerInterface.md#test_step) +- [test_end](RequiredTrainerInterface.md#test_end) - [val_dataloader](RequiredTrainerInterface.md#val_dataloader) - [test_dataloader](RequiredTrainerInterface.md#test_dataloader) - [on_save_checkpoint](RequiredTrainerInterface.md#on_save_checkpoint) @@ -63,6 +65,17 @@ class CoolModel(pl.LightningModule): avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean() return {'avg_val_loss': avg_loss} + def test_step(self, batch, batch_nb): + # OPTIONAL + x, y = batch + y_hat = self.forward(x) + return {'test_loss': F.cross_entropy(y_hat, y)} + + def test_end(self, outputs): + # OPTIONAL + avg_loss = torch.stack([x['test_loss'] for x in outputs]).mean() + return {'avg_test_loss': avg_loss} + def configure_optimizers(self): # REQUIRED return [torch.optim.Adam(self.parameters(), lr=0.02)] @@ -80,6 +93,7 @@ class CoolModel(pl.LightningModule): @pl.data_loader def test_dataloader(self): # OPTIONAL + # can also return a list of test dataloaders return DataLoader(MNIST(os.getcwd(), train=False, download=True, transform=transforms.ToTensor()), batch_size=32) ``` --- @@ -225,9 +239,10 @@ the [optimizer_step](https://williamfalcon.github.io/pytorch-lightning/Trainer/h ### validation_step ``` {.python} +# if you have one val dataloader: def validation_step(self, data_batch, batch_nb) -# if have multiple val dataloaders: +# if you have multiple val dataloaders: def validation_step(self, data_batch, batch_nb, dataloader_idx) ``` **OPTIONAL** @@ -235,7 +250,7 @@ If you don't need to validate you don't need to implement this method. In this step you'd normally generate examples or calculate anything of interest such as accuracy. -The dict you return here will be available in the validation_end method. +The dict you return here will be available in the `validation_end` method. **Params** @@ -256,11 +271,11 @@ The dict you return here will be available in the validation_end method. ``` {.python} # CASE 1: A single validation dataset def validation_step(self, data_batch, batch_nb): - x, y, z = data_batch + x, y = data_batch # implement your own out = self.forward(x) - loss = self.loss(out, x) + loss = self.loss(out, y) # log 6 example images # or generated text... or whatever @@ -338,6 +353,119 @@ def validation_end(self, outputs): return tqdm_dic ``` +### test_step + +``` {.python} +# if you have one test dataloader: +def test_step(self, data_batch, batch_nb) + +# if you have multiple test dataloaders: +def test_step(self, data_batch, batch_nb, dataloader_idx) +``` +**OPTIONAL** +If you don't need to test you don't need to implement this method. + +In this step you'd normally generate examples or calculate anything of interest such as accuracy. + +The dict you return here will be available in the `test_end` method. + +This function is used when you execute `trainer.test()`. + +**Params** + +| Param | description | +|---|---| +| data_batch | The output of your dataloader. A tensor, tuple or list | +| batch_nb | Integer displaying which batch this is | +| dataloader_i | Integer displaying which dataloader this is (only if multiple test datasets used) | + +**Return** + +| Return | description | optional | +|---|---|---| +| dict | Dict or OrderedDict with metrics to display in progress bar. All keys must be tensors. | Y | + +**Example** + +``` {.python} +# CASE 1: A single test dataset +def test_step(self, data_batch, batch_nb): + x, y = data_batch + + # implement your own + out = self.forward(x) + loss = self.loss(out, y) + + # calculate acc + labels_hat = torch.argmax(out, dim=1) + test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0) + + # all optional... + # return whatever you need for the collation function test_end + output = OrderedDict({ + 'test_loss': loss_test, + 'test_acc': torch.tensor(test_acc), # everything must be a tensor + }) + + # return an optional dict + return output +``` + +If you pass in multiple test datasets, test_step will have an additional argument. + +```python +# CASE 2: multiple test datasets +def test_step(self, data_batch, batch_nb, dataset_idx): + # dataset_idx tells you which dataset this is. +``` + +The ```dataset_idx``` corresponds to the order of datasets returned in ```test_dataloader```. + +--- +### test_end + +``` {.python} +def test_end(self, outputs) +``` +If you didn't define a test_step, this won't be called. + +Called at the end of the test step with the output of each test_step. Called once per test dataset. + +The outputs here are strictly for the progress bar. If you don't need to display anything, don't return anything. + +**Params** + +| Param | description | +|---|---| +| outputs | List of outputs you defined test_step | + +**Return** + +| Return | description | optional | +|---|---|---| +| dict | Dict of OrderedDict with metrics to display in progress bar | Y | + +**Example** + +``` {.python} +def test_end(self, outputs): + """ + Called at the end of test to aggregate outputs + :param outputs: list of individual outputs of each test step + :return: + """ + test_loss_mean = 0 + test_acc_mean = 0 + for output in outputs: + test_loss_mean += output['test_loss'] + test_acc_mean += output['test_acc'] + + test_loss_mean /= len(outputs) + test_acc_mean /= len(outputs) + tqdm_dic = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()} + return tqdm_dic +``` + --- ### on_save_checkpoint diff --git a/pytorch_lightning/models/trainer.py b/pytorch_lightning/models/trainer.py index bc075282..ed73f2e2 100644 --- a/pytorch_lightning/models/trainer.py +++ b/pytorch_lightning/models/trainer.py @@ -20,6 +20,7 @@ from pytorch_lightning.pt_overrides.override_data_parallel import ( LightningDistributedDataParallel, LightningDataParallel) from pytorch_lightning.callbacks import GradientAccumulationScheduler from pytorch_lightning.utilities.debugging import MisconfigurationException +import pdb try: from apex import amp @@ -157,6 +158,7 @@ class Trainer(TrainerIO): self.use_ddp = False self.use_dp = False self.single_gpu = False + self.testing = False # training bookeeping self.total_batch_nb = 0 @@ -370,8 +372,10 @@ class Trainer(TrainerIO): self.nb_val_batches = max(1, self.nb_val_batches) # determine number of test batches - self.nb_test_batches = len(self.test_dataloader) if self.test_dataloader is not None else 0 - self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check) + if self.test_dataloader is not None: + self.nb_test_batches = sum(len(dataloader) for dataloader in self.test_dataloader) + self.nb_test_batches = int(self.nb_test_batches * self.test_percent_check) + self.nb_test_batches = max(1, self.nb_test_batches) # determine when to check validation self.val_check_batch = int(self.nb_tng_batches * self.val_check_interval) @@ -384,40 +388,45 @@ class Trainer(TrainerIO): self.tqdm_metrics[k] = v - def __validation_forward(self, model, data_batch, batch_i, dataloader_i): + def __evaluation_forward(self, model, data_batch, batch_i, dataloader_i, test=False): # make dataloader_i arg in validation_step optional args = [data_batch, batch_i] - if len(self.val_dataloader) > 1: + + if test and len(self.test_dataloader) > 1: args.append(dataloader_i) - if self.use_ddp: + elif len(self.val_dataloader) > 1: + args.append(dataloader_i) + + # handle DP, DDP forward + if self.use_ddp or self.use_dp: output = model(*args) - elif self.use_dp: - output = model(*args) - elif self.single_gpu: - # put inputs on gpu manually + return output + + # CPU, single GPU + if self.single_gpu: + # for single GPU put inputs on gpu manually gpu_id = self.data_parallel_device_ids[0] data_batch = self.transfer_batch_to_gpu(data_batch, gpu_id) args[0] = data_batch - # do non dp, ddp step - output = model.validation_step(*args) - + if test: + output = model.test_step(*args) else: - # CPU output = model.validation_step(*args) return output - def validate(self, model, dataloader, max_batches, dataloader_i): + def evaluate(self, model, dataloader, max_batches, dataloader_i, test=False): """ - Run validation code + Run evaluation code :param model: PT model :param dataloader: PT dataloader :param max_batches: Scalar + :param dataloader_i: + :param test: boolean :return: """ - # enable eval mode model.zero_grad() model.eval() @@ -439,9 +448,10 @@ class Trainer(TrainerIO): break # ----------------- - # RUN VALIDATION STEP + # RUN EVALUATION STEP # ----------------- - output = self.__validation_forward(model, data_batch, batch_i, dataloader_i) + output = self.__evaluation_forward(model, data_batch, batch_i, dataloader_i, + test) # track outputs for collation outputs.append(output) @@ -450,13 +460,14 @@ class Trainer(TrainerIO): if self.show_progress_bar: self.progress_bar.update(1) + eval_results = {} + # give model a chance to do something with the outputs (and method defined) - val_results = {} - if self.__is_overriden('validation_end'): - if self.data_parallel: - val_results = model.module.validation_end(outputs) - else: - val_results = model.validation_end(outputs) + model = self.__get_model() + if test and self.__is_overriden('test_end'): + eval_results = model.test_end(outputs) + elif self.__is_overriden('validation_end'): + eval_results = model.validation_end(outputs) # enable train mode again model.train() @@ -464,7 +475,7 @@ class Trainer(TrainerIO): # enable gradients to save memory torch.set_grad_enabled(True) - return val_results + return eval_results def get_dataloaders(self, model): """ @@ -478,6 +489,10 @@ class Trainer(TrainerIO): self.val_dataloader = model.val_dataloader # handle returning an actual dataloader instead of a list of loaders + have_test_loaders = self.test_dataloader is not None + if have_test_loaders and not isinstance(self.test_dataloader, list): + self.test_dataloader = [self.test_dataloader] + have_val_loaders = self.val_dataloader is not None if have_val_loaders and not isinstance(self.val_dataloader, list): self.val_dataloader = [self.val_dataloader] @@ -524,11 +539,33 @@ class Trainer(TrainerIO): warnings.warn(msg) break + if self.use_ddp and self.test_dataloader is not None: + for dataloader in self.test_dataloader: + if not isinstance(dataloader, DistributedSampler): + msg = """ + Your test_dataloader(s) are not all DistributedSamplers. + You're using multiple gpus and multiple nodes without using a DistributedSampler + to assign a subset of your data to each process. To silence this warning, pass a + DistributedSampler to your DataLoader. + + ie: this: + dataset = myDataset() + dataloader = Dataloader(dataset) + + becomes: + dataset = myDataset() + dist_sampler = torch.utils.data.distributed.DistributedSampler(dataset) + dataloader = Dataloader(dataset, sampler=dist_sampler) + + If you want each process to load the full dataset, ignore this warning. + """ + warnings.warn(msg) + break + # ----------------------------- # 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 @@ -740,6 +777,7 @@ class Trainer(TrainerIO): if self.data_parallel: ref_model = model.module + # give model convenience properties ref_model.trainer = self # set local properties on the model @@ -747,6 +785,7 @@ class Trainer(TrainerIO): ref_model.use_dp = self.use_dp ref_model.use_ddp = self.use_ddp ref_model.use_amp = self.use_amp + ref_model.testing = self.testing # transfer data loaders from model self.get_dataloaders(ref_model) @@ -758,15 +797,13 @@ class Trainer(TrainerIO): if self.proc_rank == 0 and self.print_weights_summary: ref_model.summarize() - # give model convenience properties - ref_model.trainer = self - + # link up experiment object if self.experiment is not None: ref_model.experiment = self.experiment - # save exp to get started - if self.proc_rank == 0 and self.experiment is not None: - self.experiment.save() + # 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 @@ -785,7 +822,13 @@ class Trainer(TrainerIO): if self.show_progress_bar: self.progress_bar = tqdm.tqdm(0, position=self.process_position) - # run tiny validation (if validation defined) to make sure program won't crash during val + # when testing requested only run test and return + if self.testing: + self.__run_evaluation(test=True) + return + + # run tiny validation (if validation defined) + # to make sure program won't crash during val ref_model.on_sanity_check_start() if self.val_dataloader is not None and self.nb_sanity_val_steps > 0: for ds_i, dataloader in enumerate(self.val_dataloader): @@ -794,7 +837,7 @@ class Trainer(TrainerIO): if self.show_progress_bar: self.progress_bar.reset(self.nb_sanity_val_steps) - self.validate(model, dataloader, self.nb_sanity_val_steps, ds_i) + self.evaluate(model, dataloader, self.nb_sanity_val_steps, ds_i, self.testing) # --------------------------- # CORE TRAINING LOOP @@ -871,8 +914,10 @@ class Trainer(TrainerIO): # RUN VAL STEP # --------------- is_val_check_batch = (batch_nb + 1) % self.val_check_batch == 0 + can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0 if self.fast_dev_run or is_val_check_batch or early_stop_epoch: - self.__run_validation() + if can_check_epoch: + self.__run_evaluation(test=self.testing) # when batch should be saved if (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch: @@ -917,6 +962,13 @@ class Trainer(TrainerIO): model = self.__get_model() model.on_epoch_end() + def test(self, model=None): + if model is not None: + self.testing = True + self.fit(model) + else: + self.__run_evaluation(test=True) + def __metrics_to_scalars(self, metrics, blacklist=set()): new_metrics = {} for k, v in metrics.items(): @@ -1109,32 +1161,49 @@ class Trainer(TrainerIO): 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 + def __run_evaluation(self, test=False): + # when testing make sure user defined a test step + can_run_test_step = False + if test: + can_run_test_step = self.__is_overriden('test_step') and self.__is_overriden('test_end') + if not can_run_test_step: + m = '''You called .test() without defining a test step or test_end. + Please define and try again''' + raise MisconfigurationException(m) # validate only if model has validation_step defined - if self.__is_overriden('validation_step'): + # test only if test_step or validation_step are defined + run_val_step = self.__is_overriden('validation_step') + + if run_val_step or can_run_test_step: # hook - if self.__is_function_implemented('on_pre_performance_check'): - model = self.__get_model() - model.on_pre_performance_check() + model = self.__get_model() + model.on_pre_performance_check() - # use val_percent_check set on end of epoch - # use a small portion otherwise - max_batches = self.nb_val_batches if not self.fast_dev_run else 1 - for ds_i, dataloader in enumerate(self.val_dataloader): - val_out_metrics = self.validate(self.model, dataloader, max_batches, ds_i) - self.__add_tqdm_metrics(val_out_metrics) + # select dataloaders + dataloaders = self.val_dataloader + max_batches = self.nb_val_batches - # hook - if self.__is_function_implemented('on_post_performance_check'): - model = self.__get_model() + # calculate max batches to use + if test: + dataloaders = self.test_dataloader + max_batches = self.nb_test_batches + + # cap max batches to 1 when using fast_dev_run + if self.fast_dev_run: + max_batches = 1 + + for ds_i, dataloader in enumerate(dataloaders): + eval_out_metrics = self.evaluate(self.model, + dataloader, + max_batches, + ds_i, + test) + + self.__add_tqdm_metrics(eval_out_metrics) + + # hook model.on_post_performance_check() if self.show_progress_bar: @@ -1143,7 +1212,7 @@ class Trainer(TrainerIO): self.progress_bar.set_postfix(**tqdm_metrics) # model checkpointing - if self.proc_rank == 0 and self.checkpoint_callback is not None: + if self.proc_rank == 0 and self.checkpoint_callback is not None and not test: print('save callback...') self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, logs=self.__tng_tqdm_dic) diff --git a/pytorch_lightning/pt_overrides/override_data_parallel.py b/pytorch_lightning/pt_overrides/override_data_parallel.py index ab88d286..b13b7559 100644 --- a/pytorch_lightning/pt_overrides/override_data_parallel.py +++ b/pytorch_lightning/pt_overrides/override_data_parallel.py @@ -56,6 +56,8 @@ class LightningDataParallel(DataParallel): # lightning if self.module.training: return self.module.training_step(*inputs[0], **kwargs[0]) + elif self.module.testing: + return self.module.test_step(*inputs[0], **kwargs[0]) else: return self.module.validation_step(*inputs[0], **kwargs[0]) @@ -89,6 +91,8 @@ class LightningDistributedDataParallel(DistributedDataParallel): # lightning if self.module.training: output = self.module.training_step(*inputs[0], **kwargs[0]) + elif self.module.testing: + output = self.module.test_step(*inputs[0], **kwargs[0]) else: output = self.module.validation_step(*inputs[0], **kwargs[0]) else: @@ -153,6 +157,10 @@ def parallel_apply(modules, inputs, kwargs_tup=None, devices=None): # pragma: n # CHANGE if module.training: output = module.training_step(*input, **kwargs) + + elif module.testing: + output = module.test_step(*input, **kwargs) + else: output = module.validation_step(*input, **kwargs) # --------------- diff --git a/pytorch_lightning/root_module/hooks.py b/pytorch_lightning/root_module/hooks.py index 7661d344..12a187c8 100644 --- a/pytorch_lightning/root_module/hooks.py +++ b/pytorch_lightning/root_module/hooks.py @@ -5,7 +5,7 @@ class ModelHooks(torch.nn.Module): def on_sanity_check_start(self): """ - Called before starting validate + Called before starting evaluate :return: """ pass diff --git a/pytorch_lightning/root_module/root_module.py b/pytorch_lightning/root_module/root_module.py index 13433bb3..bb8be40b 100644 --- a/pytorch_lightning/root_module/root_module.py +++ b/pytorch_lightning/root_module/root_module.py @@ -55,6 +55,16 @@ class LightningModule(GradInformation, ModelIO, ModelHooks): """ pass + def test_step(self, *args, **kwargs): + """ + return whatever outputs will need to be aggregated in test_end + OPTIONAL + :param called with batch, batch_nb + additional: dataset_i if multiple val datasets used + :return: + """ + pass + def validation_end(self, outputs): """ Outputs has the appended output after each validation step @@ -64,6 +74,15 @@ class LightningModule(GradInformation, ModelIO, ModelHooks): """ pass + def test_end(self, outputs): + """ + Outputs has the appended output after each test step + OPTIONAL + :param outputs: + :return: dic_with_metrics for tqdm + """ + pass + def configure_optimizers(self): """ Return a list of optimizers and a list of schedulers (could be empty) diff --git a/pytorch_lightning/testing/lm_test_module.py b/pytorch_lightning/testing/lm_test_module.py index ac146159..16566053 100644 --- a/pytorch_lightning/testing/lm_test_module.py +++ b/pytorch_lightning/testing/lm_test_module.py @@ -20,7 +20,7 @@ class LightningTestModel(LightningModule): Sample model to show how to define a template """ - def __init__(self, hparams, force_remove_distributed_sampler=False): + def __init__(self, hparams, force_remove_distributed_sampler=False, use_two_test_sets=False): """ Pass in parsed HyperOptArgumentParser to the model :param hparams: @@ -28,6 +28,7 @@ class LightningTestModel(LightningModule): # init superclass super(LightningTestModel, self).__init__() self.hparams = hparams + self.use_two_test_sets = use_two_test_sets # for some tests regarding testing self.batch_size = hparams.batch_size @@ -167,12 +168,22 @@ class LightningTestModel(LightningModule): # 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 = output['val_loss'] + + # reduce manually when using dp + if self.trainer.use_dp: + val_loss = torch.mean(val_loss) + val_loss_mean += val_loss + + # reduce manually when using dp + val_acc = output['val_acc'] + if self.trainer.use_dp: + val_acc = torch.mean(val_acc) + + val_acc_mean += val_acc val_loss_mean /= len(outputs) val_acc_mean /= len(outputs) @@ -180,6 +191,87 @@ class LightningTestModel(LightningModule): tqdm_dic = {'val_loss': val_loss_mean.item(), 'val_acc': val_acc_mean.item()} return tqdm_dic + def test_step(self, data_batch, batch_i, dataloader_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_test = self.loss(y, y_hat) + + # acc + labels_hat = torch.argmax(y_hat, dim=1) + test_acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0) + test_acc = torch.tensor(test_acc) + + if self.on_gpu: + test_acc = test_acc.cuda(loss_test.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_test = loss_test.unsqueeze(0) + test_acc = test_acc.unsqueeze(0) + + # alternate possible outputs to test + if batch_i % 1 == 0: + output = OrderedDict({ + 'test_loss': loss_test, + 'test_acc': test_acc, + }) + return output + if batch_i % 2 == 0: + return test_acc + + if batch_i % 3 == 0: + output = OrderedDict({ + 'test_loss': loss_test, + 'test_acc': test_acc, + 'test_dic': {'test_loss_a': loss_test} + }) + return output + if batch_i % 5 == 0: + output = OrderedDict({ + f'test_loss_{dataloader_i}': loss_test, + f'test_acc_{dataloader_i}': test_acc, + }) + return output + + def test_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 test_step, outputs is a list of tensor scalars + # we return just the average in this case (if we want) + # return torch.stack(outputs).mean() + test_loss_mean = 0 + test_acc_mean = 0 + for output in outputs: + test_loss = output['test_loss'] + + # reduce manually when using dp + if self.trainer.use_dp: + test_loss = torch.mean(test_loss) + test_loss_mean += test_loss + + # reduce manually when using dp + test_acc = output['test_acc'] + if self.trainer.use_dp: + test_acc = torch.mean(test_acc) + + test_acc_mean += test_acc + + test_loss_mean /= len(outputs) + test_acc_mean /= len(outputs) + + tqdm_dic = {'test_loss': test_loss_mean.item(), 'test_acc': test_acc_mean.item()} + return tqdm_dic + def on_tng_metrics(self, logs): logs['some_tensor_to_test'] = torch.rand(1) @@ -235,6 +327,8 @@ class LightningTestModel(LightningModule): @data_loader def test_dataloader(self): + if self.use_two_test_sets: + return [self.__dataloader(train=False), self.__dataloader(train=False)] return self.__dataloader(train=False) @staticmethod diff --git a/tests/debug.py b/tests/debug.py index b480f4cf..856a1b06 100644 --- a/tests/debug.py +++ b/tests/debug.py @@ -1,5 +1,6 @@ from pytorch_lightning import Trainer from examples import LightningTemplateModel +from pytorch_lightning.testing import LightningTestModel, NoValEndTestModel, NoValModel from argparse import Namespace from test_tube import Experiment from pytorch_lightning.callbacks import ModelCheckpoint @@ -12,6 +13,7 @@ from torch.nn import functional as F from torch.utils.data import DataLoader from torchvision.datasets import MNIST import numpy as np +import pdb class CoolModel(pl.LightningModule): @@ -73,10 +75,11 @@ def get_model(): return model, hparams -def get_exp(debug=True): +def get_exp(debug=True, version=None): # 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') + save_dir = os.path.join(root_dir, 'save_dir') + exp = Experiment(debug=debug, save_dir=save_dir, name='tests_tt_dir', version=version) return exp @@ -99,7 +102,7 @@ def clear_save_dir(): shutil.rmtree(save_dir) -def load_model(exp, save_dir): +def load_model(exp, save_dir, on_gpu, map_location=None, module_class=LightningTemplateModel): # load trained model tags_path = exp.get_data_path(exp.name, exp.version) @@ -108,8 +111,10 @@ def load_model(exp, save_dir): 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) + trained_model = module_class.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' @@ -131,12 +136,10 @@ def run_prediction(dataloader, trained_model): 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, 'this model is expected to get > 0.7 in test set (it got %f)' % val_acc +# ------------------------------------------------------------------------ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True): save_dir = init_save_dir() @@ -162,7 +165,7 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True): # test model loading pretrained_model = load_model(exp, save_dir, on_gpu) - # test model preds + # test new model accuracy run_prediction(model.test_dataloader, pretrained_model) if trainer.use_ddp: @@ -177,19 +180,79 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True): clear_save_dir() +def assert_ok_val_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}' + + +def assert_ok_test_acc(trainer): + # this model should get 0.80+ acc + acc = trainer.tng_tqdm_dic['test_acc'] + assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}' + + +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 main(): + """Verify test() on fitted model""" + 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() + + # exp file to get weights + checkpoint = ModelCheckpoint(save_dir) - os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0]) - model, hparams = get_model() trainer_options = dict( + show_progress_bar=False, max_nb_epochs=1, train_percent_check=0.4, val_percent_check=0.2, + checkpoint_callback=checkpoint, + experiment=exp, gpus=[0, 1], distributed_backend='ddp' ) - run_gpu_model_test(trainer_options, model, hparams) + # fit model + trainer = Trainer(**trainer_options) + result = trainer.fit(model) + + # correct result and ok accuracy + assert result == 1, 'training failed to complete' + pretrained_model = load_model(exp, save_dir, on_gpu=True, module_class=LightningTestModel) + + new_trainer = Trainer(**trainer_options) + new_trainer.test(pretrained_model) + + # test we have good test accuracy + assert_ok_test_acc(new_trainer) + # clear_save_dir() if __name__ == '__main__': diff --git a/tests/test_models.py b/tests/test_models.py index 63af9287..df2a41c3 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -30,6 +30,177 @@ np.random.seed(SEED) # ------------------------------------------------------------------------ # TESTS # ------------------------------------------------------------------------ +def test_running_test_pretrained_model_ddp(): + """Verify test() on pretrained model""" + if not can_run_gpu_test(): + 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() + + # exp file to get weights + checkpoint = ModelCheckpoint(save_dir) + + trainer_options = dict( + show_progress_bar=False, + max_nb_epochs=1, + train_percent_check=0.4, + val_percent_check=0.2, + checkpoint_callback=checkpoint, + experiment=exp, + gpus=[0, 1], + distributed_backend='ddp' + ) + + # fit model + trainer = Trainer(**trainer_options) + result = trainer.fit(model) + + # correct result and ok accuracy + assert result == 1, 'training failed to complete' + pretrained_model = load_model(exp, save_dir, on_gpu=True, module_class=LightningTestModel) + + # run test set + new_trainer = Trainer(**trainer_options) + new_trainer.test(pretrained_model) + + run_prediction(model.test_dataloader, pretrained_model) + + # test we have good test accuracy + clear_save_dir() + + +def test_running_test_after_fitting(): + """Verify test() on fitted model""" + 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() + + # exp file to get weights + checkpoint = ModelCheckpoint(save_dir) + + trainer_options = dict( + show_progress_bar=False, + max_nb_epochs=1, + train_percent_check=0.4, + val_percent_check=0.2, + test_percent_check=0.2, + checkpoint_callback=checkpoint, + experiment=exp + ) + + # fit model + trainer = Trainer(**trainer_options) + result = trainer.fit(model) + + assert result == 1, 'training failed to complete' + + trainer.test() + + # test we have good test accuracy + assert_ok_test_acc(trainer) + + clear_save_dir() + + +def test_running_test_pretrained_model(): + """Verify test() on pretrained model""" + 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() + + # exp file to get weights + checkpoint = ModelCheckpoint(save_dir) + + trainer_options = dict( + show_progress_bar=False, + max_nb_epochs=1, + train_percent_check=0.4, + val_percent_check=0.2, + checkpoint_callback=checkpoint, + experiment=exp + ) + + # fit model + trainer = Trainer(**trainer_options) + result = trainer.fit(model) + + # correct result and ok accuracy + assert result == 1, 'training failed to complete' + pretrained_model = load_model(exp, save_dir, on_gpu=False, module_class=LightningTestModel) + + new_trainer = Trainer(**trainer_options) + new_trainer.test(pretrained_model) + + # test we have good test accuracy + assert_ok_test_acc(new_trainer) + clear_save_dir() + + +def test_running_test_pretrained_model_dp(): + """Verify test() on pretrained model""" + if not can_run_gpu_test(): + 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() + + # exp file to get weights + checkpoint = ModelCheckpoint(save_dir) + + trainer_options = dict( + show_progress_bar=False, + max_nb_epochs=1, + train_percent_check=0.4, + val_percent_check=0.2, + checkpoint_callback=checkpoint, + experiment=exp, + gpus=[0, 1], + distributed_backend='dp' + ) + + # fit model + trainer = Trainer(**trainer_options) + result = trainer.fit(model) + + # correct result and ok accuracy + assert result == 1, 'training failed to complete' + pretrained_model = load_model(exp, save_dir, on_gpu=True, module_class=LightningTestModel) + + new_trainer = Trainer(**trainer_options) + new_trainer.test(pretrained_model) + + # test we have good test accuracy + assert_ok_test_acc(new_trainer) + clear_save_dir() + + def test_gradient_accumulation_scheduling(): """ Test grad accumulation by the freq of optimizer updates @@ -107,13 +278,7 @@ 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') + if not can_run_gpu_test(): return os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0]) @@ -161,13 +326,7 @@ def test_optimizer_return_options(): def test_single_gpu_batch_parse(): - 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') + if not can_run_gpu_test(): return trainer = Trainer() @@ -261,7 +420,7 @@ def test_no_val_module(): trainer = Trainer(**trainer_options) result = trainer.fit(model) - # traning complete + # training complete assert result == 1, 'amp + ddp model failed to complete' # save model @@ -449,13 +608,7 @@ 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') + if not can_run_gpu_test(): return os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0]) @@ -667,13 +820,7 @@ 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') + if not can_run_gpu_test(): return # simulate setting slurm flags @@ -832,14 +979,9 @@ 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') + if not can_run_gpu_test(): return + model, hparams = get_model() trainer_options = dict( show_progress_bar=False, @@ -860,14 +1002,9 @@ 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') + if not can_run_gpu_test(): return + model, hparams = get_model() trainer_options = dict( max_nb_epochs=1, @@ -884,11 +1021,7 @@ 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') + if not can_run_gpu_test(): return os.environ['MASTER_PORT'] = str(np.random.randint(12000, 19000, 1)[0]) @@ -922,7 +1055,34 @@ def test_multiple_val_dataloader(): hparams = get_hparams() model = LightningTestModel(hparams) - save_dir = init_save_dir() + # exp file to get meta + trainer_options = dict( + max_nb_epochs=1, + val_percent_check=0.1, + train_percent_check=1.0, + ) + + # fit model + trainer = Trainer(**trainer_options) + result = trainer.fit(model) + + # verify tng completed + assert result == 1 + + # verify there are 2 val loaders + assert len(trainer.val_dataloader) == 2, 'Multiple val_dataloaders not initiated properly' + + # make sure predictions are good for each val set + [run_prediction(dataloader, trainer.model) for dataloader in trainer.val_dataloader] + + +def test_multiple_test_dataloader(): + """ + Verify multiple test_dataloader + :return: + """ + hparams = get_hparams() + model = LightningTestModel(hparams, use_two_test_sets=True) # exp file to get meta trainer_options = dict( @@ -939,10 +1099,10 @@ def test_multiple_val_dataloader(): assert result == 1 # verify there are 2 val loaders - assert len(trainer.val_dataloader) == 2, 'Multiple val_dataloaders not initiated properly' + assert len(trainer.test_dataloader) == 2, 'Multiple test_dataloaders not initiated properly' - # make sure predictions are good for each val set - [run_prediction(dataloader, trainer.model) for dataloader in trainer.val_dataloader] + # make sure predictions are good for each test set + [run_prediction(dataloader, trainer.model) for dataloader in trainer.test_dataloader] # ------------------------------------------------------------------------ @@ -973,7 +1133,7 @@ def run_gpu_model_test(trainer_options, model, hparams, on_gpu=True): # test model loading pretrained_model = load_model(exp, save_dir, on_gpu) - # test model preds + # test new model accuracy run_prediction(model.test_dataloader, pretrained_model) if trainer.use_ddp: @@ -1024,7 +1184,8 @@ def get_model(use_test_model=False): def get_exp(debug=True, version=None): # 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', version=version) + save_dir = os.path.join(root_dir, 'save_dir') + exp = Experiment(debug=debug, save_dir=save_dir, name='tests_tt_dir', version=version) return exp @@ -1033,7 +1194,8 @@ def init_save_dir(): save_dir = os.path.join(root_dir, 'save_dir') if os.path.exists(save_dir): - shutil.rmtree(save_dir) + n = np.random.randint(0, 10000000, 1)[0] + shutil.move(save_dir, save_dir + f'_{n}') os.makedirs(save_dir, exist_ok=True) @@ -1044,10 +1206,11 @@ 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) + n = np.random.randint(0, 10000000, 1)[0] + shutil.move(save_dir, save_dir + f'_{n}') -def load_model(exp, save_dir, on_gpu, map_location=None): +def load_model(exp, save_dir, on_gpu, map_location=None, module_class=LightningTemplateModel): # load trained model tags_path = exp.get_data_path(exp.name, exp.version) @@ -1056,10 +1219,10 @@ def load_model(exp, save_dir, on_gpu, map_location=None): 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) + trained_model = module_class.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' @@ -1078,19 +1241,35 @@ def run_prediction(dataloader, trained_model): # 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() + acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0) + acc = torch.tensor(acc) + acc = acc.item() - print(val_acc) - - assert val_acc > 0.50, 'this model is expected to get > 0.50 in test set (it got %f)' % val_acc + assert acc > 0.50, f'this model is expected to get > 0.50 in test set (it got {acc})' -def assert_ok_acc(trainer): +def assert_ok_val_acc(trainer): # this model should get 0.80+ acc acc = trainer.tng_tqdm_dic['val_acc'] - assert acc > 0.50, 'model failed to get expected 0.50 validation accuracy. Got: %f' % acc + assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}' + + +def assert_ok_test_acc(trainer): + # this model should get 0.80+ acc + acc = trainer.tng_tqdm_dic['test_acc'] + assert acc > 0.50, f'model failed to get expected 0.50 validation accuracy. Got: {acc}' + + +def can_run_gpu_test(): + 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 False + 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 False + return True if __name__ == '__main__':