refine code

Author: HydrogenSulfate

deepmd/pd/train/training.py

@@ -760,38 +760,6 @@ def step(_step_id, task_key="Default") -> None:
                     pref_lr = _lr.start_lr
                 else:
                     pref_lr = cur_lr
-                self.wrapper.load_state_dict(paddle.load("./wrapper_dict.pd"))
-                print("model loaded")
-                inp = np.load("./input_dict.npz", allow_pickle=True)
-                for k, v in inp.items():
-                    if isinstance(v, np.ndarray):
-                        # print(k, type(v), v.shape, v.dtype)
-                        try:
-                            input_dict[k] = paddle.to_tensor(v)
-                            # print(k)
-                        except Exception:
-                            pass
-                        if isinstance(input_dict[k], paddle.Tensor):
-                            input_dict[k] = input_dict[k].cuda()
-                print("input_dict loaded")
-                lab = np.load("./label_dict.npz", allow_pickle=True)
-                for k, v in lab.items():
-                    if isinstance(v, np.ndarray):
-                        # print(k, type(v), v.shape, v.dtype)
-                        try:
-                            label_dict[k] = paddle.to_tensor(v)
-                            # print(k)
-                        except Exception:
-                            pass
-                        if isinstance(label_dict[k], paddle.Tensor):
-                            label_dict[k] = label_dict[k].cuda()
-                print("label_dict loaded")
-                model_pred, loss, more_loss = self.wrapper(
-                    **input_dict, cur_lr=pref_lr, label=label_dict, task_key=task_key
-                )
-                print({k: float(v) for k, v in more_loss.items()})
-                print(f"{loss.item():.10f}")
-                exit()
 
                 # disable synchronization in forward-backward manually
                 # as derivatives exist in model forward
@@ -811,7 +779,7 @@ def step(_step_id, task_key="Default") -> None:
 
                     with nvprof_context(enable_profiling, "Backward pass"):
                         loss.backward()
-                exit()
+
                 # fuse + allreduce manually before optimization if use DDP + no_sync
                 # details in https://github.com/PaddlePaddle/Paddle/issues/48898#issuecomment-1343838622
                 if self.world_size > 1:

deepmd/pd/train/wrapper.py

