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)
| 4212 | return ckpt_files |
| 4213 | |
| 4214 | def load_checkpoint(self, |
| 4215 | load_dir, |
| 4216 | tag=None, |
| 4217 | load_module_strict=True, |
| 4218 | load_optimizer_states=True, |
| 4219 | load_lr_scheduler_states=True, |
| 4220 | load_module_only=False, |
| 4221 | custom_load_fn=None): |
| 4222 | """ |
| 4223 | Load training checkpoint |
| 4224 | |
| 4225 | Arguments: |
| 4226 | load_dir: Required. Directory to load the checkpoint from |
| 4227 | tag: Checkpoint tag used as a unique identifier for checkpoint, if not provided will attempt to load tag in 'latest' file |
| 4228 | load_module_strict: Optional. Boolean to strictly enforce that the keys in state_dict of module and checkpoint match. |
| 4229 | load_optimizer_states: Optional. Boolean to load the training optimizer states from Checkpoint. Ex. ADAM's momentum and variance |
| 4230 | load_lr_scheduler_states: Optional. Boolean to add the learning rate scheduler states from Checkpoint. |
| 4231 | load_module_only: Optional. Boolean to load only the model weights from the checkpoint. Ex. warmstarting. |
| 4232 | custom_load_fn: Optional. Custom model load function. |
| 4233 | |
| 4234 | Returns: |
| 4235 | A tuple of ``load_path`` and ``client_state``. |
| 4236 | *``load_path``: Path of the loaded checkpoint. ``None`` if loading the checkpoint failed. |
| 4237 | *``client_state``: State dictionary used for loading required training states in the client code. |
| 4238 | |
| 4239 | Important: under ZeRO3, one cannot load checkpoint with ``engine.load_checkpoint()`` right |
| 4240 | after ``engine.save_checkpoint()``. It is because ``engine.module`` is partitioned, and |
| 4241 | ``load_checkpoint()`` wants a pristine model. If insisting to do so, please reinitialize engine |
| 4242 | before ``load_checkpoint()``. |
| 4243 | |
| 4244 | """ |
| 4245 | |
| 4246 | if tag is None: |
| 4247 | latest_tag = "latest_universal" if self.load_universal_checkpoint() else "latest" |
| 4248 | latest_path = os.path.join(load_dir, latest_tag) |
| 4249 | if os.path.isfile(latest_path): |
| 4250 | with open(latest_path, "r") as fd: |
| 4251 | tag = fd.read().strip() |
| 4252 | else: |
| 4253 | if self.load_universal_checkpoint(): |
| 4254 | raise ValueError(f'Invalid for universal checkpoint: {latest_path} does not exist') |
| 4255 | else: |
| 4256 | logger.warning( |
| 4257 | f"Unable to find latest file at {latest_path}, if trying to load latest " |
| 4258 | "checkpoint please ensure this file exists or pass an explicit checkpoint tag when loading a checkpoint." |
| 4259 | ) |
| 4260 | return None, None |
| 4261 | |
| 4262 | if self._optimizer_has_ckpt_event_prologue(): |
| 4263 | # Prepare for checkpoint load by ensuring all parameters are partitioned |
| 4264 | self.optimizer.checkpoint_event_prologue() |
| 4265 | |
| 4266 | load_path, client_states = self._load_checkpoint(load_dir, |
| 4267 | tag, |
| 4268 | load_module_strict=load_module_strict, |
| 4269 | load_optimizer_states=load_optimizer_states, |
| 4270 | load_lr_scheduler_states=load_lr_scheduler_states, |
| 4271 | load_module_only=load_module_only, |