From 97794544be30b6a6eb5f4039cdc1ff3cd4f27f13 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Wed, 13 Nov 2024 13:14:49 -0500 Subject: [PATCH] docs: improve docs, examples, and error messages (#4338) 1. Refactor type embedding docs; 2. Refactor model compression docs; 3. Refactor DPA-1 docs; 4. Add error messages when type embedding is set in other backends; 5. Bump `sel` in the DPA-2 example. ## Summary by CodeRabbit ## Release Notes - **New Features** - Introduced enhanced error handling for model parameters across various backends, clarifying the usage of `type_embedding`. - Added sections on "Type embedding" and "Model compression" in multiple documentation files to improve user guidance. - New input file `input_torch_compressible.json` added for the DPA-2 model, providing a compressible configuration. - **Bug Fixes** - Updated serialization and deserialization methods to prevent unsupported configurations in model handling. - **Documentation** - Revised documentation to clarify type embedding support and model compression conditions. - Removed references to outdated system formats in the documentation. - **Configuration Updates** - Increased selection parameters for three-body interactions and representation transformers in example configuration files. --------- Signed-off-by: Jinzhe Zeng Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com> --- deepmd/dpmodel/model/model.py | 4 + deepmd/jax/model/model.py | 4 + deepmd/pt/model/model/__init__.py | 4 + deepmd/tf/model/model.py | 4 + deepmd/utils/argcheck.py | 2 +- doc/data/system.md | 2 +- doc/freeze/compress.md | 5 +- doc/model/dpa2.md | 10 ++ doc/model/train-hybrid.md | 10 ++ doc/model/train-se-a-mask.md | 8 ++ doc/model/train-se-atten.md | 39 +++--- doc/model/train-se-e2-a-tebd.md | 17 ++- doc/model/train-se-e2-a.md | 10 ++ doc/model/train-se-e2-r.md | 8 ++ doc/model/train-se-e3-tebd.md | 8 ++ doc/model/train-se-e3.md | 8 ++ examples/water/dpa2/README.md | 2 + .../water/dpa2/input_torch_compressible.json | 113 ++++++++++++++++++ examples/water/dpa2/input_torch_large.json | 4 +- examples/water/dpa2/input_torch_medium.json | 4 +- examples/water/dpa2/input_torch_small.json | 4 +- source/tests/common/test_examples.py | 1 + 22 files changed, 236 insertions(+), 35 deletions(-) create mode 100644 examples/water/dpa2/input_torch_compressible.json diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 8fa210e864..19408e58c4 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -36,6 +36,10 @@ def get_standard_model(data: dict) -> EnergyModel: data : dict The data to construct the model. """ + if "type_embedding" in data: + raise ValueError( + "In the DP backend, type_embedding is not at the model level, but within the descriptor. See type embedding documentation for details." + ) data["descriptor"]["type_map"] = data["type_map"] data["descriptor"]["ntypes"] = len(data["type_map"]) fitting_type = data["fitting_net"].pop("type") diff --git a/deepmd/jax/model/model.py b/deepmd/jax/model/model.py index 983815100c..8b7d375841 100644 --- a/deepmd/jax/model/model.py +++ b/deepmd/jax/model/model.py @@ -35,6 +35,10 @@ def get_standard_model(data: dict): The data to construct the model. """ data = deepcopy(data) + if "type_embedding" in data: + raise ValueError( + "In the JAX backend, type_embedding is not at the model level, but within the descriptor. See type embedding documentation for details." + ) descriptor_type = data["descriptor"].pop("type") data["descriptor"]["type_map"] = data["type_map"] data["descriptor"]["ntypes"] = len(data["type_map"]) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 580b83cb3d..da1edb3354 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -76,6 +76,10 @@ def _get_standard_model_components(model_params, ntypes): + if "type_embedding" in model_params: + raise ValueError( + "In the PyTorch backend, type_embedding is not at the model level, but within the descriptor. See type embedding documentation for details." + ) # descriptor model_params["descriptor"]["ntypes"] = ntypes model_params["descriptor"]["type_map"] = copy.deepcopy(model_params["type_map"]) diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 5dcbad37ca..2edfe4b651 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -842,6 +842,10 @@ def serialize(self, suffix: str = "") -> dict: Name suffix to identify this descriptor """ if self.typeebd is not None: + if not self.descrpt.explicit_ntypes: + raise RuntimeError( + "type embedding for descriptors without mixed types is not supported in other backends" + ) self.descrpt.type_embedding = self.typeebd self.fitting.tebd_dim = self.typeebd.neuron[-1] if self.spin is not None: diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index af9e2984dd..60c85123d7 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1772,7 +1772,7 @@ def model_args(exclude_hybrid=False): doc_data_stat_nbatch = "The model determines the normalization from the statistics of the data. This key specifies the number of `frames` in each `system` used for statistics." doc_data_stat_protect = "Protect parameter for atomic energy regression." doc_data_bias_nsample = "The number of training samples in a system to compute and change the energy bias." - doc_type_embedding = "The type embedding." + doc_type_embedding = "The type embedding. In other backends, the type embedding is already included in the descriptor." doc_modifier = "The modifier of model output." doc_use_srtab = "The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly." doc_smin_alpha = "The short-range tabulated interaction will be switched according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when `use_srtab` is provided." diff --git a/doc/data/system.md b/doc/data/system.md index b5db516726..b50c6fa256 100644 --- a/doc/data/system.md +++ b/doc/data/system.md @@ -1,6 +1,6 @@ # System -DeePMD-kit takes a **system** as the data structure. A snapshot of a system is called a **frame**. A system may contain multiple frames with the same atom types and numbers, i.e. the same formula (like `H2O`). To contains data with different formulas, one usually needs to divide data into multiple systems, which may sometimes result in sparse-frame systems. See a [new system format](../model/train-se-atten.md#data-format) to further combine different systems with the same atom numbers, when training with descriptor `se_atten`. +DeePMD-kit takes a **system** as the data structure. A snapshot of a system is called a **frame**. A system may contain multiple frames with the same atom types and numbers, i.e. the same formula (like `H2O`). To contains data with different formulas, one usually needs to divide data into multiple systems, which may sometimes result in sparse-frame systems. A system should contain system properties, input frame properties, and labeled frame properties. The system property contains the following property: diff --git a/doc/freeze/compress.md b/doc/freeze/compress.md index cdb12cc9e7..f7b5a39aa3 100644 --- a/doc/freeze/compress.md +++ b/doc/freeze/compress.md @@ -112,9 +112,8 @@ The model compression interface requires the version of DeePMD-kit used in the o **Acceptable descriptor type** -Descriptors with `se_e2_a`, `se_e3`, `se_e2_r` and `se_atten_v2` types are supported by the model compression feature. `Hybrid` mixed with the above descriptors is also supported. - -Notice: Model compression for the `se_atten_v2` descriptor is exclusively designed for models with the training parameter {ref}`attn_layer ` set to 0. +Not any descriptor supports model compression. +See the documentation of a specific descriptor to see whether it supports model compression. **Available activation functions for descriptor:** diff --git a/doc/model/dpa2.md b/doc/model/dpa2.md index 746f7969d7..eb641d6b01 100644 --- a/doc/model/dpa2.md +++ b/doc/model/dpa2.md @@ -31,3 +31,13 @@ See the example `examples/water/lmp/jax_dpa2.lammps`. ## Data format DPA-2 supports both the [standard data format](../data/system.md) and the [mixed type data format](../data/system.md#mixed-type). + +## Type embedding + +Type embedding is within this descriptor with the {ref}`tebd_dim ` argument. + +## Model compression + +Model compression is supported when {ref}`repinit/tebd_input_mode ` is `strip`, but only the `repinit` part is compressed. +An example is given in `examples/water/dpa2/input_torch_compressible.json`. +The performance improvement will be limited if other parts are more expensive. diff --git a/doc/model/train-hybrid.md b/doc/model/train-hybrid.md index da3b40487b..d565af5c9a 100644 --- a/doc/model/train-hybrid.md +++ b/doc/model/train-hybrid.md @@ -48,3 +48,13 @@ A complete training input script of this example can be found in the directory ```bash $deepmd_source_dir/examples/water/hybrid/input.json ``` + +## Type embedding + +Type embedding is different between the TensorFlow backend and other backends. +In the TensorFlow backend, all descriptors share the same descriptor that defined in the model level. +In other backends, each descriptor has its own type embedding and their parameters may be different. + +## Model compression + +Model compression is supported if all sub-descriptors support model compression. diff --git a/doc/model/train-se-a-mask.md b/doc/model/train-se-a-mask.md index 93edfc999e..ff1ee76c12 100644 --- a/doc/model/train-se-a-mask.md +++ b/doc/model/train-se-a-mask.md @@ -84,3 +84,11 @@ And the `loss` section in the training input script should be set as follows. "_comment": " that's all" } ``` + +## Type embedding + +Same as [`se_e2_a`](./train-se-e2-a.md). + +## Model compression + +Same as [`se_e2_a`](./train-se-e2-a.md). diff --git a/doc/model/train-se-atten.md b/doc/model/train-se-atten.md index 3e88a4e950..d8d7c3c98c 100644 --- a/doc/model/train-se-atten.md +++ b/doc/model/train-se-atten.md @@ -4,8 +4,6 @@ **Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} ::: -## DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation - ![ALT](../images/model_se_atten.png "model_se_atten") Here we propose DPA-1, a Deep Potential model with a novel attention mechanism, which is highly effective for representing the conformation and chemical spaces of atomic systems and learning the PES. @@ -68,11 +66,9 @@ Then layer normalization is added in a residual way to finally obtain the self-a [^1]: This section is built upon Jinzhe Zeng, Duo Zhang, Denghui Lu, Pinghui Mo, Zeyu Li, Yixiao Chen, Marián Rynik, Li'ang Huang, Ziyao Li, Shaochen Shi, Yingze Wang, Haotian Ye, Ping Tuo, Jiabin Yang, Ye Ding, Yifan Li, Davide Tisi, Qiyu Zeng, Han Bao, Yu Xia, Jiameng Huang, Koki Muraoka, Yibo Wang, Junhan Chang, Fengbo Yuan, Sigbjørn Løland Bore, Chun Cai, Yinnian Lin, Bo Wang, Jiayan Xu, Jia-Xin Zhu, Chenxing Luo, Yuzhi Zhang, Rhys E. A. Goodall, Wenshuo Liang, Anurag Kumar Singh, Sikai Yao, Jingchao Zhang, Renata Wentzcovitch, Jiequn Han, Jie Liu, Weile Jia, Darrin M. York, Weinan E, Roberto Car, Linfeng Zhang, Han Wang, [J. Chem. Phys. 159, 054801 (2023)](https://doi.org/10.1063/5.0155600) licensed under a [Creative Commons Attribution (CC BY) license](http://creativecommons.org/licenses/by/4.0/). -## Introduction to new features of DPA-1 - -Next, we will list the detailed settings in input.json and the data format, especially for large systems with dozens of elements. An example of DPA-1 input can be found [here](../../examples/water/se_atten/input.json). +## Descriptor `"se_atten"` -### Descriptor `"se_atten"` +Next, we will list the detailed settings in input.json and the data format, especially for large systems with dozens of elements. An example of DPA-1 input can be found in `examples/water/se_atten/input.json`. The notation of `se_atten` is short for the smooth edition of Deep Potential with an attention mechanism. This descriptor was described in detail in [the DPA-1 paper](https://arxiv.org/abs/2208.08236) and the images above. @@ -126,7 +122,7 @@ We highly recommend using the version 2.0 of the attention-based descriptor `"se "set_davg_zero": false ``` -You can use descriptor `"se_atten_v2"` and do not need to set `tebd_input_mode` and `smooth_type_embedding`. In `"se_atten_v2"`, `tebd_input_mode` is forced to be `"strip"` and `smooth_type_embedding` is forced to be `"true"`. When `tebd_input_mode` is `"strip"`, the embedding matrix $\mathcal{G}^i$ is constructed as: +You can use descriptor `"se_atten_v2"` and is not allowed to set `tebd_input_mode` and `smooth_type_embedding`. In `"se_atten_v2"`, `tebd_input_mode` is forced to be `"strip"` and `smooth_type_embedding` is forced to be `"true"`. When `tebd_input_mode` is `"strip"`, the embedding matrix $\mathcal{G}^i$ is constructed as: ```math (\mathcal{G}^i)_j = \mathcal{N}_{e,2}(s(r_{ij})) + \mathcal{N}_{e,2}(s(r_{ij})) \odot ({N}_{e,2}(\{\mathcal{A}^i, \mathcal{A}^j\}) \odot s(r_{ij})) \quad \mathrm{or} @@ -140,25 +136,28 @@ Practical evidence demonstrates that `"se_atten_v2"` offers better and more stab Notice: Model compression for the `se_atten_v2` descriptor is exclusively designed for models with the training parameter {ref}`attn_layer ` set to 0. -### Fitting `"ener"` +## Type embedding -DPA-1 only supports `"ener"` fitting type, and you can refer [here](train-energy.md) for detailed information. +DPA-1 only supports models with type embeddings. -### Type embedding - -DPA-1 only supports models with type embeddings. And the default setting is as follows: +In the TensorFlow backend, the {ref}`type_embedding ` section will be used. If it is not set, the following default parameters will be used: ```json "type_embedding":{ - "neuron": [8], - "resnet_dt": false, - "seed": 1 - } + "neuron": [8], + "resnet_dt": false, + "seed": 1 +} ``` -You can add these settings in input.json if you want to change the default ones, see [here](train-se-e2-a-tebd.md) for detailed information. +In other backends, type embedding is within this descriptor with the {ref}`tebd_dim ` argument. + +## Difference between TensorFlow and other backends -### Type map +TensorFlow and other backends have different implementations for {ref}`smooth_type_embedding `. +The results are inconsistent when `smooth_type_embedding` is `true`. + +## Type map For training large systems, especially those with dozens of elements, the {ref}`type ` determines the element index of training data: @@ -176,6 +175,10 @@ which should include all the elements in the dataset you want to train on. DPA-1 supports both the [standard data format](../data/system.md) and the [mixed type data format](../data/system.md#mixed-type). +## Model compression + +Model compression is supported only when there is no attention layer (`attn_layer` is 0) and `tebd_input_mode` is `strip`. + ## Training example Here we upload the AlMgCu example shown in the paper, you can download it here: diff --git a/doc/model/train-se-e2-a-tebd.md b/doc/model/train-se-e2-a-tebd.md index a6291bb238..00726c0d3e 100644 --- a/doc/model/train-se-e2-a-tebd.md +++ b/doc/model/train-se-e2-a-tebd.md @@ -1,7 +1,7 @@ -# Type embedding approach {{ tensorflow_icon }} +# Type embedding approach {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} :::{note} -**Supported backends**: TensorFlow {{ tensorflow_icon }} +**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} ::: We generate specific a type embedding vector for each atom type so that we can share one descriptor embedding net and one fitting net in total, which decline training complexity largely. @@ -63,8 +63,9 @@ In this way, all chemical species share the same network parameters through the [^1]: This section is built upon Jinzhe Zeng, Duo Zhang, Denghui Lu, Pinghui Mo, Zeyu Li, Yixiao Chen, Marián Rynik, Li'ang Huang, Ziyao Li, Shaochen Shi, Yingze Wang, Haotian Ye, Ping Tuo, Jiabin Yang, Ye Ding, Yifan Li, Davide Tisi, Qiyu Zeng, Han Bao, Yu Xia, Jiameng Huang, Koki Muraoka, Yibo Wang, Junhan Chang, Fengbo Yuan, Sigbjørn Løland Bore, Chun Cai, Yinnian Lin, Bo Wang, Jiayan Xu, Jia-Xin Zhu, Chenxing Luo, Yuzhi Zhang, Rhys E. A. Goodall, Wenshuo Liang, Anurag Kumar Singh, Sikai Yao, Jingchao Zhang, Renata Wentzcovitch, Jiequn Han, Jie Liu, Weile Jia, Darrin M. York, Weinan E, Roberto Car, Linfeng Zhang, Han Wang, [J. Chem. Phys. 159, 054801 (2023)](https://doi.org/10.1063/5.0155600) licensed under a [Creative Commons Attribution (CC BY) license](http://creativecommons.org/licenses/by/4.0/). -## Instructions +## Instructions for TensorFlow backend {{ tensorflow_icon }} +In the TensorFlow backend, the type embedding is at the model level. The {ref}`model ` defines how the model is constructed, adding a section of type embedding net: ```json @@ -106,6 +107,10 @@ $deepmd_source_dir/examples/water/se_e2_a_tebd/input.json See [here](../development/type-embedding.md) for further explanation of `type embedding`. -:::{note} -You can't apply the compression method while using the atom type embedding. -::: +See documentation for each descriptor for details. + +## Instructions for other backends + +In other backends, the type embedding is within the descriptor itself. + +See documentation for each descriptor for details. diff --git a/doc/model/train-se-e2-a.md b/doc/model/train-se-e2-a.md index d4a4510a31..2cc537e349 100644 --- a/doc/model/train-se-e2-a.md +++ b/doc/model/train-se-e2-a.md @@ -94,3 +94,13 @@ The construction of the descriptor is given by section {ref}`descriptor ` specifies the size of the submatrix of the embedding matrix, the axis matrix as explained in the [DeepPot-SE paper](https://arxiv.org/abs/1805.09003) - If the option {ref}`resnet_dt ` is set to `true`, then a timestep is used in the ResNet. - {ref}`seed ` gives the random seed that is used to generate random numbers when initializing the model parameters. + +## Type embedding support + +Type embdding is only supported in the TensorFlow backends. +`se_e2_a` with type embedding and [`se_atten`](./train-se-atten.md) (or its updated version) without any attention layer are mathematically equivalent, so `se_atten` can be a substitute in other backends. + +## Model compression + +Model compression is supported when type embedding is not used. +To use model compression with type embedding in the TensorFlow backend, use `se_a_tebd_v2` instead. diff --git a/doc/model/train-se-e2-r.md b/doc/model/train-se-e2-r.md index baff6d6331..dbaf8bf364 100644 --- a/doc/model/train-se-e2-r.md +++ b/doc/model/train-se-e2-r.md @@ -69,3 +69,11 @@ The training input script is very similar to that of [`se_e2_a`](train-se-e2-a.m ``` The type of the descriptor is set by the key {ref}`type `. + +## Type embedding support + +Type embdding is only supported in the TensorFlow backends. + +## Model compression + +Model compression is supported when type embedding is not used. diff --git a/doc/model/train-se-e3-tebd.md b/doc/model/train-se-e3-tebd.md index 49d0d80f42..3e0b6f9c5a 100644 --- a/doc/model/train-se-e3-tebd.md +++ b/doc/model/train-se-e3-tebd.md @@ -76,3 +76,11 @@ The training input script is very similar to that of [`se_e2_a`](train-se-e2-a.m ``` The type of the descriptor is set by the key {ref}`type `. + +## Type embedding + +Type embedding is within this descriptor with the {ref}`tebd_dim ` argument. + +## Model compression + +Model compression is not supported. diff --git a/doc/model/train-se-e3.md b/doc/model/train-se-e3.md index 714d75259a..4e7f4fa53a 100644 --- a/doc/model/train-se-e3.md +++ b/doc/model/train-se-e3.md @@ -64,3 +64,11 @@ The training input script is very similar to that of [`se_e2_a`](train-se-e2-a.m ``` The type of the descriptor is set by the key {ref}`type `. + +## Type embedding + +Use [`se_e3_tebd`](./train-se-e3-tebd.md) for type embedding support. + +## Model compression + +Model compression is supported. diff --git a/examples/water/dpa2/README.md b/examples/water/dpa2/README.md index aa37d410a8..c3c3fbb838 100644 --- a/examples/water/dpa2/README.md +++ b/examples/water/dpa2/README.md @@ -13,3 +13,5 @@ For detailed differences in their configurations, please refer to the table belo | `input_torch_small.json` | 3 | ✓ | ✗ | ✓ | Smallest DPA-2 model, optimized for speed. | | `input_torch_medium.json` | 6 | ✓ | ✓ | ✓ | Recommended well-performing DPA-2 model, balancing efficiency and precision. | | `input_torch_large.json` | 12 | ✓ | ✓ | ✓ | Most complex model with the highest precision. | + +`input_torch_compressible.json` is derived from `input_torch_small.json` and makes the `repinit` part compressible. diff --git a/examples/water/dpa2/input_torch_compressible.json b/examples/water/dpa2/input_torch_compressible.json new file mode 100644 index 0000000000..cc34571a11 --- /dev/null +++ b/examples/water/dpa2/input_torch_compressible.json @@ -0,0 +1,113 @@ +{ + "_comment": "that's all", + "model": { + "type_map": [ + "O", + "H" + ], + "descriptor": { + "type": "dpa2", + "repinit": { + "tebd_dim": 8, + "rcut": 6.0, + "rcut_smth": 0.5, + "nsel": 120, + "neuron": [ + 25, + 50, + 100 + ], + "axis_neuron": 12, + "activation_function": "tanh", + "three_body_sel": 48, + "three_body_rcut": 4.0, + "three_body_rcut_smth": 3.5, + "use_three_body": true, + "tebd_input_mode": "strip" + }, + "repformer": { + "rcut": 4.0, + "rcut_smth": 3.5, + "nsel": 48, + "nlayers": 3, + "g1_dim": 128, + "g2_dim": 32, + "attn2_hidden": 32, + "attn2_nhead": 4, + "attn1_hidden": 128, + "attn1_nhead": 4, + "axis_neuron": 4, + "update_h2": false, + "update_g1_has_conv": true, + "update_g1_has_grrg": true, + "update_g1_has_drrd": true, + "update_g1_has_attn": false, + "update_g2_has_g1g1": false, + "update_g2_has_attn": false, + "update_style": "res_residual", + "update_residual": 0.01, + "update_residual_init": "norm", + "attn2_has_gate": true, + "use_sqrt_nnei": true, + "g1_out_conv": true, + "g1_out_mlp": true + }, + "add_tebd_to_repinit_out": false + }, + "fitting_net": { + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "seed": 1, + "_comment": " that's all" + }, + "_comment": " that's all" + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-08, + "_comment": "that's all" + }, + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + "_comment": " that's all" + }, + "training": { + "stat_file": "./dpa2.hdf5", + "training_data": { + "systems": [ + "../data/data_0", + "../data/data_1", + "../data/data_2" + ], + "batch_size": 1, + "_comment": "that's all" + }, + "validation_data": { + "systems": [ + "../data/data_3" + ], + "batch_size": 1, + "_comment": "that's all" + }, + "numb_steps": 1000000, + "warmup_steps": 0, + "gradient_max_norm": 5.0, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 2000, + "_comment": "that's all" + } +} diff --git a/examples/water/dpa2/input_torch_large.json b/examples/water/dpa2/input_torch_large.json index 568cbc1a94..99eb14ec15 100644 --- a/examples/water/dpa2/input_torch_large.json +++ b/examples/water/dpa2/input_torch_large.json @@ -19,7 +19,7 @@ ], "axis_neuron": 12, "activation_function": "tanh", - "three_body_sel": 40, + "three_body_sel": 48, "three_body_rcut": 4.0, "three_body_rcut_smth": 3.5, "use_three_body": true @@ -27,7 +27,7 @@ "repformer": { "rcut": 4.0, "rcut_smth": 3.5, - "nsel": 40, + "nsel": 48, "nlayers": 12, "g1_dim": 128, "g2_dim": 32, diff --git a/examples/water/dpa2/input_torch_medium.json b/examples/water/dpa2/input_torch_medium.json index 5b739e6f27..804fb4ad68 100644 --- a/examples/water/dpa2/input_torch_medium.json +++ b/examples/water/dpa2/input_torch_medium.json @@ -19,7 +19,7 @@ ], "axis_neuron": 12, "activation_function": "tanh", - "three_body_sel": 40, + "three_body_sel": 48, "three_body_rcut": 4.0, "three_body_rcut_smth": 3.5, "use_three_body": true @@ -27,7 +27,7 @@ "repformer": { "rcut": 4.0, "rcut_smth": 3.5, - "nsel": 40, + "nsel": 48, "nlayers": 6, "g1_dim": 128, "g2_dim": 32, diff --git a/examples/water/dpa2/input_torch_small.json b/examples/water/dpa2/input_torch_small.json index 98147030b6..bf506dcb5f 100644 --- a/examples/water/dpa2/input_torch_small.json +++ b/examples/water/dpa2/input_torch_small.json @@ -19,7 +19,7 @@ ], "axis_neuron": 12, "activation_function": "tanh", - "three_body_sel": 40, + "three_body_sel": 48, "three_body_rcut": 4.0, "three_body_rcut_smth": 3.5, "use_three_body": true @@ -27,7 +27,7 @@ "repformer": { "rcut": 4.0, "rcut_smth": 3.5, - "nsel": 40, + "nsel": 48, "nlayers": 3, "g1_dim": 128, "g2_dim": 32, diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py index 8236040f27..068a91709c 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -57,6 +57,7 @@ p_examples / "water" / "dpa2" / "input_torch_small.json", p_examples / "water" / "dpa2" / "input_torch_medium.json", p_examples / "water" / "dpa2" / "input_torch_large.json", + p_examples / "water" / "dpa2" / "input_torch_compressible.json", p_examples / "property" / "train" / "input_torch.json", p_examples / "water" / "se_e3_tebd" / "input_torch.json", )