【Hackathon 7th No.19】NO.19为 Paddle 新增 load_state_dict_from_url API v1 (#68594)

zty-king · web-flow · commit 416624b51ceb · 2024-12-13T12:18:37.000+08:00
* 添加新的API:load_state_dict_from_url

* 添加新的API:load_state_dict_from_url

* 新增单侧

* 格式规范文档

* 修改了单测文件

* 上传开发的url和单测文件

* 用pre-commit修改了格式

* 修改了api代码和单测代码

* 修改相对路径

* 设置API的超时为30s

* 用pre-commit修正格式

* 用pre-commit修正格式

* CI检测卡顿，提交一个空提交刷新

* 修改了文件格式，添加了hub的__all__ list

* 添加 test_load_state_dict_from_url API

* 修正单测

* 修改单测

* 修改单测

* 添加单测的测试覆盖

* 修改单测文件保存位置

* 修改单测

* remove weights_only parameter

* 添加类型提示，规范格式
diff --git a/python/paddle/hapi/hub.py b/python/paddle/hapi/hub.py
@@ -14,22 +14,25 @@
 
 from __future__ import annotations
 
+import contextlib
+import errno
 import os
 import shutil
 import sys
+import warnings
 import zipfile
 from typing import TYPE_CHECKING, Literal
+from urllib.parse import urlparse
 
 from typing_extensions import TypeAlias
 
-from paddle.utils.download import get_path_from_url
+import paddle
+from paddle.utils.download import _download, get_path_from_url
 
 if TYPE_CHECKING:
     import builtins
     from typing import Any
 
-    import paddle
-
 __all__ = []
 
 _Source: TypeAlias = Literal["github", "gitee", "local"]
@@ -329,3 +332,128 @@ def load(
     entry = _load_entry_from_hubconf(hub_module, model)
 
     return entry(**kwargs)
+
+
+def load_state_dict_from_url(
+    url: str,
+    model_dir: str | None = None,
+    check_hash: bool = False,
+    file_name: str | None = None,
+    map_location: (
+        Literal["cpu", "gpu", "xpu", "npu", "numpy", "np"] | None
+    ) = None,
+) -> Any:
+    """Download Paddle's model weights (i.e., state_dict)
+    from the specified URL and extract the downloaded file if necessary
+
+    Args:
+        url (str): URL of the object to download.
+        model_dir (Optional[str], optional): Directory in which to save the object.
+        check_hash (bool, optional): If True, the filename part of the URL should follow the naming convention filename-<sha256>.ext where <sha256> is the first eight or more digits of the SHA256 hash of the contents of the file. The hash is used to ensure unique names and to verify the contents of the file. Default: False.
+        file_name (Optional[str], optional): Name for the downloaded file. Filename from URL will be used if not set.
+        map_location (Optional[Literal["cpu", "gpu", "xpu", "npu", "numpy", "np"]], optional): Specifies how to remap storage locations. Default: None.
+
+    Returns:
+        Any: A target object that can be used in Paddle.
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+            >>> paddle.hub.load_state_dict_from_url(url='https://paddle-hapi.bj.bcebos.com/models/resnet18.pdparams', model_dir="./paddle/test_load_from_url")
+            >>> paddle.hub.load_state_dict_from_url(url='https://x2paddle.bj.bcebos.com/resnet18.zip', model_dir="./paddle/test_file_is_zip")
+    """
+    if model_dir is None:
+        hub_dir = get_dir()
+        model_dir = os.path.join(hub_dir, 'checkpoints')
+
+    try:
+        os.makedirs(model_dir)
+    except OSError as e:
+        if e.errno == errno.EEXIST:
+            pass
+        else:
+            raise
+    parts = urlparse(url)
+    filename = os.path.basename(parts.path)
+    if file_name is not None:
+        filename = file_name
+    cached_file = os.path.join(model_dir, filename)
+    if not os.path.exists(cached_file):
+        sys.stderr.write(f'Downloading: "{url}" to {cached_file}\n')
+        hash_prefix = None
+        if check_hash:
+            hash_prefix = check_hash
+        _download(url, model_dir, hash_prefix)
+
+    if map_location:
+        assert map_location in ["cpu", "gpu", "xpu", "npu", "numpy", "np"]
+        if _is_legacy_zip_format(cached_file):
+            return _legacy_zip_load(cached_file, model_dir, map_location)
+        if map_location in ["numpy", "np"]:
+            return paddle.load(cached_file, return_numpy=True)
+        else:
+            with device_guard(map_location):
+                return paddle.load(cached_file)
+    else:
+        if _is_legacy_zip_format(cached_file):
+            return _legacy_zip_load(cached_file, model_dir)
+        return paddle.load(cached_file)
+
+
+def _is_legacy_zip_format(filename):
+    if zipfile.is_zipfile(filename):
+        infolist = zipfile.ZipFile(filename).infolist()
+        return len(infolist) == 1 and not infolist[0].is_dir()
+    return False
+
+
+def _legacy_zip_load(filename, model_dir, map_location=None):
+    with zipfile.ZipFile(filename) as f:
+        members = f.infolist()
+        if len(members) != 1:
+            raise RuntimeError(
+                'Only one file(not dir) is allowed in the zipfile'
+            )
+        f.extractall(model_dir)
+        extracted_name = members[0].filename
+        extracted_file = os.path.join(model_dir, extracted_name)
+    if map_location:
+        if map_location in ["numpy", "np"]:
+            return paddle.load(extracted_file, return_numpy=True)
+        else:
+            with device_guard(map_location):
+                return paddle.load(extracted_file)
+    else:
+        return paddle.load(extracted_file)
+
+
+def get_dir():
+    if os.getenv('PADDLE_HUB'):
+        warnings.warn(
+            'PADDLE_HUB is deprecated, please use env PADDLE_HOME instead'
+        )
+    return os.path.join(_get_paddle_home(), 'hub')
+
+
+def _get_paddle_home():
+    paddle_home = os.path.expanduser(
+        os.getenv(
+            'PADDLE_HOME',
+            os.path.join(os.getenv('XDG_CACHE_HOME', '~/.cache'), 'paddle'),
+        )
+    )
+    return paddle_home
+
+
+@contextlib.contextmanager
+def device_guard(device="cpu", dev_id=0):
+    origin_device = paddle.device.get_device()
+    if device == "cpu":
+        paddle.set_device(device)
+    elif device in ["gpu", "xpu", "npu"]:
+        paddle.set_device(f"{device}:{dev_id}")
+    try:
+        yield
+    finally:
+        paddle.set_device(origin_device)
diff --git a/python/paddle/hub.py b/python/paddle/hub.py
@@ -12,6 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .hapi.hub import help, list, load
+from .hapi.hub import help, list, load, load_state_dict_from_url
 
-__all__ = ['list', 'help', 'load']
+__all__ = ['list', 'help', 'load', 'load_state_dict_from_url']
diff --git a/test/legacy_test/CMakeLists.txt b/test/legacy_test/CMakeLists.txt
@@ -1176,3 +1176,4 @@ set_tests_properties(test_inplace_softmax_with_cross_entropy PROPERTIES TIMEOUT
 set_tests_properties(test_imperative_selected_rows_to_lod_tensor
                      PROPERTIES TIMEOUT 200)
 set_tests_properties(test_slice_op PROPERTIES TIMEOUT 120)
+set_tests_properties(test_load_state_dict_from_url PROPERTIES TIMEOUT 40)
diff --git a/test/legacy_test/test_load_state_dict_from_url.py b/test/legacy_test/test_load_state_dict_from_url.py
@@ -0,0 +1,98 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import unittest
+
+import paddle
+from paddle.vision.models import resnet18
+
+
+class TestLoadStateDictFromUrl(unittest.TestCase):
+    def setUp(self):
+        self.model = resnet18(pretrained=False)
+
+    def test_load_state_dict_from_url(self):
+        # Compare using load_state_dict_from_url to download and load an uncompressed weight file and a compressed weight file(ZIP)
+        weight1 = paddle.hub.load_state_dict_from_url(
+            url='https://paddle-hapi.bj.bcebos.com/models/resnet18.pdparams',
+            model_dir="./test/test1",
+            map_location="cpu",
+        )
+        model1 = self.model
+        model1.set_state_dict(weight1)
+        weight2 = paddle.hub.load_state_dict_from_url(
+            url='https://x2paddle.bj.bcebos.com/resnet18.zip',
+            model_dir="./test/test2",
+            map_location="cpu",
+        )
+        model2 = self.model
+        model2.set_state_dict(weight2)
+        are_parameters_equal = True
+        for (name1, param1), (name2, param2) in zip(
+            model1.named_parameters(), model2.named_parameters()
+        ):
+            if name1 != name2 or not paddle.allclose(param1, param2):
+                are_parameters_equal = False
+                break
+        assert are_parameters_equal
+
+        # Test whether the model loads properly when the model_dir is empty and Test check_hash and map_location
+        weight3 = paddle.hub.load_state_dict_from_url(
+            url='https://x2paddle.bj.bcebos.com/resnet18.zip',
+            model_dir="./test/test3",
+            file_name="resnet18.zip",
+            map_location="numpy",
+        )
+        model3 = self.model
+        model3.set_state_dict(weight3)
+        weight4 = paddle.hub.load_state_dict_from_url(
+            url='https://paddle-hapi.bj.bcebos.com/models/resnet18.pdparams',
+            file_name="resnet18.pdparams",
+            check_hash="cf548f46534aa3560945be4b95cd11c4",
+            map_location="numpy",
+        )
+        model4 = self.model
+        model4.set_state_dict(weight4)
+        are_parameters_equal = True
+        for (name1, param1), (name2, param2) in zip(
+            model3.named_parameters(), model4.named_parameters()
+        ):
+            if name1 != name2 or not paddle.allclose(param1, param2):
+                are_parameters_equal = False
+                break
+        assert are_parameters_equal
+        # Test map_location is None
+        weight5 = paddle.hub.load_state_dict_from_url(
+            url='https://x2paddle.bj.bcebos.com/resnet18.zip',
+            model_dir="./test/test4",
+        )
+        model5 = self.model
+        model5.set_state_dict(weight5)
+        weight6 = paddle.hub.load_state_dict_from_url(
+            url='https://paddle-hapi.bj.bcebos.com/models/resnet18.pdparams',
+            model_dir="./test/test5",
+        )
+        model6 = self.model
+        model6.set_state_dict(weight6)
+        are_parameters_equal = True
+        for (name1, param1), (name2, param2) in zip(
+            model5.named_parameters(), model6.named_parameters()
+        ):
+            if name1 != name2 or not paddle.allclose(param1, param2):
+                are_parameters_equal = False
+                break
+        assert are_parameters_equal
+
+
+if __name__ == '__main__':
+    unittest.main()
