1 //
2 // Copyright © 2019 Arm Ltd. All rights reserved.
3 // SPDX-License-Identifier: MIT
4 //
5
6 #pragma once
7
8 #include <armnn/Tensor.hpp>
9 #include <armnn/Types.hpp>
10
11 #include <armnnUtils/Permute.hpp>
12
13 template<typename T>
PermuteTensorNchwToNhwc(armnn::TensorInfo & tensorInfo,std::vector<T> & tensorData)14 void PermuteTensorNchwToNhwc(armnn::TensorInfo& tensorInfo, std::vector<T>& tensorData)
15 {
16 const armnn::PermutationVector nchwToNhwc = { 0, 3, 1, 2 };
17
18 tensorInfo = armnnUtils::Permuted(tensorInfo, nchwToNhwc);
19
20 std::vector<T> tmp(tensorData.size());
21 armnnUtils::Permute(tensorInfo.GetShape(), nchwToNhwc, tensorData.data(), tmp.data(), sizeof(T));
22 tensorData = tmp;
23 }
24
25 template<typename T>
PermuteTensorNhwcToNchw(armnn::TensorInfo & tensorInfo,std::vector<T> & tensorData)26 void PermuteTensorNhwcToNchw(armnn::TensorInfo& tensorInfo, std::vector<T>& tensorData)
27 {
28 const armnn::PermutationVector nhwcToNchw = { 0, 2, 3, 1 };
29
30 tensorInfo = armnnUtils::Permuted(tensorInfo, nhwcToNchw);
31
32 std::vector<T> tmp(tensorData.size());
33 armnnUtils::Permute(tensorInfo.GetShape(), nhwcToNchw, tensorData.data(), tmp.data(), sizeof(T));
34
35 tensorData = tmp;
36 }
37