Move checkpoint funtions from runner to a new sub-package

docs/en/api/checkpoint.rst

@@ -0,0 +1,34 @@
+.. role:: hidden
+    :class: hidden-section
+
+mmengine.checkpoint
+===================================
+
+.. contents:: mmengine.checkpoint
+   :depth: 2
+   :local:
+   :backlinks: top
+
+.. currentmodule:: mmengine.checkpoint
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+   :template: classtemplate.rst
+
+   CheckpointLoader
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+
+   load_checkpoint
+   save_checkpoint
+   load_state_dict
+   get_state_dict
+   weights_to_cpu
+   find_latest_checkpoint
+   get_deprecated_model_names
+   get_external_models
+   get_mmcls_models
+   get_torchvision_models

docs/en/api/runner.rst

@@ -39,6 +39,10 @@ Loop
 Checkpoints
 ----------------
 
+.. warn::
+
+   All functions and classes in this file have been moved to `mmengine.checkpoint`. Please import them from `mmengine.checkpoint`.
+
 .. autosummary::
    :toctree: generated
    :nosignatures:

docs/en/api/strategy.rst

@@ -37,5 +37,5 @@ mmengine._strategy
    :nosignatures:
    :template: classtemplate.rst
 
-   CollosalAIModelWrapper
+   ColossalAIModelWrapper
    ColossalAIOptimWrapper

docs/en/index.rst

@@ -103,6 +103,7 @@ You can switch between Chinese and English documents in the lower-left corner of
    mmengine.dataset <api/dataset>
    mmengine.infer <api/infer>
    mmengine.device <api/device>
+   mmengine.checkpoint <api/checkpoint>
    mmengine.hub <api/hub>
    mmengine.logging <api/logging>
    mmengine.visualization <api/visualization>

docs/zh_cn/api/checkpoint.rst

@@ -0,0 +1,34 @@
+.. role:: hidden
+    :class: hidden-section
+
+mmengine.checkpoint
+===================================
+
+.. contents:: mmengine.checkpoint
+   :depth: 2
+   :local:
+   :backlinks: top
+
+.. currentmodule:: mmengine.checkpoint
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+   :template: classtemplate.rst
+
+   CheckpointLoader
+
+.. autosummary::
+   :toctree: generated
+   :nosignatures:
+
+   load_checkpoint
+   save_checkpoint
+   load_state_dict
+   get_state_dict
+   weights_to_cpu
+   find_latest_checkpoint
+   get_deprecated_model_names
+   get_external_models
+   get_mmcls_models
+   get_torchvision_models

docs/zh_cn/api/runner.rst

@@ -39,6 +39,10 @@ Loop
 Checkpoints
 ----------------
 
+.. warn::
+
+   所有的函数和类在这个文件中已经被移动到 `mmengine.checkpoint`。请从 `mmengine.checkpoint` 导入它们。
+
 .. autosummary::
    :toctree: generated
    :nosignatures:

docs/zh_cn/api/strategy.rst

@@ -37,5 +37,5 @@ mmengine._strategy
    :nosignatures:
    :template: classtemplate.rst
 
-   CollosalAIModelWrapper
+   ColossalAIModelWrapper
    ColossalAIOptimWrapper

docs/zh_cn/index.rst

@@ -103,6 +103,7 @@
    mmengine.dataset <api/dataset>
    mmengine.infer <api/infer>
    mmengine.device <api/device>
+   mmengine.checkpoint <api/checkpoint>
    mmengine.hub <api/hub>
    mmengine.logging <api/logging>
    mmengine.visualization <api/visualization>

mmengine/_strategy/base.py

@@ -12,6 +12,7 @@
 from torch.optim import Optimizer
 
 import mmengine
+from mmengine.checkpoint.io import _load_checkpoint_to_model
 from mmengine.config import Config, ConfigDict
 from mmengine.dist import (broadcast, get_dist_info, infer_launcher,
                            is_distributed)
@@ -792,8 +793,6 @@ def load_model_state_dict(
         revise_keys: list = [(r'^module.', '')],
     ) -> None:
         """Load model state from dict."""
-        from mmengine.runner.checkpoint import _load_checkpoint_to_model
-
         if is_model_wrapper(self.model):
             model = self.model.module
         else:

mmengine/_strategy/colossalai.py

