random.h

Go to the documentation of this file.

/*
 * 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.
 */
 
#ifndef MSHADOW_RANDOM_H_
#define MSHADOW_RANDOM_H_
 
#include <cstdlib>
#include <algorithm>
#include <random>
#include "./base.h"
#include "./tensor.h"
#include "./tensor_container.h"
#include <random>
 
 
namespace mshadow {
template<typename Device, typename DType MSHADOW_DEFAULT_DTYPE>
class Random {};
 
template<typename DType>
class Random<cpu, DType> {
 public:
  explicit Random(int seed) {
    this->Seed(seed);
    buffer_.Resize(Shape1(kRandBufferSize));
  }
  ~Random(void) {
  }
  inline void Seed(int seed) {
    rnd_engine_.seed(seed);
    this->rseed_ = static_cast<unsigned>(seed);
  }
  inline unsigned GetSeed() const {
    return rseed_;
  }
  inline void set_stream(Stream<cpu> *stream) {
  }
 
  inline unsigned GetRandInt() {
    return rnd_engine_();
  }
 
  inline void GetRandInt(const Tensor<cpu, 1, unsigned>& dst) {
    std::generate_n(dst.dptr_, dst.size(0), [&](){ return rnd_engine_(); });
  }
 
  template<int dim, class Sampler>
  inline void SampleDistribution(Tensor<cpu, dim, DType> *dst, Sampler sampler) {
    if (dst->CheckContiguous()) {
      std::generate_n(dst->dptr_, dst->shape_.Size(), sampler);
    } else {
      Tensor<cpu, 2, DType> mat = dst->FlatTo2D();
      for (index_t i = 0; i < mat.size(0); ++i) {
        std::generate_n(mat[i].dptr_, mat.size(1), sampler);
      }
    }
  }
 
  template<int dim, typename PType>
  inline void SampleUniform(Tensor<cpu, dim, DType> *dst,
                            PType a = 0.0f , PType b = 1.0f ) {
    // Ensure that half_t is handled correctly.
    typedef typename std::conditional<std::is_floating_point<DType>::value,
                                      DType, double>::type FType;
    typedef typename std::conditional<std::is_integral<DType>::value,
                                      std::uniform_int_distribution<DType>,
                                      std::uniform_real_distribution<FType>>::type GType;
    GType dist_uniform(a, b);
    SampleDistribution(dst, [&](){ return dist_uniform(rnd_engine_);});
  }
 
  template<int dim, typename PType>
  inline void SampleGaussian(Tensor<cpu, dim, DType> *dst,
                             PType mu = 0.0f, PType sigma = 1.0f ) {
    if (sigma <= 0) {
      *dst = mu; return;
    }
    typedef typename std::conditional<std::is_floating_point<DType>::value,
                                      DType, double>::type GType;
    std::normal_distribution<GType> dist_normal(mu, sigma);
    SampleDistribution(dst, [&](){ return dist_normal(rnd_engine_);});
  }
 
  template<int dim, typename PType>
  inline void SampleGamma(Tensor<cpu, dim, DType> *dst,
                          PType alpha, PType beta) {
    typedef typename std::conditional<std::is_floating_point<DType>::value,
                                      DType, double>::type GType;
    std::gamma_distribution<GType> dist_gamma(alpha, beta);
    SampleDistribution(dst, [&](){ return dist_gamma(rnd_engine_);});
  }
 
  template<int dim, typename PType>
  inline void SampleExponential(Tensor<cpu, dim, DType> *dst, PType lambda ) {
    typedef typename std::conditional<std::is_floating_point<DType>::value,
                                      DType, double>::type GType;
    std::exponential_distribution<GType> dist_exp(lambda);
    SampleDistribution(dst, [&](){ return dist_exp(rnd_engine_);});
  }
 
  template<int dim, typename PType>
  inline void SamplePoisson(Tensor<cpu, dim, DType> *dst, PType lambda) {
    typedef typename std::conditional<std::is_integral<DType>::value, DType, int>::type GType;
    std::poisson_distribution<GType> dist_poisson(lambda);
    SampleDistribution(dst, [&](){ return static_cast<DType>(dist_poisson(rnd_engine_));});
  }
 
