diff options
author | thecha01 <theo.charalambous@arm.com> | 2020-09-07 14:50:21 +0100 |
---|---|---|
committer | Theo Charalambous <theo.charalambous@arm.com> | 2020-09-09 09:51:52 +0000 |
commit | d64444ba197c2f95dcf4d205f50a196d5a29cdeb (patch) | |
tree | 509c18d43dd6b6c614782ca2fe2d2111084ed8a7 /arm_compute/graph/frontend | |
parent | 3603aff36444e01083b93e603e8b9b0a254d4911 (diff) | |
download | ComputeLibrary-d64444ba197c2f95dcf4d205f50a196d5a29cdeb.tar.gz |
Add Reduction operation layer node to graph API
Signed-off-by: thecha01 <theo.charalambous@arm.com>
Change-Id: Ida819fb8c33790cc9da6d69eeb51e0599269197a
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/3931
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Michele Di Giorgio <michele.digiorgio@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Diffstat (limited to 'arm_compute/graph/frontend')
-rw-r--r-- | arm_compute/graph/frontend/Layers.h | 28 |
1 files changed, 28 insertions, 0 deletions
diff --git a/arm_compute/graph/frontend/Layers.h b/arm_compute/graph/frontend/Layers.h index ba303a3ef6..74c40126c8 100644 --- a/arm_compute/graph/frontend/Layers.h +++ b/arm_compute/graph/frontend/Layers.h @@ -1133,6 +1133,34 @@ private: QuantizationInfo _out_quant_info; }; +/** Reduction Layer */ +class ReductionLayer final : public ILayer +{ +public: + /** Construct a reduction layer. + * + * @param[in] op Reduction operation + * @param[in] axis Reduction axis + * @param[in] keep_dims (Optional) Whether to keep the reduced dimension after the operation. Defaults to true. + */ + ReductionLayer(ReductionOperation op, unsigned int axis, bool keep_dims) + : _op(op), _axis(axis), _keep_dims(keep_dims) + { + } + + NodeID create_layer(IStream &s) override + { + NodeParams common_params = { name(), s.hints().target_hint }; + NodeIdxPair input = { s.tail_node(), 0 }; + return GraphBuilder::add_reduction_operation_node(s.graph(), common_params, input, _op, _axis, _keep_dims); + } + +private: + ReductionOperation _op; + unsigned int _axis; + bool _keep_dims; +}; + /** Reorg Layer */ class ReorgLayer final : public ILayer { |