Skip to content

Commit

Permalink
support compressing any neuron structure
Browse files Browse the repository at this point in the history 
Fix #1370.

Signed-off-by: Jinzhe Zeng <[email protected]>
  • Loading branch information
@njzjz
njzjz committed Oct 18, 2023
1 parent 05e0d27 commit d9a68c0
Show file tree
Hide file tree
Showing 7 changed files with 77 additions and 56 deletions.
7 changes: 0 additions & 7 deletions deepmd/descriptor/se_a.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -469,13 +469,6 @@ def enable_compression(
"empty embedding-net are not supported in model compression!"
)

for ii in range(len(self.filter_neuron) - 1):
if self.filter_neuron[ii] * 2 != self.filter_neuron[ii + 1]:
raise NotImplementedError(
"Model Compression error: descriptor neuron [%s] is not supported by model compression! "
"The size of the next layer of the neural network must be twice the size of the previous layer."
% ",".join([str(item) for item in self.filter_neuron])
)
if self.stripped_type_embedding:
ret_two_side = get_pattern_nodes_from_graph_def(
graph_def, f"filter_type_all{suffix}/.+_two_side_ebd"
Expand Down
8 changes: 0 additions & 8 deletions deepmd/descriptor/se_atten.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -387,14 +387,6 @@ def enable_compression(
"empty embedding-net are not supported in model compression!"
)

for ii in range(len(self.filter_neuron) - 1):
if self.filter_neuron[ii] * 2 != self.filter_neuron[ii + 1]:
raise NotImplementedError(
"Model Compression error: descriptor neuron [%s] is not supported by model compression! "
"The size of the next layer of the neural network must be twice the size of the previous layer."
% ",".join([str(item) for item in self.filter_neuron])
)

if self.attn_layer != 0:
raise RuntimeError("can not compress model when attention layer is not 0.")

Expand Down
8 changes: 0 additions & 8 deletions deepmd/descriptor/se_r.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -334,14 +334,6 @@ def enable_compression(
not self.filter_resnet_dt
), "Model compression error: descriptor resnet_dt must be false!"

for ii in range(len(self.filter_neuron) - 1):
if self.filter_neuron[ii] * 2 != self.filter_neuron[ii + 1]:
raise NotImplementedError(
"Model Compression error: descriptor neuron [%s] is not supported by model compression! "
"The size of the next layer of the neural network must be twice the size of the previous layer."
% ",".join([str(item) for item in self.filter_neuron])
)

self.compress = True
self.table = DPTabulate(
self,
Expand Down
8 changes: 0 additions & 8 deletions deepmd/descriptor/se_t.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -349,14 +349,6 @@ def enable_compression(
not self.filter_resnet_dt
), "Model compression error: descriptor resnet_dt must be false!"

for ii in range(len(self.filter_neuron) - 1):
if self.filter_neuron[ii] * 2 != self.filter_neuron[ii + 1]:
raise NotImplementedError(
"Model Compression error: descriptor neuron [%s] is not supported by model compression! "
"The size of the next layer of the neural network must be twice the size of the previous layer."
% ",".join([str(item) for item in self.filter_neuron])
)

self.compress = True
self.table = DPTabulate(
self,
Expand Down
90 changes: 68 additions & 22 deletions deepmd/utils/tabulate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -571,13 +571,13 @@ def _make_data(self, xx, idx):
+ xx
)
dy = op_module.unaggregated_dy_dx_s(
yy,
yy - xx,
self.matrix["layer_" + str(layer + 1)][idx],
xbar,
tf.constant(self.functype),
) + tf.ones([1, 1], yy.dtype)
dy2 = op_module.unaggregated_dy2_dx_s(
yy,
yy - xx,
dy,
self.matrix["layer_" + str(layer + 1)][idx],
xbar,
Expand Down Expand Up @@ -626,26 +626,72 @@ def _make_data(self, xx, idx):
tf.matmul(yy, self.matrix["layer_" + str(layer + 1)][idx])
+ self.bias["layer_" + str(layer + 1)][idx]
)
tt, zz = self._layer_1(
yy,
self.matrix["layer_" + str(layer + 1)][idx],
self.bias["layer_" + str(layer + 1)][idx],
)
dz = op_module.unaggregated_dy_dx(
zz - tt,
self.matrix["layer_" + str(layer + 1)][idx],
dy,
ybar,
tf.constant(self.functype),
)
dy2 = op_module.unaggregated_dy2_dx(
zz - tt,
self.matrix["layer_" + str(layer + 1)][idx],
dy,
dy2,
ybar,
tf.constant(self.functype),
)
if self.neuron[layer] == self.neuron[layer - 1]:
zz = (
self._layer_0(
yy,
self.matrix["layer_" + str(layer + 1)][idx],
self.bias["layer_" + str(layer + 1)][idx],
)
+ yy
)
dz = op_module.unaggregated_dy_dx(
zz - yy,
self.matrix["layer_" + str(layer + 1)][idx],
dy,
ybar,
tf.constant(self.functype),
)
dy2 = op_module.unaggregated_dy2_dx(
zz - yy,
self.matrix["layer_" + str(layer + 1)][idx],
dy,
dy2,
ybar,
tf.constant(self.functype),
)
elif self.neuron[layer] == 2 * self.neuron[layer - 1]:
tt, zz = self._layer_1(
yy,
self.matrix["layer_" + str(layer + 1)][idx],
self.bias["layer_" + str(layer + 1)][idx],
)
dz = op_module.unaggregated_dy_dx(
zz - tt,
self.matrix["layer_" + str(layer + 1)][idx],
dy,
ybar,
tf.constant(self.functype),
)
dy2 = op_module.unaggregated_dy2_dx(
zz - tt,
self.matrix["layer_" + str(layer + 1)][idx],
dy,
dy2,
ybar,
tf.constant(self.functype),
)
else:
zz = self._layer_0(
yy,
self.matrix["layer_" + str(layer + 1)][idx],
self.bias["layer_" + str(layer + 1)][idx],
)
dz = op_module.unaggregated_dy_dx(
zz,
self.matrix["layer_" + str(layer + 1)][idx],
dy,
ybar,
tf.constant(self.functype),
)
dy2 = op_module.unaggregated_dy2_dx(
zz,
self.matrix["layer_" + str(layer + 1)][idx],
dy,
dy2,
ybar,
tf.constant(self.functype),
)
dy = dz
yy = zz

Expand Down
8 changes: 6 additions & 2 deletions source/op/unaggregated_grad.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -168,7 +168,9 @@ struct UnaggregatedDyDxFunctor {
accumulator += w[jj * width + ii] * dy_dx[kk * size + jj];
}
dz_drou *= accumulator;
dz_drou += dy_dx[kk * size + ii % size];
if (width == 2 * size || width == size) {
dz_drou += dy_dx[kk * size + ii % size];
}
dz_dx[kk * width + ii] = dz_drou;
}
}
Expand Down Expand Up @@ -256,7 +258,9 @@ struct UnaggregatedDy2DxFunctor {
dz_drou +=
grad_grad(ybar[kk * width + ii], z[kk * width + ii], functype) *
accumulator * accumulator;
dz_drou += dy2_dx[kk * size + ii % size];
if (width == 2 * size || width == size) {
dz_drou += dy2_dx[kk * size + ii % size];
}
dz2_dx[kk * width + ii] = dz_drou;
}
}
Expand Down
4 changes: 3 additions & 1 deletion source/tests/model_compression/input.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,12 @@
],
"rcut_smth": 0.50,
"rcut": 6.00,
"_comment": "N2=2N1, N2=N1, and otherwise can be tested",
"neuron": [
4,
8,
16
17,
17
],
"resnet_dt": false,
"axis_neuron": 16,
Expand Down

0 comments on commit d9a68c0

Please sign in to comment.