nce_op.cc 10.4 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
W
wanghaoshuang 已提交
2

W
wanghaoshuang 已提交
3 4 5
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
W
wanghaoshuang 已提交
6

W
wanghaoshuang 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
W
wanghaoshuang 已提交
8

W
wanghaoshuang 已提交
9 10 11 12 13
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. */
W
wanghaoshuang 已提交
14

Y
Yi Wang 已提交
15
#include "paddle/fluid/operators/nce_op.h"
W
wanghaoshuang 已提交
16

17
#include <string>
Y
Yang Yang 已提交
18 19
#include <vector>

W
wanghaoshuang 已提交
20 21 22 23 24 25 26 27 28
namespace paddle {
namespace operators {

using framework::Tensor;

class NCEOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

29
  void InferShape(framework::InferShapeContext *ctx) const override {
W
wanghaoshuang 已提交
30
    PADDLE_ENFORCE(ctx->HasInput("Input"));
W
wanghaoshuang 已提交
31
    PADDLE_ENFORCE(ctx->HasInput("Label"));
W
wanghaoshuang 已提交
32 33
    PADDLE_ENFORCE(ctx->HasInput("Weight"));
    PADDLE_ENFORCE(ctx->HasOutput("Cost"));
W
wanghaoshuang 已提交
34 35 36
    PADDLE_ENFORCE(ctx->HasOutput("SampleLogits"));
    PADDLE_ENFORCE(ctx->HasOutput("SampleLabels"));

W
wanghaoshuang 已提交
37
    auto x_dims = ctx->GetInputDim("Input");
W
wanghaoshuang 已提交
38
    auto label_dims = ctx->GetInputDim("Label");
39
    auto w_dims = ctx->GetInputDim("Weight");
W
wanghaoshuang 已提交
40
    PADDLE_ENFORCE_EQ(x_dims[0], label_dims[0]);
W
wanghaoshuang 已提交
41 42 43 44
    int num_true_classes = label_dims.size() == 2 ? label_dims[1] : 1;
    if (ctx->HasInput("Bias")) {
      PADDLE_ENFORCE_EQ(ctx->GetInputDim("Weight")[0],
                        ctx->GetInputDim("Bias")[0]);
W
wanghaoshuang 已提交
45
    }
W
wanghaoshuang 已提交
46 47
    auto num_neg_samples = ctx->Attrs().Get<int>("num_neg_samples");
    auto num_total_classes = ctx->Attrs().Get<int>("num_total_classes");
W
wanghaoshuang 已提交
48 49
    std::vector<int> custom_neg_classes =
        ctx->Attrs().Get<std::vector<int>>("custom_neg_classes");
W
wanghaoshuang 已提交
50
    PADDLE_ENFORCE_EQ(num_total_classes, ctx->GetInputDim("Weight")[0]);
W
wanghaoshuang 已提交
51 52
    if (custom_neg_classes.size() > 0) {
      PADDLE_ENFORCE_EQ(custom_neg_classes.size(),
W
wanghaoshuang 已提交
53
                        static_cast<size_t>(num_neg_samples));
W
wanghaoshuang 已提交
54
    }
W
wanghaoshuang 已提交
55
    // set dims of output(Out)
W
wanghaoshuang 已提交
56
    std::vector<int64_t> out_dims;
W
wanghaoshuang 已提交
57
    out_dims.push_back(x_dims[0]);
W
wanghaoshuang 已提交
58
    out_dims.push_back(1);
W
wanghaoshuang 已提交
59
    ctx->SetOutputDim("Cost", framework::make_ddim(out_dims));
W
wanghaoshuang 已提交
60 61

    // set dims of output(SampleOut)
W
wanghaoshuang 已提交
62
    std::vector<int64_t> sample_out_dims;
W
wanghaoshuang 已提交
63
    sample_out_dims.push_back(x_dims[0]);
W
wanghaoshuang 已提交
64
    sample_out_dims.push_back(num_neg_samples + num_true_classes);
W
wanghaoshuang 已提交
65 66 67
    ctx->SetOutputDim("SampleLogits", framework::make_ddim(sample_out_dims));
    ctx->SetOutputDim("SampleLabels", framework::make_ddim(sample_out_dims));
  }
W
wanghaoshuang 已提交
68 69

