You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Paddle/paddle/fluid/operators/inplace_abn_op.h

118 lines
4.2 KiB

// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
//
// 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.
#pragma once
#include <string>
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/operators/activation_op.h"
#include "paddle/fluid/operators/math/math_function.h"
namespace paddle {
namespace operators {
using Tensor = framework::Tensor;
template <typename T, int MajorType = Eigen::RowMajor,
typename IndexType = Eigen::DenseIndex>
using EigenMatrix = framework::EigenMatrix<T, MajorType, IndexType>;
template <typename T, int MajorType = Eigen::RowMajor,
typename IndexType = Eigen::DenseIndex>
using EigenVector = framework::EigenVector<T, MajorType, IndexType>;
enum InplaceABNActivationType { identity = 0, leakyrelu = 1, elu = 2 };
inline InplaceABNActivationType GetInplaceABNActivationType(
const std::string& type) {
if (type == "leaky_relu") {
return InplaceABNActivationType::leakyrelu;
} else if (type == "elu") {
return InplaceABNActivationType::elu;
} else if (type == "identity" || type == "") {
return InplaceABNActivationType::identity;
} else {
PADDLE_THROW(platform::errors::InvalidArgument(
"unsupported activation type %s for Op(inplace_abn)", type));
}
}
template <typename DeviceContext, typename T>
class InplaceABNActivation {
private:
template <typename Functor>
void setAttrs(const framework::ExecutionContext& ctx, Functor* functor) {
auto attrs = functor->GetAttrs();
for (auto& attr : attrs) {
*attr.second = ctx.Attr<float>(attr.first);
}
}
template <typename Functor, typename... Args>
void compute(const framework::ExecutionContext& ctx, Functor* functor,
Args... args) {
setAttrs(ctx, functor);
(*functor)(args...);
}
public:
template <typename Device, typename X, typename Y>
void Compute(const framework::ExecutionContext& ctx, const int act_type,
const Device& d, X x, Y y) {
if (act_type == InplaceABNActivationType::identity) {
y.device(d) = x;
} else if (act_type == InplaceABNActivationType::leakyrelu) {
LeakyReluFunctor<T> functor;
compute(ctx, &functor, d, x, y);
} else if (act_type == InplaceABNActivationType::elu) {
ELUFunctor<T> functor;
compute(ctx, &functor, d, x, y);
} else {
PADDLE_THROW(
platform::errors::InvalidArgument("unsupported activation type"));
}
}
template <typename Device, typename X, typename Y, typename DX, typename DY>
void GradCompute(const framework::ExecutionContext& ctx, const int act_type,
const Device& d, X x, Y y, DX dx, DY dy) {
const float alpha = ctx.Attr<float>("alpha");
if (act_type == InplaceABNActivationType::identity) {
x.device(d) = y;
dx.device(d) = dy;
} else if (act_type == InplaceABNActivationType::leakyrelu) {
auto temp1 = (y < static_cast<T>(0)).template cast<T>().eval() /
static_cast<T>(alpha);
auto temp2 = (y >= static_cast<T>(0)).template cast<T>().eval();
x.device(d) = y * (temp1 + temp2).template cast<T>();
LeakyReluGradFunctor<T> functor;
compute(ctx, &functor, d, x, y, dy, dx);
} else if (act_type == InplaceABNActivationType::elu) {
auto temp1 = (y >= static_cast<T>(0)).template cast<T>().eval();
auto temp = (y < static_cast<T>(0)).template cast<T>().eval();
auto temp2 = (y * temp / static_cast<T>(alpha) + static_cast<T>(1)).log();
x.device(d) = (y * temp1 + temp2).template cast<T>();
ELUGradFunctor<T> functor;
compute(ctx, &functor, d, x, y, dy, dx);
} else {
PADDLE_THROW(
platform::errors::InvalidArgument("unsupported activation type"));
}
}
};
} // namespace operators
} // namespace paddle