Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

support fitting net #3137

Merged
merged 2 commits into from
Jan 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions deepmd_utils/model_format/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
)
from .network import (
EmbeddingNet,
FittingNet,
NativeLayer,
NativeNet,
NetworkCollection,
Expand All @@ -31,6 +32,7 @@
"DescrptSeA",
"EnvMat",
"EmbeddingNet",
"FittingNet",
"NativeLayer",
"NativeNet",
"NetworkCollection",
Expand Down
116 changes: 113 additions & 3 deletions deepmd_utils/model_format/network.py
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,9 @@ def __init__(
self.w = w.astype(prec) if w is not None else None
self.b = b.astype(prec) if b is not None else None
self.idt = idt.astype(prec) if idt is not None else None
self.activation_function = activation_function
self.activation_function = (
activation_function if activation_function is not None else "none"
)
self.resnet = resnet
self.check_type_consistency()

Expand Down Expand Up @@ -354,6 +356,24 @@ def call(self, x: np.ndarray) -> np.ndarray:


class EmbeddingNet(NativeNet):
"""The embedding network.

Parameters
----------
in_dim
Input dimension.
neuron
The number of neurons in each layer. The output dimension
is the same as the dimension of the last layer.
activation_function
The activation function.
resnet_dt
Use time step at the resnet architecture.
precision
Floating point precision for the model paramters.

"""

def __init__(
self,
in_dim,
Expand All @@ -370,8 +390,8 @@ def __init__(
layers.append(
NativeLayer(
rng.normal(size=(i_in, i_ot)),
b=rng.normal(size=(ii)),
idt=rng.normal(size=(ii)) if resnet_dt else None,
b=rng.normal(size=(i_ot)),
idt=rng.normal(size=(i_ot)) if resnet_dt else None,
activation_function=activation_function,
resnet=True,
precision=precision,
Expand Down Expand Up @@ -417,6 +437,95 @@ def deserialize(cls, data: dict) -> "EmbeddingNet":
return obj


class FittingNet(EmbeddingNet):
"""The fitting network. It may be implemented as an embedding
net connected with a linear output layer.

Parameters
----------
in_dim
Input dimension.
out_dim
Output dimension
neuron
The number of neurons in each hidden layer.
activation_function
The activation function.
resnet_dt
Use time step at the resnet architecture.
precision
Floating point precision for the model paramters.
bias_out
The last linear layer has bias.

"""

def __init__(
self,
in_dim,
out_dim,
neuron: List[int] = [24, 48, 96],
activation_function: str = "tanh",
resnet_dt: bool = False,
precision: str = DEFAULT_PRECISION,
bias_out: bool = True,
):
super().__init__(
in_dim,
neuron=neuron,
activation_function=activation_function,
resnet_dt=resnet_dt,
precision=precision,
)
rng = np.random.default_rng()
i_in, i_ot = neuron[-1], out_dim
self.layers.append(
NativeLayer(
rng.normal(size=(i_in, i_ot)),
b=rng.normal(size=(i_ot)) if bias_out else None,
idt=None,
activation_function=None,
resnet=False,
precision=precision,
)
)
self.out_dim = out_dim
self.bias_out = bias_out

def serialize(self) -> dict:
"""Serialize the network to a dict.

Returns
-------
dict
The serialized network.
"""
return {
"in_dim": self.in_dim,
"out_dim": self.out_dim,
"neuron": self.neuron.copy(),
"activation_function": self.activation_function,
"resnet_dt": self.resnet_dt,
"precision": self.precision,
"bias_out": self.bias_out,
"layers": [layer.serialize() for layer in self.layers],
}

@classmethod
def deserialize(cls, data: dict) -> "FittingNet":
"""Deserialize the network from a dict.

Parameters
----------
data : dict
The dict to deserialize from.
"""
layers = data.pop("layers")
obj = cls(**data)
NativeNet.__init__(obj, layers)
return obj


class NetworkCollection:
"""A collection of networks for multiple elements.

Expand All @@ -439,6 +548,7 @@ class NetworkCollection:
NETWORK_TYPE_MAP: ClassVar[Dict[str, type]] = {
"network": NativeNet,
"embedding_network": EmbeddingNet,
"fitting_network": FittingNet,
}

def __init__(
Expand Down
28 changes: 28 additions & 0 deletions source/tests/test_model_format_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
DescrptSeA,
EmbeddingNet,
EnvMat,
FittingNet,
NativeLayer,
NativeNet,
NetworkCollection,
Expand Down Expand Up @@ -98,6 +99,8 @@ def test_deserialize(self):
np.testing.assert_array_equal(network[0]["resnet"], True)
np.testing.assert_array_equal(network[1]["resnet"], True)


class TestEmbeddingNet(unittest.TestCase):
def test_embedding_net(self):
for ni, act, idt, prec in itertools.product(
[1, 10],
Expand All @@ -116,6 +119,31 @@ def test_embedding_net(self):
np.testing.assert_allclose(en0.call(inp), en1.call(inp))


class TestFittingNet(unittest.TestCase):
def test_fitting_net(self):
for ni, no, act, idt, prec, bo in itertools.product(
[1, 10],
[1, 7],
["tanh", "none"],
[True, False],
["double", "single"],
[True, False],
):
en0 = FittingNet(
ni,
no,
activation_function=act,
precision=prec,
resnet_dt=idt,
bias_out=bo,
)
en1 = FittingNet.deserialize(en0.serialize())
inp = np.ones([ni])
en0.call(inp)
en1.call(inp)
np.testing.assert_allclose(en0.call(inp), en1.call(inp))


class TestNetworkCollection(unittest.TestCase):
def setUp(self) -> None:
w = np.full((2, 3), 3.0)
Expand Down
Loading