tensor_random.h

Go to the documentation of this file.

#ifndef MSHADOW_TENSOR_RANDOM_H
#define MSHADOW_TENSOR_RANDOM_H

#include <cstdlib>
#include "tensor.h"
#include "tensor_container.h"

namespace mshadow {
    template<typename Device>
    class Random {};

    template<>
    class Random<cpu> {
    public:
        Random<cpu>( int seed ){
            #if MSHADOW_USE_MKL
            int status = vslNewStream(&vStream_, VSL_BRNG_MT19937, seed);
            utils::Assert( status == VSL_STATUS_OK, "MKL VSL Random engine failed to be initialized.\n" );
            #else
            srand(seed);
            #endif
            buffer_.Resize( Shape1( kRandBufferSize ) );
        }
        ~Random<cpu>() {
            #if MSHADOW_USE_MKL
            vslDeleteStream(&vStream_);
            #endif
        }
        inline void Seed( int seed ){
            #if MSHADOW_USE_MKL
            int status = vslDeleteStream(&vStream_);
            utils::Assert(status == VSL_STATUS_OK);
            status = vslNewStream(&vStream_, VSL_BRNG_MT19937, seed);
            utils::Assert(status == VSL_STATUS_OK);
            #else
            srand( seed );
            #endif
        }
        template<int dim>
        inline void SampleUniform( Tensor<cpu, dim> &dst, real_t a=0.0f, real_t b=1.0f ) {
            Tensor<cpu, 2> mat = dst.FlatTo2D();
            for ( index_t i = 0; i < mat.shape[1]; ++i ) {
                #if MSHADOW_USE_MKL
                #if MSHADOW_SINGLE_PRECISION
                int status = vsRngUniform( 0, vStream_, mat.shape[0], mat[i].dptr, a, b );
                #else
                int status = vdRngUniform( 0, vStream_, mat.shape[0], mat[i].dptr, a, b );
                #endif
                utils::Assert(status == VSL_STATUS_OK, "Failed to generate random number by MKL.\n" );
                #else
                // use stdlib
                for ( index_t j = 0; j < mat.shape[0]; ++j ) {
                    mat[i][j] = this->RandNext()*(b-a) + a;
                }
                #endif
            }
        }
        template<int dim>
        inline void SampleGaussian( Tensor<cpu, dim> &dst, real_t mu = 0.0f, real_t sigma = 1.0f ) {
            if( sigma <= 0.0f ) {
                dst = mu; return;
            }
            Tensor<cpu, 2> mat = dst.FlatTo2D();
            for (index_t i = 0; i < mat.shape[1]; ++i) {
                #if MSHADOW_USE_MKL
                #if MSHADOW_SINGLE_PRECISION
                int status = vsRngGaussian( 0, vStream_, mat.shape[0], mat[i].dptr, mu, sigma );
                #else
                int status = vdRngGaussian( 0, vStream_, mat.shape[0], mat[i].dptr, mu, sigma );
                #endif
                utils::Assert(status == VSL_STATUS_OK, "Failed to generate random number by MKL.\n" );
                #else
                real_t g1 = 0.0f, g2 = 0.0f;
                for (index_t j = 0; j < mat.shape[0]; ++j) {
                    if( (j & 1) == 0 ){
                        this->SampleNormal2D( g1, g2 );
                        mat[i][j] = mu + g1 * sigma;
                    }else{
                        mat[i][j] = mu + g2 * sigma;
                    }
                }
                #endif
            }
        }
        template<int dim>
        inline expr::ReshapeExp<Tensor<cpu,1>,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>,dim,1> uniform( Shape<dim> shape ){
            buffer_.Resize( Shape1( shape.Size() ) );
            this->SampleUniform( buffer_, 0.0f, 1.0f );
            return expr::reshape( buffer_, shape );
        }
    private:
        inline real_t RandNext( void ){
            return static_cast<real_t>(rand()) / (static_cast<real_t>(RAND_MAX)+1.0f);
        }
        inline real_t RandNext2( void ){
            return (static_cast<real_t>( rand() ) + 1.0 ) / (static_cast<real_t>(RAND_MAX) + 2.0);
        }
        inline void SampleNormal2D( real_t &xx, real_t &yy ){
            real_t x,y,s;
            do{
                x = 2.0f * RandNext2() - 1.0f;
                y = 2.0f * RandNext2() - 1.0f;
                s = x*x + y*y;
            }while( s >= 1.0f || s == 0.0f );
            real_t t = std::sqrt( -2.0f * std::log( s ) / s ) ;
            xx = x * t; yy = y * t;
        }
    private:
        #if MSHADOW_USE_MKL

