wehub-resource-sync
211 lines
8.3 KiB
C++
211 lines
8.3 KiB
C++
|
|
//
|
|
// GridSampleGrad.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2022/09/07.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "OpGrad.hpp"
|
|
using namespace std;
|
|
namespace MNN {
|
|
using namespace MNN::Express;
|
|
|
|
class GridSampleGrad : public OpGrad {
|
|
public:
|
|
virtual std::vector<Express::VARP> onGrad(Express::EXPRP expr,
|
|
const std::vector<Express::VARP>& backwardOutput) override {
|
|
auto op = expr->get();
|
|
const auto& inputs = expr->inputs();
|
|
std::vector<VARP> res(inputs.size());
|
|
#ifndef GRIDSAMPLER_GRAD_DONT_CREATENEWOP
|
|
std::unique_ptr<MNN::OpT> gradOp(op->UnPack());
|
|
gradOp->name.clear();
|
|
gradOp->main.AsGridSample()->backward = true;
|
|
auto gradForParameterExpr = Expr::create(gradOp.get(), {backwardOutput[0], inputs[1], _Shape(inputs[0], NCHW)});
|
|
res[0] = Variable::create(gradForParameterExpr);
|
|
return res;
|
|
#else
|
|
auto param = op->main_as_GridSample();
|
|
auto sampleMode = param->mode();
|
|
auto padMode = param->paddingMode();
|
|
auto alignCorners = param->alignCorners();
|
|
|
|
auto input = inputs[0];
|
|
auto grid = inputs[1];
|
|
|
|
auto inputInfo = input->getInfo();
|
|
MNN_ASSERT(nullptr != inputInfo && (inputInfo->dim.size() == 4 || inputInfo->dim.size() == 5));
|
|
auto outputDiff = backwardOutput[0];
|
|
outputDiff = _Convert(outputDiff, NCHW);
|
|
auto diffInfo = outputDiff->getInfo();
|
|
|
|
const int inB = inputInfo->dim[0];
|
|
const int inC = inputInfo->dim[1];
|
|
const int inH = inputInfo->dim[2];
|
|
const int inW = inputInfo->dim[3];
|
|
|
|
const int outB = diffInfo->dim[0];
|
|
const int outC = diffInfo->dim[1];
|
|
const int outH = diffInfo->dim[2];
|
|
const int outW = diffInfo->dim[3];
|
|
|
|
auto one = _Scalar<float>(1.0f);
|
|
auto zeroFive = _Scalar<float>(0.5f);
|
|
std::vector<float> inShape = {float(inW), float(inH)};
|
|
auto inShapeVar = _Const((void*)inShape.data(), {1, 1, 1, 2}, NCHW);
|
|
outputDiff = _Transpose(outputDiff, {0, 2, 3, 1}); // NHWC
|
|
|
|
VARP sourceCord = nullptr;
|
|
if (alignCorners) {
|
|
sourceCord = (grid + one) * (inShapeVar - one) * zeroFive; // N, outH, outW, 2 (w, h)
|
|
} else {
|
|
sourceCord = ((grid + one) * inShapeVar - one) * zeroFive; // N, outH, outW, 2 (w, h)
|
|
}
|
|
|
|
if (sampleMode == SampleMode_NEAREST) {
|
|
auto indices = _Cast<int32_t>(_Floor(sourceCord + zeroFive)); // N, outH, outW, 2 (w, h)
|
|
|
|
{ // handle three pad modes
|
|
auto cords = _Split(indices, {1, 1}, -1);
|
|
auto wCord = cords[0];
|
|
auto hCord = cords[1];
|
|
auto zero = _Scalar<int>(0);
|
|
auto one = _Scalar<int>(1);
|
|
auto wVar = _Scalar<int>(inW);
|
|
auto hVar = _Scalar<int>(inH);
|
|
if (padMode == BorderMode_ZEROS) {
|
|
auto wMask1 = one - _Less(wCord, zero);
|
|
auto wMask2 = one - _GreaterEqual(wCord, wVar);
|
|
auto hMask1 = one - _Less(hCord, zero);
|
|
auto hMask2 = one - _GreaterEqual(hCord, hVar);
|
|
auto mask = wMask1 * wMask2 * hMask1 * hMask2;
|
|
outputDiff = outputDiff * _Cast<float>(mask);
|
|
}
|
|
wCord = Express::_Maximum(wCord, zero);
|
|
wCord = Express::_Minimum(wCord, wVar - one);
|
|
hCord = Express::_Maximum(hCord, zero);
|
|
hCord = Express::_Minimum(hCord, hVar - one);
|
|
indices = _Concat({wCord, hCord}, -1);
|
|
}
|
|
|
|
std::vector<int> bIndices;
|
|
for (int i = 0; i < outB; i++) {
|
|
for (int j = 0; j < outH*outW; j++) {
|
|
bIndices.push_back(i);
|
|
}
|
|
}
|
|
auto bIndicesVar = _Const((void*)bIndices.data(), {outB, outH, outW, 1}, NCHW, halide_type_of<int>());
|
|
|
|
indices = _Concat({bIndicesVar, indices}, -1);
|
|
auto updates = outputDiff;
|
|
std::vector<int> inputShape = {inB, inW, inH, inC};
|
|
auto shape = _Const(inputShape.data(), {4}, NCHW, halide_type_of<int>());
|
|
auto temp = _ScatterNd(indices, updates, shape, 0); // 0 for add
|
|
res[0] = _Transpose(temp, {0, 3, 2, 1}); // NCHW
|
|
} else if (sampleMode == SampleMode_BILINEAR) {
|
|
auto w0h0 = _Cast<int>(_Floor(sourceCord));
|
|
auto w1h1 = _Cast<int>(_Ceil(sourceCord));
|
|
auto factors0 = sourceCord - _Cast<float>(w0h0);
|
|
|
|
auto cords0 = _Split(w0h0, {1, 1}, -1);
|
|
auto w0 = cords0[0];
|
|
auto h0 = cords0[1];
|
|
auto cords1 = _Split(w1h1, {1, 1}, -1);
|
|
auto w1 = cords1[0];
|
|
auto h1 = cords1[1];
|
|
|
|
auto fs0 = _Split(factors0, {1, 1}, -1);
|
|
auto wf0 = fs0[0];
|
|
auto hf0 = fs0[1];
|
|
auto wf1 = _Scalar<float>(1.0f) - wf0;
|
|
auto hf1 = _Scalar<float>(1.0f) - hf0;
|
|
|
|
auto w00 = wf1 * hf1;
|
|
auto w01 = wf1 * hf0;
|
|
auto w10 = wf0 * hf1;
|
|
auto w11 = wf0 * hf0;
|
|
|
|
auto updates = outputDiff;
|
|
auto u00 = updates * w00;
|
|
auto u01 = updates * w01;
|
|
auto u10 = updates * w10;
|
|
auto u11 = updates * w11;
|
|
|
|
{ // handle three pad modes
|
|
auto zero = _Scalar<int>(0);
|
|
auto one = _Scalar<int>(1);
|
|
auto wVar = _Scalar<int>(inW);
|
|
auto hVar = _Scalar<int>(inH);
|
|
if (padMode == BorderMode_ZEROS) {
|
|
auto wMask0 = one - _Less(w0, zero);
|
|
auto wMask1 = one - _GreaterEqual(w0, wVar);
|
|
auto wMask2 = one - _Less(w1, zero);
|
|
auto wMask3 = one - _GreaterEqual(w1, wVar);
|
|
auto hMask0 = one - _Less(h0, zero);
|
|
auto hMask1 = one - _GreaterEqual(h0, hVar);
|
|
auto hMask2 = one - _Less(h1, zero);
|
|
auto hMask3 = one - _GreaterEqual(h1, hVar);
|
|
|
|
auto mask00 = wMask0 * wMask1 * hMask0 * hMask1;
|
|
auto mask01 = wMask0 * wMask1 * hMask2 * hMask3;
|
|
auto mask10 = wMask2 * wMask3 * hMask0 * hMask1;
|
|
auto mask11 = wMask2 * wMask3 * hMask2 * hMask3;
|
|
|
|
u00 = u00 * _Cast<float>(mask00);
|
|
u01 = u01 * _Cast<float>(mask01);
|
|
u10 = u10 * _Cast<float>(mask10);
|
|
u11 = u11 * _Cast<float>(mask11);
|
|
}
|
|
w0 = Express::_Maximum(w0, zero);
|
|
w0 = Express::_Minimum(w0, wVar - one);
|
|
w1 = Express::_Maximum(w1, zero);
|
|
w1 = Express::_Minimum(w1, wVar - one);
|
|
h0 = Express::_Maximum(h0, zero);
|
|
h0 = Express::_Minimum(h0, hVar - one);
|
|
h1 = Express::_Maximum(h1, zero);
|
|
h1 = Express::_Minimum(h1, hVar - one);
|
|
}
|
|
|
|
std::vector<int> bIndices;
|
|
for (int i = 0; i < outB; i++) {
|
|
for (int j = 0; j < outH*outW; j++) {
|
|
bIndices.push_back(i);
|
|
}
|
|
}
|
|
auto bIndicesVar = _Const((void*)bIndices.data(), {outB, outH, outW, 1}, NCHW, halide_type_of<int>());
|
|
|
|
auto bw0h1 = _Concat({bIndicesVar, w0, h1}, -1);
|
|
auto bw1h0 = _Concat({bIndicesVar, w1, h0}, -1);
|
|
auto bw0h0 = _Concat({bIndicesVar, w0, h0}, -1);
|
|
auto bw1h1 = _Concat({bIndicesVar, w1, h1}, -1);
|
|
|
|
std::vector<int> inputShape = {inB, inW, inH, inC};
|
|
auto shape = _Const(inputShape.data(), {4}, NCHW, halide_type_of<int>());
|
|
auto temp0 = _ScatterNd(bw0h0, u00, shape, 0); // 0 for add
|
|
auto temp1 = _ScatterNd(bw0h1, u01, shape, 0);
|
|
auto temp2 = _ScatterNd(bw1h0, u10, shape, 0);
|
|
auto temp3 = _ScatterNd(bw1h1, u11, shape, 0);
|
|
|
|
auto temp = temp0 + temp1 + temp2 + temp3;
|
|
res[0] = _Transpose(temp, {0, 3, 2, 1}); // NCHW
|
|
}
|
|
|
|
res[0] = _Convert(res[0], inputInfo->order);
|
|
return res;
|
|
#endif
|
|
}
|
|
};
|
|
|
|
static void _create() {
|
|
static GridSampleGrad _c;
|
|
OpGrad::insert((int)OpType_GridSample, &_c);
|
|
OpGrad::insert((int)OpType_Texture, &_c);
|
|
|
|
}
|
|
|
|
REGISTER_GRAD(GridSampleGrad_cpp, _create);
|
|
};
|
|
|