hodu_nn 0.1.9

hodu nn
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

use crate::compat::*;
use crate::module::Module;
use crate::state::{get_state, State};
use hodu_core::{error::HoduResult, scalar::Scalar, tensor::Tensor};

#[derive(Module, Clone)]
pub struct Dropout {
    p: Scalar,
}

impl Dropout {
    pub fn new(p: impl Into<Scalar>) -> Self {
        Dropout { p: p.into() }
    }

    pub fn forward(&self, input: &Tensor) -> HoduResult<Tensor> {
        // Only apply dropout in training mode
        if get_state() == State::Evaluation {
            return Ok(*input);
        }

        // Training mode: apply dropout
        let random = Tensor::rand_uniform_like(input, 0.0, 1.0)?;
        let mask = random.gt_scalar(self.p)?.to_dtype(input.get_dtype())?;
        let scale = 1.0 / (1.0 - self.p.to_f32());
        let scaled_mask = mask.mul_scalar(scale)?;

        input.mul(&scaled_mask)
    }
}