MKLDNNActivation.cpp 9.0 KB
Newer Older
1
/* Copyright (c) 2017 PaddlePaddle Authors. All Rights Reserved.
T
tensor-tang 已提交
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

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#include "MKLDNNActivation.h"
#include "mkldnn.hpp"
#include "paddle/utils/ClassRegistrar.h"

namespace paddle {

static ClassRegistrar<ActivationFunction> gMKLDNNActivationRegistrar;
/**
 * @def MKLDNN_ACTIVATION_CLASS_NAME
 * @note MKLDNN_ACTIVATION_CLASS_NAME(relu) relu_;
 * means mkldnn_reluActivation relu_;
 */
#define MKLDNN_ACTIVATION_CLASS_NAME(ACT_TYPE) mkldnn_##ACT_TYPE##Activation

T
tensor-tang 已提交
29
/**
T
tensor-tang 已提交
30
 * @def BEGIN_MKLDNN_ACTIVATION
T
tensor-tang 已提交
31
 */
T
tensor-tang 已提交
32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50
#define BEGIN_MKLDNN_ACTIVATION(ACT_TYPE, BASE_CLASS) \
  class MKLDNN_ACTIVATION_CLASS_NAME(ACT_TYPE) : public BASE_CLASS {
/**
 * @def END_MKLDNN_ACTIVATION
 */
#define END_MKLDNN_ACTIVATION(ACT_TYPE)                            \
private:                                                           \
  static const std::string name;                                   \
                                                                   \
public:                                                            \
  const std::string& getName() const { return name; }              \
  }                                                                \
  ;                                                                \
  const std::string MKLDNN_ACTIVATION_CLASS_NAME(ACT_TYPE)::name = \
      "mkldnn_" #ACT_TYPE;                                         \
  static InitFunction __reg_activation__mkldnn_##ACT_TYPE([] {     \
    gMKLDNNActivationRegistrar                                     \
        .registerClass<MKLDNN_ACTIVATION_CLASS_NAME(ACT_TYPE)>(    \
            "mkldnn_" #ACT_TYPE);                                  \
T
tensor-tang 已提交
51 52
  });

T
tensor-tang 已提交
53 54 55 56 57 58 59
/**
 * @def DEFINE_MKLDNN_ACTIVATION
 */
#define DEFINE_MKLDNN_ACTIVATION(ACT_TYPE, BASE_CLASS) \
  BEGIN_MKLDNN_ACTIVATION(ACT_TYPE, BASE_CLASS)        \
  END_MKLDNN_ACTIVATION(ACT_TYPE)

T
tensor-tang 已提交
60 61 62
/**
 * @def DEFINE_MKLDNN_ELTWISE_ACTIVATION
 */
T
tensor-tang 已提交
63 64 65 66 67 68 69 70 71 72 73 74 75
#define DEFINE_MKLDNN_ELTWISE_ACTIVATION(                            \
    ACT_TYPE, BASE_CLASS, ALPHA, BWD_ALPHA)                          \
  BEGIN_MKLDNN_ACTIVATION(ACT_TYPE, BASE_CLASS)                      \
private:                                                             \
  static const float alpha;                                          \
  static const float bwdAlpha;                                       \
                                                                     \
public:                                                              \
  float getAlpha() const { return alpha; }                           \
  float getBwdAlpha() const { return bwdAlpha; }                     \
  END_MKLDNN_ACTIVATION(ACT_TYPE)                                    \
  const float MKLDNN_ACTIVATION_CLASS_NAME(ACT_TYPE)::alpha = ALPHA; \
  const float MKLDNN_ACTIVATION_CLASS_NAME(ACT_TYPE)::bwdAlpha = BWD_ALPHA;
T
tensor-tang 已提交
76 77 78 79 80 81

