Skip to content

Commit

Permalink
cc: refactor DeepTensor for multiple-backend framework (#3151)
Browse files Browse the repository at this point in the history 
See #3119

---------

Signed-off-by: Jinzhe Zeng <[email protected]>
  • Loading branch information
@njzjz
njzjz authored Jan 17, 2024
1 parent 9be1ad2 commit 850575a
Show file tree
Hide file tree
Showing 4 changed files with 1,512 additions and 712 deletions.
263 changes: 179 additions & 84 deletions source/api_cc/include/DeepTensor.h
View file
Original file line number Diff line number Diff line change
@@ -1,10 +1,184 @@
// SPDX-License-Identifier: LGPL-3.0-or-later
#pragma once

#include <memory>

#include "common.h"
#include "neighbor_list.h"

namespace deepmd {
/**
* @brief Deep Tensor.
**/
class DeepTensorBase {
public:
/**
* @brief Deep Tensor constructor without initialization.
**/
DeepTensorBase(){};
virtual ~DeepTensorBase(){};
/**
* @brief Deep Tensor constructor with initialization..
* @param[in] model The name of the frozen model file.
* @param[in] gpu_rank The GPU rank. Default is 0.
* @param[in] name_scope Name scopes of operations.
**/
DeepTensorBase(const std::string& model,
const int& gpu_rank = 0,
const std::string& name_scope = "");
/**
* @brief Initialize the Deep Tensor.
* @param[in] model The name of the frozen model file.
* @param[in] gpu_rank The GPU rank. Default is 0.
* @param[in] name_scope Name scopes of operations.
**/
virtual void init(const std::string& model,
const int& gpu_rank = 0,
const std::string& name_scope = "") = 0;

/**
* @brief Evaluate the value by using this model.
* @param[out] value The value to evalute, usually would be the atomic tensor.
* @param[in] coord The coordinates of atoms. The array should be of size
*natoms x 3.
* @param[in] atype The atom types. The list should contain natoms ints.
* @param[in] box The cell of the region. The array should be of size 9.
* @{
**/
virtual void computew(std::vector<double>& value,
const std::vector<double>& coord,
const std::vector<int>& atype,
const std::vector<double>& box) = 0;
virtual void computew(std::vector<float>& value,
const std::vector<float>& coord,
const std::vector<int>& atype,
const std::vector<float>& box) = 0;
/** @} */
/**
* @brief Evaluate the value by using this model.
* @param[out] value The value to evalute, usually would be the atomic tensor.
* @param[in] coord The coordinates of atoms. The array should be of size
*natoms x 3.
* @param[in] atype The atom types. The list should contain natoms ints.
* @param[in] box The cell of the region. The array should be of size 9.
* @param[in] nghost The number of ghost atoms.
* @param[in] inlist The input neighbour list.
* @{
**/
virtual void computew(std::vector<double>& value,
const std::vector<double>& coord,
const std::vector<int>& atype,
const std::vector<double>& box,
const int nghost,
const InputNlist& inlist) = 0;
virtual void computew(std::vector<float>& value,
const std::vector<float>& coord,
const std::vector<int>& atype,
const std::vector<float>& box,
const int nghost,
const InputNlist& inlist) = 0;
/** @} */
/**
* @brief Evaluate the global tensor and component-wise force and virial.
* @param[out] global_tensor The global tensor to evalute.
* @param[out] force The component-wise force of the global tensor, size odim
*x natoms x 3.
* @param[out] virial The component-wise virial of the global tensor, size
*odim x 9.
* @param[out] atom_tensor The atomic tensor value of the model, size natoms x
*odim.
* @param[out] atom_virial The component-wise atomic virial of the global
*tensor, size odim x natoms x 9.
* @param[in] coord The coordinates of atoms. The array should be of size
*natoms x 3.
* @param[in] atype The atom types. The list should contain natoms ints.
* @param[in] box The cell of the region. The array should be of size 9.
* @{
**/
virtual void computew(std::vector<double>& global_tensor,
std::vector<double>& force,
std::vector<double>& virial,
std::vector<double>& atom_tensor,
std::vector<double>& atom_virial,
const std::vector<double>& coord,
const std::vector<int>& atype,
const std::vector<double>& box) = 0;
virtual void computew(std::vector<float>& global_tensor,
std::vector<float>& force,
std::vector<float>& virial,
std::vector<float>& atom_tensor,
std::vector<float>& atom_virial,
const std::vector<float>& coord,
const std::vector<int>& atype,
const std::vector<float>& box) = 0;
/** @} */
/**
* @brief Evaluate the global tensor and component-wise force and virial.
* @param[out] global_tensor The global tensor to evalute.
* @param[out] force The component-wise force of the global tensor, size odim
*x natoms x 3.
* @param[out] virial The component-wise virial of the global tensor, size
*odim x 9.
* @param[out] atom_tensor The atomic tensor value of the model, size natoms x
*odim.
* @param[out] atom_virial The component-wise atomic virial of the global
*tensor, size odim x natoms x 9.
* @param[in] coord The coordinates of atoms. The array should be of size
*natoms x 3.
* @param[in] atype The atom types. The list should contain natoms ints.
* @param[in] box The cell of the region. The array should be of size 9.
* @param[in] nghost The number of ghost atoms.
* @param[in] inlist The input neighbour list.
* @{
**/
virtual void computew(std::vector<double>& global_tensor,
std::vector<double>& force,
std::vector<double>& virial,
std::vector<double>& atom_tensor,
std::vector<double>& atom_virial,
const std::vector<double>& coord,
const std::vector<int>& atype,
const std::vector<double>& box,
const int nghost,
const InputNlist& inlist) = 0;
virtual void computew(std::vector<float>& global_tensor,
std::vector<float>& force,
std::vector<float>& virial,
std::vector<float>& atom_tensor,
std::vector<float>& atom_virial,
const std::vector<float>& coord,
const std::vector<int>& atype,
const std::vector<float>& box,
const int nghost,
const InputNlist& inlist) = 0;
/** @} */
/**
* @brief Get the cutoff radius.
* @return The cutoff radius.
**/
virtual double cutoff() const = 0;
/**
* @brief Get the number of types.
* @return The number of types.
**/
virtual int numb_types() const = 0;
/**
* @brief Get the output dimension.
* @return The output dimension.
**/
virtual int output_dim() const = 0;
/**
* @brief Get the list of sel types.
* @return The list of sel types.
*/
virtual const std::vector<int>& sel_types() const = 0;
/**
* @brief Get the type map (element name of the atom types) of this model.
* @param[out] type_map The type map of this model.
**/
virtual void get_type_map(std::string& type_map) = 0;
};

/**
* @brief Deep Tensor.
**/
Expand Down Expand Up @@ -169,109 +343,30 @@ class DeepTensor {
* @brief Get the cutoff radius.
* @return The cutoff radius.
**/
double cutoff() const {
assert(inited);
return rcut;
};
double cutoff() const;
/**
* @brief Get the number of types.
* @return The number of types.
**/
int numb_types() const {
assert(inited);
return ntypes;
};
int numb_types() const;
/**
* @brief Get the output dimension.
* @return The output dimension.
**/
int output_dim() const {
assert(inited);
return odim;
};
int output_dim() const;
/**
* @brief Get the list of sel types.
* @return The list of sel types.
*/
const std::vector<int>& sel_types() const {
assert(inited);
return sel_type;
};
const std::vector<int>& sel_types() const;
/**
* @brief Get the type map (element name of the atom types) of this model.
* @param[out] type_map The type map of this model.
**/
void get_type_map(std::string& type_map);

private:
tensorflow::Session* session;
std::string name_scope;
int num_intra_nthreads, num_inter_nthreads;
tensorflow::GraphDef* graph_def;
bool inited;
double rcut;
int dtype;
double cell_size;
int ntypes;
std::string model_type;
std::string model_version;
int odim;
std::vector<int> sel_type;
template <class VT>
VT get_scalar(const std::string& name) const;
template <class VT>
void get_vector(std::vector<VT>& vec, const std::string& name) const;
template <typename MODELTYPE, typename VALUETYPE>
void run_model(std::vector<VALUETYPE>& d_tensor_,
tensorflow::Session* session,
const std::vector<std::pair<std::string, tensorflow::Tensor>>&
input_tensors,
const AtomMap& atommap,
const std::vector<int>& sel_fwd,
const int nghost = 0);
template <typename MODELTYPE, typename VALUETYPE>
void run_model(std::vector<VALUETYPE>& dglobal_tensor_,
std::vector<VALUETYPE>& dforce_,
std::vector<VALUETYPE>& dvirial_,
std::vector<VALUETYPE>& datom_tensor_,
std::vector<VALUETYPE>& datom_virial_,
tensorflow::Session* session,
const std::vector<std::pair<std::string, tensorflow::Tensor>>&
input_tensors,
const AtomMap& atommap,
const std::vector<int>& sel_fwd,
const int nghost = 0);
template <typename VALUETYPE>
void compute_inner(std::vector<VALUETYPE>& value,
const std::vector<VALUETYPE>& coord,
const std::vector<int>& atype,
const std::vector<VALUETYPE>& box);
template <typename VALUETYPE>
void compute_inner(std::vector<VALUETYPE>& value,
const std::vector<VALUETYPE>& coord,
const std::vector<int>& atype,
const std::vector<VALUETYPE>& box,
const int nghost,
const InputNlist& inlist);
template <typename VALUETYPE>
void compute_inner(std::vector<VALUETYPE>& global_tensor,
std::vector<VALUETYPE>& force,
std::vector<VALUETYPE>& virial,
std::vector<VALUETYPE>& atom_tensor,
std::vector<VALUETYPE>& atom_virial,
const std::vector<VALUETYPE>& coord,
const std::vector<int>& atype,
const std::vector<VALUETYPE>& box);
template <typename VALUETYPE>
void compute_inner(std::vector<VALUETYPE>& global_tensor,
std::vector<VALUETYPE>& force,
std::vector<VALUETYPE>& virial,
std::vector<VALUETYPE>& atom_tensor,
std::vector<VALUETYPE>& atom_virial,
const std::vector<VALUETYPE>& coord,
const std::vector<int>& atype,
const std::vector<VALUETYPE>& box,
const int nghost,
const InputNlist& inlist);
std::shared_ptr<deepmd::DeepTensorBase> dt;
};
} // namespace deepmd
Loading

0 comments on commit 850575a

Please sign in to comment.