cavis/libnd4j/include/ops/declarable/generic/images/extract_image_patches.cpp

97 lines
3.8 KiB
C++
Raw Normal View History

2019-06-06 14:21:15 +02:00
/*******************************************************************************
* Copyright (c) 2015-2018 Skymind, Inc.
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://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.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// Created by GS <sgazeos@gmail.com> at 3/30/2018
//
#include <ops/declarable/CustomOperations.h>
#include <ops/declarable/helpers/extract_patches.h>
namespace sd {
2019-06-06 14:21:15 +02:00
namespace ops {
CUSTOM_OP_IMPL(extract_image_patches, 1, 1, false, 0, 7) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
int ksizeRows = INT_ARG(0);
int ksizeCols = INT_ARG(1);
int kstrideRows = INT_ARG(2);
int kstrideCols = INT_ARG(3);
int krateRows = INT_ARG(4);
int krateCols = INT_ARG(5);
bool isSame = INT_ARG(6) != 0;
REQUIRE_TRUE(input->rankOf() == 4, 0, "extract_image_patches: The rank of input array should be 4, but %i is given", input->rankOf());
//
if (output->isSameShape(input))
output->assign(input);
else {
output->nullify();
helpers::extractPatches(block.launchContext(), input, output, ksizeRows, ksizeCols, kstrideRows, kstrideCols, krateRows, krateCols, isSame);
}
return Status::OK();
}
DECLARE_TYPES(extract_image_patches) {
getOpDescriptor()
->setAllowedInputTypes(sd::DataType::ANY)
2019-06-06 14:21:15 +02:00
->setSameMode(true);
}
DECLARE_SHAPE_FN(extract_image_patches) {
auto in = inputShape->at(0);
int outRank = shape::rank(in);
Nd4jLong *outputShape = nullptr;
int ksizeRowsEffective = INT_ARG(0) + (INT_ARG(0) - 1) * (INT_ARG(4) - 1);
int ksizeColsEffective = INT_ARG(1) + (INT_ARG(1) - 1) * (INT_ARG(5) - 1);
auto batchSizeDim = shape::sizeAt(in, 0);
auto inputRowsDim = shape::sizeAt(in, 1);
auto inputColsDim = shape::sizeAt(in, 2);
auto outputDepthDim = shape::sizeAt(in, 3) * INT_ARG(0) * INT_ARG(1); // last dim * ksizeRows * ksizeCols
auto inputRowSize = inputRowsDim; //shape::sizeAt(in, inputRowsDim);
auto inputColSize = inputColsDim; //shape::sizeAt(in, inputColsDim);
Nd4jLong outRowSize;
Nd4jLong outColSize;
if (INT_ARG(6) == 0) {
// Padding is "VALID":
outRowSize = (inputRowSize - ksizeRowsEffective + INT_ARG(2)) / INT_ARG(2);
outColSize = (inputColSize - ksizeColsEffective + INT_ARG(3)) / INT_ARG(3);
} else {
// Padding is "SAME":
outRowSize = (inputRowSize + INT_ARG(2) - 1) / INT_ARG(2);
outColSize = (inputColSize + INT_ARG(3) - 1) / INT_ARG(3);
}
ALLOCATE(outputShape, block.getWorkspace(), shape::shapeInfoLength(outRank), Nd4jLong);
outputShape[0] = outRank;
outputShape[1] = batchSizeDim;
outputShape[2] = outRowSize;
outputShape[3] = outColSize;
outputShape[4] = outputDepthDim;
ShapeUtils::updateStridesAndType(outputShape, in, shape::order(in));
return SHAPELIST(CONSTANT(outputShape));
}
}
}