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

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

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

class NCEOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
78
  void Make() override {
W
wanghaoshuang 已提交
79
    AddInput("Input", "(Tensor) A tensor of shape [batch_size, dim].");
W
wanghaoshuang 已提交
80 81 82 83 84 85 86 87
    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 已提交
88 89 90
    AddInput("Weight",
             "(Tensor) A tensor of shape [num_class, dim]. 'num_class' is the "
             "total number of class.");
W
wanghaoshuang 已提交
91 92 93 94
    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 已提交
95 96
        .AsDispensable();
    AddInput("SampleWeight",
W
wanghaoshuang 已提交
97
             "(Tensor) A tensor of shape [batch_size, 1] storing a weight for "
W
wanghaoshuang 已提交
98 99 100
             "each sample. And it is a dispensable input. The default value of "
             "sample is 1.")
        .AsDispensable();
101 102

    AddInput(
103
        "CustomDistProbs",
104 105 106 107
        "(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();
108 109 110 111 112 113 114 115 116 117 118 119 120
    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 已提交
121
    AddOutput("Cost",
W
wanghaoshuang 已提交
122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138
              "(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();
139

W
wanghaoshuang 已提交
140 141 142 143
    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 已提交
144
        .SetDefault(10);
145 146 147 148 149 150 151 152
    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);
153 154
    AddAttr<bool>("is_sparse", "(boolean, default false) Sparse update.")
        .SetDefault(false);
155

T
tangwei12 已提交
156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173
    // for parameter prefetch
    AddAttr<bool>("remote_prefetch", "").SetDefault(false);
    AddAttr<int>("trainer_id", "trainer id from 0 ~ worker_num.").SetDefault(0);
    AddAttr<std::vector<int>>("height_sections",
                              "Height for each output SelectedRows.")
        .SetDefault(std::vector<int>({}));
    AddAttr<std::vector<std::string>>(
        "epmap",
        "(string vector, default 127.0.0.1:6164)"
        "Server endpoints in the order of input variables for mapping")
        .SetDefault({});
    AddAttr<std::vector<std::string>>(
        "table_names",
        "(string vector, the splited table names that will be fetched from "
        "parameter server)"
        "in the order of input variables for mapping")
        .SetDefault({});

W
wanghaoshuang 已提交
174 175 176 177
    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 已提交
178 179
                              "user should avoid setting this attribute.")
        .SetDefault({});
W
wanghaoshuang 已提交
180
    AddComment(R"DOC(
M
minqiyang 已提交
181 182 183
Compute and return the noise-contrastive estimation training loss. See
`Noise-contrastive estimation: A new estimation principle for unnormalized
statistical models
Y
Yibing Liu 已提交
184
 <http://www.jmlr.org/proceedings/papers/v9/gutmann10a/gutmann10a.pdf>`_.
W
wanghaoshuang 已提交
185
By default this operator uses a uniform distribution for sampling.
W
wanghaoshuang 已提交
186 187 188 189
)DOC");
  }
};

190 191 192 193 194 195 196 197
class NCEOpGradDescMaker : public framework::DefaultGradOpDescMaker<true> {
  using ::paddle::framework::DefaultGradOpDescMaker<
      true>::DefaultGradOpDescMaker;

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

W
wanghaoshuang 已提交
198 199 200 201
class NCEOpGrad : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

202
  void InferShape(framework::InferShapeContext *ctx) const override {
W
wanghaoshuang 已提交
203 204 205 206 207 208
    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 已提交
209
                   "The input(Out@GRAD) should not be null.");
W
wanghaoshuang 已提交
210

W
wanghaoshuang 已提交
211 212
    auto x_dims = ctx->GetInputDim("Input");
    auto x_grad_name = framework::GradVarName("Input");
W
wanghaoshuang 已提交
213 214 215 216
    if (ctx->HasOutput(x_grad_name)) {
      ctx->SetOutputDim(x_grad_name, x_dims);
    }

W
wanghaoshuang 已提交
217 218
    auto w_dims = ctx->GetInputDim("Weight");
    auto w_grad_name = framework::GradVarName("Weight");
W
wanghaoshuang 已提交
219 220 221 222
    if (ctx->HasOutput(w_grad_name)) {
      ctx->SetOutputDim(w_grad_name, w_dims);
    }

W
wanghaoshuang 已提交
223
    auto bias_grad_name = framework::GradVarName("Bias");
W
wanghaoshuang 已提交
224
    if (ctx->HasOutput(bias_grad_name)) {
W
wanghaoshuang 已提交
225
      auto bias_dims = ctx->GetInputDim("Bias");
W
wanghaoshuang 已提交
226 227 228
      ctx->SetOutputDim(bias_grad_name, bias_dims);
    }
  }
W
wanghaoshuang 已提交
229 230

 protected:
231
  framework::OpKernelType GetExpectedKernelType(
232
      const framework::ExecutionContext &ctx) const override {
Y
Yu Yang 已提交
233 234
    return framework::OpKernelType(ctx.Input<Tensor>("Input")->type(),
                                   platform::CPUPlace());
W
wanghaoshuang 已提交
235
  }
W
wanghaoshuang 已提交
236 237
};

238 239
class NCEOpGradVarTypeInference : public framework::VarTypeInference {
 public:
M
minqiyang 已提交
240 241
  void operator()(framework::InferVarTypeContext &ctx) const override {
    auto weight_grad = ctx.Output(framework::GradVarName("Weight")).front();
242

M
minqiyang 已提交
243
    auto attr = ctx.GetAttr("is_sparse");
244 245
    bool is_sparse = boost::get<bool>(attr);
    if (is_sparse) {
246
      VLOG(3) << "nce_op_grad op " << weight_grad << " and "
M
minqiyang 已提交
247
              << " is set to SelectedRows";
M
minqiyang 已提交
248
      ctx.SetType(weight_grad, framework::proto::VarType::SELECTED_ROWS);
249
    } else {
250
      VLOG(3) << "nce_op_grad op " << weight_grad << " and "
M
minqiyang 已提交
251
              << " is set to LoDTensor";
M
minqiyang 已提交
252
      ctx.SetType(weight_grad, framework::proto::VarType::LOD_TENSOR);
253
    }
M
minqiyang 已提交
254
    ctx.SetDataType(weight_grad, ctx.GetDataType(ctx.Input("Input")[0]));
255 256 257
  }
};

W
wanghaoshuang 已提交
258 259 260 261
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
262 263
REGISTER_OPERATOR(nce, ops::NCEOp, ops::NCEOpGradDescMaker, ops::NCEOpMaker);
REGISTER_OPERATOR(nce_grad, ops::NCEOpGrad, ops::NCEOpGradVarTypeInference);
W
wanghaoshuang 已提交
264 265
REGISTER_OP_CPU_KERNEL(nce, ops::NCEKernel<paddle::platform::CPUPlace, float>,
                       ops::NCEKernel<paddle::platform::CPUPlace, double>);
W
wanghaoshuang 已提交
266
REGISTER_OP_CPU_KERNEL(nce_grad,
W
wanghaoshuang 已提交
267 268
                       ops::NCEGradKernel<paddle::platform::CPUPlace, float>,
                       ops::NCEGradKernel<paddle::platform::CPUPlace, double>);