diff options
Diffstat (limited to 'src/armnn/OverrideInputRangeVisitor.hpp')
-rw-r--r-- | src/armnn/OverrideInputRangeVisitor.hpp | 45 |
1 files changed, 45 insertions, 0 deletions
diff --git a/src/armnn/OverrideInputRangeVisitor.hpp b/src/armnn/OverrideInputRangeVisitor.hpp new file mode 100644 index 0000000000..a2da6c702e --- /dev/null +++ b/src/armnn/OverrideInputRangeVisitor.hpp @@ -0,0 +1,45 @@ +// +// Copyright © 2017 Arm Ltd. All rights reserved. +// SPDX-License-Identifier: MIT +// + +#pragma once + +#include "NetworkQuantizer.hpp" +#include "LayerVisitorBase.hpp" + +#include <unordered_map> + +namespace armnn +{ + +/// Visitor object for overriding the input range of the quantized input layers in a network +class OverrideInputRangeVisitor : public LayerVisitorBase<VisitorNoThrowPolicy> +{ +private: + using MinMaxRange = std::pair<float, float>; + using MinMaxRanges = std::vector<MinMaxRange>; + +public: + OverrideInputRangeVisitor(std::unordered_map<LayerGuid, MinMaxRanges>& guidToRangesMap, + LayerBindingId layerId, + const MinMaxRange& minMaxRange); + ~OverrideInputRangeVisitor() = default; + + void VisitInputLayer(const IConnectableLayer *layer, LayerBindingId id, const char *name = nullptr) override; + +private: + /// Sets the range for the given input layer + void SetRange(const IConnectableLayer* layer); + + /// Mapping from a layer Guid to an array of ranges for outputs + std::unordered_map<LayerGuid, MinMaxRanges>& m_GuidToRangesMap; + + /// The id of the input layer of which to override the input range + LayerBindingId m_LayerId; + + /// The new input range to be applied to the input layer + MinMaxRange m_MinMaxRange; +}; + +} // namespace armnn |