wehub-resource-sync
78 lines
2.5 KiB
C++
78 lines
2.5 KiB
C++
//
|
|
// BroadcastToGrad.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2022/07/27.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "OpGrad.hpp"
|
|
#include "core/Macro.h"
|
|
using namespace std;
|
|
using namespace MNN::Express;
|
|
namespace MNN {
|
|
class BroadcastToGrad : public OpGrad {
|
|
public:
|
|
virtual std::vector<Express::VARP> onGrad(Express::EXPRP expr,
|
|
const std::vector<Express::VARP>& backwardOutput) override {
|
|
auto inputs = expr->inputs();
|
|
std::vector<VARP> res(inputs.size(), nullptr);
|
|
auto outputDiff = backwardOutput[0];
|
|
auto outputDiffInfo = outputDiff->getInfo();
|
|
|
|
auto input = inputs[0];
|
|
auto inputInfo = input->getInfo();
|
|
|
|
std::vector<int> reduceDims;
|
|
bool keepDim = true;
|
|
|
|
if (inputInfo->dim.size() < outputDiffInfo->dim.size()) {
|
|
// case like: shape(2, 3, 1) ==> shape(7, 2, 3, 3)
|
|
// will only be handled a part here
|
|
// because we need keepDim = false for dim[0] = 7
|
|
// and keepDim = true for dim[-1] = 3
|
|
auto diff = (int)outputDiffInfo->dim.size() - (int)inputInfo->dim.size();
|
|
for (int i = 0; i < diff; ++i) {
|
|
reduceDims.emplace_back(i);
|
|
}
|
|
keepDim = false;
|
|
} else {
|
|
for (int i = 0; i < outputDiffInfo->dim.size(); ++i) {
|
|
if (outputDiffInfo->dim[i] > 1 && inputInfo->dim[i] == 1) {
|
|
reduceDims.emplace_back(i);
|
|
}
|
|
}
|
|
keepDim = true;
|
|
}
|
|
|
|
if (!reduceDims.empty()) {
|
|
res[0] = _ReduceSum(outputDiff, reduceDims, keepDim);
|
|
// for case like: shape(2, 3, 1) ==> shape(7, 2, 3, 3)
|
|
if (keepDim == false) {
|
|
reduceDims.clear();
|
|
auto diff = (int)outputDiffInfo->dim.size() - (int)inputInfo->dim.size();
|
|
for (int j = 0; j < inputInfo->dim.size(); j++) {
|
|
if (outputDiffInfo->dim[j + diff] > 1 && inputInfo->dim[j] == 1) {
|
|
reduceDims.emplace_back(j);
|
|
}
|
|
}
|
|
keepDim = true;
|
|
if (!reduceDims.empty()) {
|
|
res[0] = _ReduceSum(outputDiff, reduceDims, keepDim);
|
|
}
|
|
}
|
|
} else {
|
|
res[0] = outputDiff;
|
|
}
|
|
|
|
return res;
|
|
}
|
|
};
|
|
|
|
static void _create() {
|
|
static BroadcastToGrad _c;
|
|
OpGrad::insert(OpType_BroadcastTo, &_c);
|
|
}
|
|
REGISTER_GRAD(BroadcastToGrad, _create);
|
|
};
|