wehub-resource-sync
57 lines
1.6 KiB
C++
57 lines
1.6 KiB
C++
//
|
|
// SoftmaxGrad.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/04/22.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "SoftmaxGrad.hpp"
|
|
#include "core/Macro.h"
|
|
#include <MNN/expr/ExprCreator.hpp>
|
|
using namespace std;
|
|
namespace MNN {
|
|
using namespace MNN::Express;
|
|
|
|
class SoftmaxGrad : public OpGrad {
|
|
public:
|
|
SoftmaxGrad() {
|
|
mType = NO_LINEAR;
|
|
}
|
|
virtual std::vector<Express::VARP> onGrad(Express::EXPRP expr,
|
|
const std::vector<Express::VARP>& backwardOutput) override {
|
|
MNN_ASSERT(expr->inputs().size() == 1 && backwardOutput.size() == 1);
|
|
auto input = expr->inputs()[0];
|
|
auto info = input->getInfo();
|
|
auto gradSoftmax = backwardOutput[0];
|
|
if (nullptr == info) {
|
|
return {};
|
|
}
|
|
auto axis = expr->get()->main_as_Axis()->axis();
|
|
if (axis < 0) {
|
|
axis = axis + info->dim.size();
|
|
}
|
|
auto softmax = Express::Variable::create(expr, 0);
|
|
auto originOrder = info->order;
|
|
if (originOrder == NC4HW4) {
|
|
gradSoftmax = _Convert(gradSoftmax, NCHW);
|
|
softmax = _Convert(softmax, NCHW);
|
|
}
|
|
auto sumAxis = _ReduceSum(softmax * gradSoftmax, {axis}, true);
|
|
auto inputGrad = (gradSoftmax - sumAxis) * softmax;
|
|
if (originOrder == NC4HW4) {
|
|
inputGrad = _Convert(inputGrad, NC4HW4);
|
|
}
|
|
return {inputGrad};
|
|
}
|
|
};
|
|
static void _create() {
|
|
static SoftmaxGrad _c;
|
|
OpGrad::insert(OpType_Softmax, &_c);
|
|
|
|
}
|
|
|
|
REGISTER_GRAD(SoftmaxGrad_cpp, _create);
|
|
};
|
|
|