diff options
Diffstat (limited to 'OutputShapeUtils.cpp')
-rw-r--r-- | OutputShapeUtils.cpp | 19 |
1 files changed, 19 insertions, 0 deletions
diff --git a/OutputShapeUtils.cpp b/OutputShapeUtils.cpp index a0c624c8..a1a49d70 100644 --- a/OutputShapeUtils.cpp +++ b/OutputShapeUtils.cpp @@ -6,6 +6,7 @@ #include "OutputShapeUtils.hpp" #include <algorithm> +#include <vector> namespace armnn_driver { @@ -17,6 +18,24 @@ bool IsDynamicOutput(const TensorInfo& outputInfo) return outputInfo.GetNumElements() == 0u; } +TensorShape InferPadOutputShape(const TensorShape& inputShape, + const std::vector<std::pair<unsigned int, unsigned int>>& padList) +{ + const unsigned int numDims = inputShape.GetNumDimensions(); + + std::vector<unsigned int> outputDims; + TensorShape outputShape = TensorShape(numDims); + for (unsigned int dim = 0; dim < numDims; ++dim) + { + unsigned int dimSize = inputShape[dim]; + const std::pair<unsigned int, unsigned int>& dimPadding = padList[dim]; + dimSize += dimPadding.first; + dimSize += dimPadding.second; + outputShape[dim] = dimSize; + } + return outputShape; +} + TensorShape InferPreluOutputShape(const TensorShape& inputShape, const TensorShape& alphaShape) { // NOTE: The inferred PReLU output size will be the maximum size along each dimension |