tensor-tang
8 years ago
parent
f6a940936b
commit
b2bd67133a
@ -1,222 +0,0 @@
|
||||
/* Copyright (c) 2017 PaddlePaddle Authors. All Rights Reserve.
|
||||
|
||||
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 "MkldnnLayer.h"
|
||||
|
||||
using mem = mkldnn::memory; // NOLINT
|
||||
typedef mem::format format;
|
||||
typedef mkldnn::inner_product_forward fc_fwd;
|
||||
typedef mkldnn::inner_product_backward_weights fc_bwdWgt;
|
||||
typedef mkldnn::inner_product_backward_data fc_bwdData;
|
||||
|
||||
namespace paddle {
|
||||
|
||||
bool MkldnnLayer::init(const LayerMap& layerMap,
|
||||
const ParameterMap& parameterMap) {
|
||||
if (!Layer::init(layerMap, parameterMap)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
CHECK(FLAGS_use_mkldnn) << "MkldnnLayers only support use_mkldnn."
|
||||
<< "Please set WITH_MKLDNN=ON "
|
||||
<< "and set use_mkldnn=True";
|
||||
stream_.reset(new MkldnnStream());
|
||||
engine_ = CpuEngine::Instance().getEngine();
|
||||
|
||||
// TODO(TJ): deivecId
|
||||
return true;
|
||||
}
|
||||
|
||||
void MkldnnLayer::resetForwardFC(int bs,
|
||||
int ic,
|
||||
int ih,
|
||||
int iw,
|
||||
real* botData,
|
||||
int oc,
|
||||
real* topData,
|
||||
real* wgtData,
|
||||
real* biasData) {
|
||||
bool hasSpatial = ih == 1 && iw == 1 ? false : true;
|
||||
mem::desc botMD = hasSpatial ? createMD({bs, ic, ih, iw}, format::nchw)
|
||||
: createMD({bs, ic}, format::nc);
|
||||
mem::desc wgtMD = hasSpatial ? createMD({oc, ic, ih, iw}, format::oihw)
|
||||
: createMD({oc, ic}, format::oi);
|
||||
mem::desc biasMD = biasData != NULL ? createMD({oc}, format::x)
|
||||
: createMD({}, format::format_undef);
|
||||
mem::desc topMD = createMD({bs, oc}, format::nc);
|
||||
|
||||
mem::primitive_desc botPD = mem::primitive_desc(botMD, engine_);
|
||||
if (inVal_ && inVal_->get_primitive_desc() == botPD) {
|
||||
return;
|
||||
}
|
||||
|
||||
inVal_.reset(new mem(botPD, botData));
|
||||
wgtVal_.reset(new mem(mem::primitive_desc(wgtMD, engine_), wgtData));
|
||||
outVal_.reset(new mem(mem::primitive_desc(topMD, engine_), topData));
|
||||
|
||||
mkldnn::prop_kind pk = mkldnn::prop_kind::forward;
|
||||
fc_fwd::desc fwdDesc = biasData != NULL
|
||||
? fc_fwd::desc(pk, botMD, wgtMD, biasMD, topMD)
|
||||
: fc_fwd::desc(pk, botMD, wgtMD, topMD);
|
||||
fc_fwd::primitive_desc fwdPD = fc_fwd::primitive_desc(fwdDesc, engine_);
|
||||
|
||||
if (biasData != NULL) {
|
||||
biasVal_.reset(new mem(mem::primitive_desc(biasMD, engine_), biasData));
|
||||
fwd_.reset(new fc_fwd(fwdPD, *inVal_, *wgtVal_, *biasVal_, *outVal_));
|
||||
} else {
|
||||
fwd_.reset(new fc_fwd(fwdPD, *inVal_, *wgtVal_, *outVal_));
|
||||
}
|
||||
pipelineFwd_.clear();
|
||||
pipelineFwd_.push_back(*fwd_);
|
||||
}
|
||||
|
||||
void MkldnnLayer::mkldnnForwardFC(int bs,
|
||||
int ic,
|
||||
int ih,
|
||||
int iw,
|
||||
real* botData,
|
||||
int oc,
|
||||
real* topData,
|
||||
real* wgtData,
|
||||
real* biasData) {
|
||||
// if input size changed, reset it
|
||||
resetForwardFC(bs, ic, ih, iw, botData, oc, topData, wgtData, biasData);
|
||||
|
||||
this->convertWeightsFromPaddle();
|
||||
|
||||
// update input, since the data might be changed if this is after data layer
|
||||
inVal_->set_data_handle(botData);
|
||||
|
||||
// just forward
|
||||
stream_->submit(pipelineFwd_);
|
||||
}
|
||||
|
||||
void MkldnnLayer::resetBackwardFC(int bs,
|
||||
int ic,
|
||||
int ih,
|
||||
int iw,
|
||||
real* botDiff,
|
||||
real* botData,
|
||||
int oc,
|
||||
real* topDiff,
|
||||
real* wgtDiff,
|
||||
real* wgtData,
|
||||
real* biasDiff) {
|
||||
bool hasSpatial = ih == 1 && iw == 1 ? false : true;
|
||||
|
||||
// backward weight
|
||||
mem::desc botMD = hasSpatial ? createMD({bs, ic, ih, iw}, format::nchw)
|
||||
: createMD({bs, ic}, format::nc);
|
||||
mem::desc wgtMD = hasSpatial ? createMD({oc, ic, ih, iw}, format::oihw)
|
||||
: createMD({oc, ic}, format::oi);
|
||||
mem::desc topMD = createMD({bs, oc}, format::nc);
|
||||
mem::desc biasMD = biasDiff != NULL ? createMD({oc}, format::x)
|
||||
: createMD({}, format::format_undef);
|
||||
|
||||
mem::primitive_desc topPD = mem::primitive_desc(botMD, engine_);
|
||||
if (outGrad_ && outGrad_->get_primitive_desc() == topPD) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (inVal_) {
|
||||
// update data
|
||||
inVal_->set_data_handle(botData);
|
||||
} else {
|
||||
inVal_.reset(new mem(mem::primitive_desc(botMD, engine_), botData));
|
||||
}
|
||||
wgtGrad_.reset(new mem(mem::primitive_desc(wgtMD, engine_), wgtDiff));
|
||||
outGrad_.reset(new mem(topPD, topDiff));
|
||||
|
||||
fc_fwd::desc fwdDesc =
|
||||
fc_fwd::desc(mkldnn::prop_kind::forward, botMD, wgtMD, topMD);
|
||||
fc_fwd::primitive_desc fwdPD = fc_fwd::primitive_desc(fwdDesc, engine_);
|
||||
fc_bwdWgt::desc bwdWgtDesc =
|
||||
biasDiff != NULL ? fc_bwdWgt::desc(botMD, wgtMD, biasMD, topMD)
|
||||
: fc_bwdWgt::desc(botMD, wgtMD, topMD);
|
||||
fc_bwdWgt::primitive_desc bwdWgtPD =
|
||||
fc_bwdWgt::primitive_desc(bwdWgtDesc, engine_, fwdPD);
|
||||
|
||||
if (biasDiff != NULL) {
|
||||
biasGrad_.reset(new mem(mem::primitive_desc(biasMD, engine_), biasDiff));
|
||||
bwdWgt_.reset(
|
||||
new fc_bwdWgt(bwdWgtPD, *inVal_, *outGrad_, *wgtGrad_, *biasGrad_));
|
||||
} else {
|
||||
bwdWgt_.reset(new fc_bwdWgt(bwdWgtPD, *inVal_, *outGrad_, *wgtGrad_));
|
||||
}
|
||||
pipelineBwd_.clear();
|
||||
pipelineBwd_.push_back(*bwdWgt_);
|
||||
|
||||
// backward data
|
||||
if (botDiff == NULL) {
|
||||
return;
|
||||
}
|
||||
|
||||
fc_bwdData::desc bwdDataDesc = fc_bwdData::desc(botMD, wgtMD, topMD);
|
||||
fc_bwdData::primitive_desc bwdDataPD =
|
||||
fc_bwdData::primitive_desc(bwdDataDesc, engine_, fwdPD);
|
||||
inGrad_.reset(new mem(mem::primitive_desc(botMD, engine_), botDiff));
|
||||
if (wgtVal_) {
|
||||
// update data
|
||||
wgtVal_->set_data_handle(wgtData);
|
||||
} else {
|
||||
wgtVal_.reset(new mem(mem::primitive_desc(wgtMD, engine_), wgtData));
|
||||
}
|
||||
bwdData_.reset(new fc_bwdData(bwdDataPD, *outGrad_, *wgtVal_, *inGrad_));
|
||||
pipelineBwd_.push_back(*bwdData_);
|
||||
}
|
||||
|
||||
void MkldnnLayer::mkldnnBackwardFC(int bs,
|
||||
int ic,
|
||||
int ih,
|
||||
int iw,
|
||||
real* botDiff,
|
||||
real* botData,
|
||||
int oc,
|
||||
real* topDiff,
|
||||
real* wgtDiff,
|
||||
real* wgtData,
|
||||
real* biasDiff) {
|
||||
// if input size changed, reset it
|
||||
resetBackwardFC(bs,
|
||||
ic,
|
||||
ih,
|
||||
iw,
|
||||
botDiff,
|
||||
botData,
|
||||
oc,
|
||||
topDiff,
|
||||
wgtDiff,
|
||||
wgtData,
|
||||
biasDiff);
|
||||
|
||||
// update data
|
||||
outGrad_->set_data_handle(topDiff);
|
||||
|
||||
stream_->submit(pipelineBwd_);
|
||||
}
|
||||
|
||||
void MkldnnLayer::printSizeInfo() {
|
||||
VLOG(DNN_SIZES) << getName() << ": bs: " << bs_ << ", ic: " << ic_
|
||||
<< ", ih: " << ih_ << ", iw: " << iw_ << ", oc: " << oc_
|
||||
<< ", oh: " << oh_ << ", ow: " << ow_;
|
||||
}
|
||||
|
||||
mem::desc MkldnnLayer::createMD(mem::dims dims,
|
||||
mem::format fmt,
|
||||
mem::data_type type) {
|
||||
// TODO(TJ): isFmtSuppoted(fmt)
|
||||
return mem::desc(dims, type, fmt);
|
||||
}
|
||||
|
||||
} // namespace paddle
|
||||
Loading…
Reference in new issue