[incubator-mxnet] branch master updated: Efficient MXNet sampling in the multinomial distribution (#15311)

[wkcn]

This is an automated email from the ASF dual-hosted git repository.

wkcn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git


The following commit(s) were added to refs/heads/master by this push:
     new e6fad30  Efficient MXNet sampling in the multinomial distribution 
(#15311)
e6fad30 is described below

commit e6fad30e45e6ec0ddef5c18093e8163cd2a7c62c
Author: Zixuan Wei <9527221+zixuanwe...@users.noreply.github.com>
AuthorDate: Sat Jun 22 14:26:27 2019 +0800

    Efficient MXNet sampling in the multinomial distribution (#15311)
    
    * Effective multinomial
    
    * Meaningful uniform data pointer as input
    
    * Remove beginning Zeros from CDFs
    
    * Double precision for accumulated var
---
 src/operator/random/sample_multinomial_op.h | 42 ++++++++++++++++-------------
 1 file changed, 24 insertions(+), 18 deletions(-)

diff --git a/src/operator/random/sample_multinomial_op.h 
b/src/operator/random/sample_multinomial_op.h
index 377df4f..5a0b9bb 100644
--- a/src/operator/random/sample_multinomial_op.h
+++ b/src/operator/random/sample_multinomial_op.h
@@ -122,25 +122,29 @@ inline bool SampleMultinomialOpType(const 
nnvm::NodeAttrs& attrs,
 struct SampleMultinomialKernel {
   template<typename DType, typename IType>
   MSHADOW_XINLINE static void Map(int i, index_t K, index_t M,
-                                  DType* dist, float* uniform, IType* out,
-                                  DType* prob) {
+                                  DType* dist, float* uniform, float* 
cum_table,
+                                  IType* out, DType* prob) {
+    double acc = 0.0;
+    // CDF table
+    for (index_t c = 0; c < K; ++c) {
+      acc += dist[i*K + c];
+      cum_table[i*K + c] = static_cast<float>(acc);
+    }
     for (index_t j = 0; j < M; ++j) {
+      index_t left = 0, right = K;
+      index_t middle = left + (right - left) / 2;
       DType loc = static_cast<DType>(uniform[i*M + j]);
-      DType acc = 0;
-      bool found = false;
-      for (index_t k = 0; k < K; ++k) {
-        acc += dist[i*K + k];
-        if (acc > loc) {
-          found = true;
-          out[i*M + j] = static_cast<IType>(k);
-          if (prob != nullptr) prob[i*M + j] = logf(dist[i*K + k]);
-          break;
+      while (right - left > 0) {
+        middle = left + (right - left) / 2;
+        DType cum_prob = cum_table[i*K + middle];
+        if (cum_prob < loc) {
+          left = middle + 1;
+        } else {
+          right = middle;
         }
       }
-      if (!found) {
-        out[i*M + j] = static_cast<IType>(K-1);
-        if (prob != nullptr) prob[i*M + j] = logf(dist[i*K + K - 1]);
-      }
+      out[i*M + j] = static_cast<IType>(left);
+      if (prob != nullptr) prob[i*M + j] = logf(dist[i*K + left]);
     }
   }
 };
@@ -163,12 +167,14 @@ void SampleMultinomialForward(const nnvm::NodeAttrs& 
attrs,
   Stream<xpu> *s = ctx.get_stream<xpu>();
   MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
     Random<xpu, float> *prnd = ctx.requested[0].get_random<xpu, float>(s);
-    Tensor<xpu, 1, float> uniform =
-      ctx.requested[1].get_space_typed<xpu, 1, float>(Shape1(N*M), s);
+    Tensor<xpu, 1, float> workspace =
+      ctx.requested[1].get_space_typed<xpu, 1, float>(Shape1(N*M + N*K), s);
+    Tensor<xpu, 1, float> uniform(workspace.dptr_, Shape1(N*M));
     prnd->SampleUniform(&uniform, 0, 1);
     MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, IType, {
       Kernel<SampleMultinomialKernel, xpu>::Launch(
-        s, N, K, M, inputs[0].dptr<DType>(), uniform.dptr_, 
outputs[0].dptr<IType>(),
+        s, N, K, M, inputs[0].dptr<DType>(), uniform.dptr_, workspace.dptr_ + 
N*M,
+        outputs[0].dptr<IType>(),
         param.get_prob ? outputs[1].dptr<DType>() : nullptr);
     });
   });