def restore(self, checkpoint_path: str, on_gpu: bool):
"""
Restore training state from checkpoint.
Also restores all training state like:
- epoch
- callbacks
- schedulers
- optimizer
"""
# if on_gpu:
# checkpoint = torch.load(checkpoint_path)
# else:
# load on CPU first
checkpoint = pl_load(checkpoint_path,
map_location=lambda storage, loc: storage)
# load model state
model = self.get_model()
# load the state_dict on the model automatically
model.load_state_dict(checkpoint['state_dict'])
# give model a chance to load something
model.on_load_checkpoint(checkpoint)
if on_gpu:
model.cuda(self.root_gpu)
# restore amp scaling
if self.use_amp and self.use_native_amp and 'native_amp_scaling_state' in checkpoint:
self.scaler.load_state_dict(checkpoint['native_amp_scaling_state'])
# load training state (affects trainer only)
self.restore_training_state(checkpoint)
def test_resume_from_checkpoint_epoch_restored(monkeypatch, tmpdir,
tmpdir_server, url_ckpt):
"""Verify resuming from checkpoint runs the right number of epochs"""
# set $TORCH_HOME, which determines torch hub's cache path, to tmpdir
monkeypatch.setenv('TORCH_HOME', tmpdir)
hparams = EvalModelTemplate.get_default_hparams()
def _new_model():
# Create a model that tracks epochs and batches seen
model = EvalModelTemplate(**hparams)
model.num_epochs_seen = 0
model.num_batches_seen = 0
model.num_on_load_checkpoint_called = 0
def increment_epoch(self):
self.num_epochs_seen += 1
def increment_batch(self, _):
self.num_batches_seen += 1
def increment_on_load_checkpoint(self, _):
self.num_on_load_checkpoint_called += 1
# Bind methods to keep track of epoch numbers, batch numbers it has seen
# as well as number of times it has called on_load_checkpoint()
model.on_epoch_end = types.MethodType(increment_epoch, model)
model.on_batch_start = types.MethodType(increment_batch, model)
model.on_load_checkpoint = types.MethodType(
increment_on_load_checkpoint, model)
return model
model = _new_model()
trainer_options = dict(
progress_bar_refresh_rate=0,
max_epochs=2,
train_percent_check=0.65,
val_percent_check=1,
checkpoint_callback=ModelCheckpoint(tmpdir, save_top_k=-1),
default_root_dir=tmpdir,
early_stop_callback=False,
val_check_interval=1.,
)
trainer = Trainer(**trainer_options)
# fit model
trainer.fit(model)
training_batches = trainer.num_training_batches
assert model.num_epochs_seen == 2
assert model.num_batches_seen == training_batches * 2
assert model.num_on_load_checkpoint_called == 0
# Other checkpoints can be uncommented if/when resuming mid-epoch is supported
checkpoints = sorted(
glob.glob(os.path.join(trainer.checkpoint_callback.dirpath, '*.ckpt')))
if url_ckpt:
# transform local paths into url checkpoints
ip, port = tmpdir_server
checkpoints = [
f'http://{ip}:{port}/' + os.path.basename(check)
for check in checkpoints
]
for check in checkpoints:
next_model = _new_model()
state = pl_load(check)
# Resume training
trainer_options['max_epochs'] = 2
new_trainer = Trainer(**trainer_options, resume_from_checkpoint=check)
new_trainer.fit(next_model)
assert state[
'global_step'] + next_model.num_batches_seen == training_batches * trainer_options[
'max_epochs']
assert next_model.num_on_load_checkpoint_called == 1
def load_from_checkpoint(
cls,
checkpoint_path: str,
*args,
map_location: Optional[Union[Dict[str, str], str, torch.device,
int, Callable]] = None,
hparams_file: Optional[str] = None,
tags_csv: Optional[
str] = None, # backward compatible, todo: remove in v0.9.0
**kwargs):
r"""
Primary way of loading a model from a checkpoint. When Lightning saves a checkpoint
it stores the arguments passed to `__init__` in the checkpoint under `module_arguments`
Any arguments specified through \*args and \*\*kwargs will override args stored in `hparams`.
Args:
checkpoint_path: Path to checkpoint. This can also be a URL.
args: Any positional args needed to init the model.
map_location:
If your checkpoint saved a GPU model and you now load on CPUs
or a different number of GPUs, use this to map to the new setup.
The behaviour is the same as in :func:`torch.load`.
hparams_file: Optional path to a .yaml file with hierarchical structure
as in this example::
drop_prob: 0.2
dataloader:
batch_size: 32
You most likely won't need this since Lightning will always save the hyperparameters
to the checkpoint.
However, if your checkpoint weights don't have the hyperparameters saved,
use this method to pass in a .yaml file with the hparams you'd like to use.
These will be converted into a :class:`~dict` and passed into your
:class:`LightningModule` for use.
If your model's `hparams` argument is :class:`~argparse.Namespace`
and .yaml file has hierarchical structure, you need to refactor your model to treat
`hparams` as :class:`~dict`.
.csv files are acceptable here till v0.9.0, see tags_csv argument for detailed usage.
tags_csv:
.. warning:: .. deprecated:: 0.7.6
`tags_csv` argument is deprecated in v0.7.6. Will be removed v0.9.0.
Optional path to a .csv file with two columns (key, value)
as in this example::
key,value
drop_prob,0.2
batch_size,32
Use this method to pass in a .csv file with the hparams you'd like to use.
hparam_overrides: A dictionary with keys to override in the hparams
kwargs: Any keyword args needed to init the model.
Return:
:class:`LightningModule` with loaded weights and hyperparameters (if available).
Example:
.. code-block:: python
# load weights without mapping ...
MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt')
# or load weights mapping all weights from GPU 1 to GPU 0 ...
map_location = {'cuda:1':'cuda:0'}
MyLightningModule.load_from_checkpoint(
'path/to/checkpoint.ckpt',
map_location=map_location
)
# or load weights and hyperparameters from separate files.
MyLightningModule.load_from_checkpoint(
'path/to/checkpoint.ckpt',
hparams_file='/path/to/hparams_file.yaml'
)
# override some of the params with new values
MyLightningModule.load_from_checkpoint(
PATH,
num_layers=128,
pretrained_ckpt_path: NEW_PATH,
)
# predict
pretrained_model.eval()
pretrained_model.freeze()
y_hat = pretrained_model(x)
"""
if map_location is not None:
checkpoint = pl_load(checkpoint_path, map_location=map_location)
else:
checkpoint = pl_load(checkpoint_path,
map_location=lambda storage, loc: storage)
# add the hparams from csv file to checkpoint
if tags_csv is not None:
hparams_file = tags_csv
rank_zero_warn(
'`tags_csv` argument is deprecated in v0.7.6. Will be removed v0.9.0',
DeprecationWarning)
if hparams_file is not None:
extension = hparams_file.split('.')[-1]
if extension.lower() in ('csv'):
hparams = load_hparams_from_tags_csv(hparams_file)
elif extension.lower() in ('yml', 'yaml'):
hparams = load_hparams_from_yaml(hparams_file)
else:
raise ValueError(
'.csv, .yml or .yaml is required for `hparams_file`')
hparams['on_gpu'] = False
# overwrite hparams by the given file
checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY] = hparams
# for past checkpoint need to add the new key
if cls.CHECKPOINT_HYPER_PARAMS_KEY not in checkpoint:
checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY] = {}
# override the hparams with values that were passed in
checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY].update(kwargs)
model = cls._load_model_state(checkpoint, *args, **kwargs)
return model