/**
 * @brief MKLDNN Relu Activation.
 * Actually mkldnn_relu is Leaky Relu.
 *  f(x) = x                   (x >= 0)
 *  f(x) = negative_slope * x  (x <  0)
82
 * @note the negative_slope should be -0.f in forward
T
tensor-tang 已提交
83
 */
T
tensor-tang 已提交
84
DEFINE_MKLDNN_ELTWISE_ACTIVATION(relu, MKLDNNEltwiseActivation, -0.f, 0.f)
T
tensor-tang 已提交
85 86 87 88

/**
 * @brief MKLDNN Tanh Activation.
 */
T
tensor-tang 已提交
89
DEFINE_MKLDNN_ELTWISE_ACTIVATION(tanh, MKLDNNEltwiseActivation, 0.f, 0.f)
T
tensor-tang 已提交
90 91 92 93 94 95

/**
 * @brief MKLDNN ELU(Exponential Linear Unit) Activation.
 *  f(x) = x                              (x >= 0)
 *  f(x) = negative_slope * (exp(x) - 1)  (x <  0)
 */
T
tensor-tang 已提交
96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128
DEFINE_MKLDNN_ELTWISE_ACTIVATION(elu, MKLDNNEltwiseActivation, 0.f, 0.f)

mkldnn::algorithm MKLDNNEltwiseActivation::getAlgo(std::string type) const {
  const std::map<std::string, mkldnn::algorithm> algoMap = {
      {"relu", algorithm::eltwise_relu},
      {"tanh", algorithm::eltwise_tanh},
      {"elu", algorithm::eltwise_elu}};
  type.erase(0, 7);  // remove mkldnn_
  algorithm algo = (algorithm)0;
  mapGet(type, algoMap, &algo);
  return algo;
}

void MKLDNNEltwiseActivation::resetFwd(Argument& act) {
  if (cnt_ == act.value->getElementCnt()) {
    return;
  }
  MKLDNNActivation::resetFwd(act);
  // note: alpha represents the NegativeSlope when used in relu.
  float alpha = getAlpha();
  float beta = getBeta();
  algorithm algo = getAlgo(this->getName());
  auto fwdDesc = eltwise_fwd::desc(mkldnn::prop_kind::forward_training,
                                   algo,
                                   val_->getMemoryDesc(),
                                   alpha,
                                   beta);
  fwdPD_.reset(new eltwise_fwd::primitive_desc(fwdDesc, *engine_));
  // use inplace for forward but save input value before submit
  inVal_ = val_;
  copyInVal_ = nullptr;
  if (act.grad && algo == algorithm::eltwise_tanh) {
    // tanh need save src input for backward
129
    inVal_ = MKLDNNMatrix::create(val_->getPrimitiveDesc());
T
tensor-tang 已提交
130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147
    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;
}

void MKLDNNEltwiseActivation::resetBwd(Argument& act) {
  if (!needResetBwd_) {
    return;
  }
  VLOG(MKLDNN_BASE) << getName() << " reset mkldnn backward";
  needResetBwd_ = false;
  algorithm algo = getAlgo(this->getName());
  float alpha = getBwdAlpha();
  float beta = getBeta();
148
  grad_ = MKLDNNMatrix::create(val_->getPrimitiveDesc(), act.grad);
T
tensor-tang 已提交
149 150 151 152 153 154 155 156 157
  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_);
  CHECK(inVal_);
  bwd_.reset(new eltwise_bwd(bwdPD, *inVal_, *grad_, *grad_));
  pipelineBwd_.clear();
  pipelineBwd_.push_back(*bwd_);
}
T
tensor-tang 已提交
158

T
tensor-tang 已提交
159 160 161 162 163
/**
 * @brief MKLDNN Softmax Activation
 */
DEFINE_MKLDNN_ACTIVATION(softmax, MKLDNNSoftmaxActivation)

