Merge branch 'devel' into devel-modifier-plugin

Author: ChiahsinChu

deepmd/dpmodel/descriptor/dpa1.py

@@ -899,6 +899,7 @@ def call(
         exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)
         # nfnl x nnei
         exclude_mask = xp.reshape(exclude_mask, (nf * nloc, nnei))
+        exclude_mask = xp.astype(exclude_mask, xp.bool)
         # nfnl x nnei
         nlist = xp.reshape(nlist, (nf * nloc, nnei))
         nlist = xp.where(exclude_mask, nlist, xp.full_like(nlist, -1))

deepmd/dpmodel/descriptor/repformers.py

@@ -393,6 +393,7 @@ def call(
     ):
         xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext)
         exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)
+        exclude_mask = xp.astype(exclude_mask, xp.bool)
         nlist = xp.where(exclude_mask, nlist, xp.full_like(nlist, -1))
         # nf x nloc x nnei x 4
         dmatrix, diff, sw = self.env_mat.call(

deepmd/dpmodel/descriptor/se_t_tebd.py

@@ -682,6 +682,7 @@ def call(
         exclude_mask = xp.reshape(exclude_mask, (nf * nloc, nnei))
         # nfnl x nnei
         nlist = xp.reshape(nlist, (nf * nloc, nnei))
+        exclude_mask = xp.astype(exclude_mask, xp.bool)
         nlist = xp.where(exclude_mask, nlist, xp.full_like(nlist, -1))
         # nfnl x nnei
         nlist_mask = nlist != -1

deepmd/dpmodel/fitting/general_fitting.py

@@ -488,6 +488,7 @@ def _call_common(
         )
         # nf x nloc
         exclude_mask = self.emask.build_type_exclude_mask(atype)
+        exclude_mask = xp.astype(exclude_mask, xp.bool)
         # nf x nloc x nod
         outs = xp.where(exclude_mask[:, :, None], outs, xp.zeros_like(outs))
         return {self.var_name: outs}

doc/train/tensorboard.md

@@ -26,42 +26,51 @@ Before running TensorBoard, make sure you have generated summary data in a log
 directory by modifying the input script, setting {ref}`tensorboard <training/tensorboard>` to true in the training subsection will enable the TensorBoard data analysis. eg. **water_se_a.json**.
 
 ```json
-    "training" : {
-	"systems":	["../data/"],
-	"stop_batch":	1000000,
-	"batch_size":	1,
-
-	"seed":		1,
-
-	"_comment": " display and restart",
-	"_comment": " frequencies counted in batch",
-	"disp_file":	"lcurve.out",
-	"disp_freq":	100,
-	"numb_test":	10,
-	"save_freq":	1000,
-	"save_ckpt":	"model.ckpt",
-
-	"disp_training":true,
-	"time_training":true,
-	"tensorboard":	true,
-	"tensorboard_log_dir":"log",
-	"tensorboard_freq": 1000,
-	"profiling":	false,
-	"profiling_file":"timeline.json",
-	"_comment":	"that's all"
-    }
+"training": {
+   "systems": ["../data/"],
+   "stop_batch": 1000000,
+   "batch_size": 1,
+
+   "seed": 1,
+   "_comment": " display and restart",
+   "_comment": " frequencies counted in batch",
+   "disp_file": "lcurve.out",
+   "disp_freq": 100,
+   "numb_test": 10,
+   "save_freq": 1000,
+   "save_ckpt": "model.ckpt",
+
+   "disp_training": true,
+   "time_training": true,
+   "tensorboard": true,
+   "tensorboard_log_dir": "log",
+   "tensorboard_freq": 1000,
+   "profiling": false,
+   "profiling_file": "timeline.json",
+   "_comment": "that's all"
+}
 ```
 
 Once you have event files, run TensorBoard and provide the log directory. This
 should print that TensorBoard has started. Next, connect to http://tensorboard_server_ip:6006.
 
-TensorBoard requires a logdir to read logs from. For info on configuring TensorBoard, run TensorBoard --help.
+TensorBoard requires a logdir to read logs from. For info on configuring TensorBoard, run `tensorboard --help`.
 One can easily change the log name with "tensorboard_log_dir" and the sampling frequency with "tensorboard_freq".
 
 ```bash
 tensorboard --logdir path/to/logs
 ```
 
+## PyTorch Profiler With TensorBoard {{ pytorch_icon }}
+
+DeePMD-kit has a built-in support for [PyTorch Profiler](https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html#use-profiler-to-record-execution-events).
+The profiler requires extra packages for recording and visualization:
+`pip install tensorboard torch-tb-profiler`
+Set `"enable_profiler": true` in the training section of the input script, and launch a training task with 10 steps, since the default setting of the profiler scheduler is `wait=1, warmup=1, active=3, repeat=1`.
+The profiler will generate recording files in `tensorboard_log_dir`.
+
+To [visualize the profiling data](https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html#use-tensorboard-to-view-results-and-analyze-model-performance), launch TensorBoard (see above) and navigate to the "pytorch_profiler" tab.
+
 ## Examples
 
 ### Tracking and visualizing loss metrics(red:train, blue:test)