Files
2026-07-13 13:33:03 +08:00

237 lines
9.2 KiB
C++

//
// UnaryGrad.cpp
// MNN
//
// Created by MNN on 2019/05/25.
// Copyright © 2018, Alibaba Group Holding Limited
//
#include "OpGrad.hpp"
#include "core/Macro.h"
#define MNN_PI 3.14159265358979323846
using namespace std;
namespace MNN {
using namespace MNN::Express;
class UnaryGrad : public OpGrad {
public:
virtual std::vector<Express::VARP> onGrad(Express::EXPRP expr,
const std::vector<Express::VARP>& backwardOutput) override {
std::unique_ptr<OpT> forwardOp(expr->get()->UnPack());
auto outputDiff = backwardOutput[0];
auto input = expr->inputs()[0];
std::vector<Express::VARP> res(1, nullptr);
std::vector<Express::VARP> output{Variable::create(expr, 0)};
switch (forwardOp->main.AsUnaryOp()->opType) {
case MNN::UnaryOpOperation_LOG1P: {
// d log(1+x) = 1/(1+x) * dx = dx / (1+x)
auto oneConst = _Const(1.0f, {}, NHWC);
auto addOne = _Add(input, oneConst);
res[0] = _Divide(outputDiff, addOne);
break;
}
case MNN::UnaryOpOperation_EXP: {
// d Exp(x) = Exp(x) * dx
res[0] = _Multiply(outputDiff, output[0]);
break;
}
case MNN::UnaryOpOperation_LOG: {
// d Log(x) = dx / x
res[0] = _Divide(outputDiff, input);
break;
}
case MNN::UnaryOpOperation_COS: {
// d Sin(x) = -dx * Sin(x)
res[0] = _Negative(outputDiff) * _Sin(input);
break;
}
case MNN::UnaryOpOperation_SIN: {
// d Sin(x) = dx * Cos(x)
res[0] = outputDiff * _Cos(input);
break;
}
case MNN::UnaryOpOperation_ABS: {
// d Abs(x) = dx * (x > 0 ? 1 : -1)
res[0] = outputDiff * _Sign(input);
break;
}
case MNN::UnaryOpOperation_NEG: {
// d (-x) = - dx
res[0] = _Negative(outputDiff);
break;
}
case MNN::UnaryOpOperation_SQRT: {
// d (-sqrt(x)) = 0.5 / sqrt(x) * dx
auto oneConst = _Const(0.5f, {}, NHWC);
auto mul = _Multiply(outputDiff, oneConst);
res[0] = OpGrad::divideAvoidZero(mul, output[0]);
break;
}
case MNN::UnaryOpOperation_SQUARE: {
// d (x^2) = (x*dx + x*dx)
auto mul = _Multiply(input, outputDiff);
res[0] = _Add(mul, mul);
break;
}
case MNN::UnaryOpOperation_SIGMOID: {
auto grad = OpGrad::get(OpType_Sigmoid);
res[0] = grad->onGrad(expr, backwardOutput)[0];
break;
}
case MNN::UnaryOpOperation_TANH: {
auto grad = OpGrad::get(OpType_TanH);
res[0] = grad->onGrad(expr, backwardOutput)[0];
break;
}
case MNN::UnaryOpOperation_RSQRT: {
// d (x^(-1/2)) = -1/2 * x^(-3/2) * dx
res[0] = _Scalar(-0.5f) * _Pow(input, _Scalar(-1.5f)) * outputDiff;
break;
}
case MNN::UnaryOpOperation_TAN: {
// d tan(x) = 1 / (cos(x))^2 * dx
res[0] = _Scalar(1.0f) / _Square(_Cos(input)) * outputDiff;
break;
}
case MNN::UnaryOpOperation_ASIN: {
// d asin(x) = 1 / sqrt(1 - x^2) * dx
res[0] = _Scalar(1.0f) / _Sqrt(_Scalar(1.0f) - _Square(input)) * outputDiff;
break;
}
case MNN::UnaryOpOperation_ACOS: {
// d acos(x) = -1 / sqrt(1 - x^2) * dx
res[0] = _Scalar(-1.0f) / _Sqrt(_Scalar(1.0f) - _Square(input)) * outputDiff;
break;
}
case MNN::UnaryOpOperation_ATAN: {
// d atan(x) = 1 / (1 + x^2) * dx
res[0] = _Scalar(1.0f) / (_Scalar(1.0f) + _Square(input)) * outputDiff;
break;
}
case MNN::UnaryOpOperation_RECIPROCAL: {
// d x^-1 = - x^-2 * dx
res[0] = _Negative(_Pow(input, _Scalar(-2.0f))) * outputDiff;
break;
}
case MNN::UnaryOpOperation_ACOSH: {
// d acosh(x) = 1 / sqrt(x^2 - 1) * dx
res[0] = _Scalar(1.0f) / _Sqrt(_Square(input) - _Scalar(1.0f)) * outputDiff;
break;
}
case MNN::UnaryOpOperation_SINH: {
// d sinh(x) = cosh(x) * dx
res[0] = _Cosh(input) * outputDiff;
break;
}
case MNN::UnaryOpOperation_COSH: {
// d cosh(x) = sinh(x) * dx
res[0] = _Sinh(input) * outputDiff;
break;
}
case MNN::UnaryOpOperation_ASINH: {
// d asinh(x) = 1 / sqrt(x^2 + 1) * dx
res[0] = _Scalar(1.0f) / _Sqrt(_Square(input) + _Scalar(1.0f)) * outputDiff;
break;
}
case MNN::UnaryOpOperation_ATANH: {
// d atanh(x) = 1 / (1 - x^2) * dx
res[0] = _Scalar(1.0f) / (_Scalar(1.0f) - _Square(input)) * outputDiff;
break;
}
case MNN::UnaryOpOperation_ERF: {
// d erf(x) = 2 / sqrt(pi) * exp(- x^2) * dx
res[0] = _Scalar(2.0f) / _Sqrt(_Scalar(float(MNN_PI))) * _Exp(_Negative(_Square(input))) * outputDiff;
break;
}
case MNN::UnaryOpOperation_ERFC: {
// d erfc(x) = -2 / sqrt(pi) * exp(- x^2) * dx
res[0] = _Scalar(-2.0f) / _Sqrt(_Scalar(float(MNN_PI))) * _Exp(_Negative(_Square(input))) * outputDiff;
break;
}
case MNN::UnaryOpOperation_ERFINV: {
// d erfinv(x) = sqrt(pi) / 2 * exp(erfinv(x)^2) * dx
res[0] = _Sqrt(_Scalar(float(MNN_PI))) / _Scalar(2.0f) * _Exp(_Square(_Erfinv(input))) * outputDiff;
break;
}
case MNN::UnaryOpOperation_EXPM1: {
// d expm1(x) = exp(x) * dx
res[0] = _Exp(input) * outputDiff;
break;
}
case MNN::UnaryOpOperation_HARDSWISH: {
// d hardswish(x) = (relu6(x+3) + x * relu6'(x+3)) / 6 * dx
auto inputp3 = input + _Scalar(3.0f);
auto mask0 = _Cast<float>(_Greater(inputp3, _Scalar(0.0f)));
auto mask1 = _Cast<float>(_Less(inputp3, _Scalar(6.0f)));
auto relu6Xp3Grad = mask0 * mask1;
res[0] = (_Relu6(inputp3) + input * relu6Xp3Grad) / _Scalar(6.0f) * outputDiff;
break;
}
case MNN::UnaryOpOperation_GELU:
case MNN::UnaryOpOperation_GELU_STANDARD: {
// d gelu(x) = 0.5 * ( (1 + erf(x / sqrt(2))) + x * (erf'(x / sqrt(2)) / sqrt(2)) ) * dx
auto const05 = _Scalar(0.5f);
auto inputx = input / _Sqrt(_Scalar(2.0f));
auto part1 = _Scalar(1.0f) + _Erf(inputx);
auto erfGrad = _Scalar(2.0f) / _Sqrt(_Scalar(float(MNN_PI))) * _Exp(_Negative(_Square(inputx)));
auto part2 = input * erfGrad / _Sqrt(_Scalar(2.0f));
res[0] = const05 * (part1 + part2) * outputDiff;
break;
}
default:
MNN_ERROR("Can't grad for unary: %d\n", forwardOp->main.AsUnaryOp()->opType);
return res;
}
return res;
}
};
class SigmoidGrad : public OpGrad {
public:
virtual std::vector<Express::VARP> onGrad(Express::EXPRP expr,
const std::vector<Express::VARP>& backwardOutput) override {
std::vector<Express::VARP> result(1, nullptr);
auto outputDiff = backwardOutput[0];
std::vector<Express::VARP> output{Variable::create(expr, 0)};
// y = (1/(1+e(-x))) , dy = y(1-y) * dx = (y - y*y)*dx
auto mul = _Multiply(output[0], output[0]);
auto sub = _Subtract(output[0], mul);
auto grad = _Multiply(sub, outputDiff);
result[0] = grad;
return result;
}
};
class TanhGrad : public OpGrad {
public:
virtual std::vector<Express::VARP> onGrad(Express::EXPRP expr,
const std::vector<Express::VARP>& backwardOutput) override {
std::vector<Express::VARP> result{nullptr};
std::vector<Express::VARP> output{Variable::create(expr, 0)};
auto outputDiff = backwardOutput[0];
// d tanh(x) = (1-tanh(x)^2)dx
result[0] = (_Const(1.0f, {}, NCHW) - _Square(output[0])) * outputDiff;
return result;
}
};
static void _create() {
static UnaryGrad _c;
static SigmoidGrad _s;
static TanhGrad _t;
OpGrad::insert(OpType_UnaryOp, &_c);
OpGrad::insert(OpType_Sigmoid, &_s);
OpGrad::insert(OpType_TanH, &_t);
}
REGISTER_GRAD(UnaryGrad_cpp, _create);
};