/* Copyright (c) 2016 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 #include "glog/logging.h" #include "paddle/fluid/framework/eigen.h" #include "paddle/fluid/framework/op_registry.h" namespace paddle { namespace operators { using Tensor = framework::Tensor; using DDim = framework::DDim; template using EigenTensor = framework::EigenTensor; template using EigenScalar = framework::EigenScalar; template using EigenVector = framework::EigenVector; struct SumFunctor { template void operator()(const DeviceContext& place, X* x, Y* y, const Dim& dim) { y->device(place) = x->sum(dim); } }; struct SumGradFunctor { template void operator()(const DeviceContext& place, X* x, Y* y, DX* dx, DY* dy, const Dim& dim, int size) { dx->device(place) = dy->broadcast(dim); } }; struct MeanFunctor { template void operator()(const DeviceContext& place, X* x, Y* y, const Dim& dim) { y->device(place) = x->mean(dim); } }; struct MeanGradFunctor { template void operator()(const DeviceContext& place, X* x, Y* y, DX* dx, DY* dy, const Dim& dim, int size) { dx->device(place) = dy->broadcast(dim) / dx->constant(size); } }; struct MaxFunctor { template void operator()(const DeviceContext& place, X* x, Y* y, const Dim& dim) { y->device(place) = x->maximum(dim); } }; struct MinFunctor { template void operator()(const DeviceContext& place, X* x, Y* y, const Dim& dim) { y->device(place) = x->minimum(dim); } }; struct MaxOrMinGradFunctor { template void operator()(const DeviceContext& place, X* x, Y* y, DX* dx, DY* dy, const Dim& dim, int size) { auto equals = (*x) == y->broadcast(dim); auto ones = dx->constant(1); auto zeros = dx->constant(0); // If there are multiple minimum or maximum elements, the subgradient of // each is the set [0, 1], and we pass gradient to all of them here. dx->device(place) = dy->broadcast(dim) * equals.select(ones, zeros); } }; struct ProdFunctor { template void operator()(const DeviceContext& place, X* x, Y* y, const Dim& dim) { y->device(place) = x->prod(dim); } }; struct ProdGradFunctor { template void operator()(const DeviceContext& place, X* x, Y* y, DX* dx, DY* dy, const Dim& dim, int size) { dx->device(place) = dy->broadcast(dim) * y->broadcast(dim) * x->inverse(); } }; #define HANDLE_DIM(NDIM, RDIM) \ if (ndim == NDIM && rdim == RDIM) { \ ReduceCompute(context); \ } template class ReduceKernel : public framework::OpKernel { public: void Compute(const framework::ExecutionContext& context) const override { bool reduce_all = context.Attr("reduce_all"); if (reduce_all) { // Flatten and reduce 1-D tensor auto* input = context.Input("X"); auto* output = context.Output("Out"); output->mutable_data(context.GetPlace()); auto x = EigenVector::Flatten(*input); auto out = EigenScalar::From(*output); auto& place = *context.template device_context().eigen_device(); auto reduce_dim = Eigen::array({{0}}); Functor functor; functor(place, &x, &out, reduce_dim); } else { int ndim = context.Input("X")->dims().size(); int rdim = context.Attr>("dim").size(); HANDLE_DIM(6, 5); HANDLE_DIM(6, 4); HANDLE_DIM(6, 3); HANDLE_DIM(6, 2); HANDLE_DIM(6, 1); HANDLE_DIM(5, 4); HANDLE_DIM(5, 3); HANDLE_DIM(5, 2); HANDLE_DIM(5, 1); HANDLE_DIM(4, 3); HANDLE_DIM(4, 2); HANDLE_DIM(4, 1); HANDLE_DIM(3, 2); HANDLE_DIM(3, 1); HANDLE_DIM(2, 1); HANDLE_DIM(1, 1); } } private: template void ReduceCompute(const framework::ExecutionContext& context) const { auto* input = context.Input("X"); auto* output = context.Output("Out"); output->mutable_data(context.GetPlace()); auto x = EigenTensor::From(*input); auto x_rank = static_cast(x.dimensions().size()); auto dims = context.Attr>("dim"); auto reduce_dim = Eigen::array(); for (size_t i = 0; i < dims.size(); ++i) { if (dims[i] < 0) dims[i] = x_rank + dims[i]; reduce_dim[i] = dims[i]; } // construct the squeezed output tensor bool keep_dim = context.Attr("keep_dim"); DDim out_dims = output->dims(); if (keep_dim && x_rank > 1) { const int kDelFlag = -2; auto dims_vector = vectorize(out_dims); for (size_t i = 0; i < dims.size(); ++i) { dims_vector[dims[i]] = kDelFlag; } dims_vector.erase( remove(dims_vector.begin(), dims_vector.end(), kDelFlag), dims_vector.end()); out_dims = framework::make_ddim(dims_vector); } auto& place = *context.template device_context().eigen_device(); Functor functor; if (D == 1) { auto out = EigenScalar::From(*output); functor(place, &x, &out, reduce_dim); } else { auto out = EigenTensor::From(*output, out_dims); functor(place, &x, &out, reduce_dim); } } }; template class ReduceGradKernel : public framework::OpKernel { public: void Compute(const framework::ExecutionContext& context) const override { bool reduce_all = context.Attr("reduce_all"); if (reduce_all) { auto* input0 = context.Input("X"); auto* input1 = context.Input("Out"); auto* input2 = context.Input(framework::GradVarName("Out")); auto* output = context.Output(framework::GradVarName("X")); output->mutable_data(context.GetPlace()); auto x = EigenVector::Flatten(*input0); auto x_reduce = EigenVector::From(*input1); auto x_reduce_grad = EigenVector::From(*input2); auto x_grad = EigenVector::Flatten(*output); auto& place = *context.template device_context().eigen_device(); auto broadcast_dim = Eigen::array({{static_cast(input0->numel())}}); Functor functor; functor(place, &x, &x_reduce, &x_grad, &x_reduce_grad, broadcast_dim, broadcast_dim[0]); } else { int rank = context.Input("X")->dims().size(); switch (rank) { case 1: ReduceGradCompute<1>(context); break; case 2: ReduceGradCompute<2>(context); break; case 3: ReduceGradCompute<3>(context); break; case 4: ReduceGradCompute<4>(context); break; case 5: ReduceGradCompute<5>(context); break; case 6: ReduceGradCompute<6>(context); break; } } } private: template void ReduceGradCompute(const framework::ExecutionContext& context) const { auto* input0 = context.Input("X"); auto* input1 = context.Input("Out"); auto* input2 = context.Input(framework::GradVarName("Out")); auto* output = context.Output(framework::GradVarName("X")); output->mutable_data(context.GetPlace()); auto x = EigenTensor::From(*input0); auto x_grad = EigenTensor::From(*output); auto x_rank = static_cast(x.dimensions().size()); auto dims = context.Attr>("dim"); auto x_dims = input0->dims(); auto reduced_dims_v = vectorize(x_dims); Eigen::array broadcast_dim; for (size_t i = 0; i < D; ++i) broadcast_dim[i] = 1; int broad_cats_times = 1; for (size_t i = 0; i < dims.size(); ++i) { if (dims[i] < 0) dims[i] = x_rank + dims[i]; reduced_dims_v[dims[i]] = 1; broadcast_dim[dims[i]] = x_dims[dims[i]]; broad_cats_times *= x_dims[dims[i]]; } auto reduced_dims = framework::make_ddim(reduced_dims_v); auto x_reduce = EigenTensor::From(*input1, reduced_dims); auto x_reduce_grad = EigenTensor::From(*input2, reduced_dims); auto& place = *context.template device_context().eigen_device(); Functor functor; functor(place, &x, &x_reduce, &x_grad, &x_reduce_grad, broadcast_dim, broad_cats_times); } }; } // namespace operators } // namespace paddle #define FOR_EACH_KERNEL_FUNCTOR(__macro) \ __macro(reduce_sum, SumFunctor, SumGradFunctor); \ __macro(reduce_mean, MeanFunctor, MeanGradFunctor); \ __macro(reduce_max, MaxFunctor, MaxOrMinGradFunctor); \ __macro(reduce_min, MinFunctor, MaxOrMinGradFunctor); \ __macro(reduce_prod, ProdFunctor, ProdGradFunctor);