T
tensor-tang 已提交
164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191
void MKLDNNSoftmaxActivation::resetFwd(Argument& act) {
  if (cnt_ == act.value->getElementCnt()) {
    return;
  }
  MKLDNNActivation::resetFwd(act);
  int axis = 1;
  auto fwdDesc = softmax_fwd::desc(
      mkldnn::prop_kind::forward_scoring, val_->getMemoryDesc(), axis);
  auto fwdPD = softmax_fwd::primitive_desc(fwdDesc, *engine_);
  fwd_.reset(new softmax_fwd(fwdPD, *val_, *val_));
  pipelineFwd_.push_back(*fwd_);
}

Error __must_check MKLDNNSoftmaxActivation::forward(Argument& act) {
  resetFwd(act);
  stream_->submit(pipelineFwd_);
  real* v = act.value->getData();
  real threshold = exp(-64);
#pragma omp parallel for
  for (size_t i = 0; i < act.value->getElementCnt(); ++i) {
    v[i] = v[i] < threshold ? threshold : v[i];
  }
  return Error();
}

Error __must_check MKLDNNSoftmaxActivation::backward(Argument& act) {
  MatrixPtr outputV = act.value;
  MatrixPtr outputG = act.grad;
192 193 194 195 196 197 198 199 200 201 202 203 204
  Matrix::resizeOrCreate(sftMaxDot_,
                         outputG->getHeight(),
                         outputG->getWidth(),
                         /* trans */ false,
                         /* useGpu */ false);
  Matrix::resizeOrCreate(sftMaxSum_,
                         outputG->getHeight(),
                         1,
                         /* trans */ false,
                         /* useGpu */ false);
  sftMaxDot_->dotMul(*outputG, *outputV);
  sftMaxSum_->colMerge(*sftMaxDot_);
  act.grad->softmaxDerivative(*act.value, *sftMaxSum_);
T
tensor-tang 已提交
205 206 207
  return Error();
}

T
tensor-tang 已提交
208 209 210 211 212 213 214 215 216 217 218
ActivationFunction* MKLDNNActivation::create(const std::string& type) {
  return gMKLDNNActivationRegistrar.createByType(type);
}

std::vector<std::string> MKLDNNActivation::getAllRegisteredTypes() {
  std::vector<std::string> types;
  gMKLDNNActivationRegistrar.forEachType(
      [&](const std::string& type) { types.push_back(type); });
  return types;
}

T
tensor-tang 已提交
219 220 221 222 223 224 225 226 227 228 229 230 231 232
void MKLDNNActivation::resetFwd(Argument& act) {
  VLOG(MKLDNN_BASE) << getName() << " reset mkldnn forward";
  cnt_ = act.value->getElementCnt();
  pipelineFwd_.clear();
  stream_.reset(new MKLDNNStream());
  engine_.reset(new mkldnn::engine(mkldnn::engine::cpu, 0));
  val_ = std::dynamic_pointer_cast<MKLDNNMatrix>(act.value);
  if (val_ == nullptr) {
    int bs = act.getBatchSize();
    int ih = act.getFrameHeight() > 0 ? act.getFrameHeight() : 1;
    int iw = act.getFrameWidth() > 0 ? act.getFrameWidth() : 1;
    int ic = cnt_ / bs / ih / iw;
    CHECK_EQ(cnt_, (size_t)bs * ic * ih * iw);
    val_ = MKLDNNMatrix::create(
233
        {bs, ic, ih, iw}, mkldnn::memory::format::nchw, *engine_, act.value);
T
tensor-tang 已提交
234 235 236 237 238 239 240 241 242 243 244 245 246 247 248
    CHECK(val_);
    val_->downSpatial();
  }
}

Error __must_check MKLDNNActivation::forward(Argument& act) {
  resetFwd(act);
  stream_->submit(pipelineFwd_);
  return Error();
}
Error __must_check MKLDNNActivation::backward(Argument& act) {
  resetBwd(act);
  stream_->submit(pipelineBwd_);
  return Error();
}
T
tensor-tang 已提交
249
}  // namespace paddle