@@ -26,14 +26,14 @@
 import mmengine
 from mmengine import mkdir_or_exist
 from mmengine._strategy import BaseStrategy
+from mmengine.checkpoint.io import _load_checkpoint, save_checkpoint
 from mmengine.device import get_device
 from mmengine.dist import init_dist, is_main_process
 from mmengine.fileio import join_path
 from mmengine.model import BaseDataPreprocessor
 from mmengine.optim import BaseOptimWrapper, OptimWrapper, _ParamScheduler
 from mmengine.registry import STRATEGIES, Registry
 from mmengine.registry.root import MODEL_WRAPPERS, OPTIM_WRAPPERS, OPTIMIZERS
-from mmengine.runner.checkpoint import _load_checkpoint, save_checkpoint
 from mmengine.utils import get_git_hash
 
 # Component for colossalai `plugins` and `mixed_precisions`
@@ -191,7 +191,7 @@ def __getattr__(self, name):
 class ColossalAIStrategy(BaseStrategy):
     """
     Args:
-        config: (str or dict): The colossalai config file to setup distributed
+        config (str or dict): The colossalai config file to setup distributed
             environment. See more details in the `colossalai config tutorial`_.
         mixed_precision (str or MixedPrecision): The mixed precision to run the
             training. Defaults to None. If the argument is a string, it can be

mmengine/_strategy/deepspeed.py

@@ -18,11 +18,11 @@
 import torch.nn as nn
 
 import mmengine
+from mmengine.checkpoint import save_checkpoint, weights_to_cpu
 from mmengine.dist import init_dist, is_main_process
 from mmengine.optim import BaseOptimWrapper, _ParamScheduler
 from mmengine.registry import (MODEL_WRAPPERS, OPTIM_WRAPPERS, OPTIMIZERS,
                                STRATEGIES)
-from mmengine.runner.checkpoint import save_checkpoint, weights_to_cpu
 from mmengine.utils import apply_to, digit_version, get_git_hash
 from .base import BaseStrategy
 

mmengine/_strategy/fsdp.py

@@ -19,6 +19,7 @@
 from torch.optim.lr_scheduler import LRScheduler
 
 import mmengine
+from mmengine.checkpoint import save_checkpoint
 from mmengine.config import Config, ConfigDict
 from mmengine.device import get_device
 from mmengine.dist import get_rank, is_main_process
@@ -259,8 +260,6 @@ def save_checkpoint(self,
                 checkpoint before saving the checkpoint.
                 Defaults to None.
         """
-        from mmengine.runner.checkpoint import save_checkpoint
-
         state_dict: dict = dict()
         state_dict['state_dict'] = self.model_state_dict()
 

mmengine/_strategy/single_device.py

@@ -5,6 +5,7 @@
 import torch.nn as nn
 
 import mmengine
+from mmengine.checkpoint.io import _load_checkpoint, save_checkpoint
 from mmengine.device import get_device
 from mmengine.model import revert_sync_batchnorm
 from mmengine.optim import BaseOptimWrapper, _ParamScheduler
@@ -135,8 +136,6 @@ def load_checkpoint(
                 checkpoint after loading the checkpoint.
                 Defaults to None.
         """
-        from mmengine.runner.checkpoint import _load_checkpoint
-
         self.logger.info(f'Load checkpoint from {filename}')
 
         if map_location == 'default':
@@ -255,8 +254,6 @@ def save_checkpoint(
                 checkpoint before saving the checkpoint.
                 Defaults to None.
         """
-        from mmengine.runner.checkpoint import save_checkpoint
-
         state_dict: dict = dict()
         state_dict['state_dict'] = self.model_state_dict()
 

mmengine/checkpoint/__init__.py

@@ -0,0 +1,14 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from .io import (get_state_dict, load_checkpoint, load_state_dict,
+                 save_checkpoint, weights_to_cpu)
+from .loader import CheckpointLoader
+from .utils import (find_latest_checkpoint, get_deprecated_model_names,
+                    get_external_models, get_mmcls_models,
+                    get_torchvision_models)
+
+__all__ = [
+    'CheckpointLoader', 'find_latest_checkpoint', 'get_deprecated_model_names',
+    'get_external_models', 'get_mmcls_models', 'get_state_dict',
+    'get_torchvision_models', 'load_checkpoint', 'load_state_dict',
+    'save_checkpoint', 'weights_to_cpu'
+]