 protected:
70
  framework::OpKernelType GetExpectedKernelType(
71
      const framework::ExecutionContext &ctx) const override {
W
wanghaoshuang 已提交
72 73
    return framework::OpKernelType(
        framework::ToDataType(ctx.Input<Tensor>("Input")->type()),
C
chengduo 已提交
74
        platform::CPUPlace());
W
wanghaoshuang 已提交
75
  }
W
wanghaoshuang 已提交
76 77 78 79
};

class NCEOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
80
  void Make() override {
W
wanghaoshuang 已提交
81
    AddInput("Input", "(Tensor) A tensor of shape [batch_size, dim].");
W
wanghaoshuang 已提交
82 83 84 85 86 87 88 89
    AddInput(
        "Label",
        "(Tensor) A tensor of shape [batch_size, num_true_class]. "
        "'num_true_class' is the number of target classes in each sample."
        "The number of target classes per sample should be same. "
        "If you have a variable number of target classes, "
        "you can pad them out to a constant number by either repeating them"
        " or by padding with an otherwise unused class.)");
W
wanghaoshuang 已提交
90 91 92
    AddInput("Weight",
             "(Tensor) A tensor of shape [num_class, dim]. 'num_class' is the "
             "total number of class.");
W
wanghaoshuang 已提交
93 94 95 96
    AddInput(
        "Bias",
        "(Tensor) A tensor of shape [num_class, 1]. 'num_class' is the total "
        "number of class. It is a dispensable input.")
W
wanghaoshuang 已提交
97 98
        .AsDispensable();
    AddInput("SampleWeight",
W
wanghaoshuang 已提交
99
             "(Tensor) A tensor of shape [batch_size, 1] storing a weight for "
W
wanghaoshuang 已提交
100 101 102
             "each sample. And it is a dispensable input. The default value of "
             "sample is 1.")
        .AsDispensable();
103 104

    AddInput(
105
        "CustomDistProbs",
106 107 108 109
        "(Tensor) It is used in 'CostumDist' sampler. "
        "It is a tensor with shape [num_total_classes]."
        "The i-th element is the probsbility of the i-th class being sampled.")
        .AsDispensable();
110 111 112 113 114 115 116 117 118 119 120 121 122
    AddInput(
        "CustomDistAlias",
        "(Tensor) It is used in 'CostumDist' sampler. "
        "It is a tensor with shape [num_total_classes]."
        "The i-th element is the probsbility of the i-th class being sampled.")
        .AsDispensable();
    AddInput(
        "CustomDistAliasProbs",
        "(Tensor) It is used in 'CostumDist' sampler. "
        "It is a tensor with shape [num_total_classes]."
        "The i-th element is the probsbility of the i-th class being sampled.")
        .AsDispensable();

W
wanghaoshuang 已提交
123
    AddOutput("Cost",
W
wanghaoshuang 已提交
124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140
              "(Tensor) A tensor of shape [batch_size, 1]. Cost of samples.");
    AddOutput("SampleLogits",
              "An intermediate tensor of shape[batch_size, num_neg_samples + "
              "num_pos_samples]."
              "This tensor is output of forward kernel and used in backward "
              "kernel to compute grads."
              "Given X is  the dot product of input tensor and sampled labels' "
              "weights."
              "Then 'SampleLogits' is sigmoid(X).")
        .AsIntermediate();
    AddOutput("SampleLabels",
              "An intermediate tensor of shape[batch_size, num_neg_samples + "
              "num_pos_samples]."
              "This tensor is output of forward kernel and used in backward "
              "kernel to compute grads."
              "")
        .AsIntermediate();
141

W
wanghaoshuang 已提交
142 143 144 145
    AddAttr<int>("num_total_classes",
                 "Total number of classes in all samples.");
    AddAttr<int>("num_neg_samples",
                 "The number of negative classes. The default value is 10.")
W
wanghaoshuang 已提交
146
        .SetDefault(10);
147 148 149 150 151 152 153 154
    AddAttr<int>("sampler",
                 "(int) Which sampler to be used to sample negative class."
                 "0: Uniform; 1: LogUniform; 2: CostumDist.")
        .SetDefault(0);
    AddAttr<int>("seed",
                 "(int) The seed used in sampler. If it is 0, "
                 "the sampler will generate a seed randomly.")
        .SetDefault(0);
155 156
    AddAttr<bool>("is_sparse", "(boolean, default false) Sparse update.")
        .SetDefault(false);
157

W
wanghaoshuang 已提交
158 159 160 161
    AddAttr<std::vector<int>>("custom_neg_classes",
                              "This attribute only be used in unitest. Classes "
                              "in this list wiil be used as negative classes "
                              "for every samples. Under normal conditions, "
Y
Yang Yu 已提交
162 163
                              "user should avoid setting this attribute.")
        .SetDefault({});
W
wanghaoshuang 已提交
164
    AddComment(R"DOC(
Y
Yibing Liu 已提交
165 166 167 168
Compute and return the noise-contrastive estimation training loss. See 
`Noise-contrastive estimation: A new estimation principle for unnormalized 
statistical models 
 <http://www.jmlr.org/proceedings/papers/v9/gutmann10a/gutmann10a.pdf>`_.
W
wanghaoshuang 已提交
169
By default this operator uses a uniform distribution for sampling.
W
wanghaoshuang 已提交
170 171 172 173
)DOC");
  }
};

174 175 176 177 178 179 180 181
class NCEOpGradDescMaker : public framework::DefaultGradOpDescMaker<true> {
  using ::paddle::framework::DefaultGradOpDescMaker<
      true>::DefaultGradOpDescMaker;