  template<int dim, typename PType1, typename PType2>
  inline void SampleNegativeBinomial(Tensor<cpu, dim, DType> *dst, PType1 k, PType2 p) {
    typedef typename std::conditional<std::is_integral<DType>::value, DType, int>::type GType;
    std::negative_binomial_distribution<GType> dist_negbinomial(k, p);
    SampleDistribution(dst, [&](){ return static_cast<DType>(dist_negbinomial(rnd_engine_));});
  }
 
  template<int dim, typename PType>
  inline void SampleGeneralizedNegativeBinomial(Tensor<cpu, dim, DType> *dst,
                                                PType mu, PType alpha) {
    if (alpha == PType(0)) {
      SamplePoisson(dst, mu);  // limit of Poisson
    } else {
      PType r(PType(1) / alpha);
      PType beta = mu * alpha;
      std::gamma_distribution<> dist_gamma(r, beta);
      typedef typename std::conditional<std::is_integral<DType>::value, DType, int>::type GType;
      SampleDistribution(dst,
        [&](){ std::poisson_distribution<GType> dist_poisson(dist_gamma(rnd_engine_));
               return static_cast<DType>(dist_poisson(rnd_engine_));});
    }
  }
 
  template<int dim>
  inline expr::ReshapeExp<Tensor<cpu, 1, DType>, DType, dim, 1>
  gaussian(Shape<dim> shape) {
    buffer_.Resize(Shape1(shape.Size()));
    this->SampleGaussian(&buffer_, 0.0f, 1.0f);
    return expr::reshape(buffer_, shape);
  }
  template<int dim>
  inline expr::ReshapeExp<Tensor<cpu, 1, DType>, DType, dim, 1>
  uniform(Shape<dim> shape) {
    buffer_.Resize(Shape1(shape.Size()));
    this->SampleUniform(&buffer_, 0.0f, 1.0f);
    return expr::reshape(buffer_, shape);
  }
 
  std::mt19937 &GetRndEngine() {
    return rnd_engine_;
  }
 
 private:
  std::mt19937 rnd_engine_;
  unsigned rseed_;
  TensorContainer<cpu, 1, DType> buffer_;
};  // class Random<cpu, DType>
 
// only allow GPU PRNG when cuda is enabled
#if MSHADOW_USE_CUDA
 
template<typename DType>
class Random<gpu, DType> {
 public:
  explicit Random(int seed) : gen_(NULL) {
    this->Seed(seed);
    buffer_.Resize(Shape1(kRandBufferSize));
  }
  ~Random(void) MSHADOW_THROW_EXCEPTION {
    DeleteGenerator();
  }
  inline void set_stream(Stream<gpu> *stream) {
    curandStatus_t status;
    status = curandSetStream(gen_, Stream<gpu>::GetStream(stream));
 
    CHECK_EQ(status, CURAND_STATUS_SUCCESS) << "set_stream CURAND failed";
  }
  inline void Seed(int seed) {
    // Create a new rng, either initially or if the RNG type can't reset its offset.
    if (gen_ == NULL || (curandSetGeneratorOffset(gen_, 0ULL) != CURAND_STATUS_SUCCESS))
      CreateGenerator();
    // Now set the seed.
    curandStatus_t status;
    status = curandSetPseudoRandomGeneratorSeed(gen_, static_cast<uint64_t>(seed));
    CHECK_EQ(status, CURAND_STATUS_SUCCESS) << "Set CURAND seed failed.";
  }
  inline void GetRandInt(const Tensor<gpu, 1, unsigned>& dst) {
    curandStatus_t status;
    status = curandGenerate(gen_, dst.dptr_, dst.size(0));
    CHECK_EQ(status, CURAND_STATUS_SUCCESS) << "CURAND Gen rand ints failed."
                                            << " size = " << dst.size(0);
  }
  template<int dim>
  inline void SampleUniform(Tensor<gpu, dim, DType> *dst,
                            DType a = 0.0f, DType b = 1.0f);
 
  template<int dim>
  inline void SampleGaussian(Tensor<gpu, dim, DType> *dst,
                             DType mu = 0.0f, DType sigma = 1.0f);
  template<int dim>
  inline expr::ReshapeExp<Tensor<gpu, 1, DType>, DType, dim, 1>
  gaussian(Shape<dim> shape, DType mu = 0.0f, DType sigma = 1.0f);
  template<int dim>
  inline expr::ReshapeExp<Tensor<gpu, 1, DType>, DType, dim, 1>
  uniform(Shape<dim> shape);
 