        VSLStreamStatePtr vStream_;
        #endif

        TensorContainer<cpu,1> buffer_;
    }; // class Random<cpu>

#ifdef __CUDACC__

    template<>
    class Random<gpu> {
    public:
        Random<gpu>(int seed) {
            curandStatus_t status;
            status = curandCreateGenerator(&gen_, CURAND_RNG_PSEUDO_DEFAULT);
            utils::Assert(status == CURAND_STATUS_SUCCESS, "Can not create CURAND Generator");
            this->Seed( seed );
            buffer_.Resize( Shape1(kRandBufferSize) );
        }

        ~Random<gpu>() {
            curandStatus_t status;
            status = curandDestroyGenerator(gen_);
            utils::Assert(status == CURAND_STATUS_SUCCESS, "Destory CURAND Gen failed");
        }
        inline void Seed( int seed ){
            curandStatus_t status;
            status = curandSetPseudoRandomGeneratorSeed(gen_, seed);
            utils::Assert(status == CURAND_STATUS_SUCCESS, "Set CURAND seed failed.");
        }
        template<int dim>
        inline void SampleUniform(Tensor<gpu, dim> &dst, real_t a=0.0f, real_t b=1.0f) {
            if( a == 0.0f && b == 1.0f ){
                dst = this->uniform( dst.shape );
            }else{
                dst = this->uniform( dst.shape ) *(b-a) + a;
            }
        }
        template<int dim>
        inline void SampleGaussian(Tensor<gpu, dim> &dst, real_t mu = 0.0f, real_t sigma = 1.0f) {
            dst = this->gaussian( dst.shape, mu, sigma );
        }
        template<int dim>
        inline expr::ReshapeExp<Tensor<gpu,1>,dim,1> gaussian( Shape<dim> shape, real_t mu=0.0f, real_t sigma=1.0f){
            size_t aligned_sz = ((shape.Size() + 1UL)>>1)<<1;
            // allocate alligned size
            buffer_.Resize( Shape1( aligned_sz ) );
            buffer_.Resize( Shape1( shape.Size() ) );
            curandStatus_t status;
            #if MSHADOW_SINGLE_PRECISION
            status = curandGenerateNormal(gen_, buffer_.dptr, aligned_sz , mu, sigma);
            #else
            status = curandGenerateNormalDouble(gen_, buffer_.dptr, buffer_.shape[0], mu, sigma);
            #endif
            utils::Assert(status == CURAND_STATUS_SUCCESS, "CURAND Gen Uniform failed\n");
            return expr::reshape( buffer_, shape );
        }
        template<int dim>
        inline expr::ReshapeExp<Tensor<gpu,1>,dim,1> uniform(Shape<dim> shape) {
            buffer_.Resize( Shape1( shape.Size() ) );
            curandStatus_t status;
            #if MSHADOW_SINGLE_PRECISION
            status = curandGenerateUniform(gen_, buffer_.dptr, buffer_.shape[0] );
            #else
            status = curandGenerateUniformDouble(gen_, buffer_.dptr, buffer_.shape[0] );
            #endif
            utils::Assert(status == CURAND_STATUS_SUCCESS, "CURAND Gen Uniform failed\n");
            return expr::reshape( buffer_, shape );
        }
    private:
        curandGenerator_t gen_;
        TensorContainer<gpu, 1> buffer_;
    }; // class Random<gpu>
    #endif

}; // namespace mshadow

#endif // MSHADOW_TENSOR_RANDOM_H