dropout_mask.hpp

Go to the documentation of this file.

//=======================================================================
// Copyright (c) 2014-2023 Baptiste Wicht
// Distributed under the terms of the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#pragma once

#include <chrono> //for std::time

#include "etl/impl/egblas/dropout.hpp"

namespace etl {

template <typename T>
using dropout_distribution = std::conditional_t<std::is_floating_point_v<T>, std::uniform_real_distribution<T>, std::uniform_int_distribution<T>>;

template <typename T = double>
struct dropout_mask_generator_op {
    using value_type = T; 

    const T probability;                           
    random_engine rand_engine;                     
    dropout_distribution<value_type> distribution; 

    static constexpr bool gpu_computable =
        (is_single_precision_t<T> && impl::egblas::has_sdropout_seed) || (is_double_precision_t<T> && impl::egblas::has_ddropout_seed);

    dropout_mask_generator_op(T probability) : probability(probability), rand_engine(std::time(nullptr)), distribution(T(0), T(1)) {}

    value_type operator()() {
        if (distribution(rand_engine) < probability) {
            return T(0);
        } else {
            return T(1);
        }
    }

    template <typename Y>
    auto gpu_compute_hint(Y& y) noexcept {
        std::uniform_int_distribution<long> seed_dist;

        decltype(auto) t1 = force_temporary_gpu_dim_only(y);

        T alpha(1.0);
        impl::egblas::dropout_seed(etl::size(y), probability, alpha, t1.gpu_memory(), 1, seed_dist(rand_engine));

        return t1;
    }
    template <typename Y>
    Y& gpu_compute(Y& y) noexcept {
        std::uniform_int_distribution<long> seed_dist;

        T alpha(1.0);
        impl::egblas::dropout_seed(etl::size(y), probability, alpha, y.gpu_memory(), 1, seed_dist(rand_engine));

        y.validate_gpu();
        y.invalidate_cpu();

        return y;
    }

    friend std::ostream& operator<<(std::ostream& os, const dropout_mask_generator_op& s) {
        return os << "dropout(p=" << s.probability << ")";
    }
};

template <typename G, typename T = double>
struct dropout_mask_generator_g_op {
    using value_type = T; 

    const T probability;                           
    G& rand_engine;                                
    dropout_distribution<value_type> distribution; 

    static constexpr bool gpu_computable =
        (is_single_precision_t<T> && impl::egblas::has_sdropout_seed) || (is_double_precision_t<T> && impl::egblas::has_ddropout_seed);

    dropout_mask_generator_g_op(G& g, T probability) : probability(probability), rand_engine(g), distribution(T(0), T(1)) {}

    value_type operator()() {
        if (distribution(rand_engine) < probability) {
            return T(0);
        } else {
            return T(1);
        }
    }

    template <typename Y>
    auto gpu_compute_hint(Y& y) noexcept {
        std::uniform_int_distribution<long> seed_dist;

        decltype(auto) t1 = force_temporary_gpu_dim_only(y);

        T alpha(1.0);
        impl::egblas::dropout_seed(etl::size(y), probability, alpha, t1.gpu_memory(), 1, seed_dist(rand_engine));

        return t1;
    }
    template <typename Y>
    Y& gpu_compute(Y& y) noexcept {
        std::uniform_int_distribution<long> seed_dist;

        T alpha(1.0);
        impl::egblas::dropout_seed(etl::size(y), probability, alpha, y.gpu_memory(), 1, seed_dist(rand_engine));

        y.validate_gpu();
        y.invalidate_cpu();

        return y;
    }

    friend std::ostream& operator<<(std::ostream& os, const dropout_mask_generator_g_op& s) {
        return os << "dropout(p=" << s.probability << ")";
    }
};

} //end of namespace etl