arg_min_max_op_base.h 9.8 KB
Newer Older
S
sneaxiy 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
/* Copyright (c) 2018 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
Y
yuyang18 已提交
16
#include <string>
S
sneaxiy 已提交
17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40
#include <type_traits>
#include <vector>
#include "paddle/fluid/framework/ddim.h"
#include "paddle/fluid/framework/eigen.h"
#include "paddle/fluid/framework/lod_tensor.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/operator.h"
#include "paddle/fluid/platform/enforce.h"
#include "paddle/fluid/string/printf.h"

namespace paddle {
namespace operators {

enum ArgMinMaxType { kArgMin, kArgMax };

template <typename DeviceContext, typename T, typename Tout, int64_t Rank,
          ArgMinMaxType argMinMaxValue>
struct ArgMinMaxFunctor {};

#define DECLARE_ARG_MIN_MAX_FUNCTOR(eigen_op_type, enum_argminmax_value)      \
  template <typename DeviceContext, typename T, typename Tout, int64_t Rank>  \
  struct ArgMinMaxFunctor<DeviceContext, T, Tout, Rank,                       \
                          enum_argminmax_value> {                             \
    void operator()(const DeviceContext& ctx, const framework::LoDTensor& in, \
W
wawltor 已提交
41 42 43
                    framework::LoDTensor* out, framework::DDim x_dims,        \
                    int64_t axis, bool keepdims) {                            \
      auto in_eigen = framework::EigenTensor<T, Rank>::From(in, x_dims);      \
44 45 46 47 48 49 50 51 52
      if (keepdims) {                                                         \
        auto out_eigen = framework::EigenTensor<Tout, Rank>::From(*out);      \
        out_eigen.device(*(ctx.eigen_device())) =                             \
            in_eigen.eigen_op_type(axis).template cast<Tout>();               \
      } else {                                                                \
        auto out_eigen = framework::EigenTensor<Tout, Rank - 1>::From(*out);  \
        out_eigen.device(*(ctx.eigen_device())) =                             \
            in_eigen.eigen_op_type(axis).template cast<Tout>();               \
      }                                                                       \
S
sneaxiy 已提交
53 54 55 56 57 58
    }                                                                         \
  }

DECLARE_ARG_MIN_MAX_FUNCTOR(argmin, ArgMinMaxType::kArgMin);
DECLARE_ARG_MIN_MAX_FUNCTOR(argmax, ArgMinMaxType::kArgMax);

59 60 61 62 63 64 65 66
template <typename DeviceContext, typename T, ArgMinMaxType EnumArgMinMaxValue>
struct VisitDataArgMinMaxFunctor {
  const framework::ExecutionContext& ctx;

  explicit VisitDataArgMinMaxFunctor(const framework::ExecutionContext& ctx)
      : ctx(ctx) {}
  template <typename Tout>
  void apply() const {
S
sneaxiy 已提交
67 68
    auto& x = *(ctx.Input<framework::LoDTensor>("X"));
    auto& out = *(ctx.Output<framework::LoDTensor>("Out"));
69
    out.template mutable_data<Tout>(ctx.GetPlace());
S
sneaxiy 已提交
70
    auto axis = ctx.Attr<int64_t>("axis");
71
    auto keepdims = ctx.Attr<bool>("keepdims");
W
wawltor 已提交
72
    const bool& flatten = ctx.Attr<bool>("flatten");
73 74
    // paddle do not have the scalar tensor, just return the shape [1] tensor
    if (flatten) keepdims = true;
W
wawltor 已提交
75 76 77 78 79 80 81 82 83 84 85

    // if flatten, will construct the new dims for the cacluate
    framework::DDim x_dims;
    if (flatten) {
      x_dims = framework::make_ddim({x.numel()});
      // if flatten, the axis just as 0
      axis = 0;
    } else {
      x_dims = x.dims();
      if (axis < 0) axis += x_dims.size();
    }
S
sneaxiy 已提交
86 87 88 89 90
    auto& dev_ctx = ctx.template device_context<DeviceContext>();

#define CALL_ARG_MINMAX_FUNCTOR(rank)                                \
  ArgMinMaxFunctor<DeviceContext, T, Tout, rank, EnumArgMinMaxValue> \
      functor##rank;                                                 \
W
wawltor 已提交
91
  functor##rank(dev_ctx, x, &out, x_dims, axis, keepdims)
S
sneaxiy 已提交
92

W
wawltor 已提交
93
    switch (x_dims.size()) {
S
sneaxiy 已提交
94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112
      case 1:
        CALL_ARG_MINMAX_FUNCTOR(1);
        break;
      case 2:
        CALL_ARG_MINMAX_FUNCTOR(2);
        break;
      case 3:
        CALL_ARG_MINMAX_FUNCTOR(3);
        break;
      case 4:
        CALL_ARG_MINMAX_FUNCTOR(4);
        break;
      case 5:
        CALL_ARG_MINMAX_FUNCTOR(5);
        break;
      case 6:
        CALL_ARG_MINMAX_FUNCTOR(6);
        break;
      default:
113 114 115 116 117 118
        PADDLE_ENFORCE_LE(
            x_dims.size(), 6,
            platform::errors::InvalidArgument(
                "%s operator doesn't supports tensors whose ranks are greater "
                "than 6.",
                (EnumArgMinMaxValue == kArgMin ? "argmin" : "argmax")));
S
sneaxiy 已提交
119
        break;
Y
yuyang18 已提交
120
#undef CALL_ARG_MINMAX_FUNCTOR
S
sneaxiy 已提交
121 122 123 124
    }
  }
};

125 126 127 128 129 130
template <typename DeviceContext, typename T, ArgMinMaxType EnumArgMinMaxValue>
class ArgMinMaxKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& ctx) const override {
    auto& dtype = ctx.Attr<int>("dtype");
    if (dtype < 0) {
131
      framework::VisitDataTypeTiny(
132 133 134 135 136
          static_cast<framework::proto::VarType::Type>(
              framework::proto::VarType::INT64),
          VisitDataArgMinMaxFunctor<DeviceContext, T, EnumArgMinMaxValue>(ctx));
      return;
    }
137
    framework::VisitDataTypeTiny(
138 139 140 141 142
        static_cast<framework::proto::VarType::Type>(dtype),
        VisitDataArgMinMaxFunctor<DeviceContext, T, EnumArgMinMaxValue>(ctx));
  }
};

Y
yuyang18 已提交
143
template <typename DeviceContext, typename T>
144
using ArgMinKernel = ArgMinMaxKernel<DeviceContext, T, ArgMinMaxType::kArgMin>;
S
sneaxiy 已提交
145

Y
yuyang18 已提交
146
template <typename DeviceContext, typename T>
147
using ArgMaxKernel = ArgMinMaxKernel<DeviceContext, T, ArgMinMaxType::kArgMax>;
S
sneaxiy 已提交
148

Y
yuyang18 已提交
149
class ArgMinMaxOp : public framework::OperatorWithKernel {
S
sneaxiy 已提交
150 151 152 153
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

  void InferShape(framework::InferShapeContext* ctx) const override {
154 155
    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "arg_min_max");
    OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "arg_min_max");
S
sneaxiy 已提交
156 157
    const auto& x_dims = ctx->GetInputDim("X");
    int64_t axis = ctx->Attrs().Get<int64_t>("axis");
158
    bool keepdims = ctx->Attrs().Get<bool>("keepdims");
W
wawltor 已提交
159
    const bool& flatten = ctx->Attrs().Get<bool>("flatten");
160

161 162 163 164 165 166 167 168
    PADDLE_ENFORCE_GE(axis, -x_dims.size(),
                      platform::errors::InvalidArgument(
                          "'axis'(%d) must be greater than or equal to"
                          " -Rank(X)(%d).",
                          axis, -x_dims.size()));
    PADDLE_ENFORCE_LT(
        axis, x_dims.size(),
        platform::errors::InvalidArgument(
169 170
            "'axis'(%d) must be less than Rank(X)(%d) of Input(X).", axis,
            x_dims.size()));
S
sneaxiy 已提交
171

172 173 174 175 176 177 178 179 180 181 182 183 184
    const int& dtype = ctx->Attrs().Get<int>("dtype");
    PADDLE_ENFORCE_EQ(
        (dtype < 0 || dtype == 2 || dtype == 3), true,
        platform::errors::InvalidArgument(
            "The attribute of dtype in argmin/argmax must be [%s] or [%s], but "
            "received [%s]",
            paddle::framework::DataTypeToString(
                framework::proto::VarType::INT32),
            paddle::framework::DataTypeToString(
                framework::proto::VarType::INT64),
            paddle::framework::DataTypeToString(
                static_cast<framework::proto::VarType::Type>(dtype))));

185 186 187 188 189 190 191 192 193 194 195 196 197
    auto x_rank = x_dims.size();
    if (axis < 0) axis += x_rank;
    if (ctx->IsRuntime()) {
      if (dtype == framework::proto::VarType::INT32) {
        int64_t all_element_num = 0;
        if (flatten) {
          all_element_num = framework::product(x_dims);

        } else {
          all_element_num = x_dims[axis];
        }
        PADDLE_ENFORCE_LE(
            all_element_num, INT_MAX,
198 199 200 201 202
            platform::errors::InvalidArgument(
                "The element num of the argmin/argmax input at axis is "
                "%d, is larger than int32 maximum value:%d, you must "
                "set the dtype of argmin/argmax to 'int64'.",
                all_element_num, INT_MAX));
203 204
      }
    }
S
sneaxiy 已提交
205
    std::vector<int64_t> vec;
W
wawltor 已提交
206
    if (flatten) {
207
      vec.emplace_back(static_cast<int64_t>(1));
W
wawltor 已提交
208 209 210 211 212 213
    } else {
      for (int64_t i = 0; i < axis; i++) vec.emplace_back(x_dims[i]);
      if (keepdims) {
        vec.emplace_back(static_cast<int64_t>(1));
      }
      for (int64_t i = axis + 1; i < x_rank; i++) vec.emplace_back(x_dims[i]);
214
    }
S
sneaxiy 已提交
215 216
    ctx->SetOutputDim("Out", framework::make_ddim(vec));
  }
Y
yuyang18 已提交
217
};
S
sneaxiy 已提交
218 219 220 221 222 223 224 225 226 227 228

class BaseArgMinMaxOpMaker : public framework::OpProtoAndCheckerMaker {
 protected:
  virtual const char* OpName() const = 0;
  virtual const char* Name() const = 0;

 public:
  void Make() override {
    AddInput("X", "Input tensor.");
    AddOutput("Out", "Output tensor.");
    AddAttr<int64_t>("axis", "The axis in which to compute the arg indics.");
229
    AddAttr<bool>("keepdims", "Keep the dim that to reduce.").SetDefault(false);
W
wawltor 已提交
230 231 232
    AddAttr<bool>("flatten",
                  "Flatten the input value, and search the min or max indices")
        .SetDefault(false);
233 234 235 236 237
    AddAttr<int>("dtype",
                 "(int, 3), the dtype of indices, the indices dtype must be "
                 "int32, int64."
                 "default dtype is int64, and proto value is 3.")
        .SetDefault(3);
Y
yuyang18 已提交
238 239
    AddComment(string::Sprintf(R"DOC(
      %s Operator.
S
sneaxiy 已提交
240

Y
yuyang18 已提交
241 242
      Computes the indices of the %s elements of the input tensor's element
      along the provided axis.
S
sneaxiy 已提交
243
)DOC",
Y
yuyang18 已提交
244
                               OpName(), Name()));
S
sneaxiy 已提交
245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260
  }
};

class ArgMinOpMaker : public BaseArgMinMaxOpMaker {
 protected:
  const char* OpName() const override { return "ArgMin"; }
  const char* Name() const override { return "min"; }
};

class ArgMaxOpMaker : public BaseArgMinMaxOpMaker {
 protected:
  const char* OpName() const override { return "ArgMax"; }
  const char* Name() const override { return "max"; }
};
}  // namespace operators
}  // namespace paddle