forked from njzjz/deepmd-kit
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
138 changed files
with
21,061 additions
and
25 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,133 @@ | ||
# SPDX-License-Identifier: LGPL-3.0-or-later | ||
import importlib | ||
import os | ||
import site | ||
from functools import ( | ||
lru_cache, | ||
) | ||
from importlib.machinery import ( | ||
FileFinder, | ||
) | ||
from importlib.util import ( | ||
find_spec, | ||
) | ||
from pathlib import ( | ||
Path, | ||
) | ||
from sysconfig import ( | ||
get_path, | ||
) | ||
from typing import ( | ||
Optional, | ||
Union, | ||
) | ||
|
||
|
||
@lru_cache | ||
def find_paddle() -> tuple[Optional[str], list[str]]: | ||
"""Find PaddlePadle library. | ||
Tries to find PaddlePadle in the order of: | ||
1. Environment variable `PADDLE_ROOT` if set | ||
2. The current Python environment. | ||
3. user site packages directory if enabled | ||
4. system site packages directory (purelib) | ||
Considering the default PaddlePadle package still uses old CXX11 ABI, we | ||
cannot install it automatically. | ||
Returns | ||
------- | ||
str, optional | ||
PaddlePadle library path if found. | ||
list of str | ||
Paddle requirement if not found. Empty if found. | ||
""" | ||
if os.environ.get("DP_ENABLE_PADDLE", "0") == "0": | ||
return None, [] | ||
requires = [] | ||
pd_spec = None | ||
|
||
if (pd_spec is None or not pd_spec) and os.environ.get("PADDLE_ROOT") is not None: | ||
site_packages = Path(os.environ.get("PADDLE_ROOT")).parent.absolute() | ||
pd_spec = FileFinder(str(site_packages)).find_spec("paddle") | ||
|
||
# get paddle spec | ||
# note: isolated build will not work for backend | ||
if pd_spec is None or not pd_spec: | ||
pd_spec = find_spec("paddle") | ||
|
||
if not pd_spec and site.ENABLE_USER_SITE: | ||
# first search TF from user site-packages before global site-packages | ||
site_packages = site.getusersitepackages() | ||
if site_packages: | ||
pd_spec = FileFinder(site_packages).find_spec("paddle") | ||
|
||
if not pd_spec: | ||
# purelib gets site-packages path | ||
site_packages = get_path("purelib") | ||
if site_packages: | ||
pd_spec = FileFinder(site_packages).find_spec("paddle") | ||
|
||
# get install dir from spec | ||
try: | ||
pd_install_dir = pd_spec.submodule_search_locations[0] # type: ignore | ||
# AttributeError if ft_spec is None | ||
# TypeError if submodule_search_locations are None | ||
# IndexError if submodule_search_locations is an empty list | ||
except (AttributeError, TypeError, IndexError): | ||
pd_install_dir = None | ||
requires.extend(get_pd_requirement()["paddle"]) | ||
return pd_install_dir, requires | ||
|
||
|
||
@lru_cache | ||
def get_pd_requirement(pd_version: str = "") -> dict: | ||
"""Get PaddlePadle requirement when Paddle is not installed. | ||
If pd_version is not given and the environment variable `PADDLE_VERSION` is set, use it as the requirement. | ||
Parameters | ||
---------- | ||
pd_version : str, optional | ||
Paddle version | ||
Returns | ||
------- | ||
dict | ||
PaddlePadle requirement. | ||
""" | ||
if pd_version is None: | ||
return {"paddle": []} | ||
if pd_version == "": | ||
pd_version = os.environ.get("PADDLE_VERSION", "") | ||
|
||
return { | ||
"paddle": [ | ||
"paddlepaddle>=3.0.0b1" if pd_version != "" else "paddlepaddle>=3.0.0b1", | ||
], | ||
} | ||
|
||
|
||
@lru_cache | ||
def get_pd_version(pd_path: Optional[Union[str, Path]]) -> str: | ||
"""Get Paddle version from a Paddle Python library path. | ||
Parameters | ||
---------- | ||
pd_path : str or Path | ||
Paddle Python library path, e.g. "/python3.10/site-packages/paddle/" | ||
Returns | ||
------- | ||
str | ||
version | ||
""" | ||
if pd_path is None or pd_path == "": | ||
return "" | ||
version_file = Path(pd_path) / "version" / "__init__.py" | ||
spec = importlib.util.spec_from_file_location("paddle.version", version_file) | ||
module = importlib.util.module_from_spec(spec) | ||
spec.loader.exec_module(module) | ||
return module.full_version |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,124 @@ | ||
# SPDX-License-Identifier: LGPL-3.0-or-later | ||
from importlib.util import ( | ||
find_spec, | ||
) | ||
from typing import ( | ||
TYPE_CHECKING, | ||
Callable, | ||
ClassVar, | ||
) | ||
|
||
from deepmd.backend.backend import ( | ||
Backend, | ||
) | ||
|
||
if TYPE_CHECKING: | ||
from argparse import ( | ||
Namespace, | ||
) | ||
|
||
from deepmd.infer.deep_eval import ( | ||
DeepEvalBackend, | ||
) | ||
from deepmd.utils.neighbor_stat import ( | ||
NeighborStat, | ||
) | ||
|
||
|
||
@Backend.register("pd") | ||
@Backend.register("paddle") | ||
class PaddleBackend(Backend): | ||
"""Paddle backend.""" | ||
|
||
name = "Paddle" | ||
"""The formal name of the backend.""" | ||
features: ClassVar[Backend.Feature] = ( | ||
Backend.Feature.ENTRY_POINT | ||
| Backend.Feature.DEEP_EVAL | ||
| Backend.Feature.NEIGHBOR_STAT | ||
| Backend.Feature.IO | ||
) | ||
"""The features of the backend.""" | ||
suffixes: ClassVar[list[str]] = [".json", ".pd"] | ||
"""The suffixes of the backend.""" | ||
|
||
def is_available(self) -> bool: | ||
"""Check if the backend is available. | ||
Returns | ||
------- | ||
bool | ||
Whether the backend is available. | ||
""" | ||
return find_spec("paddle") is not None | ||
|
||
@property | ||
def entry_point_hook(self) -> Callable[["Namespace"], None]: | ||
"""The entry point hook of the backend. | ||
Returns | ||
------- | ||
Callable[[Namespace], None] | ||
The entry point hook of the backend. | ||
""" | ||
from deepmd.pd.entrypoints.main import main as deepmd_main | ||
|
||
return deepmd_main | ||
|
||
@property | ||
def deep_eval(self) -> type["DeepEvalBackend"]: | ||
"""The Deep Eval backend of the backend. | ||
Returns | ||
------- | ||
type[DeepEvalBackend] | ||
The Deep Eval backend of the backend. | ||
""" | ||
from deepmd.pd.infer.deep_eval import DeepEval as DeepEvalPD | ||
|
||
return DeepEvalPD | ||
|
||
@property | ||
def neighbor_stat(self) -> type["NeighborStat"]: | ||
"""The neighbor statistics of the backend. | ||
Returns | ||
------- | ||
type[NeighborStat] | ||
The neighbor statistics of the backend. | ||
""" | ||
from deepmd.pd.utils.neighbor_stat import ( | ||
NeighborStat, | ||
) | ||
|
||
return NeighborStat | ||
|
||
@property | ||
def serialize_hook(self) -> Callable[[str], dict]: | ||
"""The serialize hook to convert the model file to a dictionary. | ||
Returns | ||
------- | ||
Callable[[str], dict] | ||
The serialize hook of the backend. | ||
""" | ||
from deepmd.pd.utils.serialization import ( | ||
serialize_from_file, | ||
) | ||
|
||
return serialize_from_file | ||
|
||
@property | ||
def deserialize_hook(self) -> Callable[[str, dict], None]: | ||
"""The deserialize hook to convert the dictionary to a model file. | ||
Returns | ||
------- | ||
Callable[[str, dict], None] | ||
The deserialize hook of the backend. | ||
""" | ||
from deepmd.pd.utils.serialization import ( | ||
deserialize_to_file, | ||
) | ||
|
||
return deserialize_to_file |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -457,7 +457,7 @@ def format_nlist( | |
Returns | ||
------- | ||
formated_nlist | ||
formatted_nlist | ||
the formatted nlist. | ||
""" | ||
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,11 @@ | ||
# SPDX-License-Identifier: LGPL-3.0-or-later | ||
|
||
# import customized OPs globally | ||
|
||
from deepmd.utils.entry_point import ( | ||
load_entry_point, | ||
) | ||
|
||
load_entry_point("deepmd.pd") | ||
|
||
__all__ = [] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
# SPDX-License-Identifier: LGPL-3.0-or-later |
Oops, something went wrong.