r"""Load training checkpoint Arguments: load_dir: Required. Directory to load the checkpoint from tag: Required. Checkpoint tag used as a unique identifier for the checkpoint. Ex. Global Step. load_module_strict: Optional. Boolean to strictly enforce that the keys in sta
(model,
load_dir,
tag,
load_module_strict=True,
load_optimizer_states=True,
load_lr_scheduler_states=True)
| 19 | |
| 20 | |
| 21 | def load_checkpoint(model, |
| 22 | load_dir, |
| 23 | tag, |
| 24 | load_module_strict=True, |
| 25 | load_optimizer_states=True, |
| 26 | load_lr_scheduler_states=True): |
| 27 | r"""Load training checkpoint |
| 28 | |
| 29 | Arguments: |
| 30 | load_dir: Required. Directory to load the checkpoint from |
| 31 | tag: Required. Checkpoint tag used as a unique identifier for the checkpoint. Ex. Global Step. |
| 32 | load_module_strict: Optional. Boolean to strictly enforce that the keys in state_dict of module and |
| 33 | checkpoint match. |
| 34 | load_optimizer_states: Optional. Boolean to load the training optimizer states from Checkpoint. |
| 35 | Ex. ADAM's momentum and variance |
| 36 | load_lr_scheduler_states: Optional. Boolean to add the learning rate scheduler states from Checkpoint. |
| 37 | Return: |
| 38 | load_path: Path of the loaded checkpoint. None if loading the checkpoint failed |
| 39 | client_state: State dictionary used for loading required training states in the client code. |
| 40 | """ |
| 41 | |
| 42 | load_path, client_states = _load_checkpoint( |
| 43 | model, |
| 44 | load_dir, |
| 45 | tag, |
| 46 | load_module_strict=load_module_strict, |
| 47 | load_optimizer_states=load_optimizer_states, |
| 48 | load_lr_scheduler_states=load_lr_scheduler_states) |
| 49 | |
| 50 | if load_optimizer_states: |
| 51 | if model.zero_optimization() and load_path is not None: |
| 52 | model._load_zero_checkpoint( |
| 53 | load_dir, tag, load_optimizer_states=load_optimizer_states) |
| 54 | |
| 55 | return load_path, client_states |
| 56 | |
| 57 | |
| 58 | def _get_ckpt_name(mp_rank, checkpoints_path, tag): |
no test coverage detected
searching dependent graphs…