 protected:
  virtual std::string GradOpType() const { return "nce_grad"; }
};

W
wanghaoshuang 已提交
182 183 184 185
class NCEOpGrad : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

186
  void InferShape(framework::InferShapeContext *ctx) const override {
W
wanghaoshuang 已提交
187 188 189 190 191 192
    PADDLE_ENFORCE(ctx->HasInput("Input"));
    PADDLE_ENFORCE(ctx->HasInput("Weight"));
    PADDLE_ENFORCE(ctx->HasInput("Cost"));
    PADDLE_ENFORCE(ctx->HasInput("SampleLogits"));
    PADDLE_ENFORCE(ctx->HasInput("SampleLabels"));
    PADDLE_ENFORCE(ctx->HasInput(framework::GradVarName("Cost")),
W
wanghaoshuang 已提交
193
                   "The input(Out@GRAD) should not be null.");
W
wanghaoshuang 已提交
194

W
wanghaoshuang 已提交
195 196
    auto x_dims = ctx->GetInputDim("Input");
    auto x_grad_name = framework::GradVarName("Input");
W
wanghaoshuang 已提交
197 198 199 200
    if (ctx->HasOutput(x_grad_name)) {
      ctx->SetOutputDim(x_grad_name, x_dims);
    }

W
wanghaoshuang 已提交
201 202
    auto w_dims = ctx->GetInputDim("Weight");
    auto w_grad_name = framework::GradVarName("Weight");
W
wanghaoshuang 已提交
203 204 205 206
    if (ctx->HasOutput(w_grad_name)) {
      ctx->SetOutputDim(w_grad_name, w_dims);
    }

W
wanghaoshuang 已提交
207
    auto bias_grad_name = framework::GradVarName("Bias");
W
wanghaoshuang 已提交
208
    if (ctx->HasOutput(bias_grad_name)) {
W
wanghaoshuang 已提交
209
      auto bias_dims = ctx->GetInputDim("Bias");
W
wanghaoshuang 已提交
210 211 212
      ctx->SetOutputDim(bias_grad_name, bias_dims);
    }
  }
W
wanghaoshuang 已提交
213 214

 protected:
215
  framework::OpKernelType GetExpectedKernelType(
216
      const framework::ExecutionContext &ctx) const override {
W
wanghaoshuang 已提交
217 218
    return framework::OpKernelType(
        framework::ToDataType(ctx.Input<Tensor>("Input")->type()),
C
chengduo 已提交
219
        platform::CPUPlace());
W
wanghaoshuang 已提交
220
  }
W
wanghaoshuang 已提交
221 222
};

223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248
class NCEOpGradVarTypeInference : public framework::VarTypeInference {
 public:
  void operator()(const framework::OpDesc &op_desc,
                  framework::BlockDesc *block) const override {
    auto weight_grad = op_desc.Output(framework::GradVarName("Weight")).front();
    auto bias_grad = op_desc.Output(framework::GradVarName("Bias")).front();

    auto attr = op_desc.GetAttr("is_sparse");
    bool is_sparse = boost::get<bool>(attr);
    if (is_sparse) {
      VLOG(30) << "nce_op_grad op " << weight_grad << " and " << bias_grad
               << " is set to SelectedRows";
      block->Var(weight_grad)
          ->SetType(framework::proto::VarType::SELECTED_ROWS);
      block->Var(bias_grad)->SetType(framework::proto::VarType::SELECTED_ROWS);
    } else {
      VLOG(30) << "nce_op_grad op " << weight_grad << " and " << bias_grad
               << " is set to LoDTensor";
      block->Var(weight_grad)->SetType(framework::proto::VarType::LOD_TENSOR);
      block->Var(bias_grad)->SetType(framework::proto::VarType::LOD_TENSOR);
    }
    block->Var(weight_grad)->SetDataType(block->Var("Input")->GetDataType());
    block->Var(bias_grad)->SetDataType(block->Var("Input")->GetDataType());
  }
};

W
wanghaoshuang 已提交
249 250 251 252
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
253 254
REGISTER_OPERATOR(nce, ops::NCEOp, ops::NCEOpGradDescMaker, ops::NCEOpMaker);
REGISTER_OPERATOR(nce_grad, ops::NCEOpGrad, ops::NCEOpGradVarTypeInference);
W
wanghaoshuang 已提交
255 256
REGISTER_OP_CPU_KERNEL(nce, ops::NCEKernel<paddle::platform::CPUPlace, float>,
                       ops::NCEKernel<paddle::platform::CPUPlace, double>);
W
wanghaoshuang 已提交
257
REGISTER_OP_CPU_KERNEL(nce_grad,
W
wanghaoshuang 已提交
258 259
                       ops::NCEGradKernel<paddle::platform::CPUPlace, float>,
                       ops::NCEGradKernel<paddle::platform::CPUPlace, double>);