chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,210 @@
|
||||
|
||||
//
|
||||
// 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);
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user