resize output vector size before copy to CPU

deepmodeling · Dec 31, 2024 · 3f7e5af · 3f7e5af
1 parent 269e96e
commit 3f7e5af
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 9 deletions.
diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt
@@ -29,14 +29,21 @@ endif()
 
 if(ENABLE_PADDLE)
   if(NOT DEFINED PADDLE_INFERENCE_DIR)
-    # message( FATAL_ERROR "Make sure PADDLE_INFERENCE_DIR is set when
+    # message(FATAL_ERROR "Make sure PADDLE_INFERENCE_DIR is set when
     # ENABLE_PADDLE=ON")
-    message(STATUS "PADDLE_INFERENCE_DIR is not defined. Downloading...")
     if(USE_CUDA_TOOLKIT)
+      message(
+        STATUS
+          "PADDLE_INFERENCE_DIR is not defined, downloading GPU infernece lib..."
+      )
       set(DOWNLOAD_URL
           "https://paddle-qa.bj.bcebos.com/paddle-pipeline/GITHUB_Docker_Compile_Test_Cuda118_cudnn860_Trt8531_D1/latest/paddle_inference.tgz"
       )
     else()
+      message(
+        STATUS
+          "PADDLE_INFERENCE_DIR is not defined, downloading CPU infernece lib..."
+      )
       set(DOWNLOAD_URL
           "https://paddle-qa.bj.bcebos.com/paddle-pipeline/GITHUB_Docker_Compile_Test_Cpu_Mkl_Avx_D1/latest/paddle_inference.tgz"
       )

diff --git a/source/api_cc/include/DeepPotPD.h b/source/api_cc/include/DeepPotPD.h
@@ -261,6 +261,36 @@ class DeepPotPD : public DeepPotBackend {
     return aparam_nall;
   };
 
+  /**
+   * @brief Print the shape of given tensor.
+   * @param[in] x Tensor x.
+   **/
+  void print_shape(const paddle_infer::Tensor& x) const {
+    std::vector<int> x_shape = x.shape();
+    std::string shape_str = "[";
+    for (int i = 0, n = x_shape.size(); i < n; ++i) {
+      if (i > 0) {
+        shape_str += ", ";
+      }
+      shape_str += std::to_string(x_shape[i]);
+    }
+    std::cout << shape_str;
+  };
+
+  /**
+   * @brief Compute the number of elements in a tensor.
+   * @param[in] x Tensor x.
+   **/
+  int numel(const paddle_infer::Tensor& x) const {
+    // TODO: There might be a overflow problem here for multiply int numbers.
+    int ret = 1;
+    std::vector<int> x_shape = x.shape();
+    for (std::size_t i = 0, n = x_shape.size(); i < n; ++i) {
+      ret *= x_shape[i];
+    }
+    return ret;
+  };
+
   // forward to template class
   void computew(std::vector<double>& ener,
                 std::vector<double>& force,

diff --git a/source/api_cc/src/DeepPotPD.cc b/source/api_cc/src/DeepPotPD.cc
@@ -77,9 +77,9 @@ void DeepPotPD::init(const std::string& model,
                                    " do not exist, please check it.");
   }
   config->SetModel(pdmodel_path, pdiparams_path);
-  config->EnableUseGpu(
-      4096, 0);  // annotate it if use cpu, default use gpu with 4G mem
-  gpu_enabled = config->use_gpu();
+  // config->EnableUseGpu(
+  //     4096, 0);  // annotate it if use cpu, default use gpu with 4G mem
+  // gpu_enabled = config->use_gpu();
   gpu_enabled = false;
   if (!gpu_enabled) {
     config->DisableGpu();
@@ -349,15 +349,23 @@ void DeepPotPD::compute(ENERGYVTYPE& ener,
   }
 
   auto output_names = predictor->GetOutputNames();
-  for (int i = 0; i < output_names.size(); ++i) {
-    printf("output_names[%d] = %s\n", i, output_names[i].c_str());
-  }
   auto energy_ = predictor->GetOutputHandle(output_names[1]);
   auto force_ = predictor->GetOutputHandle(output_names[2]);
-  auto virial_ = predictor->GetOutputHandle(output_names[3]);
+  auto virial_ = predictor->GetOutputHandle(output_names[4]);
 
+  int enery_numel = numel(*energy_);
+  assert(enery_numel > 0);
+  ener.resize(enery_numel);
   energy_->CopyToCpu(ener.data());
+
+  int forcey_numel = numel(*force_);
+  assert(forcey_numel > 0);
+  force.resize(forcey_numel);
   force_->CopyToCpu(force.data());
+
+  int virial_numel = numel(*virial_);
+  assert(virial_numel > 0);
+  virial.resize(virial_numel);
   virial_->CopyToCpu(virial.data());
 
   if (atomic) {