1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64
|
#include <caffe2/ideep/ideep_utils.h>
using namespace caffe2;
namespace {
class IDEEPSigmoidOp final : public IDEEPOperator {
public:
USE_IDEEP_DEF_ALIASES();
USE_IDEEP_OPERATOR_FUNCTIONS();
IDEEPSigmoidOp(const OperatorDef& operator_def, Workspace* ws)
: IDEEPOperator(operator_def, ws) {
}
// NOLINTNEXTLINE(modernize-use-equals-default)
~IDEEPSigmoidOp() override {}
bool RunOnDevice() override {
const auto& X = Input(INPUT);
auto* Y = Output(OUTPUT);
ideep::eltwise_forward::compute(
X, *Y, ialgo::eltwise_logistic, iprop::forward_training);
return true;
}
private:
INPUT_TAGS(INPUT);
OUTPUT_TAGS(OUTPUT);
};
class IDEEPSigmoidGradientOp final : public IDEEPOperator {
public:
USE_IDEEP_DEF_ALIASES();
USE_IDEEP_OPERATOR_FUNCTIONS();
IDEEPSigmoidGradientOp(const OperatorDef& operator_def, Workspace* ws)
: IDEEPOperator(operator_def, ws) {
}
// NOLINTNEXTLINE(modernize-use-equals-default)
~IDEEPSigmoidGradientOp() override {}
bool RunOnDevice() override {
const auto& Y = Input(OUTPUT);
const auto& dY = Input(OUTPUT_GRAD);
auto* dX = Output(INPUT_GRAD);
ideep::eltwise_backward::compute(Y, dY, *dX, ialgo::eltwise_logistic);
return true;
}
private:
INPUT_TAGS(OUTPUT, OUTPUT_GRAD);
OUTPUT_TAGS(INPUT_GRAD);
};
REGISTER_IDEEP_OPERATOR(Sigmoid, IDEEPSigmoidOp);
REGISTER_IDEEP_OPERATOR(SigmoidGradient, IDEEPSigmoidGradientOp);
} // namespace
|