Move checkpoint funtions from runner to a new sub-package

open-mmlab · Feb 18, 2024 · 478a80a · 478a80a
1 parent 4df682b
commit 478a80a
Show file tree

Hide file tree

Showing 15 changed files with 1,127 additions and 965 deletions.
diff --git a/docs/en/api/checkpoint.rst b/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
diff --git a/docs/en/api/runner.rst b/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:

diff --git a/docs/en/index.rst b/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>

diff --git a/docs/zh_cn/api/checkpoint.rst b/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
diff --git a/docs/zh_cn/api/runner.rst b/docs/zh_cn/api/runner.rst
@@ -39,6 +39,10 @@ Loop
 Checkpoints
 ----------------
 
+.. warn::
+
+   所有的函数和类在这个文件中已经被移动到 `mmengine.checkpoint`。请从 `mmengine.checkpoint` 导入它们。
+
 .. autosummary::
    :toctree: generated
    :nosignatures:

diff --git a/docs/zh_cn/index.rst b/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>

diff --git a/mmengine/checkpoint/__init__.py b/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'
+]