wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-09 11:32:07 +08:00
Code Issues Packages Projects Releases Wiki Activity

Checkpointing interval (#1272)

Browse Source 
* formatting

* formatting

* fix interval

* fix train loop

* fix test

* parametrize test

* Apply suggestions from code review

Co-Authored-By: Adrian Wälchli <adrian.waelchli@students.unibe.ch>

* fix calling

* flake8

* add types

Co-authored-by: Adrian Wälchli <adrian.waelchli@students.unibe.ch>
Co-authored-by: William Falcon <waf2107@columbia.edu>
This commit is contained in:
Jirka Borovec
2020-03-30 18:37:02 -04:00
committed by GitHub
co-authored by Adrian Wälchli William Falcon
parent 3476d2f279
commit 09167efdb5
15 changed files with 162 additions and 294 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pl_examples/multi_node_examples/multi_node_ddp_demo.py
+1 -5
View File
@@ -16,11 +16,7 @@ np.random.seed(SEED)
def main(hparams):
"""
Main training routine specific for this project
:param hparams:
:return:
"""
"""Main training routine specific for this project."""
# ------------------------
# 1 INIT LIGHTNING MODEL
# ------------------------
pytorch_lightning/callbacks/base.py
+1 -1
View File
@@ -1,6 +1,6 @@
r"""
Callback Base
==============
=============
Abstract base class used to build new callbacks.
"""
pytorch_lightning/callbacks/model_checkpoint.py
+38 -44
View File
@@ -94,7 +94,7 @@ class ModelCheckpoint(Callback):
self.save_top_k = save_top_k
self.save_weights_only = save_weights_only
self.period = period
self.epochs_since_last_check = 0
self.epoch_last_check = None
self.prefix = prefix
self.best_k_models = {}
# {filename: monitor}
@@ -139,21 +139,20 @@ class ModelCheckpoint(Callback):
def format_checkpoint_name(self, epoch, metrics, ver=None):
"""Generate a filename according define template.
Examples
--------
>>> tmpdir = os.path.dirname(__file__)
>>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{epoch}'))
>>> os.path.basename(ckpt.format_checkpoint_name(0, {}))
'epoch=0.ckpt'
>>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{epoch:03d}'))
>>> os.path.basename(ckpt.format_checkpoint_name(5, {}))
'epoch=005.ckpt'
>>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{epoch}-{val_loss:.2f}'))
>>> os.path.basename(ckpt.format_checkpoint_name(2, dict(val_loss=0.123456)))
'epoch=2-val_loss=0.12.ckpt'
>>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{missing:d}'))
>>> os.path.basename(ckpt.format_checkpoint_name(0, {}))
'missing=0.ckpt'
Examples:
>>> tmpdir = os.path.dirname(__file__)
>>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{epoch}'))
>>> os.path.basename(ckpt.format_checkpoint_name(0, {}))
'epoch=0.ckpt'
>>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{epoch:03d}'))
>>> os.path.basename(ckpt.format_checkpoint_name(5, {}))
'epoch=005.ckpt'
>>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{epoch}-{val_loss:.2f}'))
>>> os.path.basename(ckpt.format_checkpoint_name(2, dict(val_loss=0.123456)))
'epoch=2-val_loss=0.12.ckpt'
>>> ckpt = ModelCheckpoint(os.path.join(tmpdir, '{missing:d}'))
>>> os.path.basename(ckpt.format_checkpoint_name(0, {}))
'missing=0.ckpt'
"""
# check if user passed in keys to the string
groups = re.findall(r'(\{.*?)[:\}]', self.filename)
@@ -181,41 +180,36 @@ class ModelCheckpoint(Callback):
metrics = trainer.callback_metrics
epoch = trainer.current_epoch
self.epochs_since_last_check += 1
if self.save_top_k == 0:
# no models are saved
return
if self.epochs_since_last_check >= self.period:
self.epochs_since_last_check = 0
if self.epoch_last_check is not None and (epoch - self.epoch_last_check) < self.period:
# skipping in this term
return
filepath = self.format_checkpoint_name(epoch, metrics)
version_cnt = 0
while os.path.isfile(filepath):
filepath = self.format_checkpoint_name(epoch, metrics, ver=version_cnt)
# this epoch called before
version_cnt += 1
self.epoch_last_check = epoch
if self.save_top_k != -1:
current = metrics.get(self.monitor)
filepath = self.format_checkpoint_name(epoch, metrics)
version_cnt = 0
while os.path.isfile(filepath):
filepath = self.format_checkpoint_name(epoch, metrics, ver=version_cnt)
# this epoch called before
version_cnt += 1
if current is None:
warnings.warn(
f'Can save best model only with {self.monitor} available,'
' skipping.', RuntimeWarning)
else:
if self.check_monitor_top_k(current):
self._do_check_save(filepath, current, epoch)
else:
if self.verbose > 0:
log.info(
f'\nEpoch {epoch:05d}: {self.monitor}'
f' was not in top {self.save_top_k}')
if self.save_top_k != -1:
current = metrics.get(self.monitor)
else:
if self.verbose > 0:
log.info(f'\nEpoch {epoch:05d}: saving model to {filepath}')
self._save_model(filepath)
if current is None:
warnings.warn(f'Can save best model only with {self.monitor} available, skipping.', RuntimeWarning)
elif self.check_monitor_top_k(current):
self._do_check_save(filepath, current, epoch)
elif self.verbose > 0:
log.info(f'\nEpoch {epoch:05d}: {self.monitor} was not in top {self.save_top_k}')
else:
if self.verbose > 0:
log.info(f'\nEpoch {epoch:05d}: saving model to {filepath}')
self._save_model(filepath)
def _do_check_save(self, filepath, current, epoch):
# remove kth
pytorch_lightning/profiler/profiler.py
+20 -19
View File
@@ -17,15 +17,15 @@ class BaseProfiler(ABC):
"""
@abstractmethod
def start(self, action_name):
def start(self, action_name: str) -> None:
"""Defines how to start recording an action."""
@abstractmethod
def stop(self, action_name):
def stop(self, action_name: str) -> None:
"""Defines how to record the duration once an action is complete."""
@contextmanager
def profile(self, action_name):
def profile(self, action_name: str) -> None:
"""
Yields a context manager to encapsulate the scope of a profiled action.
@@ -43,7 +43,7 @@ class BaseProfiler(ABC):
finally:
self.stop(action_name)
def profile_iterable(self, iterable, action_name):
def profile_iterable(self, iterable, action_name: str) -> None:
iterator = iter(iterable)
while True:
try:
@@ -55,7 +55,7 @@ class BaseProfiler(ABC):
self.stop(action_name)
break
def describe(self):
def describe(self) -> None:
"""Logs a profile report after the conclusion of the training run."""
pass
@@ -69,10 +69,10 @@ class PassThroughProfiler(BaseProfiler):
def __init__(self):
pass
def start(self, action_name):
def start(self, action_name: str) -> None:
pass
def stop(self, action_name):
def stop(self, action_name: str) -> None:
pass
@@ -86,14 +86,14 @@ class Profiler(BaseProfiler):
self.current_actions = {}
self.recorded_durations = defaultdict(list)
def start(self, action_name):
def start(self, action_name: str) -> None:
if action_name in self.current_actions:
raise ValueError(
f"Attempted to start {action_name} which has already started."
)
self.current_actions[action_name] = time.monotonic()
def stop(self, action_name):
def stop(self, action_name: str) -> None:
end_time = time.monotonic()
if action_name not in self.current_actions:
raise ValueError(
@@ -103,7 +103,7 @@ class Profiler(BaseProfiler):
duration = end_time - start_time
self.recorded_durations[action_name].append(duration)
def describe(self):
def describe(self) -> None:
output_string = "\n\nProfiler Report\n"
def log_row(action, mean, total):
@@ -126,24 +126,25 @@ class AdvancedProfiler(BaseProfiler):
verbose and you should only use this if you want very detailed reports.
"""
def __init__(self, output_filename=None, line_count_restriction=1.0):
def __init__(self, output_filename: str = None, line_count_restriction: float = 1.0):
"""
:param output_filename (str): optionally save profile results to file instead of printing
to std out when training is finished.
:param line_count_restriction (int|float): this can be used to limit the number of functions
reported for each action. either an integer (to select a count of lines),
or a decimal fraction between 0.0 and 1.0 inclusive (to select a percentage of lines)
Args:
output_filename: optionally save profile results to file instead of printing
to std out when training is finished.
line_count_restriction: this can be used to limit the number of functions
reported for each action. either an integer (to select a count of lines),
or a decimal fraction between 0.0 and 1.0 inclusive (to select a percentage of lines)
"""
self.profiled_actions = {}
self.output_filename = output_filename
self.line_count_restriction = line_count_restriction
def start(self, action_name):
def start(self, action_name: str) -> None:
if action_name not in self.profiled_actions:
self.profiled_actions[action_name] = cProfile.Profile()
self.profiled_actions[action_name].enable()
def stop(self, action_name):
def stop(self, action_name: str) -> None:
pr = self.profiled_actions.get(action_name)
if pr is None:
raise ValueError( # pragma: no-cover
@@ -151,7 +152,7 @@ class AdvancedProfiler(BaseProfiler):
)
pr.disable()
def describe(self):
def describe(self) -> None:
self.recorded_stats = {}
for action_name, pr in self.profiled_actions.items():
s = io.StringIO()
pytorch_lightning/trainer/distrib_data_parallel.py
+3 -4
View File
@@ -199,10 +199,9 @@ class TrainerDDPMixin(ABC):
self.use_ddp2 = distributed_backend == 'ddp2'
elif distributed_backend is None:
m = 'You requested multiple GPUs but did not specify a backend' \
'Trainer(distributed_backend=dp) (or ddp, ddp2)' \
'Setting distributed_backend=dp for you'
warnings.warn(m)
warnings.warn('You requested multiple GPUs but did not specify a backend, e.g.'
' Trainer(distributed_backend=dp) (or ddp, ddp2).'
' Setting distributed_backend=dp for you.')
self.use_dp = True
self.use_ddp = False
self.use_ddp2 = False
pytorch_lightning/trainer/distrib_parts.py
+10 -14
View File
@@ -491,9 +491,8 @@ class TrainerDPMixin(ABC):
if self.precision == 16:
os.environ['XLA_USE_BF16'] = str(1)
m = f'INIT TPU local core: {self.tpu_local_core_rank}, ' \
f'global rank: {self.tpu_global_core_rank}'
log.info(m)
log.info(f'INIT TPU local core: {self.tpu_local_core_rank},'
f' global rank: {self.tpu_global_core_rank}')
# continue training routine
self.run_pretrain_routine(model)
@@ -512,12 +511,10 @@ class TrainerDPMixin(ABC):
# https://github.com/NVIDIA/apex/issues/227
if self.use_dp and self.use_amp:
if self.amp_level == 'O2':
m = f"""
Amp level {self.amp_level} with DataParallel is not supported.
See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227.
We recommend you switch to ddp if you want to use amp
"""
raise MisconfigurationException(m)
raise MisconfigurationException(
f'Amp level {self.amp_level} with DataParallel is not supported.'
f' See this note from NVIDIA for more info: https://github.com/NVIDIA/apex/issues/227.'
f' We recommend you switch to ddp if you want to use amp')
else:
model, optimizers = model.configure_apex(amp, model, self.optimizers, self.amp_level)
@@ -584,11 +581,10 @@ def sanitize_gpu_ids(gpus):
all_available_gpus = get_all_available_gpus()
for gpu in gpus:
if gpu not in all_available_gpus:
message = f"""
You requested GPUs: {gpus}
But your machine only has: {all_available_gpus}
"""
raise MisconfigurationException(message)
raise MisconfigurationException(f"""
You requested GPUs: {gpus}
But your machine only has: {all_available_gpus}
""")
return gpus
pytorch_lightning/trainer/evaluation_loop.py
+3 -3
View File
@@ -322,9 +322,9 @@ class TrainerEvaluationLoopMixin(ABC):
def run_evaluation(self, test_mode: bool = False):
# when testing make sure user defined a test step
if test_mode and not self.is_overriden('test_step'):
m = "You called `.test()` without defining model's `.test_step()`." \
" Please define and try again"
raise MisconfigurationException(m)
raise MisconfigurationException(
"You called `.test()` without defining model's `.test_step()`."
" Please define and try again")
# Validation/Test begin callbacks
if test_mode:
pytorch_lightning/trainer/trainer.py
+9 -13
View File
@@ -328,11 +328,8 @@ class Trainer(
if self.fast_dev_run:
self.num_sanity_val_steps = 1
self.max_epochs = 1
m = '''
Running in fast_dev_run mode: will run a full train,
val loop using a single batch
'''
log.info(m)
log.info('Running in fast_dev_run mode: will run a full train,'
' val loop using a single batch')
# set default save path if user didn't provide one
self.default_save_path = default_save_path
@@ -739,22 +736,22 @@ class Trainer(
# functions to overwrite with these implementations
if train_dataloader is not None:
if not self.is_overriden('training_step', model):
m = 'You called .fit() with a train_dataloader but did not define training_step()'
raise MisconfigurationException(m)
raise MisconfigurationException(
'You called `.fit()` with a `train_dataloader` but did not define `training_step()`')
model.train_dataloader = _PatchDataLoader(train_dataloader)
if val_dataloaders is not None:
if not self.is_overriden('validation_step', model):
m = 'You called .fit() with a val_dataloaders but did not define validation_step()'
raise MisconfigurationException(m)
raise MisconfigurationException(
'You called `.fit()` with a `val_dataloaders` but did not define `validation_step()`')
model.val_dataloader = _PatchDataLoader(val_dataloaders)
if test_dataloaders is not None:
if not self.is_overriden('test_step', model):
m = 'You called .fit() with a test_dataloaders but did not define test_step()'
raise MisconfigurationException(m)
raise MisconfigurationException(
'You called `.fit()` with a `test_dataloaders` but did not define `test_step()`')
model.test_dataloader = _PatchDataLoader(test_dataloaders)
@@ -855,8 +852,7 @@ class Trainer(
if self.weights_summary in ['full', 'top']:
ref_model.summarize(mode=self.weights_summary)
else:
m = "weights_summary can be None, 'full' or 'top'"
raise MisconfigurationException(m)
raise MisconfigurationException("weights_summary can be None, 'full' or 'top'")
# track model now.
# if cluster resets state, the model will update with the saved weights
pytorch_lightning/trainer/training_io.py
+7 -18
View File
@@ -143,15 +143,12 @@ class TrainerIOMixin(ABC):
# --------------------
# CHECK-POINTING
# --------------------
def restore_weights(self, model):
def restore_weights(self, model: LightningModule):
"""
We attempt to restore weights in this order:
1. HPC weights.
2. if no HPC weights restore checkpoint_path weights
3. otherwise don't restore weights
:param model:
:return:
"""
# clear cache before restore
if self.on_gpu:
@@ -230,17 +227,17 @@ class TrainerIOMixin(ABC):
# --------------------
# MODEL SAVE CHECKPOINT
# --------------------
def _atomic_save(self, checkpoint, filepath):
def _atomic_save(self, checkpoint, filepath: str):
"""Saves a checkpoint atomically, avoiding the creation of incomplete checkpoints.
This will create a temporary checkpoint with a suffix of ``.part``, then copy it to the final location once
saving is finished.
Args:
checkpoint (object): The object to save.
checkpoint: The object to save.
Built to be used with the ``dump_checkpoint`` method, but can deal with anything which ``torch.save``
accepts.
filepath (str|pathlib.Path): The path to which the checkpoint will be saved.
filepath: The path to which the checkpoint will be saved.
This points to the file that the checkpoint will be stored in.
"""
tmp_path = str(filepath) + ".part"
@@ -260,7 +257,7 @@ class TrainerIOMixin(ABC):
self._atomic_save(checkpoint, filepath)
def restore(self, checkpoint_path, on_gpu):
def restore(self, checkpoint_path: str, on_gpu: bool):
"""
Restore training state from checkpoint.
Also restores all training state like:
@@ -268,10 +265,6 @@ class TrainerIOMixin(ABC):
- callbacks
- schedulers
- optimizer
:param checkpoint_path:
:param on_gpu:
:return:
"""
# if on_gpu:
@@ -341,12 +334,8 @@ class TrainerIOMixin(ABC):
# --------------------
# HPC IO
# --------------------
def restore_hpc_weights_if_needed(self, model):
"""
If there is a set of hpc weights, use as signal to restore model
:param model:
:return:
"""
def restore_hpc_weights_if_needed(self, model: LightningModule):
"""If there is a set of hpc weights, use as signal to restore model."""
did_restore = False
# look for hpc weights
pytorch_lightning/trainer/training_loop.py
+20 -13
View File
@@ -461,8 +461,7 @@ class TrainerTrainLoopMixin(ABC):
# CHECKPOINTING, EARLY STOPPING
# ---------------
# save checkpoint even when no test or val step are defined
train_step_only = not self.is_overriden('validation_step')
if self.fast_dev_run or should_check_val or train_step_only:
if self.fast_dev_run or should_check_val:
self.call_checkpoint_callback()
if self.enable_early_stop:
@@ -483,6 +482,13 @@ class TrainerTrainLoopMixin(ABC):
if early_stop_epoch or self.fast_dev_run:
break
# in case validation step is missing and you are not running fast-dev to duplicate last batch
if not self.is_overriden('validation_step') and not (self.fast_dev_run or should_check_val):
self.call_checkpoint_callback()
if self.enable_early_stop:
self.early_stop_callback.check_metrics(self.callback_metrics)
# Epoch end events
with self.profiler.profile('on_epoch_end'):
# callbacks
@@ -709,20 +715,20 @@ class TrainerTrainLoopMixin(ABC):
with self.profiler.profile('training_end'):
output = model_ref.training_end(output)
m = 'training_end was deprecated in 0.7.0 and will be removed 1.0.0. ' \
'Use training_epoch_end instead'
warnings.warn(m, DeprecationWarning)
warnings.warn('`training_end` was deprecated in 0.7.0 and will be removed 1.0.0.'
' Use training_epoch_end instead', DeprecationWarning)
# format and reduce outputs accordingly
output = self.process_output(output, train=True)
return output
def update_learning_rates(self, interval):
''' Update learning rates
def update_learning_rates(self, interval: str):
"""Update learning rates.
Args:
interval (str): either 'epoch' or 'step'.
'''
interval: either 'epoch' or 'step'.
"""
if not self.lr_schedulers:
return
@@ -738,10 +744,11 @@ class TrainerTrainLoopMixin(ABC):
monitor_val = self.callback_metrics.get(monitor_key)
if monitor_val is None:
avail_metrics = ','.join(list(self.callback_metrics.keys()))
m = 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'
raise MisconfigurationException(m)
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'
)
lr_scheduler['scheduler'].step(monitor_val)
else:
lr_scheduler['scheduler'].step()
tests/base/__init__.py
+2 -2
View File
@@ -40,7 +40,7 @@ class LightningTestModel(LightTrainDataloader,
class LightningTestModelWithoutHyperparametersArg(LightningTestModel):
""" without hparams argument in constructor """
"""Without hparams argument in constructor """
def __init__(self):
import tests.base.utils as tutils
@@ -51,7 +51,7 @@ class LightningTestModelWithoutHyperparametersArg(LightningTestModel):
class LightningTestModelWithUnusedHyperparametersArg(LightningTestModelWithoutHyperparametersArg):
""" has hparams argument in constructor but is not used """
"""It has hparams argument in constructor but is not used."""
def __init__(self, hparams):
super().__init__()
tests/base/mixins.py
+4 -7
View File
@@ -14,11 +14,7 @@ class LightValidationStepMixin:
return self._dataloader(train=False)
def validation_step(self, batch, batch_idx, *args, **kwargs):
"""
Lightning calls this inside the validation loop
:param batch:
:return:
"""
"""Lightning calls this inside the validation loop."""
x, y = batch
x = x.view(x.size(0), -1)
y_hat = self(x)
@@ -66,8 +62,9 @@ class LightValidationMixin(LightValidationStepMixin):
def validation_epoch_end(self, outputs):
"""
Called at the end of validation to aggregate outputs
:param outputs: list of individual outputs of each validation step
:return:
Args:
outputs: list of individual outputs of each validation step
"""
# if returned a scalar from validation_step, outputs is a list of tensor scalars
# we return just the average in this case (if we want)
tests/base/models.py
+6 -23
View File
@@ -42,16 +42,10 @@ class DictHparamsModel(LightningModule):
class TestModelBase(LightningModule):
"""
Base LightningModule for testing. Implements only the required
interface
"""
"""Base LightningModule for testing. Implements only the required interface."""
def __init__(self, hparams, force_remove_distributed_sampler=False):
"""
Pass in parsed HyperOptArgumentParser to the model
:param hparams:
"""
def __init__(self, hparams, force_remove_distributed_sampler: bool = False):
"""Pass in parsed HyperOptArgumentParser to the model."""
# init superclass
super().__init__()
self.hparams = hparams
@@ -71,10 +65,7 @@ class TestModelBase(LightningModule):
# MODEL SETUP
# ---------------------
def __build_model(self):
"""
Layout model
:return:
"""
"""Layout model."""
self.c_d1 = nn.Linear(in_features=self.hparams.in_features,
out_features=self.hparams.hidden_dim)
self.c_d1_bn = nn.BatchNorm1d(self.hparams.hidden_dim)
@@ -87,11 +78,7 @@ class TestModelBase(LightningModule):
# TRAINING
# ---------------------
def forward(self, x):
"""
No special modification required for lightning, define as you normally would
:param x:
:return:
"""
"""No special modification required for lightning, define as you normally would."""
x = self.c_d1(x)
x = torch.tanh(x)
x = self.c_d1_bn(x)
@@ -107,11 +94,7 @@ class TestModelBase(LightningModule):
return nll
def training_step(self, batch, batch_idx, optimizer_idx=None):
"""
Lightning calls this inside the training loop
:param batch:
:return:
"""
"""Lightning calls this inside the training loop"""
# forward pass
x, y = batch
x = x.view(x.size(0), -1)
tests/test_profiler.py
+20 -4
View File
@@ -35,7 +35,11 @@ def advanced_profiler():
return profiler
@pytest.mark.parametrize("action,expected", [("a", [3, 1]), ("b", [2]), ("c", [1])])
@pytest.mark.parametrize(["action", "expected"], [
pytest.param("a", [3, 1]),
pytest.param("b", [2]),
pytest.param("c", [1])
])
def test_simple_profiler_durations(simple_profiler, action, expected):
"""Ensure the reported durations are reasonably accurate."""
@@ -50,7 +54,11 @@ def test_simple_profiler_durations(simple_profiler, action, expected):
)
@pytest.mark.parametrize("action,expected", [("a", [3, 1]), ("b", [2]), ("c", [1])])
@pytest.mark.parametrize(["action", "expected"], [
pytest.param("a", [3, 1]),
pytest.param("b", [2]),
pytest.param("c", [1])
])
def test_simple_profiler_iterable_durations(simple_profiler, action, expected):
"""Ensure the reported durations are reasonably accurate."""
iterable = _sleep_generator(expected)
@@ -96,7 +104,11 @@ def test_simple_profiler_value_errors(simple_profiler):
simple_profiler.stop(action)
@pytest.mark.parametrize("action,expected", [("a", [3, 1]), ("b", [2]), ("c", [1])])
@pytest.mark.parametrize(["action", "expected"], [
pytest.param("a", [3, 1]),
pytest.param("b", [2]),
pytest.param("c", [1])
])
def test_advanced_profiler_durations(advanced_profiler, action, expected):
for duration in expected:
@@ -114,7 +126,11 @@ def test_advanced_profiler_durations(advanced_profiler, action, expected):
)
@pytest.mark.parametrize("action,expected", [("a", [3, 1]), ("b", [2]), ("c", [1])])
@pytest.mark.parametrize(["action", "expected"], [
pytest.param("a", [3, 1]),
pytest.param("b", [2]),
pytest.param("c", [1])
])
def test_advanced_profiler_iterable_durations(advanced_profiler, action, expected):
"""Ensure the reported durations are reasonably accurate."""
iterable = _sleep_generator(expected)
tests/trainer/test_trainer.py
+18 -124
View File
@@ -248,7 +248,19 @@ def test_dp_output_reduce():
assert reduced['b']['c'] == out['b']['c']
def test_model_checkpoint_options(tmpdir):
@pytest.mark.parametrize(["save_top_k", "file_prefix", "expected_files"], [
pytest.param(-1, '', {'epoch=4.ckpt', 'epoch=3.ckpt', 'epoch=2.ckpt', 'epoch=1.ckpt', 'epoch=0.ckpt'},
id="CASE K=-1 (all)"),
pytest.param(1, 'test_prefix_', {'test_prefix_epoch=4.ckpt'},
id="CASE K=1 (2.5, epoch 4)"),
pytest.param(2, '', {'epoch=4.ckpt', 'epoch=2.ckpt'},
id="CASE K=2 (2.5 epoch 4, 2.8 epoch 2)"),
pytest.param(4, '', {'epoch=1.ckpt', 'epoch=4.ckpt', 'epoch=3.ckpt', 'epoch=2.ckpt'},
id="CASE K=4 (save all 4 base)"),
pytest.param(3, '', {'epoch=2.ckpt', 'epoch=3.ckpt', 'epoch=4.ckpt'},
id="CASE K=3 (save the 2nd, 3rd, 4th model)"),
])
def test_model_checkpoint_options(tmpdir, save_top_k, file_prefix, expected_files):
"""Test ModelCheckpoint options."""
def mock_save_function(filepath):
@@ -258,13 +270,9 @@ def test_model_checkpoint_options(tmpdir):
_ = LightningTestModel(hparams)
# simulated losses
save_dir = os.path.join(tmpdir, '1')
os.mkdir(save_dir)
losses = [10, 9, 2.8, 5, 2.5]
# -----------------
# CASE K=-1 (all)
checkpoint_callback = ModelCheckpoint(save_dir, save_top_k=-1, verbose=1)
checkpoint_callback = ModelCheckpoint(tmpdir, save_top_k=save_top_k, prefix=file_prefix, verbose=1)
checkpoint_callback.save_function = mock_save_function
trainer = Trainer()
@@ -274,127 +282,13 @@ def test_model_checkpoint_options(tmpdir):
trainer.callback_metrics = {'val_loss': loss}
checkpoint_callback.on_validation_end(trainer, trainer.get_model())
file_lists = set(os.listdir(save_dir))
file_lists = set(os.listdir(tmpdir))
assert len(file_lists) == len(losses), "Should save all models when save_top_k=-1"
assert len(file_lists) == len(expected_files), \
"Should save %i models when save_top_k=%i" % (len(expected_files), save_top_k)
# verify correct naming
for fname in {'epoch=4.ckpt',
'epoch=3.ckpt',
'epoch=2.ckpt',
'epoch=1.ckpt',
'epoch=0.ckpt'}:
assert fname in file_lists
save_dir = os.path.join(tmpdir, '2')
os.mkdir(save_dir)
# -----------------
# CASE K=0 (none)
checkpoint_callback = ModelCheckpoint(save_dir, save_top_k=0, verbose=1)
checkpoint_callback.save_function = mock_save_function
trainer = Trainer()
# emulate callback's calls during the training
for i, loss in enumerate(losses):
trainer.current_epoch = i
trainer.callback_metrics = {'val_loss': loss}
checkpoint_callback.on_validation_end(trainer, trainer.get_model())
file_lists = os.listdir(save_dir)
assert len(file_lists) == 0, "Should save 0 models when save_top_k=0"
save_dir = os.path.join(tmpdir, '3')
os.mkdir(save_dir)
# -----------------
# CASE K=1 (2.5, epoch 4)
checkpoint_callback = ModelCheckpoint(save_dir, save_top_k=1, verbose=1, prefix='test_prefix_')
checkpoint_callback.save_function = mock_save_function
trainer = Trainer()
# emulate callback's calls during the training
for i, loss in enumerate(losses):
trainer.current_epoch = i
trainer.callback_metrics = {'val_loss': loss}
checkpoint_callback.on_validation_end(trainer, trainer.get_model())
file_lists = set(os.listdir(save_dir))
assert len(file_lists) == 1, "Should save 1 model when save_top_k=1"
assert 'test_prefix_epoch=4.ckpt' in file_lists
save_dir = os.path.join(tmpdir, '4')
os.mkdir(save_dir)
# -----------------
# CASE K=2 (2.5 epoch 4, 2.8 epoch 2)
# make sure other files don't get deleted
checkpoint_callback = ModelCheckpoint(save_dir, save_top_k=2, verbose=1)
open(f"{save_dir}/other_file.ckpt", 'a').close()
checkpoint_callback.save_function = mock_save_function
trainer = Trainer()
# emulate callback's calls during the training
for i, loss in enumerate(losses):
trainer.current_epoch = i
trainer.callback_metrics = {'val_loss': loss}
checkpoint_callback.on_validation_end(trainer, trainer.get_model())
file_lists = set(os.listdir(save_dir))
assert len(file_lists) == 3, 'Should save 2 model when save_top_k=2'
for fname in {'epoch=4.ckpt',
'epoch=2.ckpt',
'other_file.ckpt'}:
assert fname in file_lists
save_dir = os.path.join(tmpdir, '5')
os.mkdir(save_dir)
# -----------------
# CASE K=4 (save all 4 base)
# multiple checkpoints within same epoch
checkpoint_callback = ModelCheckpoint(save_dir, save_top_k=4, verbose=1)
checkpoint_callback.save_function = mock_save_function
trainer = Trainer()
# emulate callback's calls during the training
for loss in losses:
trainer.current_epoch = 0
trainer.callback_metrics = {'val_loss': loss}
checkpoint_callback.on_validation_end(trainer, trainer.get_model())
file_lists = set(os.listdir(save_dir))
assert len(file_lists) == 4, 'Should save all 4 models when save_top_k=4 within same epoch'
save_dir = os.path.join(tmpdir, '6')
os.mkdir(save_dir)
# -----------------
# CASE K=3 (save the 2nd, 3rd, 4th model)
# multiple checkpoints within same epoch
checkpoint_callback = ModelCheckpoint(save_dir, save_top_k=3, verbose=1)
checkpoint_callback.save_function = mock_save_function
trainer = Trainer()
# emulate callback's calls during the training
for loss in losses:
trainer.current_epoch = 0
trainer.callback_metrics = {'val_loss': loss}
checkpoint_callback.on_validation_end(trainer, trainer.get_model())
file_lists = set(os.listdir(save_dir))
assert len(file_lists) == 3, 'Should save 3 models when save_top_k=3'
for fname in {'epoch=0.ckpt',
'epoch=0.ckpt',
'epoch=0.ckpt'}:
for fname in expected_files:
assert fname in file_lists