diff options
Diffstat (limited to 'src/armnnQuantizer/QuantizationDataSet.hpp')
-rw-r--r-- | src/armnnQuantizer/QuantizationDataSet.hpp | 16 |
1 files changed, 16 insertions, 0 deletions
diff --git a/src/armnnQuantizer/QuantizationDataSet.hpp b/src/armnnQuantizer/QuantizationDataSet.hpp index 3a97630ccf..47b893a7f7 100644 --- a/src/armnnQuantizer/QuantizationDataSet.hpp +++ b/src/armnnQuantizer/QuantizationDataSet.hpp @@ -43,6 +43,22 @@ private: }; /// Visitor class implementation to gather the TensorInfo for LayerBindingID for creation of ConstTensor for Refine. +class InputLayerStrategy : public armnn::IStrategy +{ +public: + virtual void ExecuteStrategy(const armnn::IConnectableLayer* layer, + const armnn::BaseDescriptor& descriptor, + const std::vector<armnn::ConstTensor>& constants, + const char* name, + const armnn::LayerBindingId id = 0) override; + + armnn::TensorInfo GetTensorInfo(armnn::LayerBindingId); +private: + std::map<armnn::LayerBindingId, armnn::TensorInfo> m_TensorInfos; +}; + + +/// Visitor class implementation to gather the TensorInfo for LayerBindingID for creation of ConstTensor for Refine. class InputLayerVisitor : public armnn::LayerVisitorBase<armnn::VisitorNoThrowPolicy> { public: |