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)
| 4379 | return ckpt_files |
| 4380 | |
| 4381 | def load_checkpoint(self, |
| 4382 | load_dir, |
| 4383 | tag=None, |
| 4384 | load_module_strict=True, |
| 4385 | load_optimizer_states=True, |
| 4386 | load_lr_scheduler_states=True, |
| 4387 | load_module_only=False, |
| 4388 | custom_load_fn=None): |
| 4389 | """ |
| 4390 | Load training checkpoint |
| 4391 | |
| 4392 | Arguments: |
| 4393 | load_dir: Required. Directory to load the checkpoint from |
| 4394 | tag: Checkpoint tag used as a unique identifier for checkpoint, if not provided will attempt to load tag in 'latest' file |
| 4395 | load_module_strict: Optional. Boolean to strictly enforce that the keys in state_dict of module and checkpoint match. |
| 4396 | load_optimizer_states: Optional. Boolean to load the training optimizer states from Checkpoint. Ex. ADAM's momentum and variance |
| 4397 | load_lr_scheduler_states: Optional. Boolean to add the learning rate scheduler states from Checkpoint. |
| 4398 | load_module_only: Optional. Boolean to load only the model weights from the checkpoint. Ex. warmstarting. |
| 4399 | custom_load_fn: Optional. Custom model load function. |
| 4400 | |
| 4401 | Returns: |
| 4402 | A tuple of ``load_path`` and ``client_state``. |
| 4403 | *``load_path``: Path of the loaded checkpoint. ``None`` if loading the checkpoint failed. |
| 4404 | *``client_state``: State dictionary used for loading required training states in the client code. |
| 4405 | |
| 4406 | Important: under ZeRO3, one cannot load checkpoint with ``engine.load_checkpoint()`` right |
| 4407 | after ``engine.save_checkpoint()``. It is because ``engine.module`` is partitioned, and |
| 4408 | ``load_checkpoint()`` wants a pristine model. If insisting to do so, please reinitialize engine |
| 4409 | before ``load_checkpoint()``. |
| 4410 | |
| 4411 | """ |
| 4412 | |
| 4413 | if tag is None: |
| 4414 | latest_tag = "latest_universal" if self.load_universal_checkpoint() else "latest" |
| 4415 | latest_path = os.path.join(load_dir, latest_tag) |
| 4416 | if os.path.isfile(latest_path): |
| 4417 | with open(latest_path, "r") as fd: |
| 4418 | tag = fd.read().strip() |
| 4419 | else: |
| 4420 | if self.load_universal_checkpoint(): |
| 4421 | raise ValueError(f'Invalid for universal checkpoint: {latest_path} does not exist') |
| 4422 | else: |
| 4423 | logger.warning( |
| 4424 | f"Unable to find latest file at {latest_path}, if trying to load latest " |
| 4425 | "checkpoint please ensure this file exists or pass an explicit checkpoint tag when loading a checkpoint." |
| 4426 | ) |
| 4427 | return None, None |
| 4428 | |
| 4429 | if self._optimizer_has_ckpt_event_prologue(): |
| 4430 | # Prepare for checkpoint load by ensuring all parameters are partitioned |
| 4431 | self.optimizer.checkpoint_event_prologue() |
| 4432 | |
| 4433 | load_path, client_states = self._load_checkpoint(load_dir, |
| 4434 | tag, |
| 4435 | load_module_strict=load_module_strict, |
| 4436 | load_optimizer_states=load_optimizer_states, |
| 4437 | load_lr_scheduler_states=load_lr_scheduler_states, |
| 4438 | load_module_only=load_module_only, |