src/operator/random/sampler.h - master - githubmirror/incubator-mxnet

src/operator/random/sampler.h
594 строки21 KB
Vincenzo Territo
[FEATURE] Add binomial sampling and fix multinomial sampling (#20734)
05 фев 2022, 19:53
05 фев 2022, 19:53e9becb9
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * 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.
 */

/*!
 * \file sampler.h
 * \brief implementations of random sampling functors.
 */

#ifndef MXNET_OPERATOR_RANDOM_SAMPLER_H_
#define MXNET_OPERATOR_RANDOM_SAMPLER_H_

#include <algorithm>

using namespace mshadow;
using namespace mxnet::op::mxnet_op;
using namespace mxnet::common::random;

namespace mxnet {
namespace op {

/*!
 * \brief Launch a generic kernel with parallel random generator.
 * \tparam gen random generator
 * \tparam N Number of iterations
 * \tparam Args Varargs type to eventually pass to the OP::Map() function
 */
template <typename OP, typename xpu, typename GType, typename... Args>
inline static void LaunchRNG(mshadow::Stream<xpu>* s,
                             common::random::RandGenerator<xpu, GType>* gen,
                             const index_t N,
                             Args... args) {
  // minimal check to avoid division by zero, below.
  // if `N` is zero the map operation is a no-op in any case.
  if (N <= 0) {
    return;
  }
  const index_t nloop = (N + RandGenerator<xpu>::kMinNumRandomPerThread - 1) /
                        RandGenerator<xpu>::kMinNumRandomPerThread;
  const index_t nthread =
      std::min(nloop, static_cast<index_t>(RandGenerator<xpu>::kNumRandomStates));
  const index_t step = (N + nthread - 1) / nthread;
  Kernel<OP, xpu>::Launch(s, nthread, *gen, N, step, args...);
}

#define RNG_KERNEL_LOOP(xpu, GType, thread_id, gen, N, step, ...)    \
  const index_t start = thread_id * step;                            \
  const index_t end   = start + step;                                \
  typename RandGenerator<xpu, GType>::Impl genImpl(&gen, thread_id); \
  for (index_t i = start; i < end && i < N; ++i) {                   \
    { __VA_ARGS__ }                                                  \
  }

template <typename xpu>
struct SampleUniformKernel {
  template <typename IType, typename OType>
  MSHADOW_XINLINE static void Map(index_t id,
                                  RandGenerator<xpu, OType> gen,
                                  const index_t N,
                                  const index_t step,
                                  index_t nParm,
                                  index_t nSample,
                                  const IType* lower,
                                  const IType* upper,
                                  OType* out) {
    RNG_KERNEL_LOOP(xpu, OType, id, gen, N, step, {
      index_t nBatch(1 + (nSample - 1) / nParm);
      out[i] =
          OType(lower[i / nBatch] + (upper[i / nBatch] - lower[i / nBatch]) * genImpl.uniform());
    });
  }
};

template <typename xpu>
struct UniformSampler {
  template <typename IType, typename OType>
  MSHADOW_FORCE_INLINE void Sample(const Tensor<xpu, 1, IType>& lower,
                                   const Tensor<xpu, 1, IType>& upper,
                                   const Tensor<xpu, 1, OType>& out,
                                   RandGenerator<xpu, OType>* pgen,
                                   Stream<xpu>* s) {
    LaunchRNG<SampleUniformKernel<xpu>, xpu>(
        s, pgen, out.size(0), lower.size(0), out.size(0), lower.dptr_, upper.dptr_, out.dptr_);
  }
};

template <typename xpu>
struct SampleRandIntKernel {
  template <typename IType, typename OType>
  MSHADOW_XINLINE static void Map(index_t id,
                                  RandGenerator<xpu, OType> gen,
                                  const index_t N,
                                  const index_t step,
                                  index_t nParm,
                                  index_t nSample,
                                  const IType* lower,
                                  const IType* upper,
                                  OType* out) {
    RNG_KERNEL_LOOP(xpu, OType, id, gen, N, step, {
      index_t nBatch(1 + (nSample - 1) / nParm);
      if (sizeof(IType) == sizeof(int64_t))
        out[i] = OType(lower[i / nBatch] +
                       genImpl.rand_int64() % (upper[i / nBatch] - lower[i / nBatch]));
      else
        out[i] =
            OType(lower[i / nBatch] + genImpl.rand() % (upper[i / nBatch] - lower[i / nBatch]));
    });
  }
};

template <typename xpu>
struct RandIntSampler {
  template <typename IType, typename OType>
  MSHADOW_FORCE_INLINE void Sample(const Tensor<xpu, 1, IType>& lower,
                                   const Tensor<xpu, 1, IType>& upper,
                                   const Tensor<xpu, 1, OType>& out,
                                   RandGenerator<xpu, OType>* pgen,
                                   Stream<xpu>* s) {
    LaunchRNG<SampleRandIntKernel<xpu>, xpu>(
        s, pgen, out.size(0), lower.size(0), out.size(0), lower.dptr_, upper.dptr_, out.dptr_);
  }
};

template <typename xpu>
struct SampleNormalKernel {
  template <typename IType, typename OType>
  MSHADOW_XINLINE static void Map(index_t id,
                                  RandGenerator<xpu, OType> gen,
                                  const index_t N,
                                  const index_t step,
                                  index_t nParm,
                                  index_t nSample,
                                  const IType* mean,
                                  const IType* std,
                                  OType* out) {
    RNG_KERNEL_LOOP(xpu, OType, id, gen, N, step, {
      index_t nBatch(1 + (nSample - 1) / nParm);
      out[i] = OType(genImpl.normal() * std[i / nBatch] + mean[i / nBatch]);
    });
  }
};

template <typename xpu>
struct NormalSampler {
  template <typename IType, typename OType>
  MSHADOW_FORCE_INLINE void Sample(const Tensor<xpu, 1, IType>& mean,
                                   const Tensor<xpu, 1, IType>& std,
                                   const Tensor<xpu, 1, OType>& out,
                                   RandGenerator<xpu, OType>* pgen,
                                   Stream<xpu>* s) {
    LaunchRNG<SampleNormalKernel<xpu>, xpu>(
        s, pgen, out.size(0), mean.size(0), out.size(0), mean.dptr_, std.dptr_, out.dptr_);
  }
};

template <typename xpu>
struct SampleExponentialKernel {
  template <typename IType, typename OType>
  MSHADOW_XINLINE static void Map(index_t id,
                                  RandGenerator<xpu, OType> gen,
                                  const index_t N,
                                  const index_t step,
                                  index_t nParm,
                                  index_t nSample,
                                  const IType* lambda,
                                  OType* out) {
    RNG_KERNEL_LOOP(xpu, OType, id, gen, N, step, {
      index_t nBatch(1 + (nSample - 1) / nParm);
      out[i] = OType(-log(1.0 - genImpl.uniform()) / lambda[i / nBatch]);
    });
  }
};

template <typename xpu>
struct ExponentialSampler {
  template <typename IType, typename OType>
  MSHADOW_FORCE_INLINE void Sample(const Tensor<xpu, 1, IType>& lambda,
                                   const Tensor<xpu, 1, OType>& out,
                                   RandGenerator<xpu, OType>* pgen,
                                   Stream<xpu>* s) {
    LaunchRNG<SampleExponentialKernel<xpu>, xpu>(
        s, pgen, out.size(0), lambda.size(0), out.size(0), lambda.dptr_, out.dptr_);
  }
};

template <typename xpu, typename IType, typename OType>
MSHADOW_XINLINE OType SampleGamma(IType a, IType b, typename RandGenerator<xpu, OType>::Impl* gen) {
  // Generate one sample of the gamma distribution
  OType sample;
  OType d = a < 1 ? a + 2.0 / 3.0 : a - 1.0 / 3.0;
  OType k = sqrt(9.0 * d);
  OType c = 1.0 / k;
  while (1) {
    OType Z = gen->normal();
    if (Z > -k) {
      OType x = 1.0 + c * Z;
      OType V = x * x * x;
      if (log(1.0 - gen->uniform()) < 0.5 * Z * Z + d * (1.0 - V + log(V))) {
        sample = d * V * b;
        break;
      }
    }
  }
  return a < 1 ? sample * pow(gen->uniform(), OType(1.0 / a)) : sample;
}

template <typename xpu>
struct SampleGammaKernel {
  template <typename IType, typename OType, typename FType>
  MSHADOW_XINLINE static void Map(index_t id,
                                  RandGenerator<xpu, FType> gen,
                                  const index_t N,
                                  const index_t step,
                                  index_t nParm,
                                  index_t nSample,
                                  const IType* alpha,
                                  const IType* beta,
                                  OType* out) {
    RNG_KERNEL_LOOP(xpu, FType, id, gen, N, step, {
      index_t nBatch(1 + (nSample - 1) / nParm);
      out[i] = OType(SampleGamma<xpu, IType, FType>(alpha[i / nBatch], beta[i / nBatch], &genImpl));
    });
  }
};

template <typename xpu>
struct GammaSampler {
  template <typename IType, typename OType>
  MSHADOW_FORCE_INLINE void Sample(const Tensor<xpu, 1, IType>& alpha,
                                   const Tensor<xpu, 1, IType>& beta,
                                   const Tensor<xpu, 1, OType>& out,
                                   RandGenerator<xpu, OType>* pgen,
                                   Stream<xpu>* s) {
    typedef
        typename std::conditional<std::is_floating_point<OType>::value, OType, float>::type FType;
    RandGenerator<xpu, FType>* gen = reinterpret_cast<RandGenerator<xpu, FType>*>(pgen);
    LaunchRNG<SampleGammaKernel<xpu>, xpu>(
        s, gen, out.size(0), alpha.size(0), out.size(0), alpha.dptr_, beta.dptr_, out.dptr_);
  }
};

template <typename xpu>
MSHADOW_XINLINE int SamplePoisson(float lambda, typename RandGenerator<xpu, float>::Impl* gen) {
  // Generate one sample of the poisson distribution. Intentionally written
  // towards a specific type (float) for internal computation which is sufficient
  // for accurate enough computation.
  if (lambda < 12.0) {
    float t = expf(-lambda);
    int x   = 0;
    for (float prod = gen->uniform(); prod > t; prod *= gen->uniform()) {
      x += 1;
    }
    return x;
  } else {
    // Approximation for high lambda according to:
    // Numerical Recipes in C: The Art of Scientific Computing
    // Cambridge University Press
    const float pi(3.1415926);
    const float sq(sqrt(2.0 * lambda));
    const float loglambda(log(lambda));
    const float g(lambda * loglambda - lgammaf(lambda + 1.0));
    float em(0), t(0), y(0);
    do {
      do {
        y  = tanf(pi * gen->uniform());
        em = sq * y + lambda;
      } while (em < 0.0);
      em = floorf(em);
      t  = 0.9 * (1.0 + y * y) * expf(em * loglambda - lgammaf(em + 1.0) - g);
    } while (gen->uniform() > t);
    return static_cast<int>(em);
  }
}

template <typename xpu>
struct SamplePoissonKernel {
  template <typename IType, typename OType>
  MSHADOW_XINLINE static void Map(index_t id,
                                  RandGenerator<xpu, float> gen,
                                  const index_t N,
                                  const index_t step,
                                  index_t nParm,
                                  index_t nSample,
                                  const IType* lambda,
                                  OType* out) {
    RNG_KERNEL_LOOP(xpu, float, id, gen, N, step, {
      index_t nBatch(1 + (nSample - 1) / nParm);
      out[i] = OType(SamplePoisson<xpu>(lambda[i / nBatch], &genImpl));
    });
  }
};

template <typename xpu>
struct PoissonSampler {
  template <typename IType, typename OType>
  MSHADOW_FORCE_INLINE void Sample(const Tensor<xpu, 1, IType>& lambda,
                                   const Tensor<xpu, 1, OType>& out,
                                   RandGenerator<xpu, OType>* pgen,
                                   Stream<xpu>* s) {
    RandGenerator<xpu, float>* gen = reinterpret_cast<RandGenerator<xpu, float>*>(pgen);
    LaunchRNG<SamplePoissonKernel<xpu>, xpu>(
        s, gen, out.size(0), lambda.size(0), out.size(0), lambda.dptr_, out.dptr_);
  }
};

MSHADOW_XINLINE double stirling_approximation(double k) {
  static const double table[] = {0.08106146679532726,
                                 0.04134069595540929,
                                 0.02767792568499834,
                                 0.02079067210376509,
                                 0.01664469118982119,
                                 0.01387612882307075,
                                 0.01189670994589177,
                                 0.01041126526197209,
                                 0.009255462182712733,
                                 0.008330563433362871};

  if (k <= 9)
    return table[static_cast<int>(k)];

  return (1.0 / 12 - (1.0 / 360 - 1.0 / 1260 / ((k + 1) * (k + 1))) / (((k + 1) * (k + 1)))) /
         (k + 1);
}

// The algorithm is explained in https://www.tandfonline.com/doi/abs/10.1080/00949659308811496
template <typename xpu, typename IType, typename OType>
MSHADOW_XINLINE OType _sample_binomial_btrd(IType N,
                                            IType p,
                                            typename RandGenerator<xpu, float>::Impl* gen) {
  OType m   = floor((N + 1) * p);
  OType r   = p / (1 - p);
  OType nr  = (N + 1) * r;
  OType npq = N * p * (1 - p);

  OType b     = 1.15 + 2.53 * sqrt(npq);
  OType a     = -0.0873 + 0.0248 * b + 0.01 * p;
  OType c     = N * p + 0.5;
  OType alpha = (2.83 + 5.1 / b) * sqrt(npq);

  OType v_r      = 0.92 - 4.2 / b;
  OType u_r__v_r = 0.86 * v_r;

  while (true) {
    OType v = gen->uniform();
    if (v <= u_r__v_r) {
      OType u = v / v_r - 0.43;

      return floor((2 * a / (0.5 - abs(u)) + b) * u + c);
    }

    OType u;
    if (v >= v_r) {
      u = gen->uniform() - 0.5;
    } else {
      u           = v / v_r - 0.93;
      OType sgn_u = ((0 < u) - (u < 0));
      u           = sgn_u * 0.5 - u;

      v = gen->uniform() * v_r;
    }

    OType us = 0.5 - abs(u);
    OType k  = floor((2 * a / us + b) * u + c);
    if (k < 0 || k > N) {
      continue;
    }

    v = v * alpha / (a / (us * us) + b);

    OType km = abs(k - m);
    if (km <= 15) {
      OType f = 1;
      for (double i = m; i < k; ++i)
        f = f * (nr / i - r);
      for (double i = k; i < m; ++i)
        v = v * (nr / i - r);

      if (v <= f) {
        return k;
      }

      continue;
    }

    v         = log(v);
    OType rho = (km / npq) * (((km / 3 + 0.625) * km + 1.0 / 6) / npq + 0.5);
    OType t   = -km * km / (2 * npq);
    if (v < t - rho) {
      return k;
    }

    if (v > t + rho) {
      continue;
    }

    OType nm = N - m + 1;
    OType h  = (m + 0.5) * log((m + 1) / (r * nm)) + stirling_approximation(m) +
              stirling_approximation(N - m);

    OType nk  = N - k + 1;
    OType tmp = h + (N + 1) * log(nm / nk) + (k + 0.5) * log(nk * r / (k + 1)) -
                stirling_approximation(k) - stirling_approximation(N - k);
    if (v <= tmp) {
      return k;
    }
  }
}

template <typename xpu, typename IType, typename OType>
MSHADOW_XINLINE OType _sample_binomial_inversion(IType n,
                                                 IType p,
                                                 typename RandGenerator<xpu, float>::Impl* gen) {
  OType N = static_cast<OType>(n);
  OType q = static_cast<OType>(p);
  if (q > 0.5)
    q = 1 - q;

  OType s = 1 - q;

  OType A = 1;
  OType B = q / s;
  OType C = (N + 1) * B;
  OType D = A;
  OType X = 0;

  OType U = gen->uniform();
  OType V = U / pow(s, N);

  do {
    if (V <= A)
      break;
    X = X + 1;
    D = D * (C / X - B);
    A = A + D;
  } while (X < N);

  if (p > 0.5)
    return N - X;

  return X;
}

template <typename xpu, typename IType, typename OType>
MSHADOW_XINLINE OType SampleBinomial(IType n,
                                     IType p,
                                     typename RandGenerator<xpu, float>::Impl* gen) {
  // Generate one sample of the binomial distribution
  if (p >= 1) {
    return static_cast<OType>(n);
  }

  if (p <= 0.5) {
    if (n * p >= 10.0) {
      return _sample_binomial_btrd<xpu, IType, OType>(n, p, gen);
    } else {
      return _sample_binomial_inversion<xpu, IType, OType>(n, p, gen);
    }
  } else {
    IType q = 1.0 - p;
    if (n * q >= 10.0) {
      return n - _sample_binomial_btrd<xpu, IType, OType>(n, q, gen);
    } else {
      return n - _sample_binomial_inversion<xpu, IType, OType>(n, q, gen);
    }
  }
}

template <typename xpu>
struct SampleBinomialKernel {
  template <typename IType, typename OType>
  MSHADOW_XINLINE static void Map(index_t id,
                                  RandGenerator<xpu, float> gen,
                                  const index_t N,
                                  const index_t step,
                                  index_t nParm,
                                  index_t nSample,
                                  const IType* n,
                                  const IType* p,
                                  OType* out) {
    RNG_KERNEL_LOOP(xpu, float, id, gen, N, step, {
      index_t nBatch(1 + (nSample - 1) / nParm);
      out[i] = SampleBinomial<xpu, IType, OType>(n[i / nBatch], p[i / nBatch], &genImpl);
    });
  }
};

template <typename xpu>
struct BinomialSampler {
  template <typename IType, typename OType>
  MSHADOW_FORCE_INLINE void Sample(const Tensor<xpu, 1, IType>& n,
                                   const Tensor<xpu, 1, IType>& p,
                                   const Tensor<xpu, 1, OType>& out,
                                   RandGenerator<xpu, OType>* pgen,
                                   Stream<xpu>* s) {
    RandGenerator<xpu, float>* gen = reinterpret_cast<RandGenerator<xpu, float>*>(pgen);
    LaunchRNG<SampleBinomialKernel<xpu>, xpu>(
        s, gen, out.size(0), n.size(0), out.size(0), n.dptr_, p.dptr_, out.dptr_);
  }
};

template <typename xpu>
struct SampleNegativeBinomialKernel {
  template <typename IType, typename OType>
  MSHADOW_XINLINE static void Map(index_t id,
                                  RandGenerator<xpu, float> gen,
                                  const index_t N,
                                  const index_t step,
                                  index_t nParm,
                                  index_t nSample,
                                  const IType* k,
                                  const IType* p,
                                  OType* out) {
    RNG_KERNEL_LOOP(xpu, float, id, gen, N, step, {
      index_t nBatch(1 + (nSample - 1) / nParm);
      float alpha  = k[i / nBatch];
      float prob   = p[i / nBatch];
      float beta   = (1.0 - prob) / prob;
      float lambda = SampleGamma<xpu, IType, float>(alpha, beta, &genImpl);
      out[i]       = OType(SamplePoisson<xpu>(lambda, &genImpl));
    });
  }
};

template <typename xpu>
struct NegativeBinomialSampler {
  template <typename IType, typename OType>
  MSHADOW_FORCE_INLINE void Sample(const Tensor<xpu, 1, IType>& k,
                                   const Tensor<xpu, 1, IType>& p,
                                   const Tensor<xpu, 1, OType>& out,
                                   RandGenerator<xpu, OType>* pgen,
                                   Stream<xpu>* s) {
    RandGenerator<xpu, float>* gen = reinterpret_cast<RandGenerator<xpu, float>*>(pgen);
    LaunchRNG<SampleNegativeBinomialKernel<xpu>, xpu>(
        s, gen, out.size(0), k.size(0), out.size(0), k.dptr_, p.dptr_, out.dptr_);
  }
};

template <typename xpu>
struct SampleGeneralizedNegativeBinomialKernel {
  template <typename IType, typename OType>
  MSHADOW_XINLINE static void Map(index_t id,
                                  RandGenerator<xpu, float> gen,
                                  const index_t N,
                                  const index_t step,
                                  index_t nParm,
                                  index_t nSample,
                                  const IType* mu,
                                  const IType* alpha,
                                  OType* out) {
    RNG_KERNEL_LOOP(xpu, float, id, gen, N, step, {
      index_t nBatch(1 + (nSample - 1) / nParm);
      float lambda =
          alpha[i / nBatch] == 0 ?
              static_cast<float>(mu[i / nBatch]) :
              SampleGamma<xpu, IType, float>(
                  IType(1) / alpha[i / nBatch], alpha[i / nBatch] * mu[i / nBatch], &genImpl);
      out[i] = OType(SamplePoisson<xpu>(lambda, &genImpl));
    });
  }
};

template <typename xpu>
struct GeneralizedNegativeBinomialSampler {
  template <typename IType, typename OType>
  MSHADOW_FORCE_INLINE void Sample(const Tensor<xpu, 1, IType>& mu,
                                   const Tensor<xpu, 1, IType>& alpha,
                                   const Tensor<xpu, 1, OType>& out,
                                   RandGenerator<xpu, OType>* pgen,
                                   Stream<xpu>* s) {
    RandGenerator<xpu, float>* gen = reinterpret_cast<RandGenerator<xpu, float>*>(pgen);
    LaunchRNG<SampleGeneralizedNegativeBinomialKernel<xpu>, xpu>(
        s, gen, out.size(0), mu.size(0), out.size(0), mu.dptr_, alpha.dptr_, out.dptr_);
  }
};

}  // namespace op
}  // namespace mxnet

#endif  // MXNET_OPERATOR_RANDOM_SAMPLER_H_