TensorFlowEngine.hpp

Go to the documentation of this file.

#pragma once
 
#include <FAST/Algorithms/NeuralNetwork/InferenceEngine.hpp>
#include <TensorFlowExport.hpp>
 
// Forward declare
namespace tensorflow {
class Session;
class SavedModelBundle;
}
 
namespace fast {
 
class TensorFlowEngine : public InferenceEngine {
    public:
        void load() override;
        void run() override;
        std::string getName() const override;
        ~TensorFlowEngine() override;
        ImageOrdering getPreferredImageOrdering() const override;
        std::vector<ModelFormat> getSupportedModelFormats() const {
            return { ModelFormat::PROTOBUF, ModelFormat::SAVEDMODEL };
        };
 
        ModelFormat getPreferredModelFormat() const {
            return ModelFormat::PROTOBUF;
        };
 
        TensorFlowEngine();
        std::vector<InferenceDeviceInfo> getDeviceList();
    protected:
        std::unique_ptr<tensorflow::Session> mSession;
        std::unique_ptr<tensorflow::SavedModelBundle> mSavedModelBundle;
        std::vector<std::string> mLearningPhaseTensors;
 
};
 
DEFINE_INFERENCE_ENGINE(TensorFlowEngine, INFERENCEENGINETENSORFLOW_EXPORT)
 
 
// Forward declare
class TensorFlowTensorWrapper;
 
class TensorFlowTensor : public Tensor {
    FAST_OBJECT(TensorFlowTensor)
    public:
        void create(TensorFlowTensorWrapper* tensorflowTensor);
        ~TensorFlowTensor();
    private:
        TensorFlowTensorWrapper* m_tensorflowTensor;
        float* getHostDataPointer() override;
        bool hasAnyData() override;
};
 
}