Skip to content

Commit

Permalink
chore(cc): merge get backend codes (#4355)
Browse files Browse the repository at this point in the history
Fix #4308.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Introduced a new function to dynamically determine the backend
framework based on the model file type.
  
- **Improvements**
- Enhanced backend detection logic in multiple classes, allowing for
more flexible model initialization.
- Simplified control flow in the initialization methods of various
components.

- **Bug Fixes**
- Improved error handling for unsupported backends and model formats
during initialization processes.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
  • Loading branch information
njzjz authored Nov 14, 2024
1 parent 6d9d8bb commit d3095cf
Show file tree
Hide file tree
Showing 6 changed files with 23 additions and 23 deletions.
7 changes: 7 additions & 0 deletions source/api_cc/include/common.h
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,13 @@ namespace deepmd {
typedef double ENERGYTYPE;
enum DPBackend { TensorFlow, PyTorch, Paddle, JAX, Unknown };

/**
* @brief Get the backend of the model.
* @param[in] model The model name.
* @return The backend of the model.
**/
DPBackend get_backend(const std::string& model);

struct NeighborListData {
/// Array stores the core region atom's index
std::vector<int> ilist;
Expand Down
3 changes: 1 addition & 2 deletions source/api_cc/src/DataModifier.cc
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,7 @@ void DipoleChargeModifier::init(const std::string& model,
<< std::endl;
return;
}
// TODO: To implement detect_backend
DPBackend backend = deepmd::DPBackend::TensorFlow;
const DPBackend backend = get_backend(model);
if (deepmd::DPBackend::TensorFlow == backend) {
#ifdef BUILD_TENSORFLOW
dcm = std::make_shared<deepmd::DipoleChargeModifierTF>(model, gpu_rank,
Expand Down
12 changes: 1 addition & 11 deletions source/api_cc/src/DeepPot.cc
Original file line number Diff line number Diff line change
Expand Up @@ -39,17 +39,7 @@ void DeepPot::init(const std::string& model,
<< std::endl;
return;
}
DPBackend backend;
if (model.length() >= 4 && model.substr(model.length() - 4) == ".pth") {
backend = deepmd::DPBackend::PyTorch;
} else if (model.length() >= 3 && model.substr(model.length() - 3) == ".pb") {
backend = deepmd::DPBackend::TensorFlow;
} else if (model.length() >= 11 &&
model.substr(model.length() - 11) == ".savedmodel") {
backend = deepmd::DPBackend::JAX;
} else {
throw deepmd::deepmd_exception("Unsupported model file format");
}
const DPBackend backend = get_backend(model);
if (deepmd::DPBackend::TensorFlow == backend) {
#ifdef BUILD_TENSORFLOW
dp = std::make_shared<deepmd::DeepPotTF>(model, gpu_rank, file_content);
Expand Down
9 changes: 1 addition & 8 deletions source/api_cc/src/DeepSpin.cc
Original file line number Diff line number Diff line change
Expand Up @@ -36,14 +36,7 @@ void DeepSpin::init(const std::string& model,
<< std::endl;
return;
}
DPBackend backend;
if (model.length() >= 4 && model.substr(model.length() - 4) == ".pth") {
backend = deepmd::DPBackend::PyTorch;
} else if (model.length() >= 3 && model.substr(model.length() - 3) == ".pb") {
backend = deepmd::DPBackend::TensorFlow;
} else {
throw deepmd::deepmd_exception("Unsupported model file format");
}
const DPBackend backend = get_backend(model);
if (deepmd::DPBackend::TensorFlow == backend) {
#ifdef BUILD_TENSORFLOW
dp = std::make_shared<deepmd::DeepSpinTF>(model, gpu_rank, file_content);
Expand Down
3 changes: 1 addition & 2 deletions source/api_cc/src/DeepTensor.cc
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,7 @@ void DeepTensor::init(const std::string &model,
<< std::endl;
return;
}
// TODO: To implement detect_backend
DPBackend backend = deepmd::DPBackend::TensorFlow;
const DPBackend backend = get_backend(model);
if (deepmd::DPBackend::TensorFlow == backend) {
#ifdef BUILD_TENSORFLOW
dt = std::make_shared<deepmd::DeepTensorTF>(model, gpu_rank, name_scope_);
Expand Down
12 changes: 12 additions & 0 deletions source/api_cc/src/common.cc
Original file line number Diff line number Diff line change
Expand Up @@ -1399,3 +1399,15 @@ void deepmd::print_summary(const std::string& pre) {
<< "set tf inter_op_parallelism_threads: " << num_inter_nthreads
<< std::endl;
}

deepmd::DPBackend deepmd::get_backend(const std::string& model) {
if (model.length() >= 4 && model.substr(model.length() - 4) == ".pth") {
return deepmd::DPBackend::PyTorch;
} else if (model.length() >= 3 && model.substr(model.length() - 3) == ".pb") {
return deepmd::DPBackend::TensorFlow;
} else if (model.length() >= 11 &&
model.substr(model.length() - 11) == ".savedmodel") {
return deepmd::DPBackend::JAX;
}
throw deepmd::deepmd_exception("Unsupported model file format");
}

0 comments on commit d3095cf

Please sign in to comment.