 private:
  inline void GenGaussian(float *dptr, size_t size, float mu, float sigma) {
    curandStatus_t status;
    status = curandGenerateNormal(gen_, dptr, size, mu, sigma);
    CHECK_EQ(status, CURAND_STATUS_SUCCESS) << "CURAND Gen Normal float failed."
                                            << " size = " << size
                                            << ",mu = " << mu
                                            << ",sigma = " << sigma;
  }
  inline void GenGaussian(double *dptr, size_t size, double mu, double sigma) {
    curandStatus_t status;
    status = curandGenerateNormalDouble(gen_, dptr, size, mu, sigma);
    CHECK_EQ(status, CURAND_STATUS_SUCCESS) << "CURAND Gen Normal double failed."
                                            << " size = " << size
                                            << ",mu = " << mu
                                            << ",sigma = " << sigma;
  }
  inline void GenUniform(float *dptr, size_t size) {
    curandStatus_t status;
    status = curandGenerateUniform(gen_, dptr, size);
    CHECK_EQ(status, CURAND_STATUS_SUCCESS) << "CURAND Gen Uniform float failed."
                                            << " size = " << size;
  }
  inline void GenUniform(double *dptr, size_t size) {
    curandStatus_t status;
    status = curandGenerateUniformDouble(gen_, dptr, size);
    CHECK_EQ(status, CURAND_STATUS_SUCCESS) << "CURAND Gen Uniform double failed."
                                            << " size = " << size;
  }
  inline void CreateGenerator() {
    if (gen_ != NULL)
      DeleteGenerator();
    curandStatus_t status;
    status = curandCreateGenerator(&gen_, CURAND_RNG_PSEUDO_DEFAULT);
    CHECK_EQ(status, CURAND_STATUS_SUCCESS) << "Cannot create CURAND Generator";
  }
  inline void DeleteGenerator() {
    if (gen_ != NULL) {
      curandStatus_t status;
      status = curandDestroyGenerator(gen_);
      CHECK_EQ(status, CURAND_STATUS_SUCCESS) << "Destory CURAND Gen failed";
      gen_ = NULL;
    }
  }
  curandGenerator_t gen_;
  TensorContainer<gpu, 1, DType> buffer_;
};  // class Random<gpu, DType>
#endif  // MSHADOW_USE_CUDA
 
#ifdef __CUDACC__
// implementations that depends on cuda kernels
template<typename DType>
template<int dim>
inline void Random<gpu, DType>::SampleUniform(
    Tensor<gpu, dim, DType> *dst, DType a, DType b) {
  if (a == 0.0f && b == 1.0f) {
    if (dst->CheckContiguous()) {
      this->GenUniform(dst->dptr_, dst->shape_.Size());
    } else {
      *dst = this->uniform(dst->shape_);
    }
  } else {
    *dst = this->uniform(dst->shape_) * (b - a) + a;
  }
}
template<typename DType>
template<int dim>
inline void Random<gpu, DType>::SampleGaussian(
    Tensor<gpu, dim, DType> *dst, DType mu, DType sigma) {
  // We need to check whether the shape size is even since CuRand supports only normal distribution
  // generation of even number of elements.
  if (dst->CheckContiguous() && (dst->shape_.Size() % 2 == 0)) {
    this->GenGaussian(dst->dptr_, dst->shape_.Size(), mu, sigma);
  } else {
    *dst = this->gaussian(dst->shape_, mu, sigma);
  }
}
 
template<typename DType>
template<int dim>
inline expr::ReshapeExp<Tensor<gpu, 1, DType>, DType, dim, 1>
Random<gpu, DType>::gaussian(Shape<dim> shape, DType mu, DType sigma) {
  size_t aligned_sz = ((shape.Size() + 1UL) >> 1) << 1;
  // allocate alligned size
  buffer_.Resize(Shape1(aligned_sz));
  buffer_.Resize(Shape1(shape.Size()));
  this->GenGaussian(buffer_.dptr_, aligned_sz, mu, sigma);
  return expr::reshape(buffer_, shape);
}
 
template<typename DType>
template<int dim>
inline expr::ReshapeExp<Tensor<gpu, 1, DType>, DType, dim, 1>
Random<gpu, DType>::uniform(Shape<dim> shape) {
  buffer_.Resize(Shape1(shape.Size()));
  this->GenUniform(buffer_.dptr_, buffer_.size(0));
  return expr::reshape(buffer_, shape);
}
#endif  // __CUDACC__
}  // namespace mshadow
#endif  // MSHADOW_RANDOM_H_