@@ -173,7 +173,6 @@ def forward(
             model_pred = self.model[task_key](**input_dict)
             return model_pred, None, None
         else:
-            print(self.loss)
             natoms = atype.shape[-1]
             model_pred, loss, more_loss = self.loss[task_key](
                 input_dict,

deepmd/pt/model/model/transform_output.py

@@ -96,9 +96,6 @@ def task_deriv_one(
         extended_virial = extended_virial.view(list(extended_virial.shape[:-2]) + [9])  # noqa:RUF005
     else:
         extended_virial = None
-    print(
-        f"extended_force: {extended_force.min().item():.10f} {extended_force.max().item():.10f} {extended_force.mean().item():.10f} {extended_force.std().item():.10f}"
-    )
     return extended_force, extended_virial
 
 

deepmd/pt/train/training.py

@@ -725,64 +725,9 @@ def step(_step_id, task_key="Default") -> None:
                     pref_lr = _lr.start_lr
                 else:
                     pref_lr = cur_lr
-
-                # save
-                # torch.save(self.wrapper.state_dict(), "wrapper_dict.pt")
-                # import paddle
-                # psd = {}
-                # for k, v in self.wrapper.state_dict().items():
-                #     if isinstance(v, torch.Tensor):
-                #         psd[k] = paddle.from_dlpack(v.detach())
-                #     else:
-                #         psd[k] = v
-                # paddle.save(psd, "wrapper_dict.pd")
-                # inp = {}
-                # for k, v in input_dict.items():
-                #     if isinstance(v, torch.Tensor):
-                #         inp[k] = v.detach().cpu().numpy()
-                #     else:
-                #         inp[k] = v
-                # np.savez("./input_dict.npz", **inp)
-                # lab = {}
-                # for k, v in label_dict.items():
-                #     if isinstance(v, torch.Tensor):
-                #         lab[k] = v.detach().cpu().numpy()
-                #     else:
-                #         lab[k] = v
-                # np.savez("./label_dict.npz", **lab)
-
-                # load
-                self.wrapper.load_state_dict(torch.load("./wrapper_dict.pt"))
-                print("model loaded")
-                inp = np.load("./input_dict.npz", allow_pickle=True)
-                for k, v in inp.items():
-                    if isinstance(v, np.ndarray):
-                        # print(k, type(v), v.shape, v.dtype)
-                        try:
-                            input_dict[k] = torch.tensor(v)
-                        except TypeError:
-                            pass
-                        if isinstance(input_dict[k], torch.Tensor):
-                            input_dict[k] = input_dict[k].cuda()
-                print("input_dict loaded")
-                lab = np.load("./label_dict.npz", allow_pickle=True)
-                for k, v in lab.items():
-                    if isinstance(v, np.ndarray):
-                        # print(k, type(v), v.shape, v.dtype)
-                        try:
-                            label_dict[k] = torch.tensor(v)
-                        except TypeError:
-                            pass
-                        if isinstance(label_dict[k], torch.Tensor):
-                            label_dict[k] = label_dict[k].cuda()
-                print("label_dict loaded")
-
                 model_pred, loss, more_loss = self.wrapper(
                     **input_dict, cur_lr=pref_lr, label=label_dict, task_key=task_key
                 )
-                print({k: float(v) for k, v in more_loss.items()})
-                print(f"{loss.item():.10f}")
-                exit()
                 loss.backward()
                 if self.gradient_max_norm > 0.0:
                     torch.nn.utils.clip_grad_norm_(

deepmd/pt/train/wrapper.py

@@ -175,7 +175,6 @@ def forward(
             return model_pred, None, None
         else:
             natoms = atype.shape[-1]
-            print(self.loss)
             model_pred, loss, more_loss = self.loss[task_key](
                 input_dict,
                 self.model[task_key],

source/api_cc/include/DeepPotPD.h

@@ -392,18 +392,18 @@ class DeepPotPD : public DeepPotBackend {
   int do_message_passing;  // 1:dpa2 model 0:others
   bool gpu_enabled;
   std::unique_ptr<paddle_infer::Tensor> firstneigh_tensor;
-  std::vector<paddle_infer::Tensor> comm_vec;
-  paddle_infer::Tensor mapping_tensor = predictor_fl->GetInputHandle("mapping");
+
   /* comm_vec flatten all the communicatoin tensors below from comm dict:
-  - send_list
-  - send_proc
-  - recv_proc
-  - send_num
-  - recv_num
-  - communicator
+  - [0] send_list
+  - [1] send_proc
+  - [2] recv_proc
+  - [3] send_num
+  - [4] recv_num
+  - [5] communicator
   */
+  std::vector<paddle_infer::Tensor> comm_vec;
 
-  // std::unordered_map<std::string, paddle::Tensor> comm_dict; # Not used yet
+  paddle_infer::Tensor mapping_tensor = predictor_fl->GetInputHandle("mapping");
 };
 
 }  // namespace deepmd

source/api_cc/src/DeepPotPD.cc

@@ -205,8 +205,6 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
     nlist_data.shuffle_exclude_empty(fwd_map);
     nlist_data.padding();
     if (do_message_passing == 1 && nghost > 0) {
-      // throw deepmd::deepmd_exception(
-      //     "(do_message_passing == 1 && nghost > 0) is not supported yet.");
       int nswap = lmp_list.nswap;
       auto sendproc_tensor = predictor_fl->GetInputHandle("sendproc");
       sendproc_tensor->Reshape({nswap});
@@ -249,10 +247,6 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
       this->mapping_tensor->Reshape({1, nall_real});
       this->mapping_tensor->CopyFromCpu(mapping.data());
     }
-    // if (do_message_passing == 1 && nghost == 0) {
-    //   throw deepmd::deepmd_exception(
-    //       "(do_message_passing == 1 && nghost == 0) is not supported yet.");
-    // }
   }
   std::vector<int> firstneigh = createNlistTensorPD(nlist_data.jlist);
   firstneigh_tensor = predictor_fl->GetInputHandle("nlist");