inference.h

Go to the documentation of this file.

 
#ifndef _NNDEPLOY_INFERENCE_INFERENCE_H_
#define _NNDEPLOY_INFERENCE_INFERENCE_H_
 
#include "nndeploy/base/any.h"
#include "nndeploy/base/common.h"
#include "nndeploy/base/glic_stl_include.h"
#include "nndeploy/base/log.h"
#include "nndeploy/base/macro.h"
#include "nndeploy/base/object.h"
#include "nndeploy/base/status.h"
#include "nndeploy/base/string.h"
#include "nndeploy/device/buffer.h"
#include "nndeploy/device/device.h"
#include "nndeploy/device/memory_pool.h"
#include "nndeploy/device/tensor.h"
#include "nndeploy/inference/inference_param.h"
 
namespace nndeploy {
namespace inference {
 
class NNDEPLOY_CC_API Inference {
 public:
  Inference(base::InferenceType type);
  virtual ~Inference();
 
  base::InferenceType getInferenceType();
 
  base::Status setParam(base::Param *param);
  base::Status setParamSharedPtr(std::shared_ptr<base::Param> param);
  
  base::Param *getParam();
  std::shared_ptr<base::Param> getParamSharedPtr();
  
  base::DeviceType getDeviceType() { return inference_param_->device_type_; }
 
  void setStream(device::Stream *stream);
  device::Stream *getStream();
 
  virtual base::Status init() = 0;
  virtual base::Status deinit() = 0;
 
  base::ShapeMap getMinShape();
  base::ShapeMap getOptShape();
  base::ShapeMap getMaxShape();
  virtual base::Status reshape(base::ShapeMap &shape_map) = 0;
 
  virtual int64_t getMemorySize();
  virtual base::Status setMemory(device::Buffer *buffer);
  virtual float getGFLOPs();
 
  virtual bool isBatch();
 
  virtual bool isShareContext();
  virtual bool isShareStream();
  virtual bool isInputDynamic();
  virtual bool isOutputDynamic();
  virtual bool canOpInput();
  virtual bool canOpOutput();
 
  virtual int getNumOfInputTensor();
  virtual int getNumOfOutputTensor();
 
  virtual std::string getInputName(int i);
  virtual std::string getOutputName(int i);
  virtual std::vector<std::string> getAllInputTensorName();
  virtual std::vector<std::string> getAllOutputTensorName();
 
  virtual base::IntVector getInputShape(const std::string &name);
  virtual base::ShapeMap getAllInputShape();
  virtual device::TensorDesc getInputTensorDesc(const std::string &name);
  virtual device::TensorDesc getOutputTensorDesc(const std::string &name);
 
  virtual device::TensorDesc getInputTensorAlignDesc(const std::string &name);
  virtual device::TensorDesc getOutputTensorAlignDesc(const std::string &name);
 
  virtual std::map<std::string, device::Tensor *> getAllInputTensorMap();
  virtual std::map<std::string, device::Tensor *> getAllOutputTensorMap();
 
  virtual std::vector<device::Tensor *> getAllInputTensorVector();
  virtual std::vector<device::Tensor *> getAllOutputTensorVector();
 
  virtual device::Tensor *getInputTensor(const std::string &name);
  virtual device::Tensor *getOutputTensor(const std::string &name);
 
  virtual base::Status setInputTensor(const std::string &name,
                                      device::Tensor *input_tensor);
 
  virtual base::Status run() = 0;
 
  virtual device::Tensor *getOutputTensorAfterRun(
      const std::string &name, base::DeviceType device_type, bool is_copy,
      base::DataFormat data_format = base::kDataFormatAuto) = 0;
 
 protected:
  base::InferenceType type_;
  std::shared_ptr<InferenceParam> inference_param_;
 
  bool is_share_context_ = true;
 
  bool is_external_stream_ = false;
  device::Stream *stream_ = nullptr;
 
  std::map<std::string, device::Tensor *> input_tensors_;
  std::map<std::string, device::Tensor *> output_tensors_;
 
  std::map<std::string, device::Tensor *> external_input_tensors_;
};
 
class InferenceCreator {
 public:
  virtual ~InferenceCreator() {};
  // virtual Inference *createInference(base::InferenceType type) = 0;
  virtual std::shared_ptr<Inference> createInference(base::InferenceType type) = 0;
};
 
template <typename T>
class TypeInferenceCreator : public InferenceCreator {
  // virtual Inference *createInference(base::InferenceType type) {
  //   return new T(type);
  // }
  virtual std::shared_ptr<Inference> createInference(base::InferenceType type) {
    return std::make_shared<T>(type);
  }
};
 
extern NNDEPLOY_CC_API std::map<base::InferenceType, std::shared_ptr<InferenceCreator>> &
getGlobalInferenceCreatorMap();
 
template <typename T>
class TypeInferenceRegister {
 public:
  explicit TypeInferenceRegister(base::InferenceType type) {
    getGlobalInferenceCreatorMap()[type] = std::shared_ptr<T>(new T());
  }
};
 
// extern NNDEPLOY_CC_API Inference *createInference(base::InferenceType type);
 
extern NNDEPLOY_CC_API std::shared_ptr<Inference> createInference(base::InferenceType type);
 
}  // namespace inference
}  // namespace nndeploy
 
#endif