diff options
Diffstat (limited to 'src/armnn/optimizations/PermuteAndBatchToSpaceAsDepthToSpace.hpp')
-rw-r--r-- | src/armnn/optimizations/PermuteAndBatchToSpaceAsDepthToSpace.hpp | 27 |
1 files changed, 27 insertions, 0 deletions
diff --git a/src/armnn/optimizations/PermuteAndBatchToSpaceAsDepthToSpace.hpp b/src/armnn/optimizations/PermuteAndBatchToSpaceAsDepthToSpace.hpp new file mode 100644 index 0000000000..4a73efca40 --- /dev/null +++ b/src/armnn/optimizations/PermuteAndBatchToSpaceAsDepthToSpace.hpp @@ -0,0 +1,27 @@ +// +// Copyright © 2019 Arm Ltd. All rights reserved. +// SPDX-License-Identifier: MIT +// +#pragma once + +#include "Optimization.hpp" + +namespace armnn +{ +namespace optimizations +{ + +/// Replaces Permute leading into BatchToSpace with a DepthToSpace +/// in the case where the Permute swaps the batch and channels dimensions +/// such that the replacement is valid. +class PermuteAndBatchToSpaceAsDepthToSpaceImpl +{ +public: + void Run(Graph& graph, InputSlot& connection) const; +}; + +using PermuteAndBatchToSpaceAsDepthToSpace = + OptimizeForConnection<PermuteLayer, BatchToSpaceNdLayer, PermuteAndBatchToSpaceAsDepthToSpaceImpl>; + +} // namespace optimizations +} // namespace armnn |