def on_save_checkpoint(self, checkpoint) -> None: results = self.trainer._results # simplify logic state_dict = results.state_dict(drop_value=False) # check device assert results['validation_step.v'].value.device.type == device assert state_dict['items']['validation_step.v'][ 'value'].device.type == device # sync fn should be kept assert results[ 'validation_step.v'].meta.sync.fn == self.trainer.training_type_plugin.reduce # sync fn dropped from the state dict assert 'fn' not in state_dict['items']['validation_step.v'][ 'meta']['_sync'] results.load_state_dict(state_dict) # check device after loading assert results['validation_step.v'].value.device.type == device # sync fn was preserved in the original result assert results[ 'validation_step.v'].meta.sync.fn == self.trainer.training_type_plugin.reduce # default sync fn new_results = ResultCollection(False, device) new_results.load_state_dict(state_dict, map_location='cpu') assert new_results['validation_step.v'].meta.sync.fn == _Sync.no_op # check map location assert new_results['validation_step.v'].value.device.type == 'cpu'
def __init__(self, min_steps: Optional[int] = 0, max_steps: int = -1) -> None: super().__init__() if max_steps is None: rank_zero_deprecation( "Setting `max_steps = None` is deprecated in v1.5 and will no longer be supported in v1.7." " Use `max_steps = -1` instead." ) max_steps = -1 elif max_steps < -1: raise MisconfigurationException( f"`max_steps` must be a non-negative integer or -1 (infinite steps). You passed in {max_steps}." ) self.min_steps = min_steps self.max_steps = max_steps self.global_step: int = 0 self.batch_progress = BatchProgress() self.scheduler_progress = SchedulerProgress() self.batch_loop: Optional[TrainingBatchLoop] = None self.val_loop: Optional["loops.EvaluationLoop"] = None self._results = ResultCollection(training=True) self._outputs: _OUTPUTS_TYPE = [] self._warning_cache = WarningCache() self._dataloader_iter: Optional[Iterator] = None # caches the loaded dataloader state until dataloader objects are available self._dataloader_state_dict: Dict[str, Any] = {}
def on_save_checkpoint(self, checkpoint) -> None: results = self.trainer._results # simplify logic state_dict = results.state_dict(drop_value=False) # check device assert results["validation_step.v"].value.device.type == device assert state_dict["items"]["validation_step.v"][ "value"].device.type == device # sync fn should be kept assert results[ "validation_step.v"].meta.sync.fn == self.trainer.training_type_plugin.reduce # sync fn dropped from the state dict assert "fn" not in state_dict["items"]["validation_step.v"][ "meta"]["_sync"] results.load_state_dict(state_dict) # check device after loading assert results["validation_step.v"].value.device.type == device # sync fn was preserved in the original result assert results[ "validation_step.v"].meta.sync.fn == self.trainer.training_type_plugin.reduce # default sync fn new_results = ResultCollection(False, device) new_results.load_state_dict(state_dict, map_location="cpu") assert new_results["validation_step.v"].meta.sync.fn is None # check map location assert new_results["validation_step.v"].value.device.type == "cpu"
def __init__(self): super().__init__() self.outputs: List[EPOCH_OUTPUT] = [] self.epoch_loop = EvaluationEpochLoop() self._results = ResultCollection(training=False) self._max_batches: Optional[Union[int, Sequence[int]]] = None self._has_run: bool = False
def __init__(self): super().__init__() self._max_batches: Optional[Union[int, Sequence[int]]] = None self.outputs = [] self.evaluation_loop = EvaluationEpochLoop() self._val_results = ResultCollection(training=False) self._test_results = ResultCollection(training=False)
def __init__(self, trainer: 'pl.Trainer'): self.trainer: 'pl.Trainer' = trainer self.outputs: EPOCH_OUTPUT = [] self.predictions: Optional[PredictionCollection] = None self.max_batches: Optional[List[Union[int, float]]] = None self.warning_cache = WarningCache() self.num_dataloaders: Optional[int] = None self._val_results = ResultCollection(training=False) self._test_results = ResultCollection(training=False)
def __init__(self, min_epochs: Optional[int] = None, max_epochs: Optional[int] = None, min_steps: Optional[int] = None, max_steps: Optional[int] = None): super().__init__() self.max_epochs = 1000 if (max_epochs is None and max_steps is None) else max_epochs self.min_epochs = 1 if (min_epochs is None and min_steps is None) else min_epochs self.training_loop = TrainingEpochLoop(min_steps, max_steps) self.results = ResultCollection(training=True)
def __init__(self, min_steps: int, max_steps: int): super().__init__() self.min_steps: int = min_steps self.max_steps: int = max_steps self.global_step: int = 0 # manually tracking which is the last batch is necessary for iterable dataset support self.is_last_batch: Optional[bool] = None self.batch_progress = Progress() self.scheduler_progress = SchedulerProgress() self.batch_loop: Optional[TrainingBatchLoop] = None self.val_loop: Optional["loops.EvaluationLoop"] = None self._results = ResultCollection(training=True) self._epoch_output: Optional[List[List[STEP_OUTPUT]]] = None
def __init__(self, min_steps: int, max_steps: int): super().__init__() self.min_steps: int = min_steps self.max_steps: int = max_steps self.global_step: int = 0 # the total batch index across all epochs self.total_batch_idx: int = 0 # the current batch index in the loop that runs over the dataloader(s) self.iteration_count: int = 0 # the current split index when the batch gets split into chunks in truncated backprop through time self.split_idx: Optional[int] = None self._dataloader_idx: Optional[int] = None self._should_stop: bool = False self.is_last_batch: Optional[bool] = None self.batches_seen: int = 0 self.warning_cache: WarningCache = WarningCache() self.epoch_output: Optional[List[List[STEP_OUTPUT]]] = None self.batch_loop: Optional[TrainingBatchLoop] = None self._results = ResultCollection(training=True)
def __init__(self, min_steps: int, max_steps: int): super().__init__() self.min_steps: int = min_steps self.max_steps: int = max_steps self.global_step: int = 0 # the total batch index across all epochs self.total_batch_idx: int = 0 self.is_last_batch: Optional[bool] = None self.batch_progress = Progress() self.scheduler_progress = SchedulerProgress() self.batch_loop: Optional[TrainingBatchLoop] = None self.val_loop: Optional["loops.EvaluationLoop"] = None self._results = ResultCollection(training=True) self._dataloader_idx: Optional[int] = None self._warning_cache: WarningCache = WarningCache() self._epoch_output: Optional[List[List[STEP_OUTPUT]]] = None
def __init__( self, trainer, max_epochs: Optional[int], min_epochs: Optional[int], max_steps: Optional[int], min_steps: Optional[int], num_sanity_val_steps: int, ): self.trainer = trainer self.accumulated_loss = None self.warning_cache = WarningCache() self.running_loss = TensorRunningAccum(window_length=20) self._skip_backward = False self._optimizer_freq_cumsum = None self._hiddens = None self.global_step = 0 self.current_epoch = 0 self.trainer.should_stop = False # the total batch index across all epochs self.total_batch_idx = 0 # the current batch index in the loop that runs over the dataloader(s) self.batch_idx = 0 # the current split index when the batch gets split into chunks in truncated backprop through time self.split_idx = None self.trainer.num_training_batches = 0 self.trainer.train_dataloader = None # If neither max_epochs or max_steps is set, then use existing default of max_epochs = 1000 self.max_epochs = 1000 if (max_epochs is None and max_steps is None) else max_epochs # If neither min_epochs or min_steps is set, then use existing default of min_epochs = 1 self.min_epochs = 1 if (min_epochs is None and min_steps is None) else min_epochs self.max_steps = max_steps self.min_steps = min_steps if num_sanity_val_steps == -1: self.trainer.num_sanity_val_steps = float("inf") else: self.trainer.num_sanity_val_steps = num_sanity_val_steps self.results = ResultCollection(training=True)
def _ddp_test_fn(rank, worldsize): _setup_ddp(rank, worldsize) torch.tensor([1.0]) metric_a = DummyMetric() metric_b = DummyMetric() metric_c = DummyMetric() metric_a = metric_a.to(f"cuda:{rank}") metric_b = metric_b.to(f"cuda:{rank}") metric_c = metric_c.to(f"cuda:{rank}") result = ResultCollection(True, torch.device(f"cuda:{rank}")) for _ in range(3): cumulative_sum = 0 for i in range(5): metric_a(i) metric_b(i) metric_c(i) cumulative_sum += i result.log("h", "a", metric_a, on_step=True, on_epoch=True) result.log("h", "b", metric_b, on_step=False, on_epoch=True) result.log("h", "c", metric_c, on_step=True, on_epoch=False) batch_log = result.metrics(True)["log"] assert batch_log == {"a_step": i, "c": i} epoch_log = result.metrics(False)["log"] result.reset() # assert metric state reset to default values assert metric_a.x == metric_a._defaults["x"], (metric_a.x, metric_a._defaults["x"]) assert metric_b.x == metric_b._defaults["x"] assert metric_c.x == metric_c._defaults["x"] assert epoch_log == { "b": cumulative_sum * worldsize, "a_epoch": cumulative_sum * worldsize }
def test_metric_result_computed_check(): """Unittest ``_get_cache`` with multielement tensors.""" sync = _Sync() metadata = _Metadata("foo", "bar", on_epoch=True, enable_graph=True) metadata.sync = sync rm = ResultMetric(metadata, is_tensor=True) computed_value = torch.tensor([1, 2, 3]) rm._computed = computed_value cache = ResultCollection._get_cache(rm, on_step=False) # `enable_graph=True` so no detach, identity works assert cache is computed_value
def __init__(self, min_steps: int, max_steps: int): super().__init__() self.min_steps: int = min_steps if max_steps and max_steps < -1: raise MisconfigurationException( f"`max_steps` must be a positive integer or -1. You passed in {max_steps}." ) self.max_steps: int = max_steps self.global_step: int = 0 # manually tracking which is the last batch is necessary for iterable dataset support self.is_last_batch: Optional[bool] = None self.batch_progress = Progress() self.scheduler_progress = SchedulerProgress() self.batch_loop: Optional[TrainingBatchLoop] = None self.val_loop: Optional["loops.EvaluationLoop"] = None self._results = ResultCollection(training=True) self._epoch_output: Optional[List[List[STEP_OUTPUT]]] = None
def test_result_metric_integration(): metric_a = DummyMetric() metric_b = DummyMetric() metric_c = DummyMetric() result = ResultCollection(True, torch.device("cpu")) for _ in range(3): cumulative_sum = 0 for i in range(5): metric_a(i) metric_b(i) metric_c(i) cumulative_sum += i result.log('h', 'a', metric_a, on_step=True, on_epoch=True) result.log('h', 'b', metric_b, on_step=False, on_epoch=True) result.log('h', 'c', metric_c, on_step=True, on_epoch=False) batch_log = result.metrics(True)[MetricSource.LOG] assert batch_log == {"a_step": i, "c": i} epoch_log = result.metrics(False)[MetricSource.LOG] result.reset() # assert metric state reset to default values assert metric_a.x == metric_a._defaults['x'] assert metric_b.x == metric_b._defaults['x'] assert metric_c.x == metric_c._defaults['x'] assert epoch_log == {"b": cumulative_sum, "a_epoch": cumulative_sum} assert str(result) == ( "ResultCollection(True, cpu, {" "'h.a': ResultMetric('a', value=DummyMetric()), " "'h.b': ResultMetric('b', value=DummyMetric()), " "'h.c': ResultMetric('c', value=DummyMetric())" "})" )
def __init__(self, min_steps: int, max_steps: int): super().__init__() self.min_steps: int = min_steps self.max_steps: int = max_steps self.global_step: int = 0 # the total batch index across all epochs self.total_batch_idx: int = 0 # the current split index when the batch gets split into chunks in truncated backprop through time self.split_idx: Optional[int] = None # the number of batches seen this run, updates immediately after batch_loop.run() # TODO: replace by progress tracking self.batches_seen: int = 0 self.is_last_batch: Optional[bool] = None self.batch_progress = Progress() self.scheduler_progress = SchedulerProgress() self.batch_loop: Optional[TrainingBatchLoop] = None self.val_loop: Optional["loops.EvaluationLoop"] = None self._results = ResultCollection(training=True) self._dataloader_idx: Optional[int] = None self._warning_cache: WarningCache = WarningCache() self._epoch_output: Optional[List[List[STEP_OUTPUT]]] = None
def test_result_collection_simple_loop(): result = ResultCollection(True, torch.device("cpu")) current_fx_name = None batch_idx = None def lightning_log(fx, *args, **kwargs): nonlocal current_fx_name if current_fx_name != fx and batch_idx in (None, 0): result.reset(metrics=False, fx=fx) result.log(fx, *args, **kwargs) current_fx_name = fx lightning_log('a0', 'a', torch.tensor(0.), on_step=True, on_epoch=True) lightning_log('a1', 'a', torch.tensor(0.), on_step=True, on_epoch=True) for epoch in range(2): lightning_log('b0', 'a', torch.tensor(1.) + epoch, on_step=True, on_epoch=True) lightning_log('b1', 'a', torch.tensor(1.) + epoch, on_step=True, on_epoch=True) for batch_idx in range(2): lightning_log('c0', 'a', torch.tensor(2.) + epoch, on_step=True, on_epoch=True) lightning_log('c1', 'a', torch.tensor(2.) + epoch, on_step=True, on_epoch=True) lightning_log('c2', 'a', torch.tensor(2.) + epoch, on_step=True, on_epoch=True) batch_idx = None lightning_log('d0', 'a', torch.tensor(3.) + epoch, on_step=False, on_epoch=True) lightning_log('d1', 'a', torch.tensor(3.) + epoch, on_step=False, on_epoch=True) for k in ('a0.a', 'a1.a'): assert result[k].value == torch.tensor(0.), k assert result[k].cumulated_batch_size == torch.tensor(1.), k for k in ('b0.a', 'b1.a'): assert result[k].value == torch.tensor(1.) + epoch, k assert result[k].cumulated_batch_size == torch.tensor(1.), k for k in ('c0.a', 'c1.a', 'c2.a'): assert result[k].value == torch.tensor(4.) + epoch * 2, k assert result[k].cumulated_batch_size == torch.tensor(2.), k for k in ('d0.a', 'd1.a'): assert result[k].value == torch.tensor(3.) + epoch, k assert result[k].cumulated_batch_size == torch.tensor(1.), k
def test_result_collection_on_tensor_with_mean_reduction(): result_collection = ResultCollection(True, torch.device("cpu")) product = [(True, True), (False, True), (True, False), (False, False)] values = torch.arange(1, 10).float( ) # need to convert to float() due to precision issues using torch 1.4 batches = values * values for i, v in enumerate(values): for prog_bar in [False, True]: for logger in [False, True]: for on_step, on_epoch in product: name = "loss" if on_step: name += "_on_step" if on_epoch: name += "_on_epoch" if prog_bar: name += "_prog_bar" if logger: name += "_logger" result_collection.log( "training_step", name, v, on_step=on_step, on_epoch=on_epoch, batch_size=batches[i], prog_bar=prog_bar, logger=logger, ) total_value = sum(values * batches) total_batches = sum(batches) assert result_collection[ "training_step.loss_on_step_on_epoch"].value == total_value assert result_collection[ "training_step.loss_on_step_on_epoch"].cumulated_batch_size == total_batches batch_metrics = result_collection.metrics(True) max_ = max(values) assert batch_metrics[MetricSource.PBAR] == { "loss_on_step_on_epoch_prog_bar_step": max_, "loss_on_step_on_epoch_prog_bar_logger_step": max_, "loss_on_step_prog_bar": max_, "loss_on_step_prog_bar_logger": max_, } assert batch_metrics[MetricSource.LOG] == { "loss_on_step_on_epoch_logger_step": max_, "loss_on_step_logger": max_, "loss_on_step_on_epoch_prog_bar_logger_step": max_, "loss_on_step_prog_bar_logger": max_, } assert batch_metrics[MetricSource.CALLBACK] == { "loss_on_step": max_, "loss_on_step_logger": max_, "loss_on_step_on_epoch": max_, "loss_on_step_on_epoch_logger": max_, "loss_on_step_on_epoch_logger_step": max_, "loss_on_step_on_epoch_prog_bar": max_, "loss_on_step_on_epoch_prog_bar_logger": max_, "loss_on_step_on_epoch_prog_bar_logger_step": max_, "loss_on_step_on_epoch_prog_bar_step": max_, "loss_on_step_on_epoch_step": max_, "loss_on_step_prog_bar": max_, "loss_on_step_prog_bar_logger": max_, } epoch_metrics = result_collection.metrics(False) mean = total_value / total_batches assert epoch_metrics[MetricSource.PBAR] == { "loss_on_epoch_prog_bar": mean, "loss_on_epoch_prog_bar_logger": mean, "loss_on_step_on_epoch_prog_bar_epoch": mean, "loss_on_step_on_epoch_prog_bar_logger_epoch": mean, } assert epoch_metrics[MetricSource.LOG] == { "loss_on_epoch_logger": mean, "loss_on_epoch_prog_bar_logger": mean, "loss_on_step_on_epoch_logger_epoch": mean, "loss_on_step_on_epoch_prog_bar_logger_epoch": mean, } assert epoch_metrics[MetricSource.CALLBACK] == { "loss_on_epoch": mean, "loss_on_epoch_logger": mean, "loss_on_epoch_prog_bar": mean, "loss_on_epoch_prog_bar_logger": mean, "loss_on_step_on_epoch": mean, "loss_on_step_on_epoch_epoch": mean, "loss_on_step_on_epoch_logger": mean, "loss_on_step_on_epoch_logger_epoch": mean, "loss_on_step_on_epoch_prog_bar": mean, "loss_on_step_on_epoch_prog_bar_epoch": mean, "loss_on_step_on_epoch_prog_bar_logger": mean, "loss_on_step_on_epoch_prog_bar_logger_epoch": mean, }
def test_result_collection_extra_reference(): """Unit-test to check that the `extra` dict reference is properly set.""" rc = ResultCollection(True) assert rc.extra is rc["_extra"]
def test_result_collection_restoration(tmpdir): """" This test make sure metrics are properly reloaded on failure. """ result = ResultCollection(True, torch.device("cpu")) metric_a = DummyMetric() metric_b = DummyMetric() metric_c = DummyMetric() metric_d = DummyMetric() current_fx_name = None batch_idx = None def lightning_log(fx, *args, **kwargs): nonlocal current_fx_name if current_fx_name != fx and batch_idx in (None, 0): result.reset(metrics=False, fx=fx) result.log(fx, *args, **kwargs, sync_dist_fn=my_sync_dist) current_fx_name = fx for _ in range(2): cumulative_sum = 0 for i in range(3): a = metric_a(i) b = metric_b(i) c = metric_c(i) metric_d(i) cumulative_sum += i metric = metric_a if i < 1 else metric_d lightning_log('training_step', 'a', metric, on_step=True, on_epoch=True) lightning_log('training_step', 'b', metric_b, on_step=False, on_epoch=True) lightning_log('training_step', 'c', metric_c, on_step=True, on_epoch=False) lightning_log('training_step', 'a_1', a, on_step=True, on_epoch=True) lightning_log('training_step', 'b_1', b, on_step=False, on_epoch=True) lightning_log('training_step', 'c_1', { '1': c, '2': c }, on_step=True, on_epoch=False) batch_log = result.metrics(on_step=True)[MetricSource.LOG] assert set(batch_log) == {"a_step", "c", "a_1_step", "c_1"} assert set(batch_log['c_1']) == {'1', '2'} result_copy = deepcopy(result) new_result = ResultCollection(True, torch.device("cpu")) state_dict = result.state_dict() # check the sync fn was dropped assert 'fn' not in state_dict['items']['training_step.a']['meta'][ '_sync'] new_result.load_state_dict(state_dict) # should match assert result_copy == new_result # the sync fn has been kept assert result_copy['training_step.a'].meta.sync.fn == new_result[ 'training_step.a'].meta.sync.fn epoch_log = result.metrics(on_step=False)[MetricSource.LOG] epoch_log_copy = result_copy.metrics(on_step=False)[MetricSource.LOG] assert epoch_log == epoch_log_copy lightning_log('train_epoch_end', 'a', metric_a, on_step=False, on_epoch=True) epoch_log = result.metrics(on_step=False)[MetricSource.LOG] assert epoch_log == { 'a_1_epoch': 1, 'a_epoch': cumulative_sum, 'a': cumulative_sum, 'b': cumulative_sum, 'b_1': 1 } # make sure can be pickled pickle.loads(pickle.dumps(result)) # make sure can be torch.loaded filepath = str(tmpdir / 'result') torch.save(result, filepath) torch.load(filepath) # assert metric state reset to default values result.reset() assert metric_a.x == metric_a._defaults['x'] assert metric_b.x == metric_b._defaults['x'] assert metric_c.x == metric_c._defaults['x'] batch_idx = None
def test_result_collection_restoration(tmpdir): """This test make sure metrics are properly reloaded on failure.""" result = ResultCollection(True, torch.device("cpu")) metric_a = DummyMetric() metric_b = DummyMetric() metric_c = DummyMetric() metric_d = DummyMetric() current_fx_name = None batch_idx = None def lightning_log(fx, *args, **kwargs): nonlocal current_fx_name if current_fx_name != fx and batch_idx in (None, 0): result.reset(metrics=False, fx=fx) result.log(fx, *args, **kwargs, sync_dist_fn=my_sync_dist) current_fx_name = fx for epoch in range(2): cumulative_sum = 0 for i in range(3): a = metric_a(i) b = metric_b(i) c = metric_c(i) metric_d(i) cumulative_sum += i metric = metric_a if i < 1 else metric_d lightning_log("training_step", "a", metric, on_step=True, on_epoch=True, metric_attribute="metric") lightning_log("training_step", "b", metric_b, on_step=False, on_epoch=True, metric_attribute="metric_b") lightning_log("training_step", "c", metric_c, on_step=True, on_epoch=False, metric_attribute="metric_c") lightning_log("training_step", "a_1", a, on_step=True, on_epoch=True) lightning_log("training_step", "b_1", b, on_step=False, on_epoch=True) lightning_log("training_step", "c_1", { "1": c, "2": c }, on_step=True, on_epoch=False) batch_log = result.metrics(on_step=True)["log"] assert set(batch_log) == {"a_step", "c", "a_1_step", "c_1"} assert set(batch_log["c_1"]) == {"1", "2"} result_copy = deepcopy(result) new_result = ResultCollection(True, torch.device("cpu")) state_dict = result.state_dict() # check the sync fn was dropped assert "fn" not in state_dict["items"]["training_step.a"]["meta"][ "_sync"] assert not new_result.result_metrics assert len(result.result_metrics) == 7 + epoch > 0 new_result.load_state_dict(state_dict, metrics={ "metric": metric, "metric_b": metric_b, "metric_c": metric_c }) # should match assert result_copy == new_result # the sync fn has been kept assert result_copy["training_step.a"].meta.sync.fn == new_result[ "training_step.a"].meta.sync.fn epoch_log = result.metrics(on_step=False)["log"] epoch_log_copy = result_copy.metrics(on_step=False)["log"] assert epoch_log == epoch_log_copy lightning_log("train_epoch_end", "a", metric_a, on_step=False, on_epoch=True) epoch_log = result.metrics(on_step=False)["log"] assert epoch_log == { "a_1_epoch": 1, "a_epoch": cumulative_sum, "a": cumulative_sum, "b": cumulative_sum, "b_1": 1, } # make sure can be pickled pickle.loads(pickle.dumps(result)) # make sure can be torch.loaded filepath = str(tmpdir / "result") torch.save(result, filepath) torch.load(filepath) # assert metric state reset to default values result.reset() assert metric_a.x == metric_a._defaults["x"] assert metric_b.x == metric_b._defaults["x"] assert metric_c.x == metric_c._defaults["x"] batch_idx = None
def test_result_collection_simple_loop(): result = ResultCollection(True, torch.device("cpu")) current_fx_name = None batch_idx = None def lightning_log(fx, *args, **kwargs): nonlocal current_fx_name if current_fx_name != fx and batch_idx in (None, 0): result.reset(metrics=False, fx=fx) result.log(fx, *args, **kwargs) current_fx_name = fx lightning_log("a0", "a", torch.tensor(0.0), on_step=True, on_epoch=True) lightning_log("a1", "a", torch.tensor(0.0), on_step=True, on_epoch=True) for epoch in range(2): lightning_log("b0", "a", torch.tensor(1.0) + epoch, on_step=True, on_epoch=True) lightning_log("b1", "a", torch.tensor(1.0) + epoch, on_step=True, on_epoch=True) for batch_idx in range(2): lightning_log("c0", "a", torch.tensor(2.0) + epoch, on_step=True, on_epoch=True) lightning_log("c1", "a", torch.tensor(2.0) + epoch, on_step=True, on_epoch=True) lightning_log("c2", "a", torch.tensor(2.0) + epoch, on_step=True, on_epoch=True) batch_idx = None lightning_log("d0", "a", torch.tensor(3.0) + epoch, on_step=False, on_epoch=True) lightning_log("d1", "a", torch.tensor(3.0) + epoch, on_step=False, on_epoch=True) for k in ("a0.a", "a1.a"): assert result[k].value == torch.tensor(0.0), k assert result[k].cumulated_batch_size == torch.tensor(1.0), k for k in ("b0.a", "b1.a"): assert result[k].value == torch.tensor(1.0) + epoch, k assert result[k].cumulated_batch_size == torch.tensor(1.0), k for k in ("c0.a", "c1.a", "c2.a"): assert result[k].value == torch.tensor(4.0) + epoch * 2, k assert result[k].cumulated_batch_size == torch.tensor(2.0), k for k in ("d0.a", "d1.a"): assert result[k].value == torch.tensor(3.0) + epoch, k assert result[k].cumulated_batch_size == torch.tensor(1.0), k
def test_result_metric_integration(): metric_a = DummyMetric() metric_b = DummyMetric() metric_c = DummyMetric() result = ResultCollection(True, torch.device("cpu")) for _ in range(3): cumulative_sum = 0 for i in range(5): metric_a(i) metric_b(i) metric_c(i) cumulative_sum += i result.log("h", "a", metric_a, on_step=True, on_epoch=True) result.log("h", "b", metric_b, on_step=False, on_epoch=True) result.log("h", "c", metric_c, on_step=True, on_epoch=False) batch_log = result.metrics(True)["log"] assert batch_log == {"a_step": i, "c": i} epoch_log = result.metrics(False)["log"] result.reset() # assert metric state reset to default values assert metric_a.x == metric_a._defaults["x"] assert metric_b.x == metric_b._defaults["x"] assert metric_c.x == metric_c._defaults["x"] assert epoch_log == {"b": cumulative_sum, "a_epoch": cumulative_sum} result.minimize = torch.tensor(1.0) result.extra = {} assert str(result) == ("ResultCollection(" "minimize=1.0, " "{" "'h.a': ResultMetric('a', value=DummyMetric()), " "'h.b': ResultMetric('b', value=DummyMetric()), " "'h.c': ResultMetric('c', value=DummyMetric())" "})") assert repr(result) == ("{" "True, " "device(type='cpu'), " "minimize=tensor(1.), " "{'h.a': ResultMetric('a', value=DummyMetric()), " "'h.b': ResultMetric('b', value=DummyMetric()), " "'h.c': ResultMetric('c', value=DummyMetric()), " "'_extra': {}}" "}")
class FitLoop(Loop): """This Loop iterates over the epochs to run the training Args: min_epochs: The minimum number of epochs max_epochs: The maximum number of epochs min_steps: The minimum number of steps max_steps: The maximum number of epoch .. note:: If neither the minimum epochs nor steps are specified the minimum number of epochs is set to 1 and if neither the maximum steps nor epochs are specified, the maximum epochs are set to 1000. """ def __init__(self, min_epochs: Optional[int] = None, max_epochs: Optional[int] = None, min_steps: Optional[int] = None, max_steps: Optional[int] = None): super().__init__() self.max_epochs = 1000 if (max_epochs is None and max_steps is None) else max_epochs self.min_epochs = 1 if (min_epochs is None and min_steps is None) else min_epochs self.training_loop = TrainingEpochLoop(min_steps, max_steps) self.results = ResultCollection(training=True) @property def current_epoch(self) -> int: """Return the current epoch""" return self.iteration_count @current_epoch.setter def current_epoch(self, value: int) -> None: """Setter for the current epoch""" self.iteration_count = value @property def global_step(self) -> int: """Returns the global step""" return self.training_loop.global_step @global_step.setter def global_step(self, value: int) -> None: """Sets the global step (forwards to training_loop)""" self.training_loop.global_step = value @property def total_batch_idx(self) -> int: """Returns the total number of batches already run (across all epochs)""" return self.training_loop.total_batch_idx @property def batch_idx(self) -> int: """Returns the number of batches already run within this epoch""" return self.training_loop.iteration_count @property def split_idx(self) -> int: """Returns the index of the current batch split (within the current batch) for bptt""" return self.training_loop.split_idx @property def min_steps(self) -> int: # TODO(@justusschock): Why aren't we using the attribute in this class? """Returns the minimum numnber of steps to run""" return self.training_loop.min_steps @property def max_steps(self) -> int: """Returns the maximum number of steps to run""" return self.training_loop.max_steps @max_steps.setter def max_steps(self, value: int) -> None: """Sets the maximum number of steps (forwards to training_loop)""" # TODO(@awaelchli): This setter is required by debugging connector (fast dev run), should be avoided self.training_loop.max_steps = value @property def running_loss(self) -> TensorRunningAccum: """Returns the running loss""" return self.training_loop.batch_loop.running_loss @property def _skip_backward(self) -> bool: """ Determines whether the loop will skip backward during automatic optimization. """ return self.training_loop.batch_loop._skip_backward @_skip_backward.setter def _skip_backward(self, value: bool) -> None: """ Determines whether the loop will skip backward during automatic optimization. """ self.training_loop.batch_loop._skip_backward = value @property def done(self) -> bool: """Evaluates when to leave the loop. Returns True if trainer.should_stop was set (e.g. by early stopping) or if the maximum number of steps or epochs is reached. """ # TODO(@awaelchli): Move track steps inside training loop and move part of these condition inside training loop stop_steps = self.max_steps is not None and self.global_step >= self.max_steps stop_epochs = self.max_epochs is not None and self.current_epoch >= self.max_epochs should_stop = False if self.trainer.should_stop: # early stopping met_min_epochs = self.current_epoch >= self.min_epochs if self.min_epochs else True met_min_steps = self.global_step >= self.min_steps if self.min_steps else True if met_min_epochs and met_min_steps: should_stop = True else: log.info( 'Trainer was signaled to stop but required minimum epochs' f' ({self.min_epochs}) or minimum steps ({self.min_steps}) has' ' not been met. Training will continue...') self.trainer.should_stop = should_stop return stop_steps or should_stop or stop_epochs @property def skip(self) -> bool: """Whether we should skip the training and immediately return from the call to :meth:`run`.""" return self.done or self.trainer.num_training_batches == 0 def connect(self, trainer: 'pl.Trainer', *args: Any, **kwargs: Any) -> None: """Connects the loop with necessary arguments like the trainer""" # TODO(@justusschock): Do we want to forward *args and **kwargs to the inner loop here? # TODO(@justusschock): Can we make the trainer a weakref/proxy? void(*args, **kwargs) self.trainer = trainer self.training_loop.connect(trainer) def reset(self) -> None: """Resets the internal state of this loop""" def on_run_start(self) -> None: """Calls the ``on_train_start`` hook.""" self.results.to(device=self.trainer.lightning_module.device) self.trainer.call_hook("on_train_start") def on_advance_start(self) -> None: """Prepares the dataloader for training and calls the hooks ``on_epoch_start`` and ``on_train_epoch_start``""" model = self.trainer.lightning_module # reset train dataloader if self.current_epoch != 0 and self.trainer.reload_dataloaders_every_epoch: self.trainer.reset_train_dataloader(model) # TODO: specify the possible exception with suppress(Exception): # set seed for distributed sampler (enables shuffling for each epoch) self.trainer.train_dataloader.sampler.set_epoch(self.current_epoch) # changing gradient according accumulation_scheduler self.trainer.accumulation_scheduler.on_train_epoch_start( self.trainer, self.trainer.lightning_module) # stores accumulated grad fractions per batch self.training_loop.batch_loop.accumulated_loss = TensorRunningAccum( window_length=self.trainer.accumulate_grad_batches) def advance(self) -> None: """Runs one whole epoch.""" train_dataloader = self.trainer.accelerator.process_dataloader( self.trainer.train_dataloader) train_dataloader = self.trainer.data_connector.get_profiled_train_dataloader( train_dataloader) with self.trainer.profiler.profile("run_training_epoch"): # run train epoch epoch_output = self.training_loop.run(train_dataloader) if epoch_output is None: return # the global step is manually decreased here due to backwards compatibility with existing loggers # as they expect that the same step is used when logging epoch end metrics even when the batch loop has # finished. this means the attribute does not exactly track the number of optimizer steps applied. # TODO(@carmocca): deprecate and rename so users don't get confused self.global_step -= 1 # log epoch metrics self.trainer.logger_connector.update_train_epoch_metrics() self.global_step += 1 def on_advance_end(self) -> None: """Updates the LR schedulers and does some internal bookkeeping""" if self.training_loop.batches_seen == 0: return self.training_loop.update_lr_schedulers('epoch', update_plateau_schedulers=True) did_train_only = self.trainer.disable_validation or self.trainer.evaluation_loop.skip if did_train_only: self.global_step -= 1 self.check_checkpoint_callback(True) self.global_step += 1 def on_run_end(self) -> None: """Runs teardown logic and calls the ``on_train_end`` hook""" # NOTE: the iteration_count/current_epoch is already incremented # Lightning today does not increment the current epoch at the last epoch run in Trainer.fit # To simulate that current behavior, we decrement here. # TODO: must be fixed by https://github.com/PyTorchLightning/pytorch-lightning/issues/5007 self.current_epoch -= 1 # trigger checkpoint check. need to temporarily decrease the global step to avoid saving duplicates # when a checkpoint was saved at the last step self.training_loop.global_step -= 1 # TODO: see discussion/rework https://github.com/PyTorchLightning/pytorch-lightning/issues/7406 self.check_checkpoint_callback(should_update=True, is_last=True) self.training_loop.global_step += 1 # hook self.trainer.call_hook("on_train_end") # todo: TPU 8 cores hangs in flush with TensorBoard. Might do for all loggers. # It might be related to xla tensors blocked when moving the cpu # kill loggers if self.trainer.logger is not None: self.trainer.logger.finalize("success") # summarize profile results self.trainer.profiler.describe() # give accelerators a chance to finish self.trainer.accelerator.on_train_end() # reset bookkeeping self.trainer._running_stage = None def should_accumulate(self) -> bool: """Whether the gradients should be accumulated""" return self.training_loop.batch_loop.should_accumulate() def get_active_optimizers( self, batch_idx: Optional[int] = None) -> List[Tuple[int, Optimizer]]: """Generates a list of active optimizers""" return self.training_loop.batch_loop.get_active_optimizers(batch_idx) def check_checkpoint_callback(self, should_update: bool, is_last: bool = False): """Checks if checkpointing needs to be done""" # TODO: bake this logic into the ModelCheckpoint callback if should_update and self.trainer.checkpoint_connector.has_trained: callbacks = self.trainer.checkpoint_callbacks if is_last and any(cb.save_last and cb.verbose for cb in callbacks): rank_zero_info("Saving latest checkpoint...") model = self.trainer.lightning_module for cb in callbacks: cb.on_validation_end(self.trainer, model)
class TrainingEpochLoop(loops.Loop[_OUTPUTS_TYPE]): """Runs over all batches in a dataloader (one epoch). Args: min_steps: The minimum number of steps (batches) to process max_steps: The maximum number of steps (batches) to process """ def __init__(self, min_steps: Optional[int] = 0, max_steps: int = -1) -> None: super().__init__() if max_steps is None: rank_zero_deprecation( "Setting `max_steps = None` is deprecated in v1.5 and will no longer be supported in v1.7." " Use `max_steps = -1` instead." ) max_steps = -1 elif max_steps < -1: raise MisconfigurationException( f"`max_steps` must be a non-negative integer or -1 (infinite steps). You passed in {max_steps}." ) self.min_steps = min_steps self.max_steps = max_steps self.global_step: int = 0 self.batch_progress = BatchProgress() self.scheduler_progress = SchedulerProgress() self.batch_loop: Optional[TrainingBatchLoop] = None self.val_loop: Optional["loops.EvaluationLoop"] = None self._results = ResultCollection(training=True) self._outputs: _OUTPUTS_TYPE = [] self._warning_cache = WarningCache() self._dataloader_iter: Optional[Iterator] = None # caches the loaded dataloader state until dataloader objects are available self._dataloader_state_dict: Dict[str, Any] = {} @property def total_batch_idx(self) -> int: """Returns the current batch index (across epochs)""" # use `ready` instead of `completed` in case this is accessed after `completed` has been increased # but before the next `ready` increase return self.batch_progress.total.ready - 1 @property def batch_idx(self) -> int: """Returns the current batch index (within this epoch)""" # use `ready` instead of `completed` in case this is accessed after `completed` has been increased # but before the next `ready` increase return self.batch_progress.current.ready - 1 @property def _is_training_done(self) -> bool: max_steps_reached = _is_max_limit_reached(self.global_step, self.max_steps) return max_steps_reached or self._num_ready_batches_reached() @property def _is_validation_done(self) -> bool: # when we are restarting we want to check whether the val loop has finished return not self.restarting or self.val_loop.done @property def done(self) -> bool: """Returns whether the training should be stopped. The criteria are that the number of steps reached the max steps, the last batch is reached or the trainer signals to stop (e.g. by early stopping). """ return (self._is_training_done and self._is_validation_done) or self.trainer.should_stop def connect( self, batch_loop: TrainingBatchLoop = None, val_loop: Optional["loops.EvaluationLoop"] = None, ) -> None: """Optionally connect a custom batch or validation loop to this training epoch loop.""" if batch_loop is not None: self.batch_loop = batch_loop if val_loop is not None: self.val_loop = val_loop def reset(self) -> None: """Resets the internal state of the loop for a new run.""" assert self.batch_loop is not None assert self.batch_loop.optimizer_loop is not None if self.restarting: self.batch_progress.reset_on_restart() self.scheduler_progress.reset_on_restart() self.batch_loop.optimizer_loop.optim_progress.reset_on_restart() else: self.batch_progress.reset_on_run() self.scheduler_progress.reset_on_run() self.batch_loop.optimizer_loop.optim_progress.reset_on_run() self._outputs = [] def on_run_start(self, data_fetcher: AbstractDataFetcher, **kwargs: Any) -> None: # hook self.trainer.logger_connector.on_epoch_start() self.trainer.call_hook("on_epoch_start") self.trainer.call_hook("on_train_epoch_start") self.trainer.fit_loop.epoch_progress.increment_started() self._reload_dataloader_state_dict(data_fetcher) self._dataloader_iter = _update_dataloader_iter(data_fetcher, self.batch_idx + 1) def advance(self, *args: Any, **kwargs: Any) -> None: """Runs a single training batch. Args: dataloader_iter: the iterator over the dataloader producing the new batch Raises: StopIteration: When the epoch is canceled by the user returning -1 """ if self.restarting and self._should_check_val_fx(self.batch_idx, self.batch_progress.is_last_batch): # skip training and run validation in `on_advance_end` return batch_idx, (batch, self.batch_progress.is_last_batch) = next(self._dataloader_iter) if not self.trainer._data_connector.train_data_fetcher.store_on_device: with self.trainer.profiler.profile("training_batch_to_device"): batch = self.trainer.accelerator.batch_to_device(batch) self.batch_progress.increment_ready() # cache the batch size value to avoid extracting it again after the batch loop runs as the value will be # different if tbptt is enabled batch_size = self.trainer.logger_connector.on_batch_start(batch_idx, batch) if batch is None: self._warning_cache.warn("train_dataloader yielded None. If this was on purpose, ignore this warning...") batch_output = [] else: # hook response = self.trainer.call_hook("on_batch_start") if response == -1: self.batch_progress.increment_processed() raise StopIteration # TODO: Update this in v1.7 (deprecation: #9816) model_fx = self.trainer.lightning_module.on_train_batch_start extra_kwargs = ( {"dataloader_idx": 0} if callable(model_fx) and is_param_in_hook_signature(model_fx, "dataloader_idx", explicit=True) else {} ) # hook response = self.trainer.call_hook("on_train_batch_start", batch, batch_idx, **extra_kwargs) if response == -1: self.batch_progress.increment_processed() raise StopIteration self.batch_progress.increment_started() with self.trainer.profiler.profile("run_training_batch"): batch_output = self.batch_loop.run(batch, batch_idx) self.trainer._results.batch_size = batch_size self.batch_progress.increment_processed() # update non-plateau LR schedulers # update epoch-interval ones only when we are at the end of training epoch self.update_lr_schedulers("step", update_plateau_schedulers=False) if self._num_ready_batches_reached(): self.update_lr_schedulers("epoch", update_plateau_schedulers=False) batch_end_outputs = self._prepare_outputs_training_batch_end( batch_output, automatic=self.trainer.lightning_module.trainer.lightning_module.automatic_optimization, num_optimizers=len(self.trainer.optimizers), ) # TODO: Update this in v1.7 (deprecation: #9816) model_fx = self.trainer.lightning_module.on_train_batch_end extra_kwargs = ( {"dataloader_idx": 0} if callable(model_fx) and is_param_in_hook_signature(model_fx, "dataloader_idx", explicit=True) else {} ) self.trainer.call_hook("on_train_batch_end", batch_end_outputs, batch, batch_idx, **extra_kwargs) self.trainer.call_hook("on_batch_end") self.trainer.logger_connector.on_batch_end() self.batch_progress.increment_completed() if is_overridden("training_epoch_end", self.trainer.lightning_module): self._outputs.append(batch_output) # ----------------------------------------- # SAVE METRICS TO LOGGERS AND PROGRESS_BAR # ----------------------------------------- self.trainer.logger_connector.update_train_step_metrics() def on_advance_end(self): """Runs validation and Checkpointing if necessary. Raises: StopIteration: if :attr:`done` evaluates to ``True`` to finish this epoch """ # ----------------------------------------- # VALIDATE IF NEEDED + CHECKPOINT CALLBACK # ----------------------------------------- should_check_val = self._should_check_val_fx(self.batch_idx, self.batch_progress.is_last_batch) if should_check_val: self.trainer.validating = True self._run_validation() self.trainer.training = True # ----------------------------------------- # SAVE LOGGERS (ie: Tensorboard, etc...) # ----------------------------------------- self._save_loggers_on_train_batch_end() # update plateau LR scheduler after metrics are logged self.update_lr_schedulers("step", update_plateau_schedulers=True) if not self._should_accumulate(): # progress global step according to grads progress self.global_step += 1 # if training finished, try to exit in `on_run_end` instead as we should have enough time # TODO: @tchaton verify this assumption is True. if not self._is_training_done: # if fault tolerant is enabled and process has been notified, exit. self.trainer._exit_gracefully_on_signal() def on_run_end(self) -> None: """Calls the on_epoch_end hook. Returns: The output of each training step for each optimizer Raises: MisconfigurationException: ``train_epoch_end`` does not return ``None`` """ # inform logger the batch loop has finished self.trainer.logger_connector.epoch_end_reached() # get the model and call model.training_epoch_end model = self.trainer.lightning_module if is_overridden("training_epoch_end", model) and self._outputs: epoch_end_outputs = self._prepare_outputs_training_epoch_end( self._outputs, automatic=model.automatic_optimization, num_optimizers=len(self.trainer.optimizers), ) # run lightning module hook training_epoch_end # refresh the result for custom logging at the epoch level model._current_fx_name = "training_epoch_end" epoch_end_outputs = model.training_epoch_end(epoch_end_outputs) if epoch_end_outputs is not None: raise MisconfigurationException( "`training_epoch_end` expects a return of None. " "HINT: remove the return statement in `training_epoch_end`." ) # free memory self._outputs = [] self.trainer.fit_loop.epoch_progress.increment_processed() # call train epoch end hooks self.trainer.call_hook("on_train_epoch_end") self.trainer.call_hook("on_epoch_end") self.trainer.logger_connector.on_epoch_end() if self._num_ready_batches_reached(): self.update_lr_schedulers("epoch", update_plateau_schedulers=True) # if fault tolerant is enabled and process has been notified, exit. self.trainer._exit_gracefully_on_signal() def teardown(self) -> None: self._results.cpu() self.batch_loop.teardown() self.val_loop.teardown() def on_save_checkpoint(self) -> Dict: state_dict = super().on_save_checkpoint() if ( self.trainer.train_dataloader is None or self._num_completed_batches_reached() # did not finish # TODO: fault-tolerance requires a minimum number of batches so probably should be > 0 or self.batch_progress.current.ready == 0 # did not start ): return state_dict state_dict["dataloader_state_dict"] = self.trainer.train_dataloader.state_dict( has_completed=self._has_completed() ) return state_dict def on_load_checkpoint(self, state_dict: Dict) -> None: # cache the dataloader state dict until the dataloader objects are available self._dataloader_state_dict = state_dict.get("dataloader_state_dict") def _run_validation(self): # reload dataloaders self.val_loop._reload_evaluation_dataloaders() with torch.no_grad(): self.val_loop.run() def _accumulated_batches_reached(self) -> bool: """Determine if accumulation will be finished by the end of the current batch.""" return self.batch_progress.current.ready % self.trainer.accumulate_grad_batches == 0 def _num_ready_batches_reached(self) -> bool: """Checks if we are in the last batch or if there are more batches to follow.""" epoch_finished_on_ready = self.batch_progress.current.ready == self.trainer.num_training_batches return epoch_finished_on_ready or self.batch_progress.is_last_batch def _num_completed_batches_reached(self) -> bool: epoch_finished_on_completed = self.batch_progress.current.completed == self.trainer.num_training_batches dataloader_consumed_successfully = self.batch_progress.is_last_batch and self._has_completed() return epoch_finished_on_completed or dataloader_consumed_successfully def _has_completed(self) -> bool: return self.batch_progress.current.ready == self.batch_progress.current.completed def _should_accumulate(self) -> bool: """Checks if the optimizer step should be performed or gradients should be accumulated for the current step.""" accumulation_done = self._accumulated_batches_reached() # Lightning steps on the final batch is_final_batch = self._num_ready_batches_reached() # but the TTP might not ttp_accumulates_on_final_batch = ( self.trainer.training_type_plugin.handles_gradient_accumulation or not is_final_batch ) return not accumulation_done and ttp_accumulates_on_final_batch @staticmethod def _prepare_outputs_training_batch_end( batch_output: _BATCH_OUTPUTS_TYPE, automatic: bool, num_optimizers: int, ) -> Union[List[List[Dict[str, Any]]], List[Dict[str, Any]]]: """Processes the outputs from the batch loop into the format passed to the ``training_batch_end`` hook. ``(tbptt_steps, n_opt) -> (n_opt, tbptt_steps)``. The optimizer dimension might have been squeezed. """ if not batch_output: return [] # convert optimizer dicts to list if automatic: batch_output = apply_to_collection( batch_output, dtype=dict, function=_convert_optim_dict, num_optimizers=num_optimizers ) array = np.array(batch_output, dtype=object) if array.ndim == 1: array = np.expand_dims(array, 1) array = array.transpose((1, 0)) array = array.squeeze() array = array.tolist() array = _recursive_unpad(array) return array @staticmethod def _prepare_outputs_training_epoch_end( batch_outputs: _OUTPUTS_TYPE, automatic: bool, num_optimizers: int, ) -> Union[List[List[List[Dict[str, Any]]]], List[List[Dict[str, Any]]], List[Dict[str, Any]]]: """Processes the outputs from the batch loop into the format passed to the ``training_epoch_end`` hook. ``(n_batches, tbptt_steps, n_opt) -> (n_opt, n_batches, tbptt_steps)``. All single-element dimensions might have been squeezed. This processing is necessary because the format of the inputs to the ``training_epoch_end`` hook does not match the loop structure and because empty dimensions are squeezed. This could break with loop customization. """ # `batch_outputs` (plural) is the same as `epoch_end_output` (singular) if not batch_outputs: return [] # convert optimizer dicts to list if automatic: batch_outputs = apply_to_collection( batch_outputs, dtype=dict, function=_convert_optim_dict, num_optimizers=num_optimizers ) array = _recursive_pad(batch_outputs) if array.ndim == 2: array = np.expand_dims(array, 2) array = array.transpose((2, 0, 1)) array = array.squeeze() array = array.tolist() array = _recursive_unpad(array) # in case we squeezed from 1-element array to a 0-dim array array = array if isinstance(array, list) else [array] # remove residual empty lists array = [item for item in array if not isinstance(item, list) or len(item)] return array def update_lr_schedulers(self, interval: str, update_plateau_schedulers: bool) -> None: """updates the lr schedulers based on the given interval.""" if interval == "step" and self._should_accumulate(): return active_optimizers = _get_active_optimizers( self.trainer.optimizers, self.trainer.optimizer_frequencies, self.total_batch_idx ) self._update_learning_rates( interval=interval, update_plateau_schedulers=update_plateau_schedulers, opt_indices=[opt_idx for opt_idx, _ in active_optimizers], ) def _update_learning_rates( self, interval: str, update_plateau_schedulers: bool, opt_indices: Optional[List[int]] = None ) -> None: """Update learning rates. Args: interval: either 'epoch' or 'step'. update_plateau_schedulers: control whether ``ReduceLROnPlateau`` or non-plateau schedulers get updated. This is used so non-plateau schedulers can be updated before running validation. Checkpoints are commonly saved during validation, however, on-plateau schedulers might monitor a validation metric so they have to be updated separately. opt_indices: indices of the optimizers to update. """ if not self.trainer.lr_schedulers or not self.trainer.lightning_module.automatic_optimization: return if opt_indices is None: opt_indices = [] for lr_scheduler in self.trainer.lr_schedulers: if isinstance(lr_scheduler["opt_idx"], int) and lr_scheduler["opt_idx"] not in opt_indices: continue if update_plateau_schedulers ^ lr_scheduler["reduce_on_plateau"]: continue current_idx = self.batch_idx if interval == "step" else self.trainer.current_epoch current_idx += 1 # account for both batch and epoch starts from 0 # Take step if call to update_learning_rates matches the interval key and # the current step modulo the schedulers frequency is zero if lr_scheduler["interval"] == interval and current_idx % lr_scheduler["frequency"] == 0: monitor_val = None if lr_scheduler["reduce_on_plateau"]: # If instance of ReduceLROnPlateau, we need a monitor monitor_key = lr_scheduler["monitor"] monitor_val = self._get_monitor_value(monitor_key) if monitor_val is None: if lr_scheduler.get("strict", True): avail_metrics = list(self.trainer.callback_metrics) raise MisconfigurationException( f"ReduceLROnPlateau conditioned on metric {monitor_key}" f" which is not available. Available metrics are: {avail_metrics}." " Condition can be set using `monitor` key in lr scheduler dict" ) rank_zero_warn( f"ReduceLROnPlateau conditioned on metric {monitor_key}" " which is not available but strict is set to `False`." " Skipping learning rate update.", RuntimeWarning, ) continue self.scheduler_progress.increment_ready() # update LR if lr_scheduler["reduce_on_plateau"]: lr_scheduler["scheduler"].step(monitor_val) else: lr_scheduler["scheduler"].step() self.scheduler_progress.increment_completed() def _get_monitor_value(self, key: str) -> Any: # this is a separate method to aid in testing return self.trainer.callback_metrics.get(key) def _should_check_val_fx(self, batch_idx: int, is_last_batch: bool) -> bool: """Decide if we should run validation.""" if not self.trainer.enable_validation: return False is_val_check_epoch = (self.trainer.current_epoch + 1) % self.trainer.check_val_every_n_epoch == 0 if not is_val_check_epoch: return False # val_check_batch is inf for iterable datasets with no length defined is_infinite_dataset = self.trainer.val_check_batch == float("inf") if is_last_batch and is_infinite_dataset: return True if self.trainer.should_stop: return True # TODO(@awaelchli): let training/eval loop handle logic around limit_*_batches and val_check_batch is_val_check_batch = is_last_batch if isinstance(self.trainer.limit_train_batches, int) and is_infinite_dataset: is_val_check_batch = (batch_idx + 1) % self.trainer.limit_train_batches == 0 elif self.trainer.val_check_batch != float("inf"): is_val_check_batch = (batch_idx + 1) % self.trainer.val_check_batch == 0 return is_val_check_batch def _save_loggers_on_train_batch_end(self) -> None: """Flushes loggers to disk.""" # when loggers should save to disk should_flush_logs = self.trainer.logger_connector.should_flush_logs if should_flush_logs and self.trainer.is_global_zero and self.trainer.logger is not None: self.trainer.logger.save() def _reload_dataloader_state_dict(self, data_fetcher: AbstractDataFetcher): if self._dataloader_state_dict: data_fetcher.dataloader.load_state_dict(self._dataloader_state_dict) self._dataloader_state_dict = None
class TrainingEpochLoop(loops.Loop): """ Runs over all batches in a dataloader (one epoch). Args: min_steps: The minimum number of steps (batches) to process max_steps: The maximum number of steps (batches) to process """ def __init__(self, min_steps: int, max_steps: int): super().__init__() self.min_steps: int = min_steps self.max_steps: int = max_steps self.global_step: int = 0 # the total batch index across all epochs self.total_batch_idx: int = 0 # the current split index when the batch gets split into chunks in truncated backprop through time self.split_idx: Optional[int] = None # the number of batches seen this run, updates immediately after batch_loop.run() # TODO: replace by progress tracking self.batches_seen: int = 0 self.is_last_batch: Optional[bool] = None self.batch_progress = Progress() self.scheduler_progress = SchedulerProgress() self.batch_loop: Optional[TrainingBatchLoop] = None self.val_loop: Optional["loops.EvaluationLoop"] = None self._results = ResultCollection(training=True) self._dataloader_idx: Optional[int] = None self._warning_cache: WarningCache = WarningCache() self._epoch_output: Optional[List[List[STEP_OUTPUT]]] = None @property def batch_idx(self) -> int: """Returns the current batch index (within this epoch)""" return self.iteration_count @property def done(self) -> bool: """Returns whether the training should be stopped. The criteria are that the number of steps reached the max steps, the last batch is reached or the trainer signals to stop (e.g. by early stopping). """ max_steps_reached = self.max_steps is not None and self.global_step >= self.max_steps return max_steps_reached or self.trainer.should_stop or self._num_training_batches_reached( self.is_last_batch) def connect(self, batch_loop: Optional[TrainingBatchLoop] = None, val_loop: Optional["loops.EvaluationLoop"] = None) -> None: """Optionally connect a custom batch or validation loop to this training epoch loop.""" if batch_loop is not None: self.batch_loop = batch_loop if val_loop is not None: self.val_loop = val_loop def reset(self) -> None: """Resets the internal state of the loop for a new run""" self.iteration_count = 0 self.batches_seen = 0 self.is_last_batch = False self._dataloader_idx = 0 # track epoch output self._epoch_output = [[] for _ in range( self.batch_loop.num_active_optimizers(self.total_batch_idx))] if self.restarting: self.iteration_count = self.batches_seen = self.batch_progress.current.completed else: self.batch_progress.current.reset() self.scheduler_progress.current.reset() self.batch_loop.optim_progress.reset_on_epoch() def on_run_start(self, *args: Any, **kwargs: Any) -> None: # hook self.trainer.logger_connector.on_epoch_start() self.trainer.call_hook("on_epoch_start") self.trainer.call_hook("on_train_epoch_start") self.trainer.fit_loop.epoch_progress.increment_started() def advance(self, dataloader_iter: Iterator, **kwargs: Any) -> None: """Runs a single training batch. Args: dataloader_iter: the iterator over the dataloader producing the new batch Raises: StopIteration: When the epoch is canceled by the user returning -1 """ _, (batch, is_last) = next(dataloader_iter) self.is_last_batch = is_last # ------------------------------------ # TRAINING_STEP + TRAINING_STEP_END # ------------------------------------ with self.trainer.profiler.profile("training_batch_to_device"): batch = self.trainer.accelerator.batch_to_device( batch, dataloader_idx=self._dataloader_idx) self.batch_progress.increment_ready() with self.trainer.profiler.profile("run_training_batch"): batch_output = self.batch_loop.run(batch, self.iteration_count, self._dataloader_idx) self.batches_seen += 1 self.batch_progress.increment_processed() # when returning -1 from train_step, we end epoch early if batch_output.signal == -1: raise StopIteration # update non-plateau LR schedulers # update epoch-interval ones only when we are at the end of training epoch self.update_lr_schedulers("step", update_plateau_schedulers=False) if self._num_training_batches_reached(is_last): self.update_lr_schedulers("epoch", update_plateau_schedulers=False) batch_end_outputs = [ opt_idx_out for opt_idx_out in batch_output.training_step_output if len(opt_idx_out) ] processed_batch_end_outputs = self._prepare_outputs(batch_end_outputs, batch_mode=True) # hook self.trainer.call_hook("on_train_batch_end", processed_batch_end_outputs, batch, self.iteration_count, self._dataloader_idx) self.trainer.call_hook("on_batch_end") self.trainer.logger_connector.on_batch_end() self.batch_progress.increment_completed() # figure out what to track for epoch end self._track_epoch_end_reduce_metrics(self._epoch_output, batch_end_outputs) # ----------------------------------------- # SAVE METRICS TO LOGGERS AND PROGRESS_BAR # ----------------------------------------- self.trainer.logger_connector.update_train_step_metrics() def on_advance_end(self): """Runs validation and Checkpointing if necessary. Raises: StopIteration: if :attr:`done` evaluates to ``True`` to finish this epoch """ # ----------------------------------------- # VALIDATE IF NEEDED + CHECKPOINT CALLBACK # ----------------------------------------- should_check_val = self._should_check_val_fx(self.iteration_count, self.is_last_batch) if should_check_val: self.trainer.validating = True self._run_validation() self.trainer.training = True # ----------------------------------------- # SAVE LOGGERS (ie: Tensorboard, etc...) # ----------------------------------------- self._save_loggers_on_train_batch_end() # update plateau LR scheduler after metrics are logged self.update_lr_schedulers("step", update_plateau_schedulers=True) self.total_batch_idx += 1 # progress global step according to grads progress self._increment_accumulated_grad_global_step() if self.done: raise StopIteration def on_run_end(self) -> List[List[STEP_OUTPUT]]: """Calls the on_epoch_end hook. Returns: The output of each training step for each optimizer Raises: MisconfigurationException: ``train_epoch_end`` does not return ``None`` """ if self.batches_seen == 0: # dataloader/iterator did not produce a batch return # inform logger the batch loop has finished self.trainer.logger_connector.epoch_end_reached() # prepare epoch output processed_outputs = self._prepare_outputs(self._epoch_output, batch_mode=False) # get the model and call model.training_epoch_end model = self.trainer.lightning_module if is_overridden("training_epoch_end", model): # run training_epoch_end # refresh the result for custom logging at the epoch level model._current_fx_name = "training_epoch_end" # lightningmodule hook training_epoch_end_output = model.training_epoch_end( processed_outputs) if training_epoch_end_output is not None: raise MisconfigurationException( "training_epoch_end expects a return of None. " "HINT: remove the return statement in training_epoch_end") self.trainer.fit_loop.epoch_progress.increment_processed() # call train epoch end hooks self._on_train_epoch_end_hook(processed_outputs) self.trainer.call_hook("on_epoch_end") self.trainer.logger_connector.on_epoch_end() self.update_lr_schedulers("epoch", update_plateau_schedulers=True) epoch_output = self._epoch_output # free memory self._epoch_output = None return epoch_output def teardown(self) -> None: self._results.cpu() self.batch_loop.teardown() self.val_loop.teardown() def _run_validation(self): # reload dataloaders self.val_loop.reload_evaluation_dataloaders() with torch.no_grad(): self.val_loop.run() def _on_train_epoch_end_hook( self, processed_epoch_output: List[List[STEP_OUTPUT]]) -> None: """Runs ``on_train_epoch_end hook``.""" # We cannot rely on Trainer.call_hook because the signatures might be different across # lightning module and callback # As a result, we need to inspect if the module accepts `outputs` in `on_train_epoch_end` # This implementation is copied from Trainer.call_hook hook_name = "on_train_epoch_end" prev_fx_name = self.trainer.lightning_module._current_fx_name self.trainer.lightning_module._current_fx_name = hook_name # always profile hooks with self.trainer.profiler.profile(hook_name): # first call trainer hook if hasattr(self.trainer, hook_name): trainer_hook = getattr(self.trainer, hook_name) trainer_hook(processed_epoch_output) # next call hook in lightningModule model_ref = self.trainer.lightning_module if is_overridden(hook_name, model_ref): hook_fx = getattr(model_ref, hook_name) if is_param_in_hook_signature(hook_fx, "outputs"): self._warning_cache.deprecation( "The signature of `ModelHooks.on_train_epoch_end` has changed in v1.3." " `outputs` parameter has been deprecated." " Support for the old signature will be removed in v1.5" ) model_ref.on_train_epoch_end(processed_epoch_output) else: model_ref.on_train_epoch_end() # call the accelerator hook if hasattr(self.trainer.accelerator, hook_name): accelerator_hook = getattr(self.trainer.accelerator, hook_name) accelerator_hook() # restore current_fx when nested context self.trainer.lightning_module._current_fx_name = prev_fx_name def _num_training_batches_reached(self, is_last_batch: bool = False) -> bool: """Checks if we are in the last batch or if there are more batches to follow.""" # TODO: Can we combine this with training_batch_loop's arg that does a similar check? return self.batches_seen == self.trainer.num_training_batches or is_last_batch def _track_epoch_end_reduce_metrics( self, epoch_output: List[List[STEP_OUTPUT]], batch_end_outputs: STEP_OUTPUT) -> None: """Adds the batch outputs to the epoch outputs and prepares reduction""" hook_overridden = self._should_add_batch_output_to_epoch_output() if not hook_overridden: return # track the outputs to reduce at the end of the epoch for opt_idx, opt_outputs in enumerate(batch_end_outputs): # with 1 step (no tbptt) don't use a sequence at epoch end if (isinstance(opt_outputs, list) and len(opt_outputs) == 1 and not isinstance(opt_outputs[0], ResultCollection)): opt_outputs = opt_outputs[0] epoch_output[opt_idx].append(opt_outputs) def _should_add_batch_output_to_epoch_output(self) -> bool: """ We add to the epoch outputs if 1. The model defines training_epoch_end OR 2. The model overrides on_train_epoch_end which has `outputs` in the signature """ # TODO: in v1.5 this only needs to check if training_epoch_end is overridden lightning_module = self.trainer.lightning_module if is_overridden("training_epoch_end", lightning_module): return True if is_overridden("on_train_epoch_end", lightning_module): model_hook_fx = getattr(lightning_module, "on_train_epoch_end") if is_param_in_hook_signature(model_hook_fx, "outputs"): return True return False @staticmethod def _prepare_outputs( outputs: List[List[List["ResultCollection"]]], batch_mode: bool ) -> Union[List[List[List[Dict]]], List[List[Dict]], List[Dict], Dict]: """ Extract required information from batch or epoch end results. Args: outputs: A 3-dimensional list of ``ResultCollection`` objects with dimensions: ``[optimizer outs][batch outs][tbptt steps]``. batch_mode: If True, ignore the batch output dimension. Returns: The cleaned outputs with ``ResultCollection`` objects converted to dictionaries. All list dimensions of size one will be collapsed. """ processed_outputs = [] for opt_outputs in outputs: # handle an edge case where an optimizer output is the empty list if len(opt_outputs) == 0: continue processed_batch_outputs = [] if batch_mode: opt_outputs = [opt_outputs] for batch_outputs in opt_outputs: processed_tbptt_outputs = [] if isinstance(batch_outputs, ResultCollection): batch_outputs = [batch_outputs] for tbptt_output in batch_outputs: out = tbptt_output.extra if tbptt_output.minimize is not None: out["loss"] = tbptt_output.minimize.detach() processed_tbptt_outputs.append(out) # if there was only one tbptt step then we can collapse that dimension if len(processed_tbptt_outputs) == 1: processed_tbptt_outputs = processed_tbptt_outputs[0] processed_batch_outputs.append(processed_tbptt_outputs) # batch_outputs should be just one dict (or a list of dicts if using tbptt) per optimizer if batch_mode: processed_batch_outputs = processed_batch_outputs[0] processed_outputs.append(processed_batch_outputs) # if there is only one optimiser then we collapse that dimension if len(processed_outputs) == 1: processed_outputs = processed_outputs[0] return processed_outputs def update_lr_schedulers(self, interval: str, update_plateau_schedulers: bool) -> None: """updates the lr schedulers based on the given interval""" if interval == "step" and self.batch_loop.should_accumulate(): return self.trainer.optimizer_connector.update_learning_rates( interval=interval, update_plateau_schedulers=update_plateau_schedulers, opt_indices=[ opt_idx for opt_idx, _ in self.batch_loop.get_active_optimizers(self.total_batch_idx) ], ) def _increment_accumulated_grad_global_step(self) -> None: """increments global step""" num_accumulated_batches_reached = self.batch_loop._accumulated_batches_reached( ) num_training_batches_reached = self._num_training_batches_reached() # progress global step according to grads progress if num_accumulated_batches_reached or num_training_batches_reached: self.global_step = self.trainer.accelerator.update_global_step( self.total_batch_idx, self.trainer.global_step) def _should_check_val_fx(self, batch_idx: int, is_last_batch: bool) -> bool: """Decide if we should run validation.""" if not self.trainer.enable_validation: return False is_val_check_epoch = (self.trainer.current_epoch + 1) % self.trainer.check_val_every_n_epoch == 0 if not is_val_check_epoch: return False # val_check_batch is inf for iterable datasets with no length defined is_infinite_dataset = self.trainer.val_check_batch == float("inf") if is_last_batch and is_infinite_dataset: return True if self.trainer.should_stop: return True # TODO(@awaelchli): let training/eval loop handle logic around limit_*_batches and val_check_batch is_val_check_batch = is_last_batch if isinstance(self.trainer.limit_train_batches, int) and is_infinite_dataset: is_val_check_batch = (batch_idx + 1) % self.trainer.limit_train_batches == 0 elif self.trainer.val_check_batch != float("inf"): is_val_check_batch = (batch_idx + 1) % self.trainer.val_check_batch == 0 return is_val_check_batch def _save_loggers_on_train_batch_end(self) -> None: """Flushes loggers to disk""" # when loggers should save to disk should_flush_logs = self.trainer.logger_connector.should_flush_logs if should_flush_logs and self.trainer.is_global_zero and self.trainer.logger is not None: self.trainer.logger.save()
class TrainingEpochLoop(loops.Loop): """ Runs over all batches in a dataloader (one epoch). Args: min_steps: The minimum number of steps (batches) to process max_steps: The maximum number of steps (batches) to process """ def __init__(self, min_steps: int, max_steps: int): super().__init__() self.min_steps: int = min_steps self.max_steps: int = max_steps self.global_step: int = 0 # manually tracking which is the last batch is necessary for iterable dataset support self.is_last_batch: Optional[bool] = None self.batch_progress = Progress() self.scheduler_progress = SchedulerProgress() self.batch_loop: Optional[TrainingBatchLoop] = None self.val_loop: Optional["loops.EvaluationLoop"] = None self._results = ResultCollection(training=True) self._epoch_output: Optional[List[List[STEP_OUTPUT]]] = None @property def total_batch_idx(self) -> int: """Returns the current batch index (across epochs)""" # use `ready` instead of `completed` in case this is accessed after `completed` has been increased # but before the next `ready` increase return self.batch_progress.total.ready - 1 @property def batch_idx(self) -> int: """Returns the current batch index (within this epoch)""" # use `ready` instead of `completed` in case this is accessed after `completed` has been increased # but before the next `ready` increase return self.batch_progress.current.ready - 1 @property def done(self) -> bool: """Returns whether the training should be stopped. The criteria are that the number of steps reached the max steps, the last batch is reached or the trainer signals to stop (e.g. by early stopping). """ max_steps_reached = self.max_steps is not None and self.global_step >= self.max_steps return max_steps_reached or self.trainer.should_stop or self._num_training_batches_reached( self.is_last_batch) def connect( self, batch_loop: TrainingBatchLoop = None, val_loop: Optional["loops.EvaluationLoop"] = None, ) -> None: """Optionally connect a custom batch or validation loop to this training epoch loop.""" if batch_loop is not None: self.batch_loop = batch_loop if val_loop is not None: self.val_loop = val_loop def reset(self) -> None: """Resets the internal state of the loop for a new run""" self.is_last_batch = False # track epoch output self._epoch_output = [[] for _ in range( self.batch_loop.num_active_optimizers(self.total_batch_idx))] if not self.restarting: self.batch_progress.current.reset() self.scheduler_progress.current.reset() self.batch_loop.optim_progress.reset_on_epoch() def on_run_start(self, dataloader_iter: Iterator, **kwargs: Any) -> None: # hook self.trainer.logger_connector.on_epoch_start() self.trainer.call_hook("on_epoch_start") self.trainer.call_hook("on_train_epoch_start") self.trainer.fit_loop.epoch_progress.increment_started() self.dataloader_iter = _prepare_dataloader_iter( dataloader_iter, self.batch_idx + 1) def advance(self, *args: Any, **kwargs: Any) -> None: """Runs a single training batch. Args: dataloader_iter: the iterator over the dataloader producing the new batch Raises: StopIteration: When the epoch is canceled by the user returning -1 """ batch_idx, (batch, is_last) = next(self.dataloader_iter) if not self.trainer.data_connector.train_data_fetcher.store_on_device: with self.trainer.profiler.profile("training_batch_to_device"): batch = self.trainer.accelerator.batch_to_device(batch) self.batch_progress.increment_ready() with self.trainer.profiler.profile("run_training_batch"): batch_output = self.batch_loop.run(batch, batch_idx) self.batch_progress.increment_processed() self.is_last_batch = is_last # when returning -1 from train_step, we end epoch early if batch_output.signal == -1: raise StopIteration # update non-plateau LR schedulers # update epoch-interval ones only when we are at the end of training epoch self.update_lr_schedulers("step", update_plateau_schedulers=False) if self._num_training_batches_reached(is_last): self.update_lr_schedulers("epoch", update_plateau_schedulers=False) batch_end_outputs = [ opt_idx_out for opt_idx_out in batch_output.training_step_output if len(opt_idx_out) ] processed_batch_end_outputs = self._prepare_outputs(batch_end_outputs, batch_mode=True) # hook self.trainer.call_hook("on_train_batch_end", processed_batch_end_outputs, batch, self.batch_idx, 0) self.trainer.call_hook("on_batch_end") self.trainer.logger_connector.on_batch_end() self.batch_progress.increment_completed() # figure out what to track for epoch end self._track_epoch_end_reduce_metrics(self._epoch_output, batch_end_outputs) # ----------------------------------------- # SAVE METRICS TO LOGGERS AND PROGRESS_BAR # ----------------------------------------- self.trainer.logger_connector.update_train_step_metrics() def on_advance_end(self): """Runs validation and Checkpointing if necessary. Raises: StopIteration: if :attr:`done` evaluates to ``True`` to finish this epoch """ # ----------------------------------------- # VALIDATE IF NEEDED + CHECKPOINT CALLBACK # ----------------------------------------- should_check_val = self._should_check_val_fx(self.batch_idx, self.is_last_batch) if should_check_val: self.trainer.validating = True self._run_validation() self.trainer.training = True # ----------------------------------------- # SAVE LOGGERS (ie: Tensorboard, etc...) # ----------------------------------------- self._save_loggers_on_train_batch_end() # update plateau LR scheduler after metrics are logged self.update_lr_schedulers("step", update_plateau_schedulers=True) # progress global step according to grads progress self._increment_accumulated_grad_global_step() def on_run_end(self) -> List[List[STEP_OUTPUT]]: """Calls the on_epoch_end hook. Returns: The output of each training step for each optimizer Raises: MisconfigurationException: ``train_epoch_end`` does not return ``None`` """ if self.batch_progress.current.ready == 0: # dataloader/iterator did not produce a batch return # inform logger the batch loop has finished self.trainer.logger_connector.epoch_end_reached() # prepare epoch output processed_outputs = self._prepare_outputs(self._epoch_output, batch_mode=False) # get the model and call model.training_epoch_end model = self.trainer.lightning_module if is_overridden("training_epoch_end", model): # run training_epoch_end # refresh the result for custom logging at the epoch level model._current_fx_name = "training_epoch_end" # lightningmodule hook training_epoch_end_output = model.training_epoch_end( processed_outputs) if training_epoch_end_output is not None: raise MisconfigurationException( "training_epoch_end expects a return of None. " "HINT: remove the return statement in training_epoch_end") self.trainer.fit_loop.epoch_progress.increment_processed() # call train epoch end hooks self.trainer.call_hook("on_train_epoch_end") self.trainer.call_hook("on_epoch_end") self.trainer.logger_connector.on_epoch_end() if self._num_training_batches_reached(self.is_last_batch): self.update_lr_schedulers("epoch", update_plateau_schedulers=True) epoch_output = self._epoch_output # free memory self._epoch_output = None return epoch_output def teardown(self) -> None: self._results.cpu() self.batch_loop.teardown() self.val_loop.teardown() def _run_validation(self): # reload dataloaders self.val_loop.reload_evaluation_dataloaders() with torch.no_grad(): self.val_loop.run() def _accumulated_batches_reached(self) -> bool: """Determine if accumulation will be finished by the end of the current batch.""" return self.batch_progress.current.ready % self.trainer.accumulate_grad_batches == 0 def _num_training_batches_reached(self, is_last_batch: bool = False) -> bool: """Checks if we are in the last batch or if there are more batches to follow. Args: is_last_batch: Whether the current batch is the last one """ return self.batch_progress.current.ready == self.trainer.num_training_batches or is_last_batch def _should_accumulate(self) -> bool: """Checks if the optimizer step should be performed or gradients should be accumulated for the current step.""" accumulation_done = self._accumulated_batches_reached() is_final_batch = self._num_training_batches_reached() return not (accumulation_done or is_final_batch) def _track_epoch_end_reduce_metrics( self, epoch_output: List[List[STEP_OUTPUT]], batch_end_outputs: STEP_OUTPUT) -> None: """Adds the batch outputs to the epoch outputs and prepares reduction""" hook_overridden = is_overridden("training_epoch_end", self.trainer.lightning_module) if not hook_overridden: return # track the outputs to reduce at the end of the epoch for opt_idx, opt_outputs in enumerate(batch_end_outputs): # with 1 step (no tbptt) don't use a sequence at epoch end if isinstance(opt_outputs, list) and len(opt_outputs) == 1: opt_outputs = opt_outputs[0] epoch_output[opt_idx].append(opt_outputs) @staticmethod def _prepare_outputs( outputs: List[List[List["ResultCollection"]]], batch_mode: bool ) -> Union[List[List[List[Dict]]], List[List[Dict]], List[Dict], Dict]: """ Extract required information from batch or epoch end results. Args: outputs: A 3-dimensional list of ``ResultCollection`` objects with dimensions: ``[optimizer outs][batch outs][tbptt steps]``. batch_mode: If True, ignore the batch output dimension. Returns: The cleaned outputs with ``ResultCollection`` objects converted to dictionaries. All list dimensions of size one will be collapsed. """ processed_outputs = [] for opt_outputs in outputs: # handle an edge case where an optimizer output is the empty list if len(opt_outputs) == 0: continue processed_batch_outputs = [] if batch_mode: opt_outputs = [opt_outputs] for batch_outputs in opt_outputs: processed_tbptt_outputs = [] if isinstance(batch_outputs, ResultCollection): batch_outputs = [batch_outputs] for tbptt_output in batch_outputs: out = {} if tbptt_output.minimize is not None: out["loss"] = tbptt_output.minimize.detach() out.update(tbptt_output.extra) processed_tbptt_outputs.append(out) # if there was only one tbptt step then we can collapse that dimension if len(processed_tbptt_outputs) == 1: processed_tbptt_outputs = processed_tbptt_outputs[0] processed_batch_outputs.append(processed_tbptt_outputs) # batch_outputs should be just one dict (or a list of dicts if using tbptt) per optimizer if batch_mode: processed_batch_outputs = processed_batch_outputs[0] processed_outputs.append(processed_batch_outputs) # if there is only one optimiser then we collapse that dimension if len(processed_outputs) == 1: processed_outputs = processed_outputs[0] return processed_outputs def update_lr_schedulers(self, interval: str, update_plateau_schedulers: bool) -> None: """updates the lr schedulers based on the given interval""" if interval == "step" and self._should_accumulate(): return self.trainer.optimizer_connector.update_learning_rates( interval=interval, update_plateau_schedulers=update_plateau_schedulers, opt_indices=[ opt_idx for opt_idx, _ in self.batch_loop.get_active_optimizers(self.total_batch_idx) ], ) def _increment_accumulated_grad_global_step(self) -> None: """Increments global step according to grads progress""" if not self._should_accumulate(): self.global_step = self.trainer.accelerator.update_global_step( self.batch_progress.current.ready, self.trainer.global_step) def _should_check_val_fx(self, batch_idx: int, is_last_batch: bool) -> bool: """Decide if we should run validation.""" if not self.trainer.enable_validation: return False is_val_check_epoch = (self.trainer.current_epoch + 1) % self.trainer.check_val_every_n_epoch == 0 if not is_val_check_epoch: return False # val_check_batch is inf for iterable datasets with no length defined is_infinite_dataset = self.trainer.val_check_batch == float("inf") if is_last_batch and is_infinite_dataset: return True if self.trainer.should_stop: return True # TODO(@awaelchli): let training/eval loop handle logic around limit_*_batches and val_check_batch is_val_check_batch = is_last_batch if isinstance(self.trainer.limit_train_batches, int) and is_infinite_dataset: is_val_check_batch = (batch_idx + 1) % self.trainer.limit_train_batches == 0 elif self.trainer.val_check_batch != float("inf"): is_val_check_batch = (batch_idx + 1) % self.trainer.val_check_batch == 0 return is_val_check_batch def _save_loggers_on_train_batch_end(self) -> None: """Flushes loggers to disk""" # when loggers should save to disk should_flush_logs = self.trainer.logger_connector.should_flush_logs if should_flush_logs and self.trainer.is_global_zero and self.trainer.logger is not None: self.trainer.logger.save()
class TrainLoop: def __init__( self, trainer, max_epochs: Optional[int], min_epochs: Optional[int], max_steps: Optional[int], min_steps: Optional[int], num_sanity_val_steps: int, ): self.trainer = trainer self.accumulated_loss = None self.warning_cache = WarningCache() self.running_loss = TensorRunningAccum(window_length=20) self._skip_backward = False self._optimizer_freq_cumsum = None self._hiddens = None self.global_step = 0 self.current_epoch = 0 self.trainer.should_stop = False # the total batch index across all epochs self.total_batch_idx = 0 # the current batch index in the loop that runs over the dataloader(s) self.batch_idx = 0 # the current split index when the batch gets split into chunks in truncated backprop through time self.split_idx = None self.trainer.num_training_batches = 0 self.trainer.train_dataloader = None # If neither max_epochs or max_steps is set, then use existing default of max_epochs = 1000 self.max_epochs = 1000 if (max_epochs is None and max_steps is None) else max_epochs # If neither min_epochs or min_steps is set, then use existing default of min_epochs = 1 self.min_epochs = 1 if (min_epochs is None and min_steps is None) else min_epochs self.max_steps = max_steps self.min_steps = min_steps if num_sanity_val_steps == -1: self.trainer.num_sanity_val_steps = float("inf") else: self.trainer.num_sanity_val_steps = num_sanity_val_steps self.results = ResultCollection(training=True) @property def num_active_optimizers(self) -> int: return len(self.get_active_optimizers()) @property def optimizer_freq_cumsum(self): if self._optimizer_freq_cumsum is None: self._optimizer_freq_cumsum = np.cumsum(self.trainer.optimizer_frequencies) return self._optimizer_freq_cumsum def should_skip_training(self) -> bool: should_by_max_steps = self.max_steps is not None and self.global_step >= self.max_steps should_by_epoch = self.max_epochs is not None and self.current_epoch >= self.max_epochs return should_by_max_steps or should_by_epoch or self.trainer.num_training_batches == 0 def on_train_start(self): self.results.to(device=self.trainer.lightning_module.device) self.trainer.call_hook("on_train_start") def on_train_end(self): # trigger checkpoint check. need to temporarily decrease the global step to avoid saving duplicates # when a checkpoint was saved at the last step self.global_step -= 1 self.check_checkpoint_callback(should_update=True, is_last=True) self.global_step += 1 # hook self.trainer.call_hook("on_train_end") # todo: TPU 8 cores hangs in flush with TensorBoard. Might do for all loggers. # It might be related to xla tensors blocked when moving the cpu # kill loggers if self.trainer.logger is not None: self.trainer.logger.finalize("success") # summarize profile results self.trainer.profiler.describe() # give accelerators a chance to finish self.trainer.accelerator.on_train_end() # reset bookkeeping self.trainer.state.stage = None def check_checkpoint_callback(self, should_update, is_last=False): # TODO bake this logic into the ModelCheckpoint callback if should_update and self.trainer.checkpoint_connector.has_trained: callbacks = self.trainer.checkpoint_callbacks if is_last and any(cb.save_last and cb.verbose for cb in callbacks): rank_zero_info("Saving latest checkpoint...") model = self.trainer.lightning_module for cb in callbacks: cb.on_validation_end(self.trainer, model) def on_train_epoch_start(self, epoch): # update training progress in trainer self.current_epoch = epoch model = self.trainer.lightning_module # reset train dataloader if epoch != 0 and self.trainer.reload_dataloaders_every_epoch: self.trainer.reset_train_dataloader(model) # todo: specify the possible exception with suppress(Exception): # set seed for distributed sampler (enables shuffling for each epoch) self.trainer.train_dataloader.sampler.set_epoch(epoch) # changing gradient according accumulation_scheduler self.trainer.accumulation_scheduler.on_train_epoch_start(self.trainer, self.trainer.lightning_module) # stores accumulated grad fractions per batch self.accumulated_loss = TensorRunningAccum(window_length=self.trainer.accumulate_grad_batches) # hook self.trainer.logger_connector.on_epoch_start() self.trainer.call_hook("on_epoch_start") self.trainer.call_hook("on_train_epoch_start") def on_train_batch_end(self, epoch_output, batch_end_outputs, batch, batch_idx, dataloader_idx): batch_end_outputs = [opt_idx_out for opt_idx_out in batch_end_outputs if len(opt_idx_out)] processed_batch_end_outputs = TrainLoop._prepare_outputs(batch_end_outputs, batch_mode=True) # hook self.trainer.call_hook('on_train_batch_end', processed_batch_end_outputs, batch, batch_idx, dataloader_idx) self.trainer.call_hook('on_batch_end') self.trainer.logger_connector.on_batch_end() # figure out what to track for epoch end self.track_epoch_end_reduce_metrics(epoch_output, batch_end_outputs) def reset_train_val_dataloaders(self, model) -> None: """ Resets train and val dataloaders if none are attached to the trainer. The val dataloader must be initialized before training loop starts, as the training loop inspects the val dataloader to determine whether to run the evaluation loop. """ if self.trainer.train_dataloader is None: self.trainer.reset_train_dataloader(model) if self.trainer.val_dataloaders is None: self.trainer.reset_val_dataloader(model) def track_epoch_end_reduce_metrics(self, epoch_output, batch_end_outputs): hook_overridden = self._should_add_batch_output_to_epoch_output() if not hook_overridden: return # track the outputs to reduce at the end of the epoch for opt_idx, opt_outputs in enumerate(batch_end_outputs): # with 1 step (no tbptt) don't use a sequence at epoch end if ( isinstance(opt_outputs, list) and len(opt_outputs) == 1 and not isinstance(opt_outputs[0], ResultCollection) ): opt_outputs = opt_outputs[0] epoch_output[opt_idx].append(opt_outputs) def _should_add_batch_output_to_epoch_output(self) -> bool: # We add to the epoch outputs if # 1. The model defines training_epoch_end OR # 2. The model overrides on_train_epoch_end which has `outputs` in the signature # TODO: in v1.5 this only needs to check if training_epoch_end is overridden lightning_module = self.trainer.lightning_module if is_overridden("training_epoch_end", lightning_module): return True if is_overridden("on_train_epoch_end", lightning_module): model_hook_fx = getattr(lightning_module, "on_train_epoch_end") if is_param_in_hook_signature(model_hook_fx, "outputs"): return True return False def get_active_optimizers(self, batch_idx: Optional[int] = None) -> List[Tuple[int, Optimizer]]: """ Returns the currently active optimizers. When multiple optimizers are used with different frequencies, only one of the optimizers is active at a time. Returns: A list of tuples (opt_idx, optimizer) of currently active optimizers. """ if not self.trainer.optimizer_frequencies: # call training_step once per optimizer return list(enumerate(self.trainer.optimizers)) batch_idx = self.total_batch_idx if batch_idx is None else batch_idx optimizers_loop_length = self.optimizer_freq_cumsum[-1] current_place_in_loop = batch_idx % optimizers_loop_length # find optimzier index by looking for the first {item > current_place} in the cumsum list opt_idx = int(np.argmax(self.optimizer_freq_cumsum > current_place_in_loop)) return [(opt_idx, self.trainer.optimizers[opt_idx])] def on_after_backward(self, batch_idx, untouched_loss): # insert after step hook self.trainer.call_hook("on_after_backward") # when in dev debugging track the losses self.trainer.dev_debugger.track_train_loss_history(batch_idx, untouched_loss.detach()) def _check_training_step_output(self, training_step_output): if isinstance(training_step_output, torch.Tensor) and not self.trainer.lightning_module.automatic_optimization: if training_step_output.grad_fn is None: # TODO: Find why - RuntimeError: Expected to mark a variable ready only once ... raise MisconfigurationException("In manual optimization, `training_step` should not return a Tensor") elif self.trainer.lightning_module.automatic_optimization: if not any(( isinstance(training_step_output, torch.Tensor), (isinstance(training_step_output, Mapping) and 'loss' in training_step_output), training_step_output is None )): raise MisconfigurationException( "In automatic optimization, `training_step` must either return a Tensor, " "a dict with key 'loss' or None (where the step will be skipped)." ) def training_step(self, split_batch, batch_idx, opt_idx, hiddens): # give the PL module a result for logging model_ref = self.trainer.lightning_module with self.trainer.profiler.profile("model_forward"): step_kwargs = self._build_kwargs(split_batch, batch_idx, opt_idx, hiddens) # manually capture logged metrics model_ref._current_fx_name = 'training_step' with self.trainer.profiler.profile("training_step"): training_step_output = self.trainer.accelerator.training_step(step_kwargs) self.trainer.accelerator.post_training_step() training_step_output = self.trainer.call_hook("training_step_end", training_step_output) self._check_training_step_output(training_step_output) training_step_output = self._process_training_step_output(training_step_output) if training_step_output is None: return closure_loss = None loss = None if self.trainer.lightning_module.automatic_optimization: # accumulate loss. if accumulate_grad_batches==1, no effect closure_loss = training_step_output.minimize / self.trainer.accumulate_grad_batches # the loss will get scaled for amp. avoid any modifications to it loss = closure_loss.detach().clone() return AttributeDict(closure_loss=closure_loss, loss=loss, training_step_output=training_step_output) def _process_training_step_output(self, training_step_output): if training_step_output is None: return None results = self.results loss = None hiddens = None results.extra = {} # handle dict return if isinstance(training_step_output, dict): loss = training_step_output.pop("loss", None) hiddens = training_step_output.pop("hiddens", None) if hiddens is not None: hiddens = hiddens.detach() results.extra = training_step_output # handle scalar return elif isinstance(training_step_output, torch.Tensor): loss = training_step_output # map to results under the hood results.minimize = loss self._hiddens = hiddens if self.trainer.move_metrics_to_cpu: results.cpu() return results @staticmethod def _prepare_outputs( outputs: List[List[List['ResultCollection']]], batch_mode: bool, ) -> Union[List[List[List[Dict]]], List[List[Dict]], List[Dict], Dict]: """ Extract required information from batch or epoch end results. Args: outputs: A 3-dimensional list of ``ResultCollection`` objects with dimensions: ``[optimizer outs][batch outs][tbptt steps]``. batch_mode: If True, ignore the batch output dimension. Returns: The cleaned outputs with ``ResultCollection`` objects converted to dictionaries. All list dimensions of size one will be collapsed. """ processed_outputs = [] for opt_outputs in outputs: # handle an edge case where an optimizer output is the empty list if len(opt_outputs) == 0: continue processed_batch_outputs = [] if batch_mode: opt_outputs = [opt_outputs] for batch_outputs in opt_outputs: processed_tbptt_outputs = [] if isinstance(batch_outputs, ResultCollection): batch_outputs = [batch_outputs] for tbptt_output in batch_outputs: out = tbptt_output.extra if tbptt_output.minimize is not None: out['loss'] = tbptt_output.minimize.detach() processed_tbptt_outputs.append(out) # if there was only one tbptt step then we can collapse that dimension if len(processed_tbptt_outputs) == 1: processed_tbptt_outputs = processed_tbptt_outputs[0] processed_batch_outputs.append(processed_tbptt_outputs) # batch_outputs should be just one dict (or a list of dicts if using tbptt) per optimizer if batch_mode: processed_batch_outputs = processed_batch_outputs[0] processed_outputs.append(processed_batch_outputs) # if there is only one optimiser then we collapse that dimension if len(processed_outputs) == 1: processed_outputs = processed_outputs[0] return processed_outputs def optimizer_step(self, optimizer, opt_idx, batch_idx, train_step_and_backward_closure): model_ref = self.trainer.lightning_module is_lbfgs = isinstance(optimizer, torch.optim.LBFGS) using_native_amp = self.trainer.amp_backend == AMPType.NATIVE # native amp + lbfgs is a no go right now if using_native_amp and is_lbfgs: raise MisconfigurationException( 'native PyTorch amp and lbfgs are not compatible.' ' To request, please file a Github issue in PyTorch and tag @mcarilli' ) # wraps into LightningOptimizer only for running step optimizer = LightningOptimizer._to_lightning_optimizer(optimizer, self.trainer, opt_idx) # model hook model_ref.optimizer_step( self.trainer.current_epoch, batch_idx, optimizer, opt_idx, train_step_and_backward_closure, on_tpu=self.trainer._device_type == DeviceType.TPU and _TPU_AVAILABLE, using_native_amp=using_native_amp, using_lbfgs=is_lbfgs, ) def on_before_zero_grad(self, optimizer): self.trainer.call_hook('on_before_zero_grad', optimizer) def optimizer_zero_grad(self, batch_idx, optimizer, opt_idx): self.trainer.accelerator.optimizer_zero_grad(self.trainer.current_epoch, batch_idx, optimizer, opt_idx) def track_and_norm_grad(self, optimizer) -> dict: # track gradient norms grad_norm_dict = {} if (self.global_step + 1) % self.trainer.log_every_n_steps == 0 and float(self.trainer.track_grad_norm) > 0: grad_norm_dict = grad_norm(self.trainer.lightning_module, self.trainer.track_grad_norm) # clip gradients self.trainer.accelerator.clip_gradients( optimizer, self.trainer.gradient_clip_val, gradient_clip_algorithm=self.trainer.gradient_clip_algorithm ) return grad_norm_dict def _tbptt_split_batch(self, batch: Any) -> List[Any]: splits = [batch] truncated_bptt_enabled = self._truncated_bptt_enabled() if truncated_bptt_enabled: model_ref = self.trainer.lightning_module with self.trainer.profiler.profile("tbptt_split_batch"): splits = model_ref.tbptt_split_batch(batch, self._truncated_bptt_steps()) return splits def run_training_epoch(self): # modify dataloader if needed (ddp, etc...) train_dataloader = self.trainer.accelerator.process_dataloader(self.trainer.train_dataloader) # track epoch output epoch_output = [[] for _ in range(self.num_active_optimizers)] train_dataloader = self.trainer.data_connector.get_profiled_train_dataloader(train_dataloader) dataloader_idx = 0 batch_idx = None for batch_idx, (batch, is_last_batch) in train_dataloader: self.batch_idx = batch_idx # ------------------------------------ # TRAINING_STEP + TRAINING_STEP_END # ------------------------------------ with self.trainer.profiler.profile("run_training_batch"): batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx) # when returning -1 from train_step, we end epoch early if batch_output.signal == -1: break # hook self.on_train_batch_end( epoch_output, batch_output.training_step_output, batch, batch_idx, dataloader_idx, ) # ----------------------------------------- # SAVE METRICS TO LOGGERS AND PROGRESS_BAR # ----------------------------------------- self.trainer.logger_connector.update_train_step_metrics() # ----------------------------------------- # VALIDATE IF NEEDED # ----------------------------------------- should_check_val = self._should_check_val_fx(batch_idx, is_last_batch) if should_check_val: self.trainer.validating = True self.trainer._run_evaluation() self.trainer.training = True # ----------------------------------------- # SAVE LOGGERS (ie: Tensorboard, etc...) # ----------------------------------------- self.save_loggers_on_train_batch_end() # update LR schedulers self.update_lr_schedulers('step') self.trainer.checkpoint_connector.has_trained = True self.total_batch_idx += 1 # progress global step according to grads progress self.increment_accumulated_grad_global_step() max_steps_reached = (self.max_steps is not None and self.max_steps <= self.global_step) if max_steps_reached or self.trainer.should_stop or self._num_training_batches_reached(is_last_batch): break if batch_idx is None: # dataloader/iterator did not produce a batch return # handle epoch_output on epoch end self.on_train_epoch_end(epoch_output) # the global step is manually decreased here due to backwards compatibility with existing loggers # as they expect that the same step is used when logging epoch end metrics even when the batch loop has # finished. this means the attribute does not exactly track the number of optimizer steps applied. # TODO(@carmocca): deprecate and rename so users don't get confused self.global_step -= 1 # log epoch metrics self.trainer.logger_connector.update_train_epoch_metrics() self.global_step += 1 self.update_lr_schedulers('epoch') did_train_only = self.trainer.disable_validation or self.trainer.evaluation_loop.should_skip_evaluation( self.trainer.num_val_batches ) if did_train_only: self.global_step -= 1 self.check_checkpoint_callback(True) self.global_step += 1 def on_train_epoch_end(self, epoch_output: List[List[List['ResultCollection']]]) -> None: # inform logger the batch loop has finished self.trainer.logger_connector.epoch_end_reached() # prepare epoch output processed_epoch_output = TrainLoop._prepare_outputs(epoch_output, batch_mode=False) # get the model and call model.training_epoch_end model = self.trainer.lightning_module if is_overridden('training_epoch_end', model): # run training_epoch_end # refresh the result for custom logging at the epoch level model._current_fx_name = 'training_epoch_end' training_epoch_end_output = model.training_epoch_end(processed_epoch_output) if training_epoch_end_output is not None: raise MisconfigurationException( 'training_epoch_end expects a return of None. ' 'HINT: remove the return statement in training_epoch_end' ) # call train epoch end hooks self._on_train_epoch_end_hook(processed_epoch_output) self.trainer.call_hook('on_epoch_end') self.trainer.logger_connector.on_epoch_end() def _on_train_epoch_end_hook(self, processed_epoch_output) -> None: # We cannot rely on Trainer.call_hook because the signatures might be different across # lightning module and callback # As a result, we need to inspect if the module accepts `outputs` in `on_train_epoch_end` # This implementation is copied from Trainer.call_hook hook_name = "on_train_epoch_end" prev_fx_name = self.trainer.lightning_module._current_fx_name self.trainer.lightning_module._current_fx_name = hook_name # always profile hooks with self.trainer.profiler.profile(hook_name): # first call trainer hook if hasattr(self.trainer, hook_name): trainer_hook = getattr(self.trainer, hook_name) trainer_hook(processed_epoch_output) # next call hook in lightningModule model_ref = self.trainer.lightning_module if is_overridden(hook_name, model_ref): hook_fx = getattr(model_ref, hook_name) if is_param_in_hook_signature(hook_fx, "outputs"): self.warning_cache.warn( "The signature of `ModelHooks.on_train_epoch_end` has changed in v1.3." " `outputs` parameter has been deprecated." " Support for the old signature will be removed in v1.5", DeprecationWarning ) model_ref.on_train_epoch_end(processed_epoch_output) else: model_ref.on_train_epoch_end() # call the accelerator hook if hasattr(self.trainer.accelerator, hook_name): accelerator_hook = getattr(self.trainer.accelerator, hook_name) accelerator_hook() # restore current_fx when nested context self.trainer.lightning_module._current_fx_name = prev_fx_name def run_training_batch(self, batch, batch_idx, dataloader_idx): # bookkeeping self._hiddens = None optimizers = list(enumerate(self.trainer.optimizers)) # track all outputs across time and num of optimizers batch_outputs = [[] for _ in range(len(optimizers))] if batch is None: self.warning_cache.warn("train_dataloader yielded None. If this was on purpose, ignore this warning...") return AttributeDict(signal=0, training_step_output=batch_outputs) # hook self.trainer.logger_connector.on_batch_start() response = self.trainer.call_hook("on_batch_start") if response == -1: return AttributeDict(signal=-1) # hook response = self.trainer.call_hook("on_train_batch_start", batch, batch_idx, dataloader_idx) if response == -1: return AttributeDict(signal=-1) # lightning module hook splits = self._tbptt_split_batch(batch) for split_idx, split_batch in enumerate(splits): self.split_idx = split_idx # let logger connector extract batch size self.trainer.logger_connector.on_train_split_start(batch_idx, split_idx, split_batch) if self.trainer.lightning_module.automatic_optimization: for opt_idx, optimizer in self.get_active_optimizers(batch_idx): result = self._run_optimization(batch_idx, split_batch, opt_idx, optimizer) if result: batch_outputs[opt_idx].append(result.training_step_output) else: # in manual optimization, there is no looping over optimizers result = self._run_optimization(batch_idx, split_batch) if result: batch_outputs[0].append(result.training_step_output) return AttributeDict(signal=0, training_step_output=batch_outputs) def _run_optimization(self, batch_idx, split_batch, opt_idx=0, optimizer=None): # TODO: In v1.5, when optimizer_idx gets removed from training_step in manual_optimization, change # opt_idx=0 to opt_idx=None in the signature here # toggle model params self.run_optimization_start(opt_idx, optimizer) result = AttributeDict() closure = self.make_closure(split_batch, batch_idx, opt_idx, optimizer, self._hiddens, result) if self.should_accumulate(): # For gradient accumulation # ------------------- # calculate loss (train step + train step end) # ------------------- # automatic_optimization=True: perform ddp sync only when performing optimizer_step # automatic_optimization=False: don't block synchronization here with self.block_ddp_sync_behaviour(): closure() # ------------------------------ # BACKWARD PASS # ------------------------------ # gradient update with accumulated gradients else: if self.trainer.lightning_module.automatic_optimization: self.optimizer_step(optimizer, opt_idx, batch_idx, closure) if len(self.trainer.optimizers) > 1: # revert back to previous state self.trainer.lightning_module.untoggle_optimizer(opt_idx) else: result = self.training_step(split_batch, batch_idx, opt_idx, self._hiddens) if not result: # user decided to skip optimization return result # update running loss + reset accumulated loss self.update_running_loss(result.loss) self._process_closure_result(result) return result def training_step_and_backward_closure( self, split_batch: Any, batch_idx: int, opt_idx: int, optimizer: Optimizer, hiddens, return_result: AttributeDict, ) -> Optional[torch.Tensor]: result = self.training_step_and_backward(split_batch, batch_idx, opt_idx, optimizer, hiddens) if result is not None: return_result.update(result) return return_result.loss def make_closure(self, *closure_args, **closure_kwargs: Any) -> Callable: """ Wraps the training step closure into a partial object which will be called within ``optimizer.step``. """ partial_func = partial(self.training_step_and_backward_closure, *closure_args, **closure_kwargs) return update_wrapper(partial_func, self.training_step_and_backward_closure) @contextmanager def block_ddp_sync_behaviour(self, should_block_sync: bool = False): """ automatic_optimization = True Blocks ddp sync gradients behaviour on backwards pass. This is useful for skipping sync when accumulating gradients, reducing communication overhead automatic_optimization = False do not block ddp gradient sync when using manual optimization as gradients are needed within the training step Returns: context manager with sync behaviour off """ if ( isinstance(self.trainer.training_type_plugin, ParallelPlugin) and (self.trainer.lightning_module.automatic_optimization or should_block_sync) ): with self.trainer.training_type_plugin.block_backward_sync(): yield None else: yield None def _process_closure_result(self, opt_closure_result: Optional[AttributeDict]) -> None: if not opt_closure_result: return # check if loss or model weights are nan if self.trainer.terminate_on_nan: self._check_finite(opt_closure_result.loss) def training_step_and_backward(self, split_batch, batch_idx, opt_idx, optimizer, hiddens): """Wrap forward, zero_grad and backward in a closure so second order methods work""" with self.trainer.profiler.profile("training_step_and_backward"): # lightning module hook result = self.training_step(split_batch, batch_idx, opt_idx, hiddens) if not self._skip_backward and self.trainer.lightning_module.automatic_optimization: is_first_batch_to_accumulate = batch_idx % self.trainer.accumulate_grad_batches == 0 if is_first_batch_to_accumulate: self.on_before_zero_grad(optimizer) self.optimizer_zero_grad(batch_idx, optimizer, opt_idx) # backward pass if result is not None: with self.trainer.profiler.profile("backward"): self.backward(result, optimizer, opt_idx) # hook - call this hook only # when gradients have finished to accumulate if not self.should_accumulate(): self.on_after_backward(batch_idx, result.loss) # check if loss or model weights are nan if self.trainer.terminate_on_nan: self._check_finite(result.loss) else: self.warning_cache.warn( "training_step returned None. If this was on purpose, ignore this warning..." ) return result def _check_finite(self, loss: torch.Tensor) -> None: if not torch.isfinite(loss).all(): raise ValueError(f'The loss returned in `training_step` is {loss}.') model = self.trainer.lightning_module detect_nan_parameters(model) def backward(self, result, optimizer, opt_idx, *args, **kwargs): self.trainer.dev_debugger.track_event("backward_call") should_accumulate = self.should_accumulate() # backward can be called manually in the training loop if isinstance(result, torch.Tensor): self.trainer.accelerator.backward(result, optimizer, opt_idx, should_accumulate, *args, **kwargs) else: result.closure_loss = self.trainer.accelerator.backward( result.closure_loss, optimizer, opt_idx, should_accumulate, *args, **kwargs ) if not self.should_accumulate(): # track gradients grad_norm_dict = self.track_and_norm_grad(optimizer=optimizer) if grad_norm_dict: self.trainer.lightning_module._current_fx_name = "on_after_backward" self.trainer.lightning_module.log_grad_norm(grad_norm_dict) def update_lr_schedulers(self, interval: str) -> None: if interval == "step": finished_accumulation = self._accumulated_batches_reached() finished_epoch = self._num_training_batches_reached() if not finished_accumulation and not finished_epoch: return self.trainer.optimizer_connector.update_learning_rates( interval=interval, opt_indices=[opt_idx for opt_idx, _ in self.get_active_optimizers()], ) def increment_accumulated_grad_global_step(self): num_accumulated_batches_reached = self._accumulated_batches_reached() num_training_batches_reached = self._num_training_batches_reached() # progress global step according to grads progress if num_accumulated_batches_reached or num_training_batches_reached: self.global_step = self.trainer.accelerator.update_global_step(self.total_batch_idx, self.global_step) def _accumulated_batches_reached(self): return (self.batch_idx + 1) % self.trainer.accumulate_grad_batches == 0 def _num_training_batches_reached(self, is_last_batch=False): return (self.batch_idx + 1) == self.trainer.num_training_batches or is_last_batch def should_accumulate(self): # checks if backward or backward + optimizer step (via closure) accumulation_done = self._accumulated_batches_reached() is_final_batch = self._num_training_batches_reached() return not (accumulation_done or is_final_batch) def _should_check_val_fx(self, batch_idx: int, is_last_batch: bool) -> bool: """ Decide if we should run validation. """ if not self.trainer.enable_validation: return False is_val_check_epoch = (self.trainer.current_epoch + 1) % self.trainer.check_val_every_n_epoch == 0 if not is_val_check_epoch: return False # val_check_batch is inf for iterable datasets with no length defined is_infinite_dataset = self.trainer.val_check_batch == float('inf') if is_last_batch and is_infinite_dataset: return True if self.trainer.should_stop: return True # TODO: let training/eval loop handle logic around limit_*_batches and val_check_batch is_val_check_batch = is_last_batch if isinstance(self.trainer.limit_train_batches, int) and is_infinite_dataset: is_val_check_batch = (batch_idx + 1) % self.trainer.limit_train_batches == 0 elif self.trainer.val_check_batch != float('inf'): is_val_check_batch = (batch_idx + 1) % self.trainer.val_check_batch == 0 return is_val_check_batch def _build_kwargs(self, batch, batch_idx, opt_idx, hiddens): # enable not needing to add opt_idx to training_step step_kwargs = OrderedDict([('batch', batch), ('batch_idx', batch_idx)]) lightning_module = self.trainer.lightning_module if len(self.trainer.optimizers) > 1: training_step_fx = getattr(lightning_module, "training_step") has_opt_idx_in_train_step = is_param_in_hook_signature(training_step_fx, "optimizer_idx") if has_opt_idx_in_train_step: if not lightning_module.automatic_optimization: self.warning_cache.warn( "`training_step` hook signature has changed in v1.3." " `optimizer_idx` argument has been removed in case of manual optimization. Support for" " the old signature will be removed in v1.5", DeprecationWarning ) step_kwargs['optimizer_idx'] = opt_idx elif not has_opt_idx_in_train_step and self.trainer.lightning_module.automatic_optimization: raise ValueError( f"Your LightningModule defines {len(self.trainer.optimizers)} optimizers but" ' `training_step` is missing the `optimizer_idx` argument.' ) # pass hiddens if using tbptt if self._truncated_bptt_enabled(): step_kwargs['hiddens'] = hiddens return step_kwargs def _truncated_bptt_enabled(self) -> bool: """ Temporary tbptt utilities until this flag is fully migrated to the lightning module. """ return self._truncated_bptt_steps() > 0 def _truncated_bptt_steps(self) -> int: lightning_module = self.trainer.lightning_module # Give precedence to the LightningModule as the Trainer flag will be removed in v1.5 if lightning_module.truncated_bptt_steps > 0: return lightning_module.truncated_bptt_steps return self.trainer.truncated_bptt_steps or 0 def save_loggers_on_train_batch_end(self): # when loggers should save to disk should_flush_logs = self.trainer.logger_connector.should_flush_logs if should_flush_logs and self.trainer.is_global_zero and self.trainer.logger is not None: self.trainer.logger.save() def run_optimization_start(self, opt_idx, optimizer): # make sure only the gradients of the current optimizer's parameters are calculated # in the training step to prevent dangling gradients in multiple-optimizer setup. if self.trainer.lightning_module.automatic_optimization and len(self.trainer.optimizers) > 1: model = self.trainer.lightning_module model.toggle_optimizer(optimizer, opt_idx) def update_running_loss(self, current_loss: torch.Tensor) -> None: if self.trainer.lightning_module.automatic_optimization: # track total loss for logging (avoid mem leaks) self.accumulated_loss.append(current_loss) accumulated_loss = self.accumulated_loss.mean() if accumulated_loss is not None: # calculate running loss for display self.running_loss.append(self.accumulated_loss.mean() * self.trainer.accumulate_grad_batches) # reset for next set of accumulated grads self.accumulated_loss.reset()
class EvaluationLoop(DataLoaderLoop): """Loops over all dataloaders for evaluation.""" def __init__(self): super().__init__() self.outputs: List[EPOCH_OUTPUT] = [] self.epoch_loop = EvaluationEpochLoop() self._results = ResultCollection(training=False) self._max_batches: Optional[Union[int, Sequence[int]]] = None self._has_run: bool = False @property def num_dataloaders(self) -> int: """Returns the total number of dataloaders.""" # case where user does: # return dl1, dl2 dataloaders = self.dataloaders if dataloaders is None: return 0 length = len(dataloaders) if length > 0 and isinstance(dataloaders[0], (list, tuple)): length = len(dataloaders[0]) return length @property def dataloaders(self) -> Sequence[DataLoader]: """Returns the validation or test dataloaders.""" if self.trainer.testing: return self.trainer.test_dataloaders return self.trainer.val_dataloaders def connect(self, epoch_loop: EvaluationEpochLoop): """Connect the evaluation epoch loop with this loop.""" self.epoch_loop = epoch_loop @property def done(self) -> bool: """Returns whether all dataloaders are processed or evaluation should be skipped altogether.""" return super().done or self.skip @property def skip(self) -> bool: """Returns whether the evaluation should be skipped.""" max_batches = self.get_max_batches() return sum(max_batches) == 0 def reset(self) -> None: """Resets the internal state of the loop.""" self._max_batches = self.get_max_batches() # bookkeeping self.outputs = [] if isinstance(self._max_batches, int): self._max_batches = [self._max_batches] * len(self.dataloaders) super().reset() def on_skip(self) -> List: return [] def on_run_start(self, *args: Any, **kwargs: Any) -> None: """Runs the ``on_evaluation_model_eval``, ``on_evaluation_start`` and ``on_evaluation_epoch_start`` hooks.""" void(*args, **kwargs) # hook self.on_evaluation_model_eval() self.trainer.lightning_module.zero_grad() self.on_evaluation_start() self.on_evaluation_epoch_start() def advance(self, *args: Any, **kwargs: Any) -> None: """Performs evaluation on one single dataloader.""" void(*args, **kwargs) dataloader_idx: int = self.current_dataloader_idx dataloader = self.trainer.accelerator.process_dataloader( self.current_dataloader) dataloader = self.trainer.data_connector.get_profiled_dataloader( dataloader, dataloader_idx=dataloader_idx) dl_max_batches = self._max_batches[dataloader_idx] dl_outputs = self.epoch_loop.run(dataloader, dataloader_idx, dl_max_batches, self.num_dataloaders) # store batch level output per dataloader self.outputs.append(dl_outputs) if not self.trainer.sanity_checking: # indicate the loop has run self._has_run = True def on_run_end(self) -> Any: """Runs the ``on_evaluation_epoch_end`` hook.""" outputs = self.outputs # free memory self.outputs = [] # with a single dataloader don't pass a 2D list if len(outputs) > 0 and self.num_dataloaders == 1: outputs = outputs[0] # lightning module method self.evaluation_epoch_end(outputs) # hook self.on_evaluation_epoch_end() # log epoch metrics eval_loop_results = self.trainer.logger_connector.update_eval_epoch_metrics( ) # hook self.on_evaluation_end() # enable train mode again self.on_evaluation_model_train() return eval_loop_results def get_max_batches(self) -> List[Union[int, float]]: """Returns the max number of batches for each dataloader.""" if self.trainer.testing: max_batches = self.trainer.num_test_batches else: if self.trainer.sanity_checking: self.trainer.num_sanity_val_batches = [ min(self.trainer.num_sanity_val_steps, val_batches) for val_batches in self.trainer.num_val_batches ] max_batches = self.trainer.num_sanity_val_batches else: max_batches = self.trainer.num_val_batches return max_batches def reload_evaluation_dataloaders(self) -> None: """Reloads dataloaders if necessary.""" if self.trainer.testing: self.trainer.reset_test_dataloader() elif self.trainer.val_dataloaders is None or self.trainer._should_reload_dl_epoch: self.trainer.reset_val_dataloader() def on_evaluation_start(self, *args: Any, **kwargs: Any) -> None: """Runs ``on_{validation/test}_start`` hooks.""" assert self._results is not None self._results.to(device=self.trainer.lightning_module.device) if self.trainer.testing: self.trainer.call_hook("on_test_start", *args, **kwargs) else: self.trainer.call_hook("on_validation_start", *args, **kwargs) def on_evaluation_model_eval(self) -> None: """Sets model to eval mode.""" if self.trainer.testing: self.trainer.call_hook("on_test_model_eval") else: self.trainer.call_hook("on_validation_model_eval") def on_evaluation_model_train(self) -> None: """Sets model to train mode.""" model_ref = self.trainer.lightning_module if self.trainer.testing: model_ref.on_test_model_train() else: model_ref.on_validation_model_train() def on_evaluation_end(self, *args: Any, **kwargs: Any) -> None: """Runs ``on_{validation/test}_end`` hook.""" if self.trainer.testing: self.trainer.call_hook("on_test_end", *args, **kwargs) else: self.trainer.call_hook("on_validation_end", *args, **kwargs) # reset any `torchmetrics.Metric` and the logger connector state self.trainer.logger_connector.reset(metrics=True) def on_evaluation_epoch_start(self, *args: Any, **kwargs: Any) -> None: """Runs ``on_epoch_start`` and ``on_{validation/test}_epoch_start`` hooks.""" self.trainer.logger_connector.on_epoch_start() self.trainer.call_hook("on_epoch_start", *args, **kwargs) if self.trainer.testing: self.trainer.call_hook("on_test_epoch_start", *args, **kwargs) else: self.trainer.call_hook("on_validation_epoch_start", *args, **kwargs) def evaluation_epoch_end(self, outputs: EPOCH_OUTPUT) -> None: """Runs ``{validation/test}_epoch_end``""" # inform logger the batch loop has finished self.trainer.logger_connector.epoch_end_reached() # call the model epoch end model = self.trainer.lightning_module # unset dataloader_idx in model model._current_dataloader_idx = None if self.trainer.testing: if is_overridden("test_epoch_end", model): model._current_fx_name = "test_epoch_end" model.test_epoch_end(outputs) else: if is_overridden("validation_epoch_end", model): model._current_fx_name = "validation_epoch_end" model.validation_epoch_end(outputs) def on_evaluation_epoch_end(self) -> None: """Runs ``on_{validation/test}_epoch_end`` hook.""" hook_name = "on_test_epoch_end" if self.trainer.testing else "on_validation_epoch_end" self.trainer.call_hook(hook_name) self.trainer.call_hook("on_epoch_end") self.trainer.logger_connector.on_epoch_end() def teardown(self) -> None: self._results.cpu() self.epoch_loop.teardown()
def test_result_collection_on_tensor_with_mean_reduction(): result_collection = ResultCollection(True, torch.device("cpu")) product = [(True, True), (False, True), (True, False), (False, False)] values = torch.arange(1, 10).float( ) # need to convert to float() due to precision issues using torch 1.4 batches = values * values for i, v in enumerate(values): for prog_bar in [False, True]: for logger in [False, True]: for on_step, on_epoch in product: name = "loss" if on_step: name += "_on_step" if on_epoch: name += "_on_epoch" if prog_bar: name += "_prog_bar" if logger: name += "_logger" log_kwargs = dict( fx="training_step", name=name, value=v, on_step=on_step, on_epoch=on_epoch, batch_size=batches[i], prog_bar=prog_bar, logger=logger, ) if not on_step and not on_epoch: with pytest.raises( MisconfigurationException, match="on_step=False, on_epoch=False"): result_collection.log(**log_kwargs) else: result_collection.log(**log_kwargs) total_value = sum(values * batches) total_batches = sum(batches) assert result_collection[ "training_step.loss_on_step_on_epoch"].value == total_value assert result_collection[ "training_step.loss_on_step_on_epoch"].cumulated_batch_size == total_batches batch_metrics = result_collection.metrics(True) max_ = max(values) assert batch_metrics["pbar"] == { "loss_on_step_on_epoch_prog_bar_step": max_, "loss_on_step_on_epoch_prog_bar_logger_step": max_, "loss_on_step_prog_bar": max_, "loss_on_step_prog_bar_logger": max_, } assert batch_metrics["log"] == { "loss_on_step_on_epoch_logger_step": max_, "loss_on_step_logger": max_, "loss_on_step_on_epoch_prog_bar_logger_step": max_, "loss_on_step_prog_bar_logger": max_, } assert batch_metrics["callback"] == { "loss_on_step": max_, "loss_on_step_logger": max_, "loss_on_step_on_epoch": max_, "loss_on_step_on_epoch_logger": max_, "loss_on_step_on_epoch_logger_step": max_, "loss_on_step_on_epoch_prog_bar": max_, "loss_on_step_on_epoch_prog_bar_logger": max_, "loss_on_step_on_epoch_prog_bar_logger_step": max_, "loss_on_step_on_epoch_prog_bar_step": max_, "loss_on_step_on_epoch_step": max_, "loss_on_step_prog_bar": max_, "loss_on_step_prog_bar_logger": max_, } epoch_metrics = result_collection.metrics(False) mean = total_value / total_batches assert epoch_metrics["pbar"] == { "loss_on_epoch_prog_bar": mean, "loss_on_epoch_prog_bar_logger": mean, "loss_on_step_on_epoch_prog_bar_epoch": mean, "loss_on_step_on_epoch_prog_bar_logger_epoch": mean, } assert epoch_metrics["log"] == { "loss_on_epoch_logger": mean, "loss_on_epoch_prog_bar_logger": mean, "loss_on_step_on_epoch_logger_epoch": mean, "loss_on_step_on_epoch_prog_bar_logger_epoch": mean, } assert epoch_metrics["callback"] == { "loss_on_epoch": mean, "loss_on_epoch_logger": mean, "loss_on_epoch_prog_bar": mean, "loss_on_epoch_prog_bar_logger": mean, "loss_on_step_on_epoch": mean, "loss_on_step_on_epoch_epoch": mean, "loss_on_step_on_epoch_logger": mean, "loss_on_step_on_epoch_logger_epoch": mean, "loss_on_step_on_epoch_prog_bar": mean, "loss_on_step_on_epoch_prog_bar_epoch": mean, "loss_on_step_on_epoch_prog_bar_logger": mean, "loss_on_step_on_epoch_prog_bar_logger_epoch": mean, }