Tomasz Patejko
7 years ago
parent
a431f984f9
commit
192cc5dd32
@ -0,0 +1,189 @@
|
|||||||
|
/* Copyright (c) 2018 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 "paddle/fluid/framework/tensor.h"
|
||||||
|
#include "paddle/fluid/operators/lrn_op.h"
|
||||||
|
#include "paddle/fluid/platform/mkldnn_helper.h"
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace operators {
|
||||||
|
|
||||||
|
using paddle::framework::Tensor;
|
||||||
|
using paddle::platform::MKLDNNDeviceContext;
|
||||||
|
|
||||||
|
namespace {
|
||||||
|
mkldnn::algorithm LRNAlgorithm(const paddle::framework::ExecutionContext& ctx) {
|
||||||
|
mkldnn::algorithm algorithm = mkldnn::lrn_across_channels;
|
||||||
|
|
||||||
|
std::string algorithm_str = ctx.Attr<std::string>("algorithm");
|
||||||
|
if (algorithm_str == "WITHIN_CHANNEL") {
|
||||||
|
algorithm = mkldnn::lrn_within_channel;
|
||||||
|
}
|
||||||
|
return algorithm;
|
||||||
|
}
|
||||||
|
} // namespace
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
class LRNMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
|
||||||
|
public:
|
||||||
|
void Compute(const paddle::framework::ExecutionContext& ctx) const override {
|
||||||
|
PADDLE_ENFORCE(std::is_same<T, float>::value,
|
||||||
|
"MKLDNN LRN must use float data.");
|
||||||
|
PADDLE_ENFORCE(paddle::platform::is_cpu_place(ctx.GetPlace()),
|
||||||
|
"MKLDNN LRN must use CPUPlace.");
|
||||||
|
|
||||||
|
auto& dev_ctx = ctx.template device_context<MKLDNNDeviceContext>();
|
||||||
|
const auto& mkldnn_engine = dev_ctx.GetEngine();
|
||||||
|
|
||||||
|
auto x = ctx.Input<Tensor>("X");
|
||||||
|
auto out = ctx.Output<Tensor>("Out");
|
||||||
|
auto mid = ctx.Output<Tensor>("MidOut");
|
||||||
|
|
||||||
|
auto input_data = x->data<T>();
|
||||||
|
auto output_data = out->mutable_data<T>(ctx.GetPlace());
|
||||||
|
mid->mutable_data<T>(ctx.GetPlace());
|
||||||
|
|
||||||
|
const std::string key = ctx.op().Output("Out");
|
||||||
|
const std::string key_src_memory = key + "@lrn_src_memory";
|
||||||
|
const std::string key_pd = key + "@lrn_pd";
|
||||||
|
const std::string key_workspace_memory = key + "@lrn_workspace_memory";
|
||||||
|
|
||||||
|
const int n = ctx.Attr<int>("n");
|
||||||
|
const float alpha = ctx.Attr<float>("alpha");
|
||||||
|
const float beta = ctx.Attr<float>("beta");
|
||||||
|
const float k = ctx.Attr<float>("k");
|
||||||
|
|
||||||
|
auto algorithm = LRNAlgorithm(ctx);
|
||||||
|
|
||||||
|
auto e_mid = framework::EigenTensor<T, 4>::From(*mid);
|
||||||
|
e_mid = e_mid.constant(k);
|
||||||
|
|
||||||
|
auto dims = paddle::framework::vectorize2int(x->dims());
|
||||||
|
|
||||||
|
auto src_md = paddle::platform::MKLDNNMemDesc(
|
||||||
|
dims, mkldnn::memory::data_type::f32, mkldnn::memory::format::nchw);
|
||||||
|
|
||||||
|
auto dst_md = paddle::platform::MKLDNNMemDesc(
|
||||||
|
dims, mkldnn::memory::data_type::f32, mkldnn::memory::format::nchw);
|
||||||
|
|
||||||
|
auto forward_desc = mkldnn::lrn_forward::desc{
|
||||||
|
mkldnn::prop_kind::forward, algorithm, src_md, n, alpha, beta, k};
|
||||||
|
|
||||||
|
auto forward_pd = std::make_shared<mkldnn::lrn_forward::primitive_desc>(
|
||||||
|
forward_desc, mkldnn_engine);
|
||||||
|
|
||||||
|
dev_ctx.SetBlob(key_pd, forward_pd);
|
||||||
|
|
||||||
|
auto src_memory_pd = mkldnn::memory::primitive_desc{src_md, mkldnn_engine};
|
||||||
|
auto src_memory = std::make_shared<mkldnn::memory>(
|
||||||
|
src_memory_pd, static_cast<void*>(const_cast<float*>(input_data)));
|
||||||
|
|
||||||
|
dev_ctx.SetBlob(key_src_memory, src_memory);
|
||||||
|
auto dst_memory = mkldnn::memory{{dst_md, mkldnn_engine},
|
||||||
|
static_cast<void*>(output_data)};
|
||||||
|
|
||||||
|
auto workspace_md = forward_pd->workspace_primitive_desc();
|
||||||
|
auto workspace_memory = std::make_shared<mkldnn::memory>(workspace_md);
|
||||||
|
|
||||||
|
dev_ctx.SetBlob(key_workspace_memory, workspace_memory);
|
||||||
|
|
||||||
|
auto forward_op = mkldnn::lrn_forward{*forward_pd, *src_memory,
|
||||||
|
*workspace_memory, dst_memory};
|
||||||
|
|
||||||
|
std::vector<mkldnn::primitive> pipeline = {forward_op};
|
||||||
|
mkldnn::stream(mkldnn::stream::kind::eager).submit(pipeline).wait();
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
class LRNMKLDNNGradOpKernel : public paddle::framework::OpKernel<T> {
|
||||||
|
public:
|
||||||
|
void Compute(const paddle::framework::ExecutionContext& ctx) const override {
|
||||||
|
PADDLE_ENFORCE(std::is_same<T, float>::value,
|
||||||
|
"MKLDNN LRN must use float data.");
|
||||||
|
PADDLE_ENFORCE(paddle::platform::is_cpu_place(ctx.GetPlace()),
|
||||||
|
"MKLDNN LRN must use CPUPlace.");
|
||||||
|
|
||||||
|
auto x = ctx.Input<Tensor>("X");
|
||||||
|
|
||||||
|
auto out_grad = ctx.Input<Tensor>(framework::GradVarName("Out"));
|
||||||
|
auto x_grad = ctx.Output<Tensor>(framework::GradVarName("X"));
|
||||||
|
|
||||||
|
const std::string key = ctx.op().Input("Out");
|
||||||
|
const std::string key_src_memory = key + "@lrn_src_memory";
|
||||||
|
const std::string key_pd = key + "@lrn_pd";
|
||||||
|
const std::string key_workspace_memory = key + "@lrn_workspace_memory";
|
||||||
|
|
||||||
|
const int n = ctx.Attr<int>("n");
|
||||||
|
const float alpha = ctx.Attr<float>("alpha");
|
||||||
|
const float beta = ctx.Attr<float>("beta");
|
||||||
|
const float k = ctx.Attr<float>("k");
|
||||||
|
|
||||||
|
auto& dev_ctx = ctx.template device_context<MKLDNNDeviceContext>();
|
||||||
|
const auto& mkldnn_engine = dev_ctx.GetEngine();
|
||||||
|
|
||||||
|
auto x_grad_data = x_grad->mutable_data<T>(ctx.GetPlace());
|
||||||
|
auto out_grad_data = out_grad->data<T>();
|
||||||
|
|
||||||
|
auto dims = paddle::framework::vectorize2int(x->dims());
|
||||||
|
|
||||||
|
auto src_md = paddle::platform::MKLDNNMemDesc(
|
||||||
|
dims, mkldnn::memory::data_type::f32, mkldnn::memory::format::nchw);
|
||||||
|
|
||||||
|
auto diff_src_md = paddle::platform::MKLDNNMemDesc(
|
||||||
|
dims, mkldnn::memory::data_type::f32, mkldnn::memory::format::nchw);
|
||||||
|
|
||||||
|
auto diff_dst_md = paddle::platform::MKLDNNMemDesc(
|
||||||
|
dims, mkldnn::memory::data_type::f32, mkldnn::memory::format::nchw);
|
||||||
|
|
||||||
|
auto diff_dst_memory =
|
||||||
|
mkldnn::memory{{diff_dst_md, mkldnn_engine},
|
||||||
|
static_cast<void*>(const_cast<float*>(out_grad_data))};
|
||||||
|
|
||||||
|
auto diff_src_memory = mkldnn::memory{{diff_src_md, mkldnn_engine},
|
||||||
|
static_cast<void*>(x_grad_data)};
|
||||||
|
|
||||||
|
auto algorithm = LRNAlgorithm(ctx);
|
||||||
|
|
||||||
|
auto backward_desc = mkldnn::lrn_backward::desc{
|
||||||
|
algorithm, src_md, diff_src_md, n, alpha, beta, k};
|
||||||
|
|
||||||
|
auto forward_pd = dev_ctx.GetBlob(key_pd);
|
||||||
|
|
||||||
|
auto backward_pd = mkldnn::lrn_backward::primitive_desc{
|
||||||
|
backward_desc, mkldnn_engine,
|
||||||
|
*static_cast<mkldnn::lrn_forward::primitive_desc*>(forward_pd.get())};
|
||||||
|
|
||||||
|
std::shared_ptr<void> workspace_memory =
|
||||||
|
dev_ctx.GetBlob(key_workspace_memory);
|
||||||
|
|
||||||
|
auto src_memory = dev_ctx.GetBlob(key_src_memory);
|
||||||
|
auto backward_op = mkldnn::lrn_backward{
|
||||||
|
backward_pd, *static_cast<mkldnn::memory*>(src_memory.get()),
|
||||||
|
diff_dst_memory, *static_cast<mkldnn::memory*>(workspace_memory.get()),
|
||||||
|
diff_src_memory};
|
||||||
|
|
||||||
|
std::vector<mkldnn::primitive> pipeline = {backward_op};
|
||||||
|
mkldnn::stream(mkldnn::stream::kind::eager).submit(pipeline).wait();
|
||||||
|
}
|
||||||
|
};
|
||||||
|
} // namespace operators
|
||||||
|
} // namespace paddle
|
||||||
|
|
||||||
|
namespace ops = paddle::operators;
|
||||||
|
|
||||||
|
REGISTER_OP_KERNEL(lrn, MKLDNN, paddle::platform::CPUPlace,
|
||||||
|
ops::LRNMKLDNNOpKernel<float>);
|
||||||
|
REGISTER_OP_KERNEL(lrn_grad, MKLDNN, paddle::platform::CPUPlace,
|
||||||
|
ops::LRNMKLDNNGradOpKernel<float>);
|
||||||
Loading…
Reference in new issue