mirrors/mmengine
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmengine.git synced 2025-06-03 21:54:44 +08:00
Code Issues Projects Releases Wiki Activity

[Fix] Support exclude_frozen_parameters for DeepSpeedStrategy's resume (#1424)

Browse Source
This commit is contained in:
Zhihao Lin 2023-11-08 23:35:12 +08:00 committed by GitHub
parent 46784185cf
commit 26f22ed283
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
mmengine/_strategy/deepspeed.py
View File

@ -463,6 +463,13 @@ class DeepSpeedStrategy(BaseStrategy):
self.logger.info(f'Resume checkpoint from {filename}') self.logger.info(f'Resume checkpoint from {filename}')
dirname, basename = osp.split(filename) dirname, basename = osp.split(filename)
if digit_version(deepspeed.__version__) >= digit_version('0.10.1'):
_, extra_ckpt = self.model.load_checkpoint(
dirname,
tag=basename,
load_optimizer_states=resume_optimizer,
load_module_strict=not self.exclude_frozen_parameters)
else:
_, extra_ckpt = self.model.load_checkpoint( _, extra_ckpt = self.model.load_checkpoint(
dirname, tag=basename, load_optimizer_states=resume_optimizer) dirname, tag=basename, load_optimizer_states=resume_optimizer)