mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-09-09 11:32:07 +08:00
No auto load weights (#985)
* remove autoload * remove autoload * added weights loading docs * checkpoint loading saving docs * checkpoint loading saving docs * checkpoint loading saving docs * docs (#1010) * remove autoload * remove autoload * added weights loading docs * checkpoint loading saving docs * checkpoint loading saving docs * checkpoint loading saving docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs * docs
This commit is contained in:
William Falcon
@@ -1,80 +0,0 @@
|
||||
Checkpointing
|
||||
==============
|
||||
|
||||
.. _model-saving:
|
||||
|
||||
Model saving
|
||||
-------------------
|
||||
To save a LightningModule, provide a :meth:`pytorch_lightning.callbacks.ModelCheckpoint` callback.
|
||||
|
||||
The Lightning checkpoint also saves the hparams (hyperparams) passed into the LightningModule init.
|
||||
|
||||
.. note:: hparams is a `Namespace <https://docs.python.org/2/library/argparse.html#argparse.Namespace>`_ or dictionary.
|
||||
|
||||
.. code-block:: python
|
||||
:emphasize-lines: 8
|
||||
|
||||
from argparse import Namespace
|
||||
|
||||
# usually these come from command line args
|
||||
args = Namespace(**{'learning_rate':0.001})
|
||||
|
||||
# define you module to have hparams as the first arg
|
||||
# this means your checkpoint will have everything that went into making
|
||||
# this model (in this case, learning rate)
|
||||
class MyLightningModule(pl.LightningModule):
|
||||
|
||||
def __init__(self, hparams, ...):
|
||||
self.hparams = hparams
|
||||
|
||||
my_model = MyLightningModule(args)
|
||||
|
||||
# auto-saves checkpoint
|
||||
checkpoint_callback = ModelCheckpoint(filepath='my_path')
|
||||
Trainer(checkpoint_callback=checkpoint_callback)
|
||||
|
||||
|
||||
Model loading
|
||||
-----------------------------------
|
||||
|
||||
To load a model, use :meth:`pytorch_lightning.core.LightningModule.load_from_checkpoint`
|
||||
|
||||
.. note:: If lightning created your checkpoint, your model will receive all the hyperparameters used
|
||||
to create the checkpoint. (See: :ref:`model-saving`).
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# load weights without mapping
|
||||
MyLightningModule.load_from_checkpoint('path/to/checkpoint.ckpt')
|
||||
|
||||
# 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)
|
||||
|
||||
Restoring training session
|
||||
-----------------------------------
|
||||
|
||||
If you want to pick up training from where you left off, you have a few options.
|
||||
|
||||
1. Pass in a logger with the same experiment version to continue training.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# train the first time and set the version number
|
||||
logger = TensorboardLogger(version=10)
|
||||
trainer = Trainer(logger=logger)
|
||||
trainer.fit(model)
|
||||
|
||||
# when you init another logger with that same version, the model
|
||||
# will continue where it left off
|
||||
logger = TensorboardLogger(version=10)
|
||||
trainer = Trainer(logger=logger)
|
||||
trainer.fit(model)
|
||||
|
||||
2. A second option is to pass in a path to a checkpoint (see: :ref:`pytorch_lightning.trainer.trainer.Trainer`).
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# train the first time and set the version number
|
||||
trainer = Trainer(resume_from_checkpoint='some/path/to/my_checkpoint.ckpt')
|
||||
trainer.fit(model)
|
||||
@@ -45,7 +45,6 @@ PyTorch-Lightning Documentation
|
||||
:caption: Common Use Cases
|
||||
|
||||
apex
|
||||
checkpointing
|
||||
slurm
|
||||
debugging
|
||||
experiment_logging
|
||||
@@ -54,6 +53,7 @@ PyTorch-Lightning Documentation
|
||||
fast_training
|
||||
hooks
|
||||
multi_gpu
|
||||
weights_loading
|
||||
single_gpu
|
||||
sequences
|
||||
training_tricks
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
Saving and loading weights
|
||||
==========================
|
||||
|
||||
Lightning can automate saving and loading checkpoints.
|
||||
|
||||
Checkpoint saving
|
||||
-----------------
|
||||
|
||||
Checkpointing is enabled by default to the current working directory.
|
||||
To change the checkpoint path pass in:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
Trainer(default_save_path='/your/path/to/save/checkpoints')
|
||||
|
||||
To modify the behavior of checkpointing pass in your own callback.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from pytorch_lightning.callbacks import ModelCheckpoint
|
||||
|
||||
# DEFAULTS used by the Trainer
|
||||
checkpoint_callback = ModelCheckpoint(
|
||||
filepath=os.getcwd(),
|
||||
save_best_only=True,
|
||||
verbose=True,
|
||||
monitor='val_loss',
|
||||
mode='min',
|
||||
prefix=''
|
||||
)
|
||||
|
||||
trainer = Trainer(checkpoint_callback=checkpoint_callback)
|
||||
|
||||
|
||||
Or disable it by passing
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
trainer = Trainer(checkpoint_callback=False)
|
||||
|
||||
|
||||
The Lightning checkpoint also saves the hparams (hyperparams) passed into the LightningModule init.
|
||||
|
||||
.. note:: hparams is a `Namespace <https://docs.python.org/2/library/argparse.html#argparse.Namespace>`_.
|
||||
|
||||
.. code-block:: python
|
||||
:emphasize-lines: 8
|
||||
|
||||
from argparse import Namespace
|
||||
|
||||
# usually these come from command line args
|
||||
args = Namespace(**{'learning_rate':0.001})
|
||||
|
||||
# define you module to have hparams as the first arg
|
||||
# this means your checkpoint will have everything that went into making
|
||||
# this model (in this case, learning rate)
|
||||
class MyLightningModule(pl.LightningModule):
|
||||
|
||||
def __init__(self, hparams, ...):
|
||||
self.hparams = hparams
|
||||
|
||||
Checkpoint Loading
|
||||
------------------
|
||||
|
||||
You might want to not only load a model but also continue training it. Use this method to
|
||||
restore the trainer state as well. This will continue from the epoch and global step you last left off.
|
||||
However, the dataloaders will start from the first batch again (if you shuffled it shouldn't matter).
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
model = MyLightingModule.load_from_checkpoint(PATH)
|
||||
model.eval()
|
||||
y_hat = model(x)
|
||||
|
||||
A LightningModule is no different than a nn.Module. This means you can load it and use it for
|
||||
predictions as you would a nn.Module.
|
||||
@@ -62,7 +62,7 @@ class TrainerCallbackConfigMixin(ABC):
|
||||
self.weights_save_path = self.default_save_path
|
||||
|
||||
def configure_early_stopping(self, early_stop_callback):
|
||||
if early_stop_callback is True:
|
||||
if early_stop_callback is True or None:
|
||||
self.early_stop_callback = EarlyStopping(
|
||||
monitor='val_loss',
|
||||
patience=3,
|
||||
@@ -71,15 +71,6 @@ class TrainerCallbackConfigMixin(ABC):
|
||||
mode='min'
|
||||
)
|
||||
self.enable_early_stop = True
|
||||
elif early_stop_callback is None:
|
||||
self.early_stop_callback = EarlyStopping(
|
||||
monitor='val_loss',
|
||||
patience=3,
|
||||
strict=False,
|
||||
verbose=False,
|
||||
mode='min'
|
||||
)
|
||||
self.enable_early_stop = True
|
||||
elif not early_stop_callback:
|
||||
self.early_stop_callback = None
|
||||
self.enable_early_stop = False
|
||||
|
||||
@@ -347,6 +347,10 @@ class TrainerEvaluationLoopMixin(ABC):
|
||||
# add metrics to prog bar
|
||||
self.add_tqdm_metrics(prog_bar_metrics)
|
||||
|
||||
# log results of test
|
||||
if test_mode:
|
||||
model.print(prog_bar_metrics)
|
||||
|
||||
# log metrics
|
||||
self.log_metrics(log_metrics, {})
|
||||
|
||||
|
||||
@@ -1096,7 +1096,8 @@ class Trainer(TrainerIOMixin,
|
||||
self.register_slurm_signal_handlers()
|
||||
|
||||
# print model summary
|
||||
if self.proc_rank == 0 and self.weights_summary is not None:
|
||||
# TODO: remove self.testing condition because model.summarize() is wiping out the weights
|
||||
if self.proc_rank == 0 and self.weights_summary is not None and not self.testing:
|
||||
if self.weights_summary in ['full', 'top']:
|
||||
ref_model.summarize(mode=self.weights_summary)
|
||||
else:
|
||||
@@ -1116,7 +1117,7 @@ class Trainer(TrainerIOMixin,
|
||||
# when testing requested only run test and return
|
||||
if self.testing:
|
||||
# only load test dataloader for testing
|
||||
self.reset_test_dataloader(ref_model)
|
||||
# self.reset_test_dataloader(ref_model)
|
||||
self.run_evaluation(test_mode=True)
|
||||
return
|
||||
|
||||
@@ -1189,8 +1190,10 @@ class Trainer(TrainerIOMixin,
|
||||
"""
|
||||
self.testing = True
|
||||
if model is not None:
|
||||
self.model = model
|
||||
self.fit(model)
|
||||
self.run_evaluation(test_mode=True)
|
||||
else:
|
||||
self.run_evaluation(test_mode=True)
|
||||
|
||||
|
||||
class _PatchDataLoader(object):
|
||||
|
||||
@@ -1,94 +1,3 @@
|
||||
"""
|
||||
Lightning can automate saving and loading checkpoints
|
||||
=====================================================
|
||||
|
||||
Checkpointing is enabled by default to the current working directory.
|
||||
To change the checkpoint path pass in::
|
||||
|
||||
Trainer(default_save_path='/your/path/to/save/checkpoints')
|
||||
|
||||
|
||||
To modify the behavior of checkpointing pass in your own callback.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from pytorch_lightning.callbacks import ModelCheckpoint
|
||||
|
||||
# DEFAULTS used by the Trainer
|
||||
checkpoint_callback = ModelCheckpoint(
|
||||
filepath=os.getcwd(),
|
||||
save_best_only=True,
|
||||
verbose=True,
|
||||
monitor='val_loss',
|
||||
mode='min',
|
||||
prefix=''
|
||||
)
|
||||
|
||||
trainer = Trainer(checkpoint_callback=checkpoint_callback)
|
||||
|
||||
|
||||
Restoring training session
|
||||
--------------------------
|
||||
|
||||
You might want to not only load a model but also continue training it. Use this method to
|
||||
restore the trainer state as well. This will continue from the epoch and global step you last left off.
|
||||
However, the dataloaders will start from the first batch again (if you shuffled it shouldn't matter).
|
||||
|
||||
Lightning will restore the session if you pass a logger with the same version and there's a saved checkpoint.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from pytorch_lightning import Trainer
|
||||
from pytorch_lightning.loggers import TestTubeLogger
|
||||
|
||||
logger = TestTubeLogger(
|
||||
save_dir='./savepath',
|
||||
version=1 # An existing version with a saved checkpoint
|
||||
)
|
||||
trainer = Trainer(
|
||||
logger=logger,
|
||||
default_save_path='./savepath'
|
||||
)
|
||||
|
||||
# this fit call loads model weights and trainer state
|
||||
# the trainer continues seamlessly from where you left off
|
||||
# without having to do anything else.
|
||||
trainer.fit(model)
|
||||
|
||||
|
||||
The trainer restores:
|
||||
|
||||
- global_step
|
||||
- current_epoch
|
||||
- All optimizers
|
||||
- All lr_schedulers
|
||||
- Model weights
|
||||
|
||||
You can even change the logic of your model as long as the weights and "architecture" of
|
||||
the system isn't different. If you add a layer, for instance, it might not work.
|
||||
|
||||
At a rough level, here's what happens inside Trainer :py:mod:`pytorch_lightning.base_module.model_saving.py`:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
self.global_step = checkpoint['global_step']
|
||||
self.current_epoch = checkpoint['epoch']
|
||||
|
||||
# restore the optimizers
|
||||
optimizer_states = checkpoint['optimizer_states']
|
||||
for optimizer, opt_state in zip(self.optimizers, optimizer_states):
|
||||
optimizer.load_state_dict(opt_state)
|
||||
|
||||
# restore the lr schedulers
|
||||
lr_schedulers = checkpoint['lr_schedulers']
|
||||
for scheduler, lrs_state in zip(self.lr_schedulers, lr_schedulers):
|
||||
scheduler.load_state_dict(lrs_state)
|
||||
|
||||
# uses the model you passed into trainer
|
||||
model.load_state_dict(checkpoint['state_dict'])
|
||||
|
||||
"""
|
||||
|
||||
import logging as log
|
||||
import os
|
||||
import re
|
||||
@@ -150,11 +59,11 @@ class TrainerIOMixin(ABC):
|
||||
# --------------------
|
||||
def restore_weights(self, model):
|
||||
"""
|
||||
To restore weights we have two cases.
|
||||
First, attempt to restore hpc weights. If successful, don't restore
|
||||
other weights.
|
||||
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
|
||||
|
||||
Otherwise, try to restore actual weights
|
||||
:param model:
|
||||
:return:
|
||||
"""
|
||||
@@ -172,9 +81,6 @@ class TrainerIOMixin(ABC):
|
||||
if not did_restore_hpc_weights:
|
||||
if self.resume_from_checkpoint is not None:
|
||||
self.restore(self.resume_from_checkpoint, on_gpu=self.on_gpu)
|
||||
else:
|
||||
# restore weights if same exp version
|
||||
self.restore_state_if_checkpoint_exists(model)
|
||||
|
||||
# wait for all models to restore weights
|
||||
if self.use_ddp or self.use_ddp2:
|
||||
@@ -190,42 +96,6 @@ class TrainerIOMixin(ABC):
|
||||
if self.on_gpu:
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def restore_state_if_checkpoint_exists(self, model):
|
||||
did_restore = False
|
||||
|
||||
# do nothing if there's not dir or callback
|
||||
no_ckpt_callback = (self.checkpoint_callback is None) or (not self.checkpoint_callback)
|
||||
if no_ckpt_callback or not os.path.exists(self.checkpoint_callback.filepath):
|
||||
return did_restore
|
||||
|
||||
# restore trainer state and model if there is a weight for this experiment
|
||||
last_epoch = -1
|
||||
last_ckpt_name = None
|
||||
|
||||
# find last epoch
|
||||
checkpoints = os.listdir(self.checkpoint_callback.filepath)
|
||||
for name in checkpoints:
|
||||
# ignore hpc ckpts
|
||||
if 'hpc_' in name:
|
||||
continue
|
||||
|
||||
if '.ckpt' in name:
|
||||
epoch = name.split('epoch_')[1]
|
||||
epoch = int(re.sub('[^0-9]', '', epoch))
|
||||
|
||||
if epoch > last_epoch:
|
||||
last_epoch = epoch
|
||||
last_ckpt_name = name
|
||||
|
||||
# restore last checkpoint
|
||||
if last_ckpt_name is not None:
|
||||
last_ckpt_path = os.path.join(self.checkpoint_callback.filepath, last_ckpt_name)
|
||||
self.restore(last_ckpt_path, self.on_gpu)
|
||||
log.info(f'Model and Trainer restored from checkpoint: {last_ckpt_path}')
|
||||
did_restore = True
|
||||
|
||||
return did_restore
|
||||
|
||||
# --------------------
|
||||
# HPC SIGNAL HANDLING
|
||||
# --------------------
|
||||
@@ -304,6 +174,18 @@ class TrainerIOMixin(ABC):
|
||||
self._atomic_save(checkpoint, filepath)
|
||||
|
||||
def restore(self, checkpoint_path, on_gpu):
|
||||
"""
|
||||
Restore training state from checkpoint.
|
||||
Also restores all training state like:
|
||||
- epoch
|
||||
- callbacks
|
||||
- schedulers
|
||||
- optimizer
|
||||
:param checkpoint_path:
|
||||
:param on_gpu:
|
||||
|
||||
:return:
|
||||
"""
|
||||
|
||||
# if on_gpu:
|
||||
# checkpoint = torch.load(checkpoint_path)
|
||||
|
||||
@@ -25,37 +25,6 @@ It can be useful to force training for a minimum number of epochs or limit to a
|
||||
# DEFAULT
|
||||
trainer = Trainer(min_epochs=1, max_epochs=1000)
|
||||
|
||||
Early stopping
|
||||
--------------
|
||||
|
||||
The trainer already sets up default early stopping for you.
|
||||
To modify this behavior, pass in your own EarlyStopping callback.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from pytorch_lightning.callbacks import EarlyStopping
|
||||
|
||||
# DEFAULTS used by Trainer
|
||||
early_stop_callback = EarlyStopping(
|
||||
monitor='val_loss',
|
||||
min_delta=0.00,
|
||||
patience=3,
|
||||
verbose=False,
|
||||
mode='min'
|
||||
)
|
||||
|
||||
# without passing anything in, uses the default callback above
|
||||
trainer = Trainer()
|
||||
|
||||
# pass in your own to override the default callback
|
||||
trainer = Trainer(early_stop_callback=early_stop_callback)
|
||||
|
||||
# pass in min_epochs to enable the callback after min_epochs have run
|
||||
trainer = Trainer(early_stop_callback=early_stop_callback, min_epochs=5)
|
||||
|
||||
# pass in None to disable it
|
||||
trainer = Trainer(early_stop_callback=None)
|
||||
|
||||
Force disable early stop
|
||||
------------------------
|
||||
|
||||
|
||||
@@ -158,8 +158,7 @@ def load_model(exp, root_weights_dir, module_class=LightningTemplateModel, path_
|
||||
checkpoints = [x for x in os.listdir(root_weights_dir) if '.ckpt' in x]
|
||||
weights_dir = os.path.join(root_weights_dir, checkpoints[0])
|
||||
|
||||
trained_model = module_class.load_from_metrics(weights_path=weights_dir,
|
||||
tags_csv=tags_path)
|
||||
trained_model = module_class.load_from_checkpoint(weights_dir)
|
||||
|
||||
assert trained_model is not None, 'loading model failed'
|
||||
|
||||
|
||||
@@ -126,6 +126,7 @@ def test_load_model_from_checkpoint(tmpdir):
|
||||
# fit model
|
||||
trainer = Trainer(**trainer_options)
|
||||
result = trainer.fit(model)
|
||||
trainer.test()
|
||||
|
||||
# correct result and ok accuracy
|
||||
assert result == 1, 'training failed to complete'
|
||||
@@ -140,6 +141,10 @@ def test_load_model_from_checkpoint(tmpdir):
|
||||
for k, v in vars(hparams).items():
|
||||
assert getattr(pretrained_model.hparams, k) == v
|
||||
|
||||
# assert weights are the same
|
||||
for (old_name, old_p), (new_name, new_p) in zip(model.named_parameters(), pretrained_model.named_parameters()):
|
||||
assert torch.all(torch.eq(old_p, new_p)), 'loaded weights are not the same as the saved weights'
|
||||
|
||||
new_trainer = Trainer(**trainer_options)
|
||||
new_trainer.test(pretrained_model)
|
||||
|
||||
@@ -203,7 +208,7 @@ def test_dp_resume(tmpdir):
|
||||
|
||||
trainer_options = dict(
|
||||
show_progress_bar=True,
|
||||
max_epochs=2,
|
||||
max_epochs=3,
|
||||
gpus=2,
|
||||
distributed_backend='dp',
|
||||
)
|
||||
@@ -240,7 +245,7 @@ def test_dp_resume(tmpdir):
|
||||
new_logger = tutils.get_test_tube_logger(tmpdir, version=logger.version)
|
||||
trainer_options['logger'] = new_logger
|
||||
trainer_options['checkpoint_callback'] = ModelCheckpoint(tmpdir)
|
||||
trainer_options['train_percent_check'] = 0.2
|
||||
trainer_options['train_percent_check'] = 0.5
|
||||
trainer_options['val_percent_check'] = 0.2
|
||||
trainer_options['max_epochs'] = 1
|
||||
new_trainer = Trainer(**trainer_options)
|
||||
@@ -269,68 +274,6 @@ def test_dp_resume(tmpdir):
|
||||
model.unfreeze()
|
||||
|
||||
|
||||
def test_cpu_restore_training(tmpdir):
|
||||
"""Verify continue training session on CPU."""
|
||||
tutils.reset_seed()
|
||||
|
||||
hparams = tutils.get_hparams()
|
||||
model = LightningTestModel(hparams)
|
||||
|
||||
# logger file to get meta
|
||||
test_logger_version = 10
|
||||
logger = tutils.get_test_tube_logger(tmpdir, False, version=test_logger_version)
|
||||
|
||||
trainer_options = dict(
|
||||
max_epochs=8,
|
||||
val_check_interval=0.50,
|
||||
val_percent_check=0.2,
|
||||
train_percent_check=0.2,
|
||||
logger=logger,
|
||||
checkpoint_callback=ModelCheckpoint(tmpdir, save_top_k=-1)
|
||||
)
|
||||
|
||||
# fit model
|
||||
trainer = Trainer(**trainer_options)
|
||||
result = trainer.fit(model)
|
||||
# Increment since we've finished the current epoch, don't want to rerun
|
||||
real_global_epoch = trainer.current_epoch + 1
|
||||
|
||||
# traning complete
|
||||
assert result == 1, 'amp + ddp model failed to complete'
|
||||
|
||||
# wipe-out trainer and model
|
||||
# retrain with not much data... this simulates picking training back up after slurm
|
||||
# we want to see if the weights come back correctly
|
||||
new_logger = tutils.get_test_tube_logger(tmpdir, False, version=test_logger_version)
|
||||
trainer_options = dict(
|
||||
max_epochs=2,
|
||||
val_check_interval=0.50,
|
||||
val_percent_check=0.2,
|
||||
train_percent_check=0.2,
|
||||
logger=new_logger,
|
||||
checkpoint_callback=ModelCheckpoint(tmpdir),
|
||||
)
|
||||
trainer = Trainer(**trainer_options)
|
||||
model = LightningTestModel(hparams)
|
||||
|
||||
# set the epoch start hook so we can predict before the model does the full training
|
||||
def assert_good_acc():
|
||||
assert trainer.current_epoch == real_global_epoch
|
||||
assert trainer.current_epoch >= 0
|
||||
|
||||
# if model and state loaded correctly, predictions will be good even though we
|
||||
# haven't trained with the new loaded model
|
||||
trainer.model.eval()
|
||||
for dataloader in trainer.val_dataloaders:
|
||||
tutils.run_prediction(dataloader, trainer.model)
|
||||
|
||||
model.on_train_start = assert_good_acc
|
||||
|
||||
# by calling fit again, we trigger training, loading weights from the cluster
|
||||
# and our hook to predict using current model before any more weight updates
|
||||
trainer.fit(model)
|
||||
|
||||
|
||||
def test_model_saving_loading(tmpdir):
|
||||
"""Tests use case where trainer saves the model, and user loads it from tags independently."""
|
||||
tutils.reset_seed()
|
||||
|
||||
Reference in New Issue
Block a user