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

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 15 16 17 18 19 20 21 22 23 24 25 26

#include "paddle/operators/nce_op.h"

namespace paddle {
namespace operators {

using framework::Tensor;

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

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

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

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

 protected:
66
  framework::OpKernelType GetExpectedKernelType(
W
wanghaoshuang 已提交
67 68 69
      const framework::ExecutionContext& ctx) const override {
    return framework::OpKernelType(
        framework::ToDataType(ctx.Input<Tensor>("Input")->type()),
Q
QI JUN 已提交
70
        ctx.GetPlace());
W
wanghaoshuang 已提交
71
  }
W
wanghaoshuang 已提交
72 73 74 75
};

class NCEOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
76
  NCEOpMaker(OpProto* proto, OpAttrChecker* op_checker)
W
wanghaoshuang 已提交
77
      : OpProtoAndCheckerMaker(proto, op_checker) {
W
wanghaoshuang 已提交
78
    AddInput("Input", "(Tensor) A tensor of shape [batch_size, dim].");
W
wanghaoshuang 已提交
79 80 81 82 83 84 85 86
    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 已提交
87 88 89
    AddInput("Weight",
             "(Tensor) A tensor of shape [num_class, dim]. 'num_class' is the "
             "total number of class.");
W
wanghaoshuang 已提交
90 91 92 93
    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 已提交
94 95
        .AsDispensable();
    AddInput("SampleWeight",
W
wanghaoshuang 已提交
96
             "(Tensor) A tensor of shape [batch_size, 1] storing a weight for "
W
wanghaoshuang 已提交
97 98 99 100
             "each sample. And it is a dispensable input. The default value of "
             "sample is 1.")
        .AsDispensable();
    AddOutput("Cost",
W
wanghaoshuang 已提交
101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121
              "(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();
    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 已提交
122
        .SetDefault(10);
W
wanghaoshuang 已提交
123 124 125 126
    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 已提交
127 128
                              "user should avoid setting this attribute.")
        .SetDefault({});
W
wanghaoshuang 已提交
129
    AddComment(R"DOC(
W
wanghaoshuang 已提交
130
Compute and return the noise-contrastive estimation training loss.
W
wanghaoshuang 已提交
131
See [Noise-contrastive estimation: A new estimation principle for unnormalized statistical models](http://www.jmlr.org/proceedings/papers/v9/gutmann10a/gutmann10a.pdf).
W
wanghaoshuang 已提交
132
By default this operator uses a uniform distribution for sampling.
W
wanghaoshuang 已提交
133 134 135 136 137 138 139 140 141
)DOC");
  }
};

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

  void InferShape(framework::InferShapeContext* ctx) const override {
W
wanghaoshuang 已提交
142 143 144 145 146 147
    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 已提交
148
                   "The input(Out@GRAD) should not be null.");
W
wanghaoshuang 已提交
149

W
wanghaoshuang 已提交
150 151
    auto x_dims = ctx->GetInputDim("Input");
    auto x_grad_name = framework::GradVarName("Input");
W
wanghaoshuang 已提交
152 153 154 155
    if (ctx->HasOutput(x_grad_name)) {
      ctx->SetOutputDim(x_grad_name, x_dims);
    }

W
wanghaoshuang 已提交
156 157
    auto w_dims = ctx->GetInputDim("Weight");
    auto w_grad_name = framework::GradVarName("Weight");
W
wanghaoshuang 已提交
158 159 160 161
    if (ctx->HasOutput(w_grad_name)) {
      ctx->SetOutputDim(w_grad_name, w_dims);
    }

W
wanghaoshuang 已提交
162
    auto bias_grad_name = framework::GradVarName("Bias");
W
wanghaoshuang 已提交
163
    if (ctx->HasOutput(bias_grad_name)) {
W
wanghaoshuang 已提交
164
      auto bias_dims = ctx->GetInputDim("Bias");
W
wanghaoshuang 已提交
165 166 167
      ctx->SetOutputDim(bias_grad_name, bias_dims);
    }
  }
W
wanghaoshuang 已提交
168 169

 protected:
170
  framework::OpKernelType GetExpectedKernelType(
W
wanghaoshuang 已提交
171 172 173
      const framework::ExecutionContext& ctx) const override {
    return framework::OpKernelType(
        framework::ToDataType(ctx.Input<Tensor>("Input")->type()),
Q
QI JUN 已提交
174
        ctx.GetPlace());
W
wanghaoshuang 已提交
175
  }
W
wanghaoshuang 已提交
176 177 178 179 180 181 182
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
REGISTER_OP(nce, ops::NCEOp, ops::NCEOpMaker, nce_grad, ops::NCEOpGrad);
W
wanghaoshuang 已提交
183 184
REGISTER_OP_CPU_KERNEL(nce, ops::NCEKernel<paddle::platform::CPUPlace, float>,
                       ops::NCEKernel<paddle::platform::CPUPlace, double>);
W
wanghaoshuang 已提交
185
REGISTER_OP_CPU_KERNEL(nce_grad,
W
wanghaoshuang 已提交
186 187
                       ops::NCEGradKernel<paddle::platform::CPUPlace, float>,
                       ops::NCEGradKernel<paddle::platform::CPUPlace, double>);