MNN/source/backend/cpu/CPUROIPooling.cpp

204 lines
9.6 KiB
C++

//
// CPUROIPooling.cpp
// MNN
//
// Created by MNN on 2018/07/19.
// Copyright © 2018, Alibaba Group Holding Limited
//
#include "backend/cpu/CPUROIPooling.hpp"
#include <float.h>
#include <math.h>
#include "CPUTensorConvert.hpp"
#include "backend/cpu/CPUBackend.hpp"
#include "core/TensorUtils.hpp"
namespace MNN {
CPUROIPooling::CPUROIPooling(Backend *backend, int pooledWidth, int pooledHeight, float spatialScale, bool outputGrad)
: CPUROIAlign(backend, pooledWidth, pooledHeight, 0, spatialScale, false, PoolType_MAX, outputGrad) {
// nothing to do
}
static inline int max(int a, int b) { return a > b ? a : b; }
static inline int min(int a, int b) { return a < b ? a : b; }
ErrorCode CPUROIPooling::onExecute(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs) {
auto &input = inputs[0];
auto &output = outputs[0];
auto core = static_cast<CPUBackend*>(backend())->functions();
CPUTensorConverter::convert(inputs[1], &mROI, core);
// dataType of ROI must be float32.
Tensor *roiTensor = &mROI;
auto roiPtrSrc = roiTensor->host<float>();
if (core->bytes != 4) {
core->MNNLowpToFp32(mROI.host<int16_t>(), mROITemp->host<float>(), mROI.elementSize());
roiPtrSrc = mROITemp->host<float>();
}
if (mOutputGrad == false) {
// get params
auto iw = input->width(), ih = input->height(), is = iw * ih * core->pack;
auto ow = output->width(), oh = output->height(), os = ow * oh * core->pack;
auto slice = UP_DIV(input->channel(), core->pack);
auto numROI = inputs[1]->batch();
for (int n = 0; n < numROI; ++n) {
auto batchOutput = output->host<uint8_t>() + os * n * core->bytes;
auto roiPtr = roiPtrSrc + roiTensor->buffer().dim[0].stride * n;
int roi = roiPtr[0];
int x1 = round(roiPtr[1] * mSpatialScale);
int y1 = round(roiPtr[2] * mSpatialScale);
int x2 = round(roiPtr[3] * mSpatialScale);
int y2 = round(roiPtr[4] * mSpatialScale);
MNN_ASSERT(roi < input->batch());
int roiW = max(x2 - x1 + 1, 1);
int roiH = max(y2 - y1 + 1, 1);
float binSizeW = (float)roiW / (float)mPooledWidth;
float binSizeH = (float)roiH / (float)mPooledHeight;
auto batchInput = input->host<uint8_t>() + is * roi * core->bytes;
for (int s = 0; s < slice; s++) {
auto sliceInput = batchInput + is * input->batch() * s * core->bytes;
auto rowOutput = batchOutput + os * output->batch() * s * core->bytes;
float binPh = 0;
for (int ph = 0; ph < mPooledHeight; ph++, rowOutput += mPooledWidth * core->pack * core->bytes) {
// Compute pooling region for this output unit:
// start (included) = floor(ph * roiHeight / pooledHeight)
// end (excluded) = ceil((ph + 1) * roiHeight / pooledHeight)
int hStart = min(max(y1 + (int)floorf(binPh), 0), ih);
binPh += binSizeH;
int hEnd = min(max(y1 + (int)ceilf(binPh), 0), ih);
int hLen = hEnd - hStart;
if (hLen <= 0) {
memset(rowOutput, 0, mPooledWidth * core->pack * core->bytes * sizeof(uint8_t));
continue;
}
float binPw = 0;
for (int pw = 0; pw < mPooledWidth; pw++) {
int wStart = min(max(x1 + (int)floorf(binPw), 0), iw);
binPw += binSizeW;
int wEnd = min(max(x1 + (int)ceilf(binPw), 0), iw);
int wLen = wEnd - wStart;
if (wLen <= 0) {
memset(rowOutput + pw * core->pack * core->bytes, 0, core->pack * core->bytes * sizeof(uint8_t));
continue;
}
core->MNNRoiPoolingMax((float *)(rowOutput + pw * core->pack * core->bytes), (float *)(sliceInput + (hStart * iw + wStart) * core->pack * core->bytes), hLen, wLen, iw);
}
}
}
}
} else {
#ifndef MNN_REDUCE_SIZE
// get params
auto iw = input->width(), ih = input->height(), is = iw * ih * core->pack;
// backward mode, output shape is the same with input[0] shape
auto& bwDiff = inputs[2];
auto ow = bwDiff->width(), oh = bwDiff->height(), os = ow * oh * core->pack;
auto slice = UP_DIV(input->channel(), core->pack);
auto numROI = inputs[1]->batch();
::memset(output->host<uint8_t>(), 0, static_cast<CPUBackend*>(backend())->getTensorSize(output, true));
for (int n = 0; n < numROI; ++n) {
auto batchBwDiff = inputs[2]->host<uint8_t>() + os * n * core->bytes;
auto roiPtr = roiPtrSrc + roiTensor->buffer().dim[0].stride * n;
int roi = roiPtr[0];
int x1 = round(roiPtr[1] * mSpatialScale);
int y1 = round(roiPtr[2] * mSpatialScale);
int x2 = round(roiPtr[3] * mSpatialScale);
int y2 = round(roiPtr[4] * mSpatialScale);
MNN_ASSERT(roi < input->batch());
int roiW = max(x2 - x1 + 1, 1);
int roiH = max(y2 - y1 + 1, 1);
float binSizeW = (float)roiW / (float)mPooledWidth;
float binSizeH = (float)roiH / (float)mPooledHeight;
auto batchInput = input->host<uint8_t>() + is * roi * core->bytes;
auto batchOutput = output->host<uint8_t>() + is * roi * core->bytes;
for (int s = 0; s < slice; s++) {
auto sliceInput = batchInput + is * input->batch() * s * core->bytes;
auto sliceOutput = batchOutput + is * input->batch() * s * core->bytes;
auto rowBwDiff = batchBwDiff + os * bwDiff->batch() * s * core->bytes;
float binPh = 0;
for (int ph = 0; ph < mPooledHeight; ph++, rowBwDiff += mPooledWidth * core->pack * core->bytes) {
// Compute pooling region for this output unit:
// start (included) = floor(ph * roiHeight / pooledHeight)
// end (excluded) = ceil((ph + 1) * roiHeight / pooledHeight)
int hStart = min(max(y1 + (int)floorf(binPh), 0), ih);
binPh += binSizeH;
int hEnd = min(max(y1 + (int)ceilf(binPh), 0), ih);
int hLen = hEnd - hStart;
if (hLen <= 0) {
continue;
}
float binPw = 0;
for (int pw = 0; pw < mPooledWidth; pw++) {
int wStart = min(max(x1 + (int)floorf(binPw), 0), iw);
binPw += binSizeW;
int wEnd = min(max(x1 + (int)ceilf(binPw), 0), iw);
int wLen = wEnd - wStart;
if (wLen <= 0) {
continue;
}
{
std::vector<int> indices(core->pack);
std::vector<float> maxes(core->pack, -FLT_MAX);
float* src = (float *)(sliceInput + (hStart * iw + wStart) * core->pack * core->bytes);
float* diff = (float *)(rowBwDiff + pw * core->pack * core->bytes);
for (int h = 0; h < hLen; h++, src += iw * core->pack) {
for (int w = 0; w < wLen; w++) {
int spatialIndex = (h + hStart) * iw + (wStart + w);
float* srcPtr = src + w * core->pack;
std::vector<float*> pre(core->pack, nullptr);
for (int k = 0; k < core->pack; k++) {
if (srcPtr[k] > maxes[k]) {
maxes[k] = srcPtr[k];
indices[k] = spatialIndex;
}
}
}
}
for (int k = 0; k < core->pack; k++) {
int h = indices[k] / iw;
int w = indices[k] % iw;
float* out = (float *)(sliceOutput + (h * iw + w) * core->pack * core->bytes);
out[k] += diff[k];
}
}
}
}
}
}
#endif
}
return NO_ERROR;
}
class CPUROIPoolingCreator : public CPUBackend::Creator {
public:
virtual Execution *onCreate(const std::vector<Tensor *> &inputs, const std::vector<Tensor *> &outputs,
const MNN::Op *op, Backend *backend) const {
auto roi = op->main_as_RoiParameters();
auto core = static_cast<CPUBackend*>(backend)->functions();
if (core->MNNRoiPoolingMax == nullptr) {
MNN_ERROR("Don't have function for CPUROIPooling\n");
return nullptr;
}
if (core->bytes < 4 && roi->outputGrad()) {
return nullptr;
}
return new CPUROIPooling(backend, roi->pooledWidth(), roi->pooledHeight(), roi->spatialScale(), roi->outputGrad());
}
};
REGISTER_CPU_OP_CREATOR(CPUROIPoolingCreator, OpType_ROIPooling);
} // namespace MNN