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> sampled_labels =
        ctx->Attrs().Get<std::vector<int>>("sampled_labels");
W
wanghaoshuang 已提交
46
    PADDLE_ENFORCE_EQ(num_total_classes, ctx->GetInputDim("Weight")[0]);
W
wanghaoshuang 已提交
47 48
    if (sampled_labels.size() > 0) {
      PADDLE_ENFORCE_EQ(sampled_labels.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 66 67 68 69 70 71

 protected:
  framework::OpKernelType GetKernelType(
      const framework::ExecutionContext& ctx) const override {
    return framework::OpKernelType(
        framework::ToDataType(ctx.Input<Tensor>("Input")->type()),
        ctx.device_context());
  }
W
wanghaoshuang 已提交
72 73 74 75 76 77
};

class NCEOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
  NCEOpMaker(framework::OpProto* proto, framework::OpAttrChecker* op_checker)
      : 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 127
    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, "
                              "user should avoid setting this attribute.");
W
wanghaoshuang 已提交
128
    AddComment(R"DOC(
W
wanghaoshuang 已提交
129
Compute and return the noise-contrastive estimation training loss.
W
wanghaoshuang 已提交
130
See [Noise-contrastive estimation: A new estimation principle for unnormalized statistical models](http://www.jmlr.org/proceedings/papers/v9/gutmann10a/gutmann10a.pdf).
W
wanghaoshuang 已提交
131
By default this operator uses a uniform distribution for sampling.
W
wanghaoshuang 已提交
132 133 134 135 136 137 138 139 140
)DOC");
  }
};

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

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

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

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

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

 protected:
  framework::OpKernelType GetKernelType(
      const framework::ExecutionContext& ctx) const override {
    return framework::OpKernelType(
        framework::ToDataType(ctx.Input<Tensor>("Input")->type()),
        ctx.device_context());
  }
W
wanghaoshuang 已提交
175 176 177 178 179 180 181
};

}  // namespace operators
}  // namespace paddle

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