Make nce support more distribution. (#13549)

wanghaoshuang · web-flow · commit 1722678258fa · 2018-11-16T16:08:22.000+08:00
* Fix truncated normal.

* Fix.

* Make nce support more distribution.

* Fix API.spec.

* Fix python API.

* Fix.
test=develop

* Fix API.spec
test=develop

* Fix sampler.

* Fix order of arguments in python API.
test=develop
diff --git a/paddle/fluid/API.spec b/paddle/fluid/API.spec
@@ -97,7 +97,7 @@ paddle.fluid.layers.warpctc ArgSpec(args=['input', 'label', 'blank', 'norm_by_ti
 paddle.fluid.layers.sequence_reshape ArgSpec(args=['input', 'new_dim'], varargs=None, keywords=None, defaults=None)
 paddle.fluid.layers.transpose ArgSpec(args=['x', 'perm', 'name'], varargs=None, keywords=None, defaults=(None,))
 paddle.fluid.layers.im2sequence ArgSpec(args=['input', 'filter_size', 'stride', 'padding', 'input_image_size', 'out_stride', 'name'], varargs=None, keywords=None, defaults=(1, 1, 0, None, 1, None))
-paddle.fluid.layers.nce ArgSpec(args=['input', 'label', 'num_total_classes', 'sample_weight', 'param_attr', 'bias_attr', 'num_neg_samples', 'name'], varargs=None, keywords=None, defaults=(None, None, None, None, None))
+paddle.fluid.layers.nce ArgSpec(args=['input', 'label', 'num_total_classes', 'sample_weight', 'param_attr', 'bias_attr', 'num_neg_samples', 'name', 'sampler', 'custom_dist', 'seed'], varargs=None, keywords=None, defaults=(None, None, None, None, None, 'uniform', None, 0))
 paddle.fluid.layers.hsigmoid ArgSpec(args=['input', 'label', 'num_classes', 'param_attr', 'bias_attr', 'name'], varargs=None, keywords=None, defaults=(None, None, None))
 paddle.fluid.layers.beam_search ArgSpec(args=['pre_ids', 'pre_scores', 'ids', 'scores', 'beam_size', 'end_id', 'level', 'name'], varargs=None, keywords=None, defaults=(0, None))
 paddle.fluid.layers.row_conv ArgSpec(args=['input', 'future_context_size', 'param_attr', 'act'], varargs=None, keywords=None, defaults=(None, None))
diff --git a/paddle/fluid/operators/CMakeLists.txt b/paddle/fluid/operators/CMakeLists.txt
@@ -308,6 +308,7 @@ op_library(flatten_op DEPS reshape_op)
 op_library(sequence_pad_op DEPS sequence_padding)
 op_library(unstack_op DEPS stack_op)
 op_library(fake_quantize_op DEPS memory)
+op_library(nce_op DEPS sampler)
 if (NOT WIN32)
 op_library(crf_decoding_op DEPS jit_kernel)
 op_library(fusion_lstm_op DEPS jit_kernel)
diff --git a/paddle/fluid/operators/math/CMakeLists.txt b/paddle/fluid/operators/math/CMakeLists.txt
@@ -41,6 +41,7 @@ math_library(cross_entropy)
 math_library(cos_sim_functor)
 math_library(depthwise_conv)
 math_library(im2col)
+math_library(sampler)
 
 if (NOT WIN32) # windows do not support avx functions yet.
     math_library(gru_compute DEPS activation_functions math_function)
diff --git a/paddle/fluid/operators/math/sampler.cc b/paddle/fluid/operators/math/sampler.cc
@@ -1,4 +1,4 @@
-/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -13,58 +13,123 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/operators/math/sampler.h"
+#include <iostream>
+#include <queue>
+#include <utility>
+#include <vector>
 
 namespace paddle {
-namespace random {
+namespace operators {
+namespace math {
 
 Sampler::~Sampler() {}
 
-UniformSampler::UniformSampler(int64 range)
-    : Sampler(range), inv_range_(1.0 / range) {
-  random_engine_ = std::make_shared<std::mt19937>(seed_);
+UniformSampler::UniformSampler(int64_t range, unsigned int seed)
+    : Sampler(range, seed), inv_range_(1.0 / (range + 1)) {
+  random_engine_ = std::make_shared<std::mt19937_64>(seed_);
   dist_ = std::make_shared<std::uniform_int_distribution<>>(0, range);
 }
 
-UniformSampler::UniformSampler(int64 range, unsigned int seed)
-    : Sampler(range, seed), inv_range_(1.0 / range) {
-  random_engine_ = std::make_shared<std::mt19937>(seed_);
-  dist_ = std::make_shared<std::uniform_int_distribution<>>(0, range);
-}
-
-int64 UniformSampler::Sample() const { return (*dist_)(*random_engine_); }
+int64_t UniformSampler::Sample() const { return (*dist_)(*random_engine_); }
 
-float UniformSampler::Probability(int64 value) const { return inv_range_; }
+float UniformSampler::Probability(int64_t value) const { return inv_range_; }
 
-LogUniformSampler::LogUniformSampler(int64 range)
-    : Sampler(range), log_range_(log(range + 1)) {
-  random_engine_ = std::make_shared<std::mt19937>(seed_);
-  dist_ = std::make_shared<std::uniform_real_distribution<>>(0, 1);
-}
-
-LogUniformSampler::LogUniformSampler(int64 range, unsigned int seed)
+LogUniformSampler::LogUniformSampler(int64_t range, unsigned int seed)
     : Sampler(range, seed), log_range_(log(range + 1)) {
-  random_engine_ = std::make_shared<std::mt19937>(seed_);
+  random_engine_ = std::make_shared<std::mt19937_64>(seed_);
   dist_ = std::make_shared<std::uniform_real_distribution<>>(0, 1);
 }
-int64 LogUniformSampler::Sample() const {
+
+int64_t LogUniformSampler::Sample() const {
   // Got Log Uniform distribution from uniform distribution by
   // inverse_transform_sampling method
   // More details:
   // https://wanghaoshuang.github.io/2017/11/Log-uniform-distribution-sampler/
-  const int64 value =
-      static_cast<int64>(exp((*dist_)(*random_engine_) * log_range_)) - 1;
+  const int64_t value =
+      static_cast<int64_t>(exp((*dist_)(*random_engine_) * log_range_)) - 1;
   // Mathematically, value should be <= range_, but might not be due to some
   // floating point roundoff, so we mod by range_.
   return value % range_;
 }
 
-float LogUniformSampler::Probability(int64 value) const {
+float LogUniformSampler::Probability(int64_t value) const {
   // Given f(x) = 1/[(x+1) * log_range_]
   // The value's  probability  is integral of f(x) from value to (value + 1)
   // More details:
   // https://wanghaoshuang.github.io/2017/11/Log-uniform-distribution-sampler
   return (log((value + 2.0) / (value + 1.0))) / log_range_;
 }
 
-}  // namespace random
+CustomSampler::CustomSampler(int64_t range, const float* probabilities,
+                             unsigned int seed)
+    : Sampler(range, seed) {
+  random_engine_ = std::make_shared<std::mt19937_64>(seed_);
+  real_dist_ = std::make_shared<std::uniform_real_distribution<>>(0, 1);
+  int_dist_ = std::make_shared<std::uniform_int_distribution<>>(0, range);
+  alias_probs_ = std::make_shared<std::vector<float>>(range + 1);
+  alias_ = std::make_shared<std::vector<int64_t>>(range + 1);
+  probs_ = std::make_shared<std::vector<float>>(range + 1);
+
+  std::queue<std::pair<int64_t, float>> bigs;
+  std::queue<std::pair<int64_t, float>> littles;
+  for (int64_t i = 0; i <= range; ++i) {
+    (*probs_)[i] = probabilities[i];
+    float normal_prob = probabilities[i] * (range + 1);
+    if (normal_prob - 1.0 > 1e-4) {
+      bigs.emplace(i, normal_prob);
+    } else if (1.0 - normal_prob > 1e-4) {
+      littles.emplace(i, normal_prob);
+    } else {
+      (*alias_probs_)[i] = normal_prob;
+      (*alias_)[i] = -1;
+    }
+  }
+
+  while ((!littles.empty()) && (!bigs.empty())) {
+    auto big = bigs.front();
+    auto little = littles.front();
+    bigs.pop();
+    littles.pop();
+    (*alias_probs_)[little.first] = little.second;
+    (*alias_)[little.first] = big.first;
+    auto big_left = big.second - (1 - little.second);
+    if (big_left - 1.0 > 1e-4) {
+      bigs.emplace(big.first, big_left);
+    } else if (1.0 - big_left > 1e-4) {
+      littles.emplace(big.first, big_left);
+    } else {
+      (*alias_probs_)[big.first] = big_left;
+      (*alias_)[big.first] = -1;
+    }
+  }
+
+  if (!littles.empty()) {  // littles.second is close to 1.0
+    auto little = littles.front();
+    (*alias_probs_)[little.first] = 1.0;
+    (*alias_)[little.first] = -1;
+  }
+
+  if (!bigs.empty()) {  // bigs.second is close to 1.0
+    auto big = bigs.front();
+    (*alias_probs_)[big.first] = 1.0;
+    (*alias_)[big.first] = -1;
+  }
+}
+
+int64_t CustomSampler::Sample() const {
+  auto index = (*int_dist_)(*random_engine_);
+  auto p = (*real_dist_)(*random_engine_);
+  if (p > (*alias_probs_)[index]) {
+    return (*alias_)[index];
+  } else {
+    return index;
+  }
+}
+
+float CustomSampler::Probability(int64_t value) const {
+  return (*probs_)[value];
+}
+
+}  // namespace math
+}  // namespace operators
 }  // namespace paddle
diff --git a/paddle/fluid/operators/math/sampler.h b/paddle/fluid/operators/math/sampler.h
@@ -16,6 +16,8 @@ limitations under the License. */
 #include <cstdint>
 #include <memory>
 #include <random>
+#include <vector>
+
 namespace paddle {
 namespace operators {
 namespace math {
@@ -27,22 +29,22 @@ namespace math {
 */
 class Sampler {
  public:
-  explicit Sampler(int64_t range) : range_(range) {
-    PADDLE_ENFORCE_GT(range, 0);
-    std::random_device r;
-    seed_ = r();
-  }
-  explicit Sampler(int64_t range, unsigned int seed)
-      : range_(range), seed_(seed) {
-    PADDLE_ENFORCE_GT(range, 0);
+  explicit Sampler(int64_t range, unsigned int seed = 0UL) : range_(range) {
+    //    PADDLE_ENFORCE_GT(range, 0, "Range should be greater than 0.");
+    if (seed == 0) {
+      std::random_device r;
+      seed_ = r();
+    } else {
+      seed_ = seed;
+    }
   }
   virtual ~Sampler();
   // Sample a single value
   virtual int64_t Sample() const = 0;
   // The probability that a single call to Sample() returns the given value.
   virtual float Probability(int64_t value) const = 0;
 
-  int64 range() { return range_; }
+  int64_t range() { return range_; }
 
  protected:
   const int64_t range_;
@@ -56,13 +58,11 @@ class Sampler {
  */
 class UniformSampler : public Sampler {
  public:
-  explicit UniformSampler(int64_t range);
-
-  explicit UniformSampler(int64_t range, unsigned int seed);
+  explicit UniformSampler(int64_t range, unsigned int seed = 0UL);
 
   ~UniformSampler() override {}
 
-  int64 Sample() const override;
+  int64_t Sample() const override;
 
   float Probability(int64_t value) const override;
 
@@ -79,13 +79,11 @@ class UniformSampler : public Sampler {
  */
 class LogUniformSampler : public Sampler {
  public:
-  explicit LogUniformSampler(int64_t range);
-
-  explicit LogUniformSampler(int64_t range, unsigned int seed);
+  explicit LogUniformSampler(int64_t range, unsigned int seed = 0UL);
 
   ~LogUniformSampler() override {}
 
-  int64 Sample() const override;
+  int64_t Sample() const override;
 
   float Probability(int64_t value) const override;
 
@@ -95,6 +93,29 @@ class LogUniformSampler : public Sampler {
   std::shared_ptr<std::uniform_real_distribution<>> dist_;
 };
 
+/**
+ * Sample integers from [0, range) from custom distribution.
+ */
+class CustomSampler : public Sampler {
+ public:
+  explicit CustomSampler(int64_t range, const float* probabilities,
+                         unsigned int seed = 0UL);
+
+  ~CustomSampler() override {}
+
+  int64_t Sample() const override;
+
+  float Probability(int64_t value) const override;
+
+ private:
+  std::shared_ptr<std::vector<float>> alias_probs_;
+  std::shared_ptr<std::vector<int64_t>> alias_;
+  std::shared_ptr<std::vector<float>> probs_;
+  std::shared_ptr<std::mt19937_64> random_engine_;
+  std::shared_ptr<std::uniform_real_distribution<>> real_dist_;
+  std::shared_ptr<std::uniform_int_distribution<>> int_dist_;
+};
+
 }  // namespace math
 }  // namespace operators
 }  // namespace paddle
diff --git a/paddle/fluid/operators/nce_op.cc b/paddle/fluid/operators/nce_op.cc
@@ -35,6 +35,7 @@ class NCEOp : public framework::OperatorWithKernel {
 
     auto x_dims = ctx->GetInputDim("Input");
     auto label_dims = ctx->GetInputDim("Label");
+    auto w_dims = ctx->GetInputDim("Weight");
     PADDLE_ENFORCE_EQ(x_dims[0], label_dims[0]);
     int num_true_classes = label_dims.size() == 2 ? label_dims[1] : 1;
     if (ctx->HasInput("Bias")) {
@@ -98,6 +99,13 @@ class NCEOpMaker : public framework::OpProtoAndCheckerMaker {
              "each sample. And it is a dispensable input. The default value of "
              "sample is 1.")
         .AsDispensable();
+
+    AddInput(
+        "CustomDistribution",
+        "(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();
     AddOutput("Cost",
               "(Tensor) A tensor of shape [batch_size, 1]. Cost of samples.");
     AddOutput("SampleLogits",
@@ -121,6 +129,17 @@ class NCEOpMaker : public framework::OpProtoAndCheckerMaker {
     AddAttr<int>("num_neg_samples",
                  "The number of negative classes. The default value is 10.")
         .SetDefault(10);
+
+    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);
+
     AddAttr<std::vector<int>>("custom_neg_classes",
                               "This attribute only be used in unitest. Classes "
                               "in this list wiil be used as negative classes "
diff --git a/paddle/fluid/operators/nce_op.h b/paddle/fluid/operators/nce_op.h
diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py
diff --git a/python/paddle/fluid/tests/unittests/test_nce.py b/python/paddle/fluid/tests/unittests/test_nce.py
