1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
|
//
// Copyright © 2017 Arm Ltd. All rights reserved.
// SPDX-License-Identifier: MIT
//
#include "SpaceToDepth.hpp"
#include <DataLayoutIndexed.hpp>
using namespace armnnUtils;
namespace {
unsigned int GetOffset(const armnn::TensorShape& shape,
unsigned int c,
unsigned int h,
unsigned int w,
unsigned int b,
const DataLayoutIndexed& dataLayout)
{
if (dataLayout.GetDataLayout() == armnn::DataLayout::NHWC)
{
return ((b * shape[dataLayout.GetHeightIndex()] + h) * shape[dataLayout.GetWidthIndex()] + w) *
shape[dataLayout.GetChannelsIndex()] + c;
}
else
{
return ((b * shape[dataLayout.GetChannelsIndex()] + c) * shape[dataLayout.GetHeightIndex()] + h) *
shape[dataLayout.GetWidthIndex()] + w;
}
}
}
namespace armnn
{
void SpaceToDepth(const TensorInfo& inputInfo,
const TensorInfo& outputInfo,
const SpaceToDepthDescriptor& params,
Decoder<float>& inputData,
Encoder<float>& outputData)
{
DataLayoutIndexed dataLayout = params.m_DataLayout;
const TensorShape& inputShape = inputInfo.GetShape();
const TensorShape& outputShape = outputInfo.GetShape();
const unsigned int inputBatchSize = inputShape[0];
const unsigned int inputChannels = inputShape[dataLayout.GetChannelsIndex()];
const unsigned int outputHeight = outputShape[dataLayout.GetHeightIndex()];
const unsigned int outputWidth = outputShape[dataLayout.GetWidthIndex()];
const unsigned int outputChannels = outputShape[dataLayout.GetChannelsIndex()];
const unsigned int blockSize = params.m_BlockSize;
if (blockSize == 0)
{
throw InvalidArgumentException(
"Input shape must be divisible by block size in all spatial dimensions: Block size is"
" equal to zero");
}
for (unsigned int outChannelIndex = 0; outChannelIndex < outputChannels; outChannelIndex++)
{
unsigned int inChannelIndex = outChannelIndex % inputChannels;
unsigned int shiftW = (outChannelIndex / inputChannels) % blockSize;
unsigned int shiftH = (outChannelIndex / inputChannels) / blockSize;
for (unsigned int outH = 0; outH < outputHeight; outH++)
{
for (unsigned int outW = 0; outW < outputWidth; outW++)
{
for (unsigned int inBatchIndex = 0; inBatchIndex < inputBatchSize; inBatchIndex++)
{
unsigned int inOffset = GetOffset(inputShape,
inChannelIndex,
(outH * blockSize + shiftH),
(outW * blockSize + shiftW),
inBatchIndex,
dataLayout);
unsigned int outOffset = GetOffset(outputShape,
outChannelIndex,
outH,
outW,
inBatchIndex,
dataLayout);
outputData += outOffset;
inputData += inOffset;
outputData.Set(inputData.Get());
inputData -= inOffset;
outputData -= outOffset;
}
}
}
}
}
void SpaceToDepth(const TensorInfo& inputInfo,
const TensorInfo& outputInfo,
const SpaceToDepthDescriptor& params,
Decoder<float>& inputData,
Encoder<float>& outData);
} //namespace armnn
|
