|
|
|
@ -30,6 +30,9 @@ class MKLDNNActivation : public ActivationFunction {
|
|
|
|
|
protected:
|
|
|
|
|
// input value element count
|
|
|
|
|
size_t cnt_;
|
|
|
|
|
// should not merge the resetBwd into resetFwd,
|
|
|
|
|
// because the grad data would be changing before backward.
|
|
|
|
|
bool needResetBwd_;
|
|
|
|
|
// mkldnn matrix, primitive, stream and pipeline
|
|
|
|
|
MKLDNNMatrixPtr val_;
|
|
|
|
|
MKLDNNMatrixPtr grad_;
|
|
|
|
@ -40,7 +43,7 @@ protected:
|
|
|
|
|
std::vector<mkldnn::primitive> pipelineBwd_;
|
|
|
|
|
|
|
|
|
|
public:
|
|
|
|
|
MKLDNNActivation() : cnt_(0) {}
|
|
|
|
|
MKLDNNActivation() : cnt_(0), needResetBwd_(true) {}
|
|
|
|
|
~MKLDNNActivation() {}
|
|
|
|
|
static ActivationFunction* create(const std::string& type);
|
|
|
|
|
static std::vector<std::string> getAllRegisteredTypes();
|
|
|
|
@ -57,19 +60,43 @@ class MKLDNNEltwiseActivation : public MKLDNNActivation {
|
|
|
|
|
typedef mkldnn::eltwise_forward eltwise_fwd;
|
|
|
|
|
typedef mkldnn::eltwise_backward eltwise_bwd;
|
|
|
|
|
|
|
|
|
|
protected:
|
|
|
|
|
// save the forward primitive desc, which can be used backward
|
|
|
|
|
std::shared_ptr<eltwise_fwd::primitive_desc> fwdPD_;
|
|
|
|
|
// eltwise_bwd need src input value
|
|
|
|
|
MKLDNNMatrixPtr inVal_;
|
|
|
|
|
// use for copy data
|
|
|
|
|
std::shared_ptr<mkldnn::reorder> copyInVal_;
|
|
|
|
|
|
|
|
|
|
public:
|
|
|
|
|
MKLDNNEltwiseActivation() {}
|
|
|
|
|
|
|
|
|
|
~MKLDNNEltwiseActivation() {}
|
|
|
|
|
|
|
|
|
|
virtual const std::string& getName() const = 0;
|
|
|
|
|
|
|
|
|
|
// in common, the alpha of forward and backward should be equal.
|
|
|
|
|
// but for relu, to avoid negative value, they should be opposite
|
|
|
|
|
virtual float getAlpha() const = 0;
|
|
|
|
|
virtual float getBwdAlpha() const = 0;
|
|
|
|
|
virtual float getBeta() const { return 0.f; }
|
|
|
|
|
virtual mkldnn::algorithm getAlgo(const std::string& type) const {
|
|
|
|
|
if (type == "mkldnn_relu") {
|
|
|
|
|
return mkldnn::algorithm::eltwise_relu;
|
|
|
|
|
} else if (type == "mkldnn_tanh") {
|
|
|
|
|
return mkldnn::algorithm::eltwise_tanh;
|
|
|
|
|
} else if (type == "mkldnn_elu") {
|
|
|
|
|
return mkldnn::algorithm::eltwise_elu;
|
|
|
|
|
} else {
|
|
|
|
|
LOG(FATAL) << "Unkown eltwise activation type: " << type;
|
|
|
|
|
}
|
|
|
|
|
return (mkldnn::algorithm)0;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/**
|
|
|
|
|
* reshape and reset the forward and backward primitives
|
|
|
|
|
* reshape and reset the forward primitives
|
|
|
|
|
*/
|
|
|
|
|
void resetPrimitives(Argument& act) {
|
|
|
|
|
void resetFwd(Argument& act) {
|
|
|
|
|
if (cnt_ == act.value->getElementCnt()) {
|
|
|
|
|
return;
|
|
|
|
|
}
|
|
|
|
@ -78,21 +105,13 @@ public:
|
|
|
|
|
auto eng = CPUEngine::Instance().getEngine();
|
|
|
|
|
|
|
|
|
|
// get algo setting
|
|
|
|
|
mkldnn::algorithm algo;
|
|
|
|
|
if (this->getName() == "mkldnn_relu") {
|
|
|
|
|
algo = mkldnn::algorithm::eltwise_relu;
|
|
|
|
|
} else if (this->getName() == "mkldnn_tanh") {
|
|
|
|
|
algo = mkldnn::algorithm::eltwise_tanh;
|
|
|
|
|
} else if (this->getName() == "mkldnn_elu") {
|
|
|
|
|
algo = mkldnn::algorithm::eltwise_elu;
|
|
|
|
|
} else {
|
|
|
|
|
LOG(FATAL) << "Unkown eltwise activation type: " << this->getName();
|
|
|
|
|
}
|
|
|
|
|
mkldnn::algorithm algo = getAlgo(this->getName());
|
|
|
|
|
// note: alpha represents the NegativeSlope when used in relu.
|
|
|
|
|
float alpha = getAlpha();
|
|
|
|
|
float beta = getBeta();
|
|
|
|
|
|
|
|
|
|
/// forward
|
|
|
|
|
pipelineFwd_.clear();
|
|
|
|
|
val_ = std::dynamic_pointer_cast<MKLDNNMatrix>(act.value);
|
|
|
|
|
if (val_ == nullptr) {
|
|
|
|
|
int bs = act.getBatchSize();
|
|
|
|
@ -109,33 +128,52 @@ public:
|
|
|
|
|
val_->getMemoryDesc(),
|
|
|
|
|
alpha,
|
|
|
|
|
beta);
|
|
|
|
|
auto fwdPD = eltwise_fwd::primitive_desc(fwdDesc, eng);
|
|
|
|
|
// inplace buffer, dst = src
|
|
|
|
|
fwd_.reset(new eltwise_fwd(fwdPD, *val_, *val_));
|
|
|
|
|
pipelineFwd_.clear();
|
|
|
|
|
fwdPD_.reset(new eltwise_fwd::primitive_desc(fwdDesc, eng));
|
|
|
|
|
// use inplace for forward but save input value before submit
|
|
|
|
|
inVal_ = val_;
|
|
|
|
|
if (act.grad) {
|
|
|
|
|
// only copy when need do backward
|
|
|
|
|
inVal_ = MKLDNNMatrix::create(nullptr, val_->getPrimitiveDesc());
|
|
|
|
|
copyInVal_ = std::make_shared<mkldnn::reorder>(*val_, *inVal_);
|
|
|
|
|
CHECK(copyInVal_) << "should not be emptry";
|
|
|
|
|
pipelineFwd_.push_back(*copyInVal_);
|
|
|
|
|
}
|
|
|
|
|
fwd_.reset(new eltwise_fwd(*fwdPD_, *val_, *val_));
|
|
|
|
|
pipelineFwd_.push_back(*fwd_);
|
|
|
|
|
needResetBwd_ = true;
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// backward
|
|
|
|
|
if (act.grad == nullptr) {
|
|
|
|
|
grad_ = nullptr;
|
|
|
|
|
/**
|
|
|
|
|
* reset the backward primitives, can not merge into resetFwd as the grad data
|
|
|
|
|
* would be changing before backward.
|
|
|
|
|
*/
|
|
|
|
|
void resetBwd(Argument& act) {
|
|
|
|
|
if (!needResetBwd_) {
|
|
|
|
|
return;
|
|
|
|
|
}
|
|
|
|
|
needResetBwd_ = false;
|
|
|
|
|
mkldnn::algorithm algo = getAlgo(this->getName());
|
|
|
|
|
float alpha = getBwdAlpha();
|
|
|
|
|
float beta = getBeta();
|
|
|
|
|
grad_ = MKLDNNMatrix::create(act.grad, val_->getPrimitiveDesc());
|
|
|
|
|
auto eng = CPUEngine::Instance().getEngine();
|
|
|
|
|
auto bwdDesc = eltwise_bwd::desc(
|
|
|
|
|
algo, grad_->getMemoryDesc(), val_->getMemoryDesc(), alpha, beta);
|
|
|
|
|
auto bwdPD = eltwise_bwd::primitive_desc(bwdDesc, eng, fwdPD);
|
|
|
|
|
bwd_.reset(new eltwise_bwd(bwdPD, *val_, *grad_, *grad_));
|
|
|
|
|
auto bwdPD = eltwise_bwd::primitive_desc(bwdDesc, eng, *fwdPD_);
|
|
|
|
|
CHECK(inVal_);
|
|
|
|
|
bwd_.reset(new eltwise_bwd(bwdPD, *inVal_, *grad_, *grad_));
|
|
|
|
|
pipelineBwd_.clear();
|
|
|
|
|
pipelineBwd_.push_back(*bwd_);
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
Error __must_check forward(Argument& act) {
|
|
|
|
|
resetPrimitives(act);
|
|
|
|
|
resetFwd(act);
|
|
|
|
|
stream_->submit(pipelineFwd_);
|
|
|
|
|
return Error();
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
Error __must_check backward(Argument& act) {
|
|
|
|
|
resetBwd(act);
|
|
|
|
|
stream_->submit(pipelineBwd_);
|
|
|
|
|
return Error();
|
|
|
|
|
}
|
|
|
|
|