nce_op.cc 7.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

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

W
wanghaoshuang 已提交
19 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;

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

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

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

 protected:
68
  framework::OpKernelType GetExpectedKernelType(
W
wanghaoshuang 已提交
69 70 71
      const framework::ExecutionContext& ctx) const override {
    return framework::OpKernelType(
        framework::ToDataType(ctx.Input<Tensor>("Input")->type()),
Q
QI JUN 已提交
72
        ctx.GetPlace());
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 101
             "each sample. And it is a dispensable input. The default value of "
             "sample is 1.")
        .AsDispensable();
    AddOutput("Cost",
W
wanghaoshuang 已提交
102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122
              "(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 已提交
123
        .SetDefault(10);
W
wanghaoshuang 已提交
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, "
Y
Yang Yu 已提交
128 129
                              "user should avoid setting this attribute.")
        .SetDefault({});
W
wanghaoshuang 已提交
130
    AddComment(R"DOC(
Y
Yibing Liu 已提交
131 132 133 134
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 已提交
135
By default this operator uses a uniform distribution for sampling.
W
wanghaoshuang 已提交
136 137 138 139 140 141 142 143 144
)DOC");
  }
};

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

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

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

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

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

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

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
Y
Yang Yang 已提交
185
REGISTER_OPERATOR(nce, ops::NCEOp, ops::NCEOpMaker,
186 187
                  paddle::framework::DefaultGradOpDescMaker<true>);
REGISTER_OPERATOR(nce_grad, ops::NCEOpGrad);
W
wanghaoshuang 已提交
188 189
REGISTER_OP_CPU_KERNEL(nce, ops::NCEKernel<paddle::platform::CPUPlace, float>,
                       ops::NCEKernel<paddle::platform::CPUPlace, double>);
W
wanghaoshuang 已提交
190
REGISTER_OP_CPU_KERNEL(nce_grad,
W
wanghaoshuang 已提交
191 192
                       ops::NCEGradKernel<paddle::platform::CPUPlace, float>,
                       ops::NCEGradKernel<paddle::platform::CPUPlace, double>);