dropout_module.hpp

Go to the documentation of this file.

 
#pragma once
 
 
#include <optional>
#include <cstdint>
 
#include "pulsatrix/module.hpp"
 
namespace pulsatrix {
 
class DropoutModule : public Module {
public:
    DropoutModule(float p, DeviceBackend* backend, uint64_t seed);
 
    DropoutModule(float p, DeviceBackend* backend);
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override;
 
    [[nodiscard]] OpType op_type() const override { return OpType::Elementwise; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] bool supports_lrp_rule(LRPRule) const override { return true; }
 
 
    [[nodiscard]] std::optional<DeviceType> compute_device() const override { return backend_->device(); }
 
protected:
    [[nodiscard]] Tensor forward_impl(const Tensor& input) override;
 
private:
    float p_;
    float scale_;
    DeviceBackend* backend_;
    uint64_t seed_;
    uint64_t draws_ = 0;  // elements consumed from the (seed_, k) stream so far
    Shape last_shape_ = Shape({0});
    Tensor last_mask_;    // 1.0 (kept) or 0.0 (dropped); meaningful only when !last_was_identity_
    bool last_was_identity_ = true;
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix