blob: f89325829ed5c8c7b9d1a374fdcdfe17a32d0c4b [file] [log] [blame]
Aron Virginas-Tarf97f6da2019-10-01 18:35:44 +01001//
2// Copyright © 2019 Arm Ltd. All rights reserved.
3// SPDX-License-Identifier: MIT
4//
5
6#pragma once
7
8#include <Permute.hpp>
9
10#include <armnn/Tensor.hpp>
11#include <armnn/Types.hpp>
12
13template<typename T>
14void 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
25template<typename T>
26void 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}