Create a model in TensorFlow #
If you’d like to create a new model that isn’t covered by the existing DeePMD-kit library, but reuse DeePMD-kit’s other efficient modules such as data processing, trainer, etc, you may want to read this section.
To incorporate your custom model you’ll need to:
Register and implement new components (e.g. descriptor) in a Python file. You may also want to register new TensorFlow OPs if necessary.
Register new arguments for user inputs.
Package new codes into a Python package.
Test new models.
Design a new component#
When creating a new component, take descriptor as the example, one should inherit from the deepmd.tf.descriptor.descriptor.Descriptor
class and override several methods. Abstract methods such as deepmd.tf.descriptor.descriptor.Descriptor.build
must be implemented and others are not. You should keep arguments of these methods unchanged.
After implementation, you need to register the component with a key:
from deepmd.tf.descriptor import Descriptor
@Descriptor.register("some_descrpt")
class SomeDescript(Descriptor):
def __init__(self, arg1: bool, arg2: float) -> None:
pass
Register new arguments#
To let someone uses your new component in their input file, you need to create a new method that returns some Argument
of your new component, and then register new arguments. For example, the code below
from typing import List
from dargs import Argument
from deepmd.utils.argcheck import descrpt_args_plugin
@descrpt_args_plugin.register("some_descrpt")
def descrpt_some_args() -> list[Argument]:
return [
Argument("arg1", bool, optional=False, doc="balabala"),
Argument("arg2", float, optional=True, default=6.0, doc="haha"),
]
allows one to use your new descriptor as below:
"descriptor" :{
"type": "some_descrpt",
"arg1": true,
"arg2": 6.0
}
The arguments here should be consistent with the class arguments of your new component.
Package new codes#
You may package new codes into a new Python package if you don’t want to contribute it to the main DeePMD-kit repository. It’s crucial to add your new component to project.entry-points."deepmd"
in pyproject.toml
:
[project.entry-points."deepmd"]
some_descrpt = "deepmd_some_descrtpt:SomeDescript"
where deepmd_some_descrtpt
is the module of your codes. It is equivalent to from deepmd_some_descrtpt import SomeDescript
.
If you place SomeDescript
and descrpt_some_args
into different modules, you are also expected to add descrpt_some_args
to entry_points
.
After you install your new package, you can now use dp train
to run your new model.
Package customized C++ OPs#
You may need to use customized TensorFlow C++ OPs in the new model. Follow TensorFlow documentation to create one library.
When using your customized C++ OPs in the Python interface, use tf.load_op_library()
to load the OP library in the module defined in entry_points
.
When using your customized C++ OPs in the C++ library, define the environment variable DP_PLUGIN_PATH
to load the OP library.