From 14241577314c5c9adffdf14b107100e6656f5332 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Tue, 22 Oct 2019 04:16:51 +0300 Subject: [PATCH] Refactor (#407) * moved dp, ddp outside of trainer * added main mixins * finished major mixin refactor * flake8 * finished major mixin refactor * finished major mixin refactor * finished major mixin refactor * finished major mixin refactor * finished major mixin refactor * finished major mixin refactor * finished major mixin refactor --- pytorch_lightning/trainer/amp_mixin.py | 23 + .../trainer/callback_config_mixin.py | 70 + .../trainer/data_loading_mixin.py | 141 ++ pytorch_lightning/trainer/ddp_mixin.py | 236 ++++ pytorch_lightning/trainer/dp_mixin.py | 111 ++ .../trainer/evaluation_loop_mixin.py | 171 +++ pytorch_lightning/trainer/logging_mixin.py | 157 +++ .../trainer/model_hooks_mixin.py | 17 + pytorch_lightning/trainer/train_loop_mixin.py | 266 ++++ pytorch_lightning/trainer/trainer.py | 1223 +---------------- pytorch_lightning/trainer/trainer_io.py | 10 +- .../trainer/training_tricks_mixin.py | 28 + tests/test_models.py | 10 +- 13 files changed, 1276 insertions(+), 1187 deletions(-) create mode 100644 pytorch_lightning/trainer/amp_mixin.py create mode 100644 pytorch_lightning/trainer/callback_config_mixin.py create mode 100644 pytorch_lightning/trainer/data_loading_mixin.py create mode 100644 pytorch_lightning/trainer/ddp_mixin.py create mode 100644 pytorch_lightning/trainer/dp_mixin.py create mode 100644 pytorch_lightning/trainer/evaluation_loop_mixin.py create mode 100644 pytorch_lightning/trainer/logging_mixin.py create mode 100644 pytorch_lightning/trainer/model_hooks_mixin.py create mode 100644 pytorch_lightning/trainer/train_loop_mixin.py create mode 100644 pytorch_lightning/trainer/training_tricks_mixin.py diff --git a/pytorch_lightning/trainer/amp_mixin.py b/pytorch_lightning/trainer/amp_mixin.py new file mode 100644 index 00000000..85f2248d --- /dev/null +++ b/pytorch_lightning/trainer/amp_mixin.py @@ -0,0 +1,23 @@ +try: + from apex import amp + APEX_AVAILABLE = True +except ImportError: + APEX_AVAILABLE = False + + +class TrainerAMPMixin(object): + + def init_amp(self, use_amp): + self.use_amp = use_amp and APEX_AVAILABLE + if self.use_amp: + print('using 16bit precision') + + if use_amp and not APEX_AVAILABLE: # pragma: no cover + msg = """ + You set use_amp=True but do not have apex installed. + Install apex first using this guide and rerun with use_amp=True: + https://github.com/NVIDIA/apex#linux + + this run will NOT use 16 bit precision + """ + raise ModuleNotFoundError(msg) diff --git a/pytorch_lightning/trainer/callback_config_mixin.py b/pytorch_lightning/trainer/callback_config_mixin.py new file mode 100644 index 00000000..1b36bd46 --- /dev/null +++ b/pytorch_lightning/trainer/callback_config_mixin.py @@ -0,0 +1,70 @@ +from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping +from pytorch_lightning.logging import TestTubeLogger + + +class TrainerCallbackConfigMixin(object): + def configure_checkpoint_callback(self): + """ + Weight path set in this priority: + Checkpoint_callback's path (if passed in). + User provided weights_saved_path + Otherwise use os.getcwd() + """ + if self.checkpoint_callback is True: + # init a default one + if isinstance(self.logger, TestTubeLogger): + ckpt_path = '{}/{}/version_{}/{}'.format( + self.default_save_path, + self.logger.experiment.name, + self.logger.experiment.version, + 'checkpoints') + else: + ckpt_path = self.default_save_path + + self.checkpoint_callback = ModelCheckpoint( + filepath=ckpt_path + ) + elif self.checkpoint_callback is False: + self.checkpoint_callback = None + + 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 weights_save_path is still none here, set to current working dir + if self.weights_save_path is None: + self.weights_save_path = self.default_save_path + + def configure_early_stopping(self, early_stop_callback, logger): + 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 + 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 + elif logger is False: + self.logger = None + else: + self.logger = logger + self.logger.rank = 0 diff --git a/pytorch_lightning/trainer/data_loading_mixin.py b/pytorch_lightning/trainer/data_loading_mixin.py new file mode 100644 index 00000000..5f34ff87 --- /dev/null +++ b/pytorch_lightning/trainer/data_loading_mixin.py @@ -0,0 +1,141 @@ +import warnings + +from torch.utils.data.distributed import DistributedSampler +import torch.distributed as dist + +try: + from apex import amp + APEX_AVAILABLE = True +except ImportError: + APEX_AVAILABLE = False + + +class TrainerDataLoadingMixin(object): + + def layout_bookeeping(self): + + # determine number of training batches + self.nb_training_batches = len(self.get_train_dataloader()) + self.nb_training_batches = int(self.nb_training_batches * self.train_percent_check) + + # determine number of validation batches + # val datasets could be none, 1 or 2+ + if self.get_val_dataloaders() is not None: + self.nb_val_batches = sum(len(dataloader) for dataloader in self.get_val_dataloaders()) + self.nb_val_batches = int(self.nb_val_batches * self.val_percent_check) + self.nb_val_batches = max(1, self.nb_val_batches) + + # determine number of test batches + if self.get_test_dataloaders() is not None: + self.nb_test_batches = sum( + len(dataloader) for dataloader in self.get_test_dataloaders() + ) + 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_training_batches * self.val_check_interval) + self.val_check_batch = max(1, self.val_check_batch) + + def get_dataloaders(self, model): + """ + Dataloaders are provided by the model + :param model: + :return: + """ + self.get_train_dataloader = model.train_dataloader + self.get_test_dataloaders = model.test_dataloader + self.get_val_dataloaders = model.val_dataloader + + # call warnings from proc zero only which triggers dataloaders + # if those have to download data it will only happen on proc 0 + if self.proc_rank == 0: + on_ddp = self.use_ddp or self.use_ddp2 + if on_ddp and not isinstance(self.get_train_dataloader().sampler, DistributedSampler): + msg = """ + 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) + + if on_ddp and self.get_val_dataloaders() is not None: + for dataloader in self.get_val_dataloaders(): + if not isinstance(dataloader.sampler, DistributedSampler): + msg = """ + Your val_dataloader(s) don't use DistributedSampler. + + 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 + + if on_ddp and self.get_test_dataloaders() is not None: + for dataloader in self.get_test_dataloaders(): + if not isinstance(dataloader.sampler, DistributedSampler): + msg = """ + Your test_dataloader(s) don't use DistributedSampler. + + 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 + + if self.use_ddp or self.use_ddp2: + # wait for all processes to catch up + dist.barrier() + + # load each dataloader + self.get_train_dataloader() + self.get_test_dataloaders() + self.get_val_dataloaders() + + def determine_data_use_amount(self, train_percent_check, val_percent_check, + test_percent_check, overfit_pct): + """ + Use less data for debugging purposes + """ + self.train_percent_check = train_percent_check + self.val_percent_check = val_percent_check + self.test_percent_check = test_percent_check + if overfit_pct > 0: + self.train_percent_check = overfit_pct + self.val_percent_check = overfit_pct + self.test_percent_check = overfit_pct diff --git a/pytorch_lightning/trainer/ddp_mixin.py b/pytorch_lightning/trainer/ddp_mixin.py new file mode 100644 index 00000000..1b9cac56 --- /dev/null +++ b/pytorch_lightning/trainer/ddp_mixin.py @@ -0,0 +1,236 @@ +import os +import re +import warnings + +import torch +import torch.distributed as dist +from pytorch_lightning.pt_overrides.override_data_parallel import LightningDistributedDataParallel +from pytorch_lightning.utilities.debugging import MisconfigurationException + + +try: + from apex import amp + APEX_AVAILABLE = True +except ImportError: + APEX_AVAILABLE = False + + +class TrainerDDPMixin(object): + def set_distributed_mode(self, distributed_backend, nb_gpu_nodes): + # skip for CPU + if self.num_gpus == 0: + return + + # single GPU case + # in single gpu case we allow ddp so we can train on multiple + # nodes, 1 gpu per node + if self.num_gpus == 1: + self.single_gpu = True + + if distributed_backend is not None: + self.use_dp = distributed_backend == 'dp' + self.use_ddp = distributed_backend == 'ddp' + self.use_ddp2 = distributed_backend == 'ddp2' + + # disable single gpu when using ddp2 + if self.use_ddp2: + self.single_gpu = False + + # multiple GPU case + elif self.num_gpus > 1: + if distributed_backend is not None: + # DP, DDP case + self.use_dp = distributed_backend == 'dp' + self.use_ddp = distributed_backend == 'ddp' + self.use_ddp2 = distributed_backend == 'ddp2' + + elif distributed_backend is None: + m = 'When using multiple GPUs set ' \ + 'Trainer(distributed_backend=dp) (or ddp)' + raise MisconfigurationException(m) + + # use ddp automatically if nb_gpu_nodes > 1 + if nb_gpu_nodes > 1 and self.use_dp: # pragma: no cover + self.use_ddp = True + self.use_dp = False + w = 'DataParallel does not support nb_gpu_nodes > 1. ' \ + 'Switching to DistributedDataParallel for you. ' \ + 'To silence this warning set distributed_backend=ddp' + warnings.warn(w) + + print('gpu available: {}, used: {}'.format(torch.cuda.is_available(), self.on_gpu)) + + def configure_slurm_ddp(self, nb_gpu_nodes): + self.is_slurm_managing_tasks = False + + # extract SLURM flag vars + # whenever we have the correct number of tasks, we let slurm manage processes + # otherwise we launch the required number of processes + if self.use_ddp: + self.nb_requested_gpus = self.num_gpus * nb_gpu_nodes + self.nb_slurm_tasks = 0 + try: + self.nb_slurm_tasks = int(os.environ['SLURM_NTASKS']) + self.is_slurm_managing_tasks = self.nb_slurm_tasks == self.nb_requested_gpus + + # in interactive mode we don't manage tasks + job_name = os.environ['SLURM_JOB_NAME'] + if job_name == 'bash': + self.is_slurm_managing_tasks = False + + except Exception: + # likely not on slurm, so set the slurm managed flag to false + self.is_slurm_managing_tasks = False + + # used for tests only, set this flag to simulate slurm managing a task + try: + should_fake = int(os.environ['FAKE_SLURM_MANAGING_TASKS']) + if should_fake: + self.is_slurm_managing_tasks = True + except Exception as e: + pass + + def set_nvidia_flags(self, is_slurm_managing_tasks, data_parallel_device_ids): + if data_parallel_device_ids is None: + return + + # set the correct cuda visible devices (using pci order) + os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + + # when slurm is managing the task it sets the visible devices + if not is_slurm_managing_tasks: + if type(data_parallel_device_ids) is int: + id_str = ','.join(str(x) for x in list(range(data_parallel_device_ids))) + os.environ["CUDA_VISIBLE_DEVICES"] = id_str + else: + gpu_str = ','.join([str(x) for x in data_parallel_device_ids]) + os.environ["CUDA_VISIBLE_DEVICES"] = gpu_str + + print(f'VISIBLE GPUS: {os.environ["CUDA_VISIBLE_DEVICES"]}') + + def ddp_train(self, gpu_nb, model): + """ + Entry point into a DP thread + :param gpu_nb: + :param model: + :param cluster_obj: + :return: + """ + # node rank using relative slurm id + # otherwise default to node rank 0 + try: + node_id = os.environ['SLURM_NODEID'] + self.node_rank = int(node_id) + except Exception: + self.node_rank = 0 + + # show progressbar only on progress_rank 0 + self.show_progress_bar = self.show_progress_bar and self.node_rank == 0 and gpu_nb == 0 + + # determine which process we are and world size + if self.use_ddp: + self.proc_rank = self.node_rank * self.num_gpus + gpu_nb + self.world_size = self.nb_gpu_nodes * self.num_gpus + + elif self.use_ddp2: + self.proc_rank = self.node_rank + self.world_size = self.nb_gpu_nodes + + # let the exp know the rank to avoid overwriting logs + if self.logger is not None: + self.logger.rank = self.proc_rank + + # set up server using proc 0's ip address + # try to init for 20 times at max in case ports are taken + # where to store ip_table + self.__init_tcp_connection() + + # CHOOSE OPTIMIZER + # allow for lr schedulers as well + self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers()) + + # MODEL + # copy model to each gpu + if self.distributed_backend == 'ddp': + torch.cuda.set_device(gpu_nb) + model.cuda(gpu_nb) + + # set model properties before going into wrapper + self.copy_trainer_model_properties(model) + + # override root GPU + self.root_gpu = gpu_nb + + # AMP + # run through amp wrapper before going to distributed DP + if self.use_amp: + # An example + model, optimizers = amp.initialize( + model, self.optimizers, opt_level=self.amp_level, + ) + self.optimizers = optimizers + + # DDP2 uses all GPUs on the machine + if self.distributed_backend == 'ddp': + device_ids = [gpu_nb] + elif self.use_ddp2: + device_ids = None + + model = LightningDistributedDataParallel( + model, + device_ids=device_ids, + find_unused_parameters=True + ) + + # continue training routine + self.run_pretrain_routine(model) + + def __init_tcp_connection(self): + """ + Connect all procs in the world using the env:// init + Use the first node as the root address + :param port: + :param tries: + :return: + """ + + # use slurm job id for the port number + # guarantees unique ports across jobs from same grid search + try: + # use the last 4 numbers in the job id as the id + default_port = os.environ['SLURM_JOB_ID'] + default_port = default_port[-4:] + + # all ports should be in the 10k+ range + default_port = int(default_port) + 15000 + + except Exception as e: + default_port = 12910 + + # if user gave a port number, use that one instead + try: + default_port = os.environ['MASTER_PORT'] + except Exception: + os.environ['MASTER_PORT'] = str(default_port) + + # figure out the root node addr + try: + root_node = os.environ['SLURM_NODELIST'].split(' ')[0] + except Exception: + root_node = '127.0.0.2' + + root_node = self.resolve_root_node_address(root_node) + os.environ['MASTER_ADDR'] = root_node + dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size) + + def resolve_root_node_address(self, root_node): + if '[' in root_node: + name = root_node.split('[')[0] + number = root_node.split(',')[0] + if '-' in number: + number = number.split('-')[0] + + number = re.sub('[^0-9]', '', number) + root_node = name + number + + return root_node diff --git a/pytorch_lightning/trainer/dp_mixin.py b/pytorch_lightning/trainer/dp_mixin.py new file mode 100644 index 00000000..76403164 --- /dev/null +++ b/pytorch_lightning/trainer/dp_mixin.py @@ -0,0 +1,111 @@ +import os +import re +import signal +import pdb +from subprocess import call + +import torch +import torch.distributed as dist +from pytorch_lightning.pt_overrides.override_data_parallel import ( + LightningDistributedDataParallel, LightningDataParallel) +from pytorch_lightning.utilities.debugging import MisconfigurationException + +try: + from apex import amp + APEX_AVAILABLE = True +except ImportError: + APEX_AVAILABLE = False + + +class TrainerDPMixin(object): + def copy_trainer_model_properties(self, model): + if isinstance(model, LightningDataParallel): + ref_model = model.module + elif isinstance(model, LightningDistributedDataParallel): + ref_model = model.module + else: + ref_model = model + + for m in [model, ref_model]: + m.trainer = self + m.on_gpu = self.on_gpu + m.use_dp = self.use_dp + m.use_ddp2 = self.use_ddp2 + m.use_ddp = self.use_ddp + m.use_amp = self.use_amp + m.testing = self.testing + m.single_gpu = self.single_gpu + + def transfer_batch_to_gpu(self, batch, gpu_id): + # base case: object can be directly moved using `cuda` or `to` + if callable(getattr(batch, 'cuda', None)): + return batch.cuda(gpu_id) + + elif callable(getattr(batch, 'to', None)): + return batch.to(torch.device('cuda', gpu_id)) + + # when list + elif isinstance(batch, list): + for i, x in enumerate(batch): + batch[i] = self.transfer_batch_to_gpu(x, gpu_id) + return batch + + # when tuple + elif isinstance(batch, tuple): + batch = list(batch) + for i, x in enumerate(batch): + batch[i] = self.transfer_batch_to_gpu(x, gpu_id) + return tuple(batch) + + # when dict + elif isinstance(batch, dict): + for k, v in batch.items(): + batch[k] = self.transfer_batch_to_gpu(v, gpu_id) + + return batch + + # nothing matches, return the value as is without transform + return batch + + def single_gpu_train(self, model): + # CHOOSE OPTIMIZER + # allow for lr schedulers as well + self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers()) + + model.cuda(self.root_gpu) + + if self.use_amp: + # An example + model, optimizers = amp.initialize( + model, self.optimizers, opt_level=self.amp_level, + ) + self.optimizers = optimizers + + self.run_pretrain_routine(model) + + def dp_train(self, model): + + # CHOOSE OPTIMIZER + # allow for lr schedulers as well + self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers()) + + model.cuda(self.root_gpu) + + # check for this bug (amp + dp + !01 doesn't work) + # https://github.com/NVIDIA/apex/issues/227 + if self.use_dp and self.use_amp: + m = f""" + Amp level {self.amp_level} with DataParallel is not supported. + See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227. + We recommend you switch to ddp if you want to use amp + """ + raise MisconfigurationException(m) + + # create list of device ids + device_ids = self.data_parallel_device_ids + if type(device_ids) is int: + device_ids = list(range(device_ids)) + + model = LightningDataParallel(model, device_ids=device_ids) + + self.run_pretrain_routine(model) diff --git a/pytorch_lightning/trainer/evaluation_loop_mixin.py b/pytorch_lightning/trainer/evaluation_loop_mixin.py new file mode 100644 index 00000000..0373c188 --- /dev/null +++ b/pytorch_lightning/trainer/evaluation_loop_mixin.py @@ -0,0 +1,171 @@ +import torch +from pytorch_lightning.utilities.debugging import MisconfigurationException + + +class TrainerEvaluationLoopMixin(object): + + def evaluate(self, model, dataloaders, max_batches, test=False): + """ + Run evaluation code + :param model: PT model + :param dataloaders: list of PT dataloaders + :param max_batches: Scalar + :param test: boolean + :return: + """ + # enable eval mode + model.zero_grad() + model.eval() + + # copy properties for forward overrides + self.copy_trainer_model_properties(model) + + # disable gradients to save memory + torch.set_grad_enabled(False) + + # bookkeeping + outputs = [] + + # run training + for dataloader_idx, dataloader in enumerate(dataloaders): + dl_outputs = [] + for batch_idx, batch in enumerate(dataloader): + + if batch is None: # pragma: no cover + continue + + # stop short when on fast_dev_run (sets max_batch=1) + if batch_idx >= max_batches: + break + + # ----------------- + # RUN EVALUATION STEP + # ----------------- + output = self.evaluation_forward(model, + batch, + batch_idx, + dataloader_idx, + test) + + # track outputs for collation + dl_outputs.append(output) + + # batch done + if self.show_progress_bar: + self.progress_bar.update(1) + outputs.append(dl_outputs) + + eval_results = {} + + # with a single dataloader don't pass an array + if len(dataloaders) == 1: + outputs = outputs[0] + + # give model a chance to do something with the outputs (and method defined) + 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() + + # enable gradients to save memory + torch.set_grad_enabled(True) + + return eval_results + + 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 + # 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 + model = self.get_model() + model.on_pre_performance_check() + + # select dataloaders + if test: + dataloaders = self.get_test_dataloaders() + max_batches = self.nb_test_batches + else: + # val + dataloaders = self.get_val_dataloaders() + max_batches = self.nb_val_batches + + # cap max batches to 1 when using fast_dev_run + if self.fast_dev_run: + max_batches = 1 + + # run evaluation + eval_results = self.evaluate(self.model, + dataloaders, + max_batches, + test) + _, prog_bar_metrics, log_metrics, callback_metrics = self.process_output(eval_results) + + # add metrics to prog bar + self.add_tqdm_metrics(prog_bar_metrics) + + # log metrics + self.log_metrics(log_metrics, {}) + + # track metrics for callbacks + self.callback_metrics = callback_metrics + + # hook + model.on_post_performance_check() + + if self.show_progress_bar: + # add model specific metrics + tqdm_metrics = self.training_tqdm_dict + self.progress_bar.set_postfix(**tqdm_metrics) + + # model checkpointing + if self.proc_rank == 0 and self.checkpoint_callback is not None and not test: + self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, + logs=self.callback_metrics) + + def evaluation_forward(self, model, batch, batch_idx, dataloader_idx, test=False): + # make dataloader_idx arg in validation_step optional + args = [batch, batch_idx] + + if test and len(self.get_test_dataloaders()) > 1: + args.append(dataloader_idx) + + elif not test and len(self.get_val_dataloaders()) > 1: + args.append(dataloader_idx) + + # handle DP, DDP forward + if self.use_ddp or self.use_dp or self.use_ddp2: + output = model(*args) + return output + + # single GPU + if self.single_gpu: + # for single GPU put inputs on gpu manually + root_gpu = 0 + if type(self.data_parallel_device_ids) is list: + root_gpu = self.data_parallel_device_ids[0] + batch = self.transfer_batch_to_gpu(batch, root_gpu) + args[0] = batch + + # CPU + if test: + output = model.test_step(*args) + else: + output = model.validation_step(*args) + + return output diff --git a/pytorch_lightning/trainer/logging_mixin.py b/pytorch_lightning/trainer/logging_mixin.py new file mode 100644 index 00000000..a66541bf --- /dev/null +++ b/pytorch_lightning/trainer/logging_mixin.py @@ -0,0 +1,157 @@ +import torch +from pytorch_lightning.root_module import memory + + +class TrainerLoggingMixin(object): + + def log_metrics(self, metrics, grad_norm_dic): + """ + Logs the metric dict passed in + :param metrics: + :param grad_norm_dic: + :return: + """ + # added metrics by Lightning for convenience + metrics['epoch'] = self.current_epoch + + # add gpu memory + if self.on_gpu and self.log_gpu_memory: + mem_map = memory.get_memory_profile(self.log_gpu_memory) + metrics.update(mem_map) + + # add norms + metrics.update(grad_norm_dic) + + # turn all tensors to scalars + scalar_metrics = self.metrics_to_scalars(metrics) + + # log actual metrics + if self.proc_rank == 0 and self.logger is not None: + self.logger.log_metrics(scalar_metrics, step_num=self.global_step) + self.logger.save() + + def add_tqdm_metrics(self, metrics): + for k, v in metrics.items(): + if type(v) is torch.Tensor: + v = v.item() + + self.tqdm_metrics[k] = v + + def metrics_to_scalars(self, metrics): + new_metrics = {} + for k, v in metrics.items(): + if isinstance(v, torch.Tensor): + v = v.item() + + if type(v) is dict: + v = self.metrics_to_scalars(v) + + new_metrics[k] = v + + return new_metrics + + def process_output(self, output, train=False): + """ + Reduces output according to the training mode. + Separates loss from logging and tqdm metrics + :param output: + :return: + """ + # --------------- + # EXTRACT CALLBACK KEYS + # --------------- + # all keys not progress_bar or log are candidates for callbacks + callback_metrics = {} + for k, v in output.items(): + if k not in ['progress_bar', 'log']: + callback_metrics[k] = v + + if train and (self.use_dp or self.use_ddp2): + nb_gpus = self.num_gpus + callback_metrics = self.reduce_distributed_output(callback_metrics, nb_gpus) + + for k, v in callback_metrics.items(): + callback_metrics[k] = v.item() + + # --------------- + # EXTRACT PROGRESS BAR KEYS + # --------------- + try: + progress_output = output['progress_bar'] + + # reduce progress metrics for tqdm when using dp + if train and (self.use_dp or self.use_ddp2): + nb_gpus = self.num_gpus + progress_output = self.reduce_distributed_output(progress_output, nb_gpus) + + progress_bar_metrics = progress_output + except Exception: + progress_bar_metrics = {} + + # --------------- + # EXTRACT LOGGING KEYS + # --------------- + # extract metrics to log to experiment + try: + log_output = output['log'] + + # reduce progress metrics for tqdm when using dp + if train and (self.use_dp or self.use_ddp2): + nb_gpus = self.num_gpus + log_output = self.reduce_distributed_output(log_output, nb_gpus) + + log_metrics = log_output + except Exception: + log_metrics = {} + + # --------------- + # EXTRACT LOSS + # --------------- + # if output dict doesn't have the keyword loss + # then assume the output=loss if scalar + loss = None + if train: + try: + loss = output['loss'] + except Exception: + if type(output) is torch.Tensor: + loss = output + else: + raise RuntimeError( + 'No `loss` value in the dictionary returned from `model.training_step()`.' + ) + + # when using dp need to reduce the loss + if self.use_dp or self.use_ddp2: + loss = self.reduce_distributed_output(loss, self.num_gpus) + + # use every metric passed in as a candidate for callback + callback_metrics.update(progress_bar_metrics) + callback_metrics.update(log_metrics) + + # convert tensors to numpy + for k, v in callback_metrics.items(): + if isinstance(v, torch.Tensor): + callback_metrics[k] = v.item() + + return loss, progress_bar_metrics, log_metrics, callback_metrics + + def reduce_distributed_output(self, output, nb_gpus): + if nb_gpus <= 1: + return output + + # when using DP, we get one output per gpu + # average outputs and return + if type(output) is torch.Tensor: + return output.mean() + + for k, v in output.items(): + # recurse on nested dics + if isinstance(output[k], dict): + output[k] = self.reduce_distributed_output(output[k], nb_gpus) + + # reduce only metrics that have the same nb of gpus + elif output[k].size(0) == nb_gpus: + reduced = torch.mean(output[k]) + output[k] = reduced + return output diff --git a/pytorch_lightning/trainer/model_hooks_mixin.py b/pytorch_lightning/trainer/model_hooks_mixin.py new file mode 100644 index 00000000..537d342d --- /dev/null +++ b/pytorch_lightning/trainer/model_hooks_mixin.py @@ -0,0 +1,17 @@ +from pytorch_lightning.root_module.root_module import LightningModule + + +class TrainerModelHooksMixin(object): + + def is_function_implemented(self, f_name): + model = self.get_model() + f_op = getattr(model, f_name, None) + return callable(f_op) + + def is_overriden(self, f_name): + model = self.get_model() + super_object = LightningModule + + # when code pointers are different, it was overriden + is_overriden = getattr(model, f_name).__code__ is not getattr(super_object, f_name).__code__ + return is_overriden diff --git a/pytorch_lightning/trainer/train_loop_mixin.py b/pytorch_lightning/trainer/train_loop_mixin.py new file mode 100644 index 00000000..92bb65e1 --- /dev/null +++ b/pytorch_lightning/trainer/train_loop_mixin.py @@ -0,0 +1,266 @@ +import numpy as np + +try: + from apex import amp + APEX_AVAILABLE = True +except ImportError: + APEX_AVAILABLE = False + + +class TrainerTrainLoopMixin(object): + + def train(self): + # run all epochs + for epoch_nb in range(self.current_epoch, self.max_nb_epochs): + # set seed for distributed sampler (enables shuffling for each epoch) + if self.use_ddp and hasattr(self.get_train_dataloader().sampler, 'set_epoch'): + self.get_train_dataloader().sampler.set_epoch(epoch_nb) + + # get model + model = self.get_model() + + # update training progress in trainer and model + model.current_epoch = epoch_nb + self.current_epoch = epoch_nb + self.total_batches = self.nb_training_batches + self.nb_val_batches + self.batch_loss_value = 0 # accumulated grads + + # limit the number of batches to 1 in fast_dev_run + if self.fast_dev_run: + self.total_batches = 1 + + # init progress_bar when requested + if self.show_progress_bar: + self.progress_bar.reset(self.total_batches) + + # changing gradient according accumulation_scheduler + self.accumulation_scheduler.on_epoch_begin(epoch_nb, self) + + # ----------------- + # RUN TNG EPOCH + # ----------------- + self.run_training_epoch() + + # update LR schedulers + if self.lr_schedulers is not None: + for lr_scheduler in self.lr_schedulers: + lr_scheduler.step(self.current_epoch) + + # early stopping + met_min_epochs = epoch_nb > self.min_nb_epochs + if self.enable_early_stop and (met_min_epochs or self.fast_dev_run): + should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb, + logs=self.callback_metrics) + # stop training + stop = should_stop and met_min_epochs + if stop: + return + + if self.logger is not None: + self.logger.finalize("success") + + def run_training_epoch(self): + # before epoch hook + if self.is_function_implemented('on_epoch_start'): + model = self.get_model() + model.on_epoch_start() + + # run epoch + for batch_nb, batch in enumerate(self.get_train_dataloader()): + self.batch_nb = batch_nb + self.global_step += 1 + + model = self.get_model() + model.global_step = self.global_step + + # stop when the flag is changed or we've gone past the amount + # requested in the batches + self.total_batch_nb += 1 + met_batch_limit = batch_nb >= self.nb_training_batches + if met_batch_limit: + break + + # --------------- + # RUN TRAIN STEP + # --------------- + output = self.run_training_batch(batch, batch_nb) + batch_result, grad_norm_dic, batch_step_metrics = output + early_stop_epoch = batch_result == -1 + + # --------------- + # 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 + should_check_val = ((is_val_check_batch or early_stop_epoch) and can_check_epoch) + + # fast_dev_run always forces val checking after train batch + if self.fast_dev_run or should_check_val: + self.run_evaluation(test=self.testing) + + # when logs should be saved + should_save_log = (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch + if should_save_log or self.fast_dev_run: + if self.proc_rank == 0 and self.logger is not None: + self.logger.save() + + # when metrics should be logged + should_log_metrics = batch_nb % self.row_log_interval == 0 or early_stop_epoch + if should_log_metrics or self.fast_dev_run: + + # logs user requested information to logger + self.log_metrics(batch_step_metrics, grad_norm_dic) + + # end epoch early + if early_stop_epoch or self.fast_dev_run: + break + + # epoch end hook + if self.is_function_implemented('on_epoch_end'): + model = self.get_model() + model.on_epoch_end() + + def run_training_batch(self, batch, batch_nb): + # track grad norms + grad_norm_dic = {} + + # track all metrics for callbacks + all_callback_metrics = [] + + # track metrics to log + all_log_metrics = [] + + if batch is None: + return 0, grad_norm_dic + + # hook + if self.is_function_implemented('on_batch_start'): + model_ref = self.get_model() + response = model_ref.on_batch_start(batch) + + if response == -1: + return -1, grad_norm_dic + + if self.show_progress_bar: + self.progress_bar.update(1) + + # call training_step once per optimizer + for opt_idx, optimizer in enumerate(self.optimizers): + + # wrap the forward step in a closure so second order methods work + def optimizer_closure(): + # forward pass + output = self.training_forward(batch, batch_nb, opt_idx) + closure_loss, progress_bar_metrics, log_metrics, callback_metrics = output + + # track metrics for callbacks + all_callback_metrics.append(callback_metrics) + + # track progress bar metrics + self.add_tqdm_metrics(progress_bar_metrics) + all_log_metrics.append(log_metrics) + + # accumulate loss + # (if accumulate_grad_batches = 1 no effect) + closure_loss = closure_loss / self.accumulate_grad_batches + + # backward pass + if self.use_amp: + with amp.scale_loss(closure_loss, optimizer) as scaled_loss: + scaled_loss.backward() + else: + closure_loss.backward() + + # insert after step hook + if self.is_function_implemented('on_after_backward'): + model_ref = self.get_model() + model_ref.on_after_backward() + + return closure_loss + + # calculate loss + loss = optimizer_closure() + + # nan grads + if self.print_nan_grads: + self.print_nan_gradients() + + # track total loss for logging (avoid mem leaks) + self.batch_loss_value += loss.item() + + # gradient update with accumulated gradients + if (self.batch_nb + 1) % self.accumulate_grad_batches == 0: + + # track gradient norms when requested + if batch_nb % self.row_log_interval == 0: + if self.track_grad_norm > 0: + model = self.get_model() + grad_norm_dic = model.grad_norm(self.track_grad_norm) + + # clip gradients + self.clip_gradients() + + # calls .step(), .zero_grad() + # override function to modify this behavior + model = self.get_model() + model.optimizer_step(self.current_epoch, batch_nb, + optimizer, opt_idx, optimizer_closure) + + # calculate running loss for display + self.running_loss.append(self.batch_loss_value) + self.batch_loss_value = 0 + self.avg_loss = np.mean(self.running_loss[-100:]) + + # update progress bar + if self.show_progress_bar: + # add model specific metrics + tqdm_metrics = self.training_tqdm_dict + self.progress_bar.set_postfix(**tqdm_metrics) + + # activate batch end hook + if self.is_function_implemented('on_batch_end'): + model = self.get_model() + model.on_batch_end() + + # collapse all metrics into one dict + all_log_metrics = {k: v for d in all_log_metrics for k, v in d.items()} + + # track all metrics for callbacks + self.callback_metrics = {k: v for d in all_callback_metrics for k, v in d.items()} + + return 0, grad_norm_dic, all_log_metrics + + def training_forward(self, batch, batch_nb, opt_idx): + """ + Handle forward for each training case (distributed, single gpu, etc...) + :param batch: + :param batch_nb: + :return: + """ + # --------------- + # FORWARD + # --------------- + # enable not needing to add opt_idx to training_step + args = [batch, batch_nb] + if len(self.optimizers) > 1: + args.append(opt_idx) + + if self.use_ddp or self.use_ddp2: + output = self.model(*args) + elif self.use_dp: + output = self.model(*args) + elif self.single_gpu: + gpu_id = 0 + if type(self.data_parallel_device_ids) is list: + gpu_id = self.data_parallel_device_ids[0] + batch = self.transfer_batch_to_gpu(batch, gpu_id) + args[0] = batch + output = self.model.training_step(*args) + + else: + output = self.model.training_step(*args) + + # format and reduce outputs accordingly + output = self.process_output(output, train=True) + loss, progress_bar_metrics, log_metrics, callback_metrics = output + return loss, progress_bar_metrics, log_metrics, callback_metrics diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py index d51a6de8..844b7ea4 100644 --- a/pytorch_lightning/trainer/trainer.py +++ b/pytorch_lightning/trainer/trainer.py @@ -3,29 +3,28 @@ The trainer handles all the logic for running a val loop, training loop, distrib """ import os -import re import warnings -import numpy as np import tqdm import torch -from torch.utils.data.distributed import DistributedSampler import torch.multiprocessing as mp import torch.distributed as dist 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 TrainerIOMixin -from pytorch_lightning.pt_overrides.override_data_parallel import ( - LightningDistributedDataParallel, LightningDataParallel) -from pytorch_lightning.callbacks import GradientAccumulationScheduler, \ - ModelCheckpoint, EarlyStopping +from pytorch_lightning.trainer.ddp_mixin import TrainerDDPMixin +from pytorch_lightning.trainer.dp_mixin import TrainerDPMixin +from pytorch_lightning.trainer.amp_mixin import TrainerAMPMixin +from pytorch_lightning.trainer.data_loading_mixin import TrainerDataLoadingMixin +from pytorch_lightning.trainer.evaluation_loop_mixin import TrainerEvaluationLoopMixin +from pytorch_lightning.trainer.train_loop_mixin import TrainerTrainLoopMixin +from pytorch_lightning.trainer.logging_mixin import TrainerLoggingMixin +from pytorch_lightning.trainer.training_tricks_mixin import TrainerTrainingTricksMixin +from pytorch_lightning.trainer.callback_config_mixin import TrainerCallbackConfigMixin +from pytorch_lightning.trainer.model_hooks_mixin import TrainerModelHooksMixin + from pytorch_lightning.utilities.debugging import MisconfigurationException import pdb -from pytorch_lightning.trainer import ignored_warnings - try: from apex import amp @@ -34,28 +33,17 @@ except ImportError: APEX_AVAILABLE = False -def reduce_distributed_output(output, nb_gpus): - if nb_gpus <= 1: - return output - - # when using DP, we get one output per gpu - # average outputs and return - if type(output) is torch.Tensor: - return output.mean() - - for k, v in output.items(): - # recurse on nested dics - if isinstance(output[k], dict): - output[k] = reduce_distributed_output(output[k], nb_gpus) - - # reduce only metrics that have the same nb of gpus - elif output[k].size(0) == nb_gpus: - reduced = torch.mean(output[k]) - output[k] = reduced - return output - - -class Trainer(TrainerIOMixin): +class Trainer(TrainerIOMixin, + TrainerDDPMixin, + TrainerDPMixin, + TrainerDataLoadingMixin, + TrainerAMPMixin, + TrainerEvaluationLoopMixin, + TrainerTrainLoopMixin, + TrainerLoggingMixin, + TrainerTrainingTricksMixin, + TrainerCallbackConfigMixin, + TrainerModelHooksMixin): def __init__(self, logger=True, @@ -185,43 +173,14 @@ class Trainer(TrainerIOMixin): # configure early stop callback # creates a default one if none passed in 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 - 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 - elif logger is False: - self.logger = None - else: - self.logger = logger - self.logger.rank = 0 + self.configure_early_stopping(early_stop_callback, logger) # configure checkpoint callback self.checkpoint_callback = checkpoint_callback - self.weights_save_path = weights_save_path # accumulated grads - self.__configure_accumulated_gradients(accumulate_grad_batches) + self.configure_accumulated_gradients(accumulate_grad_batches) # allow int, string and gpu list self.data_parallel_device_ids = self.__parse_gpu_ids(gpus) @@ -233,16 +192,16 @@ class Trainer(TrainerIOMixin): self.use_dp = False self.single_gpu = False self.distributed_backend = distributed_backend - self.__set_distributed_mode(distributed_backend, nb_gpu_nodes) + self.set_distributed_mode(distributed_backend, nb_gpu_nodes) # init flags for SLURM+ddp to work self.proc_rank = 0 self.world_size = 1 self.node_rank = 0 - self.__configure_slurm_ddp(nb_gpu_nodes) + self.configure_slurm_ddp(nb_gpu_nodes) # nvidia setup - self.__set_nvidia_flags(self.is_slurm_managing_tasks, self.data_parallel_device_ids) + self.set_nvidia_flags(self.is_slurm_managing_tasks, self.data_parallel_device_ids) # can't init progress bar here because starting a new process # means the progress_bar won't survive pickling @@ -259,12 +218,12 @@ class Trainer(TrainerIOMixin): self.row_log_interval = row_log_interval # how much of the data to use - self.__determine_data_use_amount(train_percent_check, val_percent_check, - test_percent_check, overfit_pct) + self.determine_data_use_amount(train_percent_check, val_percent_check, + test_percent_check, overfit_pct) # 16 bit mixed precision training using apex self.amp_level = amp_level - self.__init_amp(use_amp) + self.init_amp(use_amp) @property def slurm_job_id(self): @@ -275,67 +234,6 @@ class Trainer(TrainerIOMixin): job_id = None return job_id - def __configure_checkpoint_callback(self): - """ - Weight path set in this priority: - Checkpoint_callback's path (if passed in). - User provided weights_saved_path - Otherwise use os.getcwd() - """ - if self.checkpoint_callback is True: - # init a default one - if isinstance(self.logger, TestTubeLogger): - ckpt_path = '{}/{}/version_{}/{}'.format( - self.default_save_path, - self.logger.experiment.name, - self.logger.experiment.version, - 'checkpoints') - else: - ckpt_path = self.default_save_path - - self.checkpoint_callback = ModelCheckpoint( - filepath=ckpt_path - ) - elif self.checkpoint_callback is False: - self.checkpoint_callback = None - - 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 weights_save_path is still none here, set to current working dir - if self.weights_save_path is None: - self.weights_save_path = self.default_save_path - - def __init_amp(self, use_amp): - self.use_amp = use_amp and APEX_AVAILABLE - if self.use_amp: - print('using 16bit precision') - - if use_amp and not APEX_AVAILABLE: # pragma: no cover - msg = """ - You set use_amp=True but do not have apex installed. - Install apex first using this guide and rerun with use_amp=True: - https://github.com/NVIDIA/apex#linux - - this run will NOT use 16 bit precision - """ - raise ModuleNotFoundError(msg) - - def __configure_accumulated_gradients(self, accumulate_grad_batches): - self.accumulate_grad_batches = None - - if isinstance(accumulate_grad_batches, dict): - self.accumulation_scheduler = GradientAccumulationScheduler(accumulate_grad_batches) - elif isinstance(accumulate_grad_batches, int): - schedule = {1: accumulate_grad_batches} - self.accumulation_scheduler = GradientAccumulationScheduler(schedule) - else: - raise TypeError("Gradient accumulation supports only int and dict types") - def __parse_gpu_ids(self, gpus): """ :param gpus: Int, string or list of ids @@ -383,133 +281,16 @@ class Trainer(TrainerIOMixin): m = 'gpus must be int, none or list of ints' raise MisconfigurationException(m) - def __set_distributed_mode(self, distributed_backend, nb_gpu_nodes): - # skip for CPU - if self.num_gpus == 0: - return - - # single GPU case - # in single gpu case we allow ddp so we can train on multiple - # nodes, 1 gpu per node - if self.num_gpus == 1: - self.single_gpu = True - - if distributed_backend is not None: - self.use_dp = distributed_backend == 'dp' - self.use_ddp = distributed_backend == 'ddp' - self.use_ddp2 = distributed_backend == 'ddp2' - - # disable single gpu when using ddp2 - if self.use_ddp2: - self.single_gpu = False - - # multiple GPU case - elif self.num_gpus > 1: - if distributed_backend is not None: - # DP, DDP case - self.use_dp = distributed_backend == 'dp' - self.use_ddp = distributed_backend == 'ddp' - self.use_ddp2 = distributed_backend == 'ddp2' - - elif distributed_backend is None: - m = 'When using multiple GPUs set ' \ - 'Trainer(distributed_backend=dp) (or ddp)' - raise MisconfigurationException(m) - - # use ddp automatically if nb_gpu_nodes > 1 - if nb_gpu_nodes > 1 and self.use_dp: # pragma: no cover - self.use_ddp = True - self.use_dp = False - w = 'DataParallel does not support nb_gpu_nodes > 1. ' \ - 'Switching to DistributedDataParallel for you. ' \ - 'To silence this warning set distributed_backend=ddp' - warnings.warn(w) - - print('gpu available: {}, used: {}'.format(torch.cuda.is_available(), self.on_gpu)) - - def __configure_slurm_ddp(self, nb_gpu_nodes): - self.is_slurm_managing_tasks = False - - # extract SLURM flag vars - # whenever we have the correct number of tasks, we let slurm manage processes - # otherwise we launch the required number of processes - if self.use_ddp: - self.nb_requested_gpus = self.num_gpus * nb_gpu_nodes - self.nb_slurm_tasks = 0 - try: - self.nb_slurm_tasks = int(os.environ['SLURM_NTASKS']) - self.is_slurm_managing_tasks = self.nb_slurm_tasks == self.nb_requested_gpus - - # in interactive mode we don't manage tasks - job_name = os.environ['SLURM_JOB_NAME'] - if job_name == 'bash': - self.is_slurm_managing_tasks = False - - except Exception: - # likely not on slurm, so set the slurm managed flag to false - self.is_slurm_managing_tasks = False - - # used for tests only, set this flag to simulate slurm managing a task - try: - should_fake = int(os.environ['FAKE_SLURM_MANAGING_TASKS']) - if should_fake: - self.is_slurm_managing_tasks = True - except Exception as e: - pass - - def __set_nvidia_flags(self, is_slurm_managing_tasks, data_parallel_device_ids): - if data_parallel_device_ids is None: - return - - # set the correct cuda visible devices (using pci order) - os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" - - # when slurm is managing the task it sets the visible devices - if not is_slurm_managing_tasks: - if type(data_parallel_device_ids) is int: - id_str = ','.join(str(x) for x in list(range(data_parallel_device_ids))) - os.environ["CUDA_VISIBLE_DEVICES"] = id_str - else: - gpu_str = ','.join([str(x) for x in data_parallel_device_ids]) - os.environ["CUDA_VISIBLE_DEVICES"] = gpu_str - - print(f'VISIBLE GPUS: {os.environ["CUDA_VISIBLE_DEVICES"]}') - @property def data_parallel(self): return self.use_dp or self.use_ddp or self.use_ddp2 - def __determine_data_use_amount(self, train_percent_check, val_percent_check, - test_percent_check, overfit_pct): - """ - Use less data for debugging purposes - """ - self.train_percent_check = train_percent_check - self.val_percent_check = val_percent_check - self.test_percent_check = test_percent_check - if overfit_pct > 0: - self.train_percent_check = overfit_pct - self.val_percent_check = overfit_pct - self.test_percent_check = overfit_pct - - def __get_model(self): - return self.model.module if self.data_parallel else self.model - - def __is_function_implemented(self, f_name): - model = self.__get_model() - f_op = getattr(model, f_name, None) - return callable(f_op) - - def __is_overriden(self, f_name): - model = self.__get_model() - super_object = LightningModule - - # when code pointers are different, it was overriden - is_overriden = getattr(model, f_name).__code__ is not getattr(super_object, f_name).__code__ - return is_overriden - @property - def __training_tqdm_dict(self): + def training_tqdm_dict(self): + """ + Read-only for tqdm metrics + :return: + """ tqdm_dict = { 'loss': '{0:.3f}'.format(self.avg_loss), 'epoch': '{}'.format(self.current_epoch), @@ -526,14 +307,6 @@ class Trainer(TrainerIOMixin): return tqdm_dict - @property - def training_tqdm_dict(self): - """ - Read-only for tqdm metrics - :return: - """ - return self.__training_tqdm_dict - @property def tng_tqdm_dic(self): """ @@ -544,232 +317,6 @@ class Trainer(TrainerIOMixin): DeprecationWarning) return self.training_tqdm_dict - def __layout_bookeeping(self): - - # determine number of training batches - self.nb_training_batches = len(self.get_train_dataloader()) - self.nb_training_batches = int(self.nb_training_batches * self.train_percent_check) - - # determine number of validation batches - # val datasets could be none, 1 or 2+ - if self.get_val_dataloaders() is not None: - self.nb_val_batches = sum(len(dataloader) for dataloader in self.get_val_dataloaders()) - self.nb_val_batches = int(self.nb_val_batches * self.val_percent_check) - self.nb_val_batches = max(1, self.nb_val_batches) - - # determine number of test batches - if self.get_test_dataloaders() is not None: - self.nb_test_batches = sum( - len(dataloader) for dataloader in self.get_test_dataloaders() - ) - 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_training_batches * self.val_check_interval) - self.val_check_batch = max(1, self.val_check_batch) - - def __add_tqdm_metrics(self, metrics): - for k, v in metrics.items(): - if type(v) is torch.Tensor: - v = v.item() - - self.tqdm_metrics[k] = v - - def __evaluation_forward(self, model, batch, batch_idx, dataloader_idx, test=False): - # make dataloader_idx arg in validation_step optional - args = [batch, batch_idx] - - if test and len(self.get_test_dataloaders()) > 1: - args.append(dataloader_idx) - - elif not test and len(self.get_val_dataloaders()) > 1: - args.append(dataloader_idx) - - # handle DP, DDP forward - if self.use_ddp or self.use_dp or self.use_ddp2: - output = model(*args) - return output - - # single GPU - if self.single_gpu: - # for single GPU put inputs on gpu manually - root_gpu = 0 - if type(self.data_parallel_device_ids) is list: - root_gpu = self.data_parallel_device_ids[0] - batch = self.transfer_batch_to_gpu(batch, root_gpu) - args[0] = batch - - # CPU - if test: - output = model.test_step(*args) - else: - output = model.validation_step(*args) - - return output - - def evaluate(self, model, dataloaders, max_batches, test=False): - """ - Run evaluation code - :param model: PT model - :param dataloaders: list of PT dataloaders - :param max_batches: Scalar - :param test: boolean - :return: - """ - # enable eval mode - model.zero_grad() - model.eval() - - # copy properties for forward overrides - self.__copy_trainer_model_properties(model) - - # disable gradients to save memory - torch.set_grad_enabled(False) - - # bookkeeping - outputs = [] - - # run training - for dataloader_idx, dataloader in enumerate(dataloaders): - dl_outputs = [] - for batch_idx, batch in enumerate(dataloader): - - if batch is None: # pragma: no cover - continue - - # stop short when on fast_dev_run (sets max_batch=1) - if batch_idx >= max_batches: - break - - # ----------------- - # RUN EVALUATION STEP - # ----------------- - output = self.__evaluation_forward(model, - batch, - batch_idx, - dataloader_idx, - test) - - # track outputs for collation - dl_outputs.append(output) - - # batch done - if self.show_progress_bar: - self.progress_bar.update(1) - outputs.append(dl_outputs) - - eval_results = {} - - # with a single dataloader don't pass an array - if len(dataloaders) == 1: - outputs = outputs[0] - - # give model a chance to do something with the outputs (and method defined) - 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() - - # enable gradients to save memory - torch.set_grad_enabled(True) - - return eval_results - - def get_dataloaders(self, model): - """ - Dataloaders are provided by the model - :param model: - :return: - """ - self.get_train_dataloader = model.train_dataloader - self.get_test_dataloaders = model.test_dataloader - self.get_val_dataloaders = model.val_dataloader - - # call warnings from proc zero only which triggers dataloaders - # if those have to download data it will only happen on proc 0 - if self.proc_rank == 0: - on_ddp = self.use_ddp or self.use_ddp2 - if on_ddp and not isinstance(self.get_train_dataloader().sampler, DistributedSampler): - msg = """ - 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) - - if on_ddp and self.get_val_dataloaders() is not None: - for dataloader in self.get_val_dataloaders(): - if not isinstance(dataloader.sampler, DistributedSampler): - msg = """ - Your val_dataloader(s) don't use DistributedSampler. - - 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 - - if on_ddp and self.get_test_dataloaders() is not None: - for dataloader in self.get_test_dataloaders(): - if not isinstance(dataloader.sampler, DistributedSampler): - msg = """ - Your test_dataloader(s) don't use DistributedSampler. - - 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 - - if self.use_ddp or self.use_ddp2: - # wait for all processes to catch up - dist.barrier() - - # load each dataloader - self.get_train_dataloader() - self.get_test_dataloaders() - self.get_val_dataloaders() - # ----------------------------- # MODEL TRAINING # ----------------------------- @@ -789,10 +336,10 @@ class Trainer(TrainerIOMixin): # 1 gpu or dp option triggers training using DP module # easier to avoid NCCL issues elif self.use_dp: - self.__dp_train(model) + self.dp_train(model) elif self.single_gpu: - self.__single_gpu_train(model) + self.single_gpu_train(model) # ON CPU else: @@ -805,7 +352,7 @@ class Trainer(TrainerIOMixin): # allow for lr schedulers as well self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers()) - self.__run_pretrain_routine(model) + self.run_pretrain_routine(model) # return 1 when finished # used for testing or when we need to know that training succeeded @@ -826,195 +373,7 @@ class Trainer(TrainerIOMixin): elif isinstance(optimizers, list) or isinstance(optimizers, tuple): return optimizers, [] - def __single_gpu_train(self, model): - # CHOOSE OPTIMIZER - # allow for lr schedulers as well - self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers()) - - model.cuda(self.root_gpu) - - if self.use_amp: - # An example - model, optimizers = amp.initialize( - model, self.optimizers, opt_level=self.amp_level, - ) - self.optimizers = optimizers - - self.__run_pretrain_routine(model) - - def __dp_train(self, model): - - # CHOOSE OPTIMIZER - # allow for lr schedulers as well - self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers()) - - model.cuda(self.root_gpu) - - # check for this bug (amp + dp + !01 doesn't work) - # https://github.com/NVIDIA/apex/issues/227 - if self.use_dp and self.use_amp: - m = f""" - Amp level {self.amp_level} with DataParallel is not supported. - See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227. - We recommend you switch to ddp if you want to use amp - """ - raise MisconfigurationException(m) - - # create list of device ids - device_ids = self.data_parallel_device_ids - if type(device_ids) is int: - device_ids = list(range(device_ids)) - - model = LightningDataParallel(model, device_ids=device_ids) - - self.__run_pretrain_routine(model) - - def __copy_trainer_model_properties(self, model): - if isinstance(model, LightningDataParallel): - ref_model = model.module - elif isinstance(model, LightningDistributedDataParallel): - ref_model = model.module - else: - ref_model = model - - for m in [model, ref_model]: - m.trainer = self - m.on_gpu = self.on_gpu - m.use_dp = self.use_dp - m.use_ddp2 = self.use_ddp2 - m.use_ddp = self.use_ddp - m.use_amp = self.use_amp - m.testing = self.testing - m.single_gpu = self.single_gpu - - def ddp_train(self, gpu_nb, model): - """ - Entry point into a DP thread - :param gpu_nb: - :param model: - :param cluster_obj: - :return: - """ - # node rank using relative slurm id - # otherwise default to node rank 0 - try: - node_id = os.environ['SLURM_NODEID'] - self.node_rank = int(node_id) - except Exception: - self.node_rank = 0 - - # show progressbar only on progress_rank 0 - self.show_progress_bar = self.show_progress_bar and self.node_rank == 0 and gpu_nb == 0 - - # determine which process we are and world size - if self.use_ddp: - self.proc_rank = self.node_rank * self.num_gpus + gpu_nb - self.world_size = self.nb_gpu_nodes * self.num_gpus - - elif self.use_ddp2: - self.proc_rank = self.node_rank - self.world_size = self.nb_gpu_nodes - - # let the exp know the rank to avoid overwriting logs - if self.logger is not None: - self.logger.rank = self.proc_rank - - # set up server using proc 0's ip address - # try to init for 20 times at max in case ports are taken - # where to store ip_table - self.__init_tcp_connection() - - # CHOOSE OPTIMIZER - # allow for lr schedulers as well - self.optimizers, self.lr_schedulers = self.init_optimizers(model.configure_optimizers()) - - # MODEL - # copy model to each gpu - if self.distributed_backend == 'ddp': - torch.cuda.set_device(gpu_nb) - model.cuda(gpu_nb) - - # set model properties before going into wrapper - self.__copy_trainer_model_properties(model) - - # override root GPU - self.root_gpu = gpu_nb - - # AMP - # run through amp wrapper before going to distributed DP - if self.use_amp: - # An example - model, optimizers = amp.initialize( - model, self.optimizers, opt_level=self.amp_level, - ) - self.optimizers = optimizers - - # DDP2 uses all GPUs on the machine - if self.distributed_backend == 'ddp': - device_ids = [gpu_nb] - elif self.use_ddp2: - device_ids = None - - model = LightningDistributedDataParallel( - model, - device_ids=device_ids, - find_unused_parameters=True - ) - - # continue training routine - self.__run_pretrain_routine(model) - - def __init_tcp_connection(self): - """ - Connect all procs in the world using the env:// init - Use the first node as the root address - :param port: - :param tries: - :return: - """ - - # use slurm job id for the port number - # guarantees unique ports across jobs from same grid search - try: - # use the last 4 numbers in the job id as the id - default_port = os.environ['SLURM_JOB_ID'] - default_port = default_port[-4:] - - # all ports should be in the 10k+ range - default_port = int(default_port) + 15000 - - except Exception as e: - default_port = 12910 - - # if user gave a port number, use that one instead - try: - default_port = os.environ['MASTER_PORT'] - except Exception: - os.environ['MASTER_PORT'] = str(default_port) - - # figure out the root node addr - try: - root_node = os.environ['SLURM_NODELIST'].split(' ')[0] - except Exception: - root_node = '127.0.0.2' - - root_node = self.resolve_root_node_address(root_node) - os.environ['MASTER_ADDR'] = root_node - dist.init_process_group("nccl", rank=self.proc_rank, world_size=self.world_size) - - def resolve_root_node_address(self, root_node): - if '[' in root_node: - name = root_node.split('[')[0] - number = root_node.split(',')[0] - if '-' in number: - number = number.split('-')[0] - - number = re.sub('[^0-9]', '', number) - root_node = name + number - - return root_node - - def __run_pretrain_routine(self, model): + def run_pretrain_routine(self, model): """ Sanity check a few things before starting actual training :param model: @@ -1028,7 +387,7 @@ class Trainer(TrainerIOMixin): ref_model.trainer = self # set local properties on the model - self.__copy_trainer_model_properties(ref_model) + self.copy_trainer_model_properties(ref_model) # link up experiment object if self.logger is not None: @@ -1044,7 +403,7 @@ class Trainer(TrainerIOMixin): dist.barrier() # set up checkpoint callback - self.__configure_checkpoint_callback() + self.configure_checkpoint_callback() # register auto-resubmit when on SLURM self.register_slurm_signal_handlers() @@ -1053,7 +412,7 @@ class Trainer(TrainerIOMixin): self.get_dataloaders(ref_model) # init training constants - self.__layout_bookeeping() + self.layout_bookeeping() # print model summary if self.proc_rank == 0 and self.weights_summary is not None: @@ -1076,7 +435,7 @@ class Trainer(TrainerIOMixin): # when testing requested only run test and return if self.testing: - self.__run_evaluation(test=True) + self.run_evaluation(test=True) return # run tiny validation (if validation defined) @@ -1089,504 +448,12 @@ class Trainer(TrainerIOMixin): self.evaluate(model, self.get_val_dataloaders(), self.nb_sanity_val_steps, self.testing) - # --------------------------- # CORE TRAINING LOOP - # --------------------------- - self.__train() - - def __train(self): - # run all epochs - for epoch_nb in range(self.current_epoch, self.max_nb_epochs): - # set seed for distributed sampler (enables shuffling for each epoch) - if self.use_ddp and hasattr(self.get_train_dataloader().sampler, 'set_epoch'): - self.get_train_dataloader().sampler.set_epoch(epoch_nb) - - # get model - model = self.__get_model() - - # update training progress in trainer and model - model.current_epoch = epoch_nb - self.current_epoch = epoch_nb - self.total_batches = self.nb_training_batches + self.nb_val_batches - self.batch_loss_value = 0 # accumulated grads - - # limit the number of batches to 1 in fast_dev_run - if self.fast_dev_run: - self.total_batches = 1 - - # init progress_bar when requested - if self.show_progress_bar: - self.progress_bar.reset(self.total_batches) - - # changing gradient according accumulation_scheduler - self.accumulation_scheduler.on_epoch_begin(epoch_nb, self) - - # ----------------- - # RUN TNG EPOCH - # ----------------- - self.run_training_epoch() - - # update LR schedulers - if self.lr_schedulers is not None: - for lr_scheduler in self.lr_schedulers: - lr_scheduler.step(self.current_epoch) - - # early stopping - met_min_epochs = epoch_nb > self.min_nb_epochs - if self.enable_early_stop and (met_min_epochs or self.fast_dev_run): - should_stop = self.early_stop_callback.on_epoch_end(epoch=epoch_nb, - logs=self.callback_metrics) - # stop training - stop = should_stop and met_min_epochs - if stop: - return - - if self.logger is not None: - self.logger.finalize("success") - - def run_training_epoch(self): - # before epoch hook - if self.__is_function_implemented('on_epoch_start'): - model = self.__get_model() - model.on_epoch_start() - - # run epoch - for batch_nb, batch in enumerate(self.get_train_dataloader()): - self.batch_nb = batch_nb - self.global_step += 1 - - model = self.__get_model() - model.global_step = self.global_step - - # stop when the flag is changed or we've gone past the amount - # requested in the batches - self.total_batch_nb += 1 - met_batch_limit = batch_nb >= self.nb_training_batches - if met_batch_limit: - break - - # --------------- - # RUN TRAIN STEP - # --------------- - output = self.__run_training_batch(batch, batch_nb) - batch_result, grad_norm_dic, batch_step_metrics = output - early_stop_epoch = batch_result == -1 - - # --------------- - # 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 - should_check_val = ((is_val_check_batch or early_stop_epoch) and can_check_epoch) - - # fast_dev_run always forces val checking after train batch - if self.fast_dev_run or should_check_val: - self.__run_evaluation(test=self.testing) - - # when logs should be saved - should_save_log = (batch_nb + 1) % self.log_save_interval == 0 or early_stop_epoch - if should_save_log or self.fast_dev_run: - if self.proc_rank == 0 and self.logger is not None: - self.logger.save() - - # when metrics should be logged - should_log_metrics = batch_nb % self.row_log_interval == 0 or early_stop_epoch - if should_log_metrics or self.fast_dev_run: - - # logs user requested information to logger - self.__log_metrics(batch_step_metrics, grad_norm_dic) - - # end epoch early - if early_stop_epoch or self.fast_dev_run: - break - - # epoch end hook - if self.__is_function_implemented('on_epoch_end'): - model = self.__get_model() - model.on_epoch_end() - - def __log_metrics(self, metrics, grad_norm_dic): - """ - Logs the metric dict passed in - :param metrics: - :param grad_norm_dic: - :return: - """ - # added metrics by Lightning for convenience - metrics['epoch'] = self.current_epoch - - # add gpu memory - if self.on_gpu and self.log_gpu_memory: - mem_map = memory.get_memory_profile(self.log_gpu_memory) - metrics.update(mem_map) - - # add norms - metrics.update(grad_norm_dic) - - # turn all tensors to scalars - scalar_metrics = self.__metrics_to_scalars(metrics) - - # log actual metrics - if self.proc_rank == 0 and self.logger is not None: - self.logger.log_metrics(scalar_metrics, step_num=self.global_step) - self.logger.save() + self.train() def test(self, model=None): self.testing = True if model is not None: self.fit(model) else: - self.__run_evaluation(test=True) - - def __metrics_to_scalars(self, metrics): - new_metrics = {} - for k, v in metrics.items(): - if isinstance(v, torch.Tensor): - v = v.item() - - if type(v) is dict: - v = self.__metrics_to_scalars(v) - - new_metrics[k] = v - - return new_metrics - - def __log_vals_blacklist(self): - """avoid logging some vals lightning uses to maintain state""" - blacklist = {'batch_nb', 'v_nb', 'gpu'} - return blacklist - - def transfer_batch_to_gpu(self, batch, gpu_id): - # base case: object can be directly moved using `cuda` or `to` - if callable(getattr(batch, 'cuda', None)): - return batch.cuda(gpu_id) - - elif callable(getattr(batch, 'to', None)): - return batch.to(torch.device('cuda', gpu_id)) - - # when list - elif isinstance(batch, list): - for i, x in enumerate(batch): - batch[i] = self.transfer_batch_to_gpu(x, gpu_id) - return batch - - # when tuple - elif isinstance(batch, tuple): - batch = list(batch) - for i, x in enumerate(batch): - batch[i] = self.transfer_batch_to_gpu(x, gpu_id) - return tuple(batch) - - # when dict - elif isinstance(batch, dict): - for k, v in batch.items(): - batch[k] = self.transfer_batch_to_gpu(v, gpu_id) - - return batch - - # nothing matches, return the value as is without transform - return batch - - def __training_forward(self, batch, batch_nb, opt_idx): - """ - Handle forward for each training case (distributed, single gpu, etc...) - :param batch: - :param batch_nb: - :return: - """ - # --------------- - # FORWARD - # --------------- - # enable not needing to add opt_idx to training_step - args = [batch, batch_nb] - if len(self.optimizers) > 1: - args.append(opt_idx) - - if self.use_ddp or self.use_ddp2: - output = self.model(*args) - elif self.use_dp: - output = self.model(*args) - elif self.single_gpu: - gpu_id = 0 - if type(self.data_parallel_device_ids) is list: - gpu_id = self.data_parallel_device_ids[0] - batch = self.transfer_batch_to_gpu(batch, gpu_id) - args[0] = batch - output = self.model.training_step(*args) - - else: - output = self.model.training_step(*args) - - # format and reduce outputs accordingly - output = self.__process_output(output, train=True) - loss, progress_bar_metrics, log_metrics, callback_metrics = output - return loss, progress_bar_metrics, log_metrics, callback_metrics - - def __process_output(self, output, train=False): - """ - Reduces output according to the training mode. - Separates loss from logging and tqdm metrics - :param output: - :return: - """ - # --------------- - # EXTRACT CALLBACK KEYS - # --------------- - # all keys not progress_bar or log are candidates for callbacks - callback_metrics = {} - for k, v in output.items(): - if k not in ['progress_bar', 'log']: - callback_metrics[k] = v - - if train and (self.use_dp or self.use_ddp2): - nb_gpus = self.num_gpus - callback_metrics = reduce_distributed_output(callback_metrics, nb_gpus) - - for k, v in callback_metrics.items(): - callback_metrics[k] = v.item() - - # --------------- - # EXTRACT PROGRESS BAR KEYS - # --------------- - try: - progress_output = output['progress_bar'] - - # reduce progress metrics for tqdm when using dp - if train and (self.use_dp or self.use_ddp2): - nb_gpus = self.num_gpus - progress_output = reduce_distributed_output(progress_output, nb_gpus) - - progress_bar_metrics = progress_output - except Exception: - progress_bar_metrics = {} - - # --------------- - # EXTRACT LOGGING KEYS - # --------------- - # extract metrics to log to experiment - try: - log_output = output['log'] - - # reduce progress metrics for tqdm when using dp - if train and (self.use_dp or self.use_ddp2): - nb_gpus = self.num_gpus - log_output = reduce_distributed_output(log_output, nb_gpus) - - log_metrics = log_output - except Exception: - log_metrics = {} - - # --------------- - # EXTRACT LOSS - # --------------- - # if output dict doesn't have the keyword loss - # then assume the output=loss if scalar - loss = None - if train: - try: - loss = output['loss'] - except Exception: - if type(output) is torch.Tensor: - loss = output - else: - raise RuntimeError( - 'No `loss` value in the dictionary returned from `model.training_step()`.' - ) - - # when using dp need to reduce the loss - if self.use_dp or self.use_ddp2: - loss = reduce_distributed_output(loss, self.num_gpus) - - # use every metric passed in as a candidate for callback - callback_metrics.update(progress_bar_metrics) - callback_metrics.update(log_metrics) - - # convert tensors to numpy - for k, v in callback_metrics.items(): - if isinstance(v, torch.Tensor): - callback_metrics[k] = v.item() - - return loss, progress_bar_metrics, log_metrics, callback_metrics - - def __clip_gradients(self): - if self.gradient_clip_val > 0: - model = self.__get_model() - torch.nn.utils.clip_grad_norm_(model.parameters(), self.gradient_clip_val) - - def __print_nan_grads(self): - model = self.__get_model() - for param in model.parameters(): - if torch.isnan(param.grad.float()).any(): - print(param, param.grad) - - def __run_training_batch(self, batch, batch_nb): - # track grad norms - grad_norm_dic = {} - - # track all metrics for callbacks - all_callback_metrics = [] - - # track metrics to log - all_log_metrics = [] - - if batch is None: - return 0, grad_norm_dic - - # hook - if self.__is_function_implemented('on_batch_start'): - model_ref = self.__get_model() - response = model_ref.on_batch_start(batch) - - if response == -1: - return -1, grad_norm_dic - - if self.show_progress_bar: - self.progress_bar.update(1) - - # call training_step once per optimizer - for opt_idx, optimizer in enumerate(self.optimizers): - - # wrap the forward step in a closure so second order methods work - def optimizer_closure(): - # forward pass - output = self.__training_forward(batch, batch_nb, opt_idx) - closure_loss, progress_bar_metrics, log_metrics, callback_metrics = output - - # track metrics for callbacks - all_callback_metrics.append(callback_metrics) - - # track progress bar metrics - self.__add_tqdm_metrics(progress_bar_metrics) - all_log_metrics.append(log_metrics) - - # accumulate loss - # (if accumulate_grad_batches = 1 no effect) - closure_loss = closure_loss / self.accumulate_grad_batches - - # backward pass - if self.use_amp: - with amp.scale_loss(closure_loss, optimizer) as scaled_loss: - scaled_loss.backward() - else: - closure_loss.backward() - - # insert after step hook - if self.__is_function_implemented('on_after_backward'): - model_ref = self.__get_model() - model_ref.on_after_backward() - - return closure_loss - - # calculate loss - loss = optimizer_closure() - - # nan grads - if self.print_nan_grads: - self.__print_nan_grads() - - # track total loss for logging (avoid mem leaks) - self.batch_loss_value += loss.item() - - # gradient update with accumulated gradients - if (self.batch_nb + 1) % self.accumulate_grad_batches == 0: - - # track gradient norms when requested - if batch_nb % self.row_log_interval == 0: - if self.track_grad_norm > 0: - model = self.__get_model() - grad_norm_dic = model.grad_norm(self.track_grad_norm) - - # clip gradients - self.__clip_gradients() - - # calls .step(), .zero_grad() - # override function to modify this behavior - model = self.__get_model() - model.optimizer_step(self.current_epoch, batch_nb, - optimizer, opt_idx, optimizer_closure) - - # calculate running loss for display - self.running_loss.append(self.batch_loss_value) - self.batch_loss_value = 0 - self.avg_loss = np.mean(self.running_loss[-100:]) - - # update progress bar - if self.show_progress_bar: - # add model specific metrics - tqdm_metrics = self.__training_tqdm_dict - self.progress_bar.set_postfix(**tqdm_metrics) - - # activate batch end hook - if self.__is_function_implemented('on_batch_end'): - model = self.__get_model() - model.on_batch_end() - - # collapse all metrics into one dict - all_log_metrics = {k: v for d in all_log_metrics for k, v in d.items()} - - # track all metrics for callbacks - self.callback_metrics = {k: v for d in all_callback_metrics for k, v in d.items()} - - return 0, grad_norm_dic, all_log_metrics - - 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 - # 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 - model = self.__get_model() - model.on_pre_performance_check() - - # select dataloaders - if test: - dataloaders = self.get_test_dataloaders() - max_batches = self.nb_test_batches - else: - # val - dataloaders = self.get_val_dataloaders() - max_batches = self.nb_val_batches - - # cap max batches to 1 when using fast_dev_run - if self.fast_dev_run: - max_batches = 1 - - # run evaluation - eval_results = self.evaluate(self.model, - dataloaders, - max_batches, - test) - _, prog_bar_metrics, log_metrics, callback_metrics = self.__process_output(eval_results) - - # add metrics to prog bar - self.__add_tqdm_metrics(prog_bar_metrics) - - # log metrics - self.__log_metrics(log_metrics, {}) - - # track metrics for callbacks - self.callback_metrics = callback_metrics - - # hook - model.on_post_performance_check() - - if self.show_progress_bar: - # add model specific metrics - tqdm_metrics = self.__training_tqdm_dict - self.progress_bar.set_postfix(**tqdm_metrics) - - # model checkpointing - if self.proc_rank == 0 and self.checkpoint_callback is not None and not test: - self.checkpoint_callback.on_epoch_end(epoch=self.current_epoch, - logs=self.callback_metrics) + self.run_evaluation(test=True) diff --git a/pytorch_lightning/trainer/trainer_io.py b/pytorch_lightning/trainer/trainer_io.py index 8852e7d8..65934837 100644 --- a/pytorch_lightning/trainer/trainer_io.py +++ b/pytorch_lightning/trainer/trainer_io.py @@ -12,7 +12,7 @@ from pytorch_lightning.pt_overrides.override_data_parallel import ( class TrainerIOMixin(object): - def __get_model(self): + def get_model(self): is_dp_module = isinstance(self.model, (LightningDistributedDataParallel, LightningDataParallel)) model = self.model.module if is_dp_module else self.model @@ -134,7 +134,7 @@ class TrainerIOMixin(object): checkpoint = torch.load(checkpoint_path, map_location=lambda storage, loc: storage) # load model state - model = self.__get_model() + model = self.get_model() # load the state_dict on the model automatically model.load_state_dict(checkpoint['state_dict']) @@ -173,7 +173,7 @@ class TrainerIOMixin(object): checkpoint['lr_schedulers'] = lr_schedulers # add the state_dict from the model - model = self.__get_model() + model = self.get_model() checkpoint['state_dict'] = model.state_dict() # give the model a chance to add a few things @@ -252,7 +252,7 @@ class TrainerIOMixin(object): filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, ckpt_number) # give model a chance to do something on hpc_save - model = self.__get_model() + model = self.get_model() checkpoint = self.dump_checkpoint() model.on_hpc_save(checkpoint) @@ -269,7 +269,7 @@ class TrainerIOMixin(object): checkpoint = torch.load(filepath, map_location=lambda storage, loc: storage) # load model state - model = self.__get_model() + model = self.get_model() # load the state_dict on the model automatically model.load_state_dict(checkpoint['state_dict']) diff --git a/pytorch_lightning/trainer/training_tricks_mixin.py b/pytorch_lightning/trainer/training_tricks_mixin.py new file mode 100644 index 00000000..1efc9a17 --- /dev/null +++ b/pytorch_lightning/trainer/training_tricks_mixin.py @@ -0,0 +1,28 @@ +import torch + +from pytorch_lightning.callbacks import GradientAccumulationScheduler + + +class TrainerTrainingTricksMixin(object): + + def clip_gradients(self): + if self.gradient_clip_val > 0: + model = self.get_model() + torch.nn.utils.clip_grad_norm_(model.parameters(), self.gradient_clip_val) + + def print_nan_gradients(self): + model = self.get_model() + for param in model.parameters(): + if torch.isnan(param.grad.float()).any(): + print(param, param.grad) + + def configure_accumulated_gradients(self, accumulate_grad_batches): + self.accumulate_grad_batches = None + + if isinstance(accumulate_grad_batches, dict): + self.accumulation_scheduler = GradientAccumulationScheduler(accumulate_grad_batches) + elif isinstance(accumulate_grad_batches, int): + schedule = {1: accumulate_grad_batches} + self.accumulation_scheduler = GradientAccumulationScheduler(schedule) + else: + raise TypeError("Gradient accumulation supports only int and dict types") diff --git a/tests/test_models.py b/tests/test_models.py index 24fca323..a500f224 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -25,11 +25,12 @@ from pytorch_lightning.callbacks import ( ) from pytorch_lightning.utilities.debugging import MisconfigurationException from pytorch_lightning.root_module import memory -from pytorch_lightning.trainer.trainer import reduce_distributed_output from pytorch_lightning.root_module import model_saving from pytorch_lightning.trainer import trainer_io from pytorch_lightning.logging import TestTubeLogger from pl_examples import LightningTemplateModel +from pytorch_lightning.trainer.logging_mixin import TrainerLoggingMixin + # generate a list of random seeds for each test RANDOM_FILE_PATHS = list(np.random.randint(12000, 19000, 1000)) @@ -996,14 +997,15 @@ def test_loading_meta_tags(): def test_dp_output_reduce(): + mixin = TrainerLoggingMixin() reset_seed() # test identity when we have a single gpu out = torch.rand(3, 1) - assert reduce_distributed_output(out, nb_gpus=1) is out + assert mixin.reduce_distributed_output(out, nb_gpus=1) is out # average when we have multiples - assert reduce_distributed_output(out, nb_gpus=2) == out.mean() + assert mixin.reduce_distributed_output(out, nb_gpus=2) == out.mean() # when we have a dict of vals out = { @@ -1012,7 +1014,7 @@ def test_dp_output_reduce(): 'c': out } } - reduced = reduce_distributed_output(out, nb_gpus=3) + reduced = mixin.reduce_distributed_output(out, nb_gpus=3) assert reduced['a'] == out['a'] assert reduced['b']['c'] == out['b']['c']