| /* |
| * 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_); |
| } |
| }; |
| |
| 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_ |
