Load training checkpoint Arguments: load_dir: Required. Directory to load the checkpoint from tag: Checkpoint tag used as a unique identifier for checkpoint, if not provided will attempt to load tag in 'latest' file load_module_strict: Optional.
(self,
load_dir,
tag=None,
load_module_strict=True,
load_optimizer_states=True,
load_lr_scheduler_states=True,
load_module_only=False,
custom_load_fn=None)
| 4077 | return ckpt_files |
| 4078 | |
| 4079 | def load_checkpoint(self, |
| 4080 | load_dir, |
| 4081 | tag=None, |
| 4082 | load_module_strict=True, |
| 4083 | load_optimizer_states=True, |
| 4084 | load_lr_scheduler_states=True, |
| 4085 | load_module_only=False, |
| 4086 | custom_load_fn=None): |
| 4087 | """ |
| 4088 | Load training checkpoint |
| 4089 | |
| 4090 | Arguments: |
| 4091 | load_dir: Required. Directory to load the checkpoint from |
| 4092 | tag: Checkpoint tag used as a unique identifier for checkpoint, if not provided will attempt to load tag in 'latest' file |
| 4093 | load_module_strict: Optional. Boolean to strictly enforce that the keys in state_dict of module and checkpoint match. |
| 4094 | load_optimizer_states: Optional. Boolean to load the training optimizer states from Checkpoint. Ex. ADAM's momentum and variance |
| 4095 | load_lr_scheduler_states: Optional. Boolean to add the learning rate scheduler states from Checkpoint. |
| 4096 | load_module_only: Optional. Boolean to load only the model weights from the checkpoint. Ex. warmstarting. |
| 4097 | custom_load_fn: Optional. Custom model load function. |
| 4098 | |
| 4099 | Returns: |
| 4100 | A tuple of ``load_path`` and ``client_state``. |
| 4101 | *``load_path``: Path of the loaded checkpoint. ``None`` if loading the checkpoint failed. |
| 4102 | *``client_state``: State dictionary used for loading required training states in the client code. |
| 4103 | |
| 4104 | Important: under ZeRO3, one cannot load checkpoint with ``engine.load_checkpoint()`` right |
| 4105 | after ``engine.save_checkpoint()``. It is because ``engine.module`` is partitioned, and |
| 4106 | ``load_checkpoint()`` wants a pristine model. If insisting to do so, please reinitialize engine |
| 4107 | before ``load_checkpoint()``. |
| 4108 | |
| 4109 | """ |
| 4110 | |
| 4111 | if tag is None: |
| 4112 | latest_tag = "latest_universal" if self.load_universal_checkpoint() else "latest" |
| 4113 | latest_path = os.path.join(load_dir, latest_tag) |
| 4114 | if os.path.isfile(latest_path): |
| 4115 | with open(latest_path, "r") as fd: |
| 4116 | tag = fd.read().strip() |
| 4117 | else: |
| 4118 | if self.load_universal_checkpoint(): |
| 4119 | raise ValueError(f'Invalid for universal checkpoint: {latest_path} does not exist') |
| 4120 | else: |
| 4121 | logger.warning( |
| 4122 | f"Unable to find latest file at {latest_path}, if trying to load latest " |
| 4123 | "checkpoint please ensure this file exists or pass an explicit checkpoint tag when loading a checkpoint." |
| 4124 | ) |
| 4125 | return None, None |
| 4126 | |
| 4127 | if self._optimizer_has_ckpt_event_prologue(): |
| 4128 | # Prepare for checkpoint load by ensuring all parameters are partitioned |
| 4129 | self.optimizer.checkpoint_event_prologue() |
| 4130 | |
| 4131 | load_path, client_states = self._load_checkpoint(load_dir, |
| 4132 | tag, |
| 4133 | load_module_strict=load_module_strict, |
| 4134 | load_optimizer_states=load_optimizer_states, |
| 4135 | load_lr_scheduler_states=load_lr_scheduler_states, |
| 4136 | load_module_only=load_module_only, |