summaryrefslogtreecommitdiff
path: root/runtimes/libs/ARMComputeEx/arm_compute/core/utils/misc/ShapeCalculatorEx.h
blob: bacb1140c9e08a48a4f21d01457a6289309abcae (plain)
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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
/*
 * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
 * Copyright (c) 2016-2018 ARM Limited.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef __ARM_COMPUTE_MISC_SHAPE_CALCULATOR_EX_H__
#define __ARM_COMPUTE_MISC_SHAPE_CALCULATOR_EX_H__

#include "arm_compute/core/Helpers.h"
#include "arm_compute/core/ITensorInfo.h"
#include "arm_compute/core/Utils.h"

#include "arm_compute/core/utils/helpers/tensor_transform.h"

#include <cmath>

namespace arm_compute
{
namespace misc
{
namespace shape_calculator
{

/** Calculate the upsampled output shape used for transpose convolution
 *
 * @param[in] input              Input tensor info
 * @param[in] weights            Weights tensor shape
 * @param[in] info               Padding and stride info
 * @param[in] out_dims           Output shape dimensions
 * @param[in] invalid_right      The number of zeros added to right edge of the output.
 * @param[in] invalid_bottom     The number of zeros added to bottom edge of the output.
 * @param[out] pad_left          Padding on left
 * @param[out] pad_right         Padding on right
 * @param[out] pad_top           Padding on top
 * @param[out] pad_bottom        Padding on bottom
 *
 * @return the calculated shape
 */
inline TensorShape compute_transposeconv_upsampled_shape(
    const ITensorInfo &input, const ITensorInfo &weights, const PadStrideInfo &info,
    std::pair<unsigned int, unsigned int> &out_dims, unsigned int invalid_right,
    unsigned int invalid_bottom, unsigned int &pad_left, unsigned int &pad_right,
    unsigned int &pad_top, unsigned int &pad_bottom)
{
  unsigned int sx = info.stride().first;
  unsigned int sy = info.stride().second;
  const DataLayout data_layout = input.data_layout();
  const size_t idx_w = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
  const size_t idx_h = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);

  // Find the upsampled dimensions
  // transpose conv out:
  //    tconv_out + pad = 1 + (in - 1) * stride + invalid
  //    tconv_out = 1 + (in - 1) * stride + invalid - pad
  // upsample out:
  //    upsample_out = 1 + (in - 1) * stride
  unsigned int out_x = (input.dimension(idx_w) - 1) * sx + 1;
  unsigned int out_y = (input.dimension(idx_h) - 1) * sy + 1;

  // Find the padding needed for the convolution with stride 1 in order to match output shape
  // upsample+pad out:
  //    upsample_out + pad = tconv_out + kernel - 1
  //    pad = tconv_out + kernel - 1 - upsample_out
  unsigned int padx = out_dims.first - (out_x - weights.dimension(idx_w) + 1);
  unsigned int pady = out_dims.second - (out_y - weights.dimension(idx_h) + 1);
  out_x += padx;
  out_y += pady;

  unsigned int padx_all_except_invallid = padx + info.pad_left() + info.pad_right() - invalid_right;
  unsigned int pady_all_except_invallid =
      pady + info.pad_top() + info.pad_bottom() - invalid_bottom;
  pad_left = (padx_all_except_invallid + 1) / 2 - info.pad_left();
  pad_right = pady_all_except_invallid / 2 - info.pad_right() + invalid_right;
  pad_top = (padx_all_except_invallid + 1) / 2 - info.pad_top();
  pad_bottom = pady_all_except_invallid / 2 - info.pad_bottom() + invalid_bottom;

  TensorShape scale_out_shape(input.tensor_shape());
  scale_out_shape.set(idx_w, out_x);
  scale_out_shape.set(idx_h, out_y);

  return scale_out_shape;
}

/** Calculate the output shape of the transpose convolution layer
 *
 * @param[in] out_dims Output x and y shape dimensions
 * @param[in] input    Input tensor info
 * @param[in] weights  Weights tensor shape
 *
 * @return the calculated shape
 */
inline TensorShape
compute_transposeconv_output_shape(const std::pair<unsigned int, unsigned int> &out_dims,
                                   const ITensorInfo &input, const ITensorInfo &weights)
{
  const TensorShape input_shape{input.tensor_shape()};
  const TensorShape weights_shape{weights.tensor_shape()};

  const DataLayout data_layout = input.data_layout();
  const int width_idx = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
  const int height_idx = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
  const int channel_idx =
      get_data_layout_dimension_index(data_layout, DataLayoutDimension::CHANNEL);
  const int batch_idx = get_data_layout_dimension_index(data_layout, DataLayoutDimension::BATCHES);

  TensorShape out_shape{input_shape};
  out_shape.set(width_idx, out_dims.first);
  out_shape.set(height_idx, out_dims.second);
  out_shape.set(channel_idx, weights_shape[batch_idx]);
  return out_shape;
}

} // namespace shape_calculator
} // namespace misc
} // namespace arm_compute

#endif // __ARM_COMPUTE_MISC_SHAPE_CALCULATOR_EX_H__