hodu_nn 0.2.1

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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

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

#[derive(Module, Clone)]
#[module(inputs = 2)]
pub struct BCELoss {
    epsilon: Scalar,
}

impl Default for BCELoss {
    fn default() -> Self {
        Self {
            epsilon: Scalar::F32(1e-7), // for numerical stability
        }
    }
}

impl BCELoss {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_epsilon(epsilon: impl Into<Scalar>) -> Self {
        Self {
            epsilon: epsilon.into(),
        }
    }

    pub fn forward(&self, (pred, target): (&Tensor, &Tensor)) -> HoduResult<Tensor> {
        // Clamp predictions to avoid log(0)
        // pred_clamped = clamp(pred, epsilon, 1 - epsilon)
        let one_minus_eps_scalar = Scalar::one(self.epsilon.dtype()) - self.epsilon;
        let pred_clamped = pred.clamp(self.epsilon, one_minus_eps_scalar)?;

        // BCE = -[target * log(pred) + (1 - target) * log(1 - pred)]

        // First term: target * log(pred)
        let log_pred = pred_clamped.ln()?;
        let first_term = target.mul(&log_pred)?;

        // Second term: (1 - target) * log(1 - pred)
        let one = Tensor::ones_like(target)?;
        let one_minus_target = one.sub(target)?;
        let one_minus_pred = Tensor::ones_like(&pred_clamped)?.sub(&pred_clamped)?;
        let log_one_minus_pred = one_minus_pred.ln()?;
        let second_term = one_minus_target.mul(&log_one_minus_pred)?;

        // Combine
        let bce = first_term.add(&second_term)?.neg()?;

        // Mean over all elements
        bce.mean_all()
    }
}

#[derive(Module, Clone, Default)]
#[module(inputs = 2)]
pub struct BCEWithLogitsLoss;

impl BCEWithLogitsLoss {
    pub fn new() -> Self {
        Self
    }

    pub fn forward(&self, (logits, target): (&Tensor, &Tensor)) -> HoduResult<Tensor> {
        // More numerically stable: BCE(sigmoid(x)) = log(1 + exp(-x)) if target=1
        //                                            = log(1 + exp(x))  if target=0
        // This can be simplified to: max(x, 0) - x * target + log(1 + exp(-|x|))

        let zeros = Tensor::zeros_like(logits)?;
        let max_val = logits.maximum(&zeros)?; // max(x, 0)

        let neg_abs = logits.abs()?.neg()?; // -|x|
        let log_term = neg_abs.exp()?.add_scalar(Scalar::one(neg_abs.dtype()))?.ln()?; // log(1 + exp(-|x|))

        let target_term = logits.mul(target)?; // x * target

        // max(x, 0) - x * target + log(1 + exp(-|x|))
        let loss = max_val.sub(&target_term)?.add(&log_term)?;

        // Mean over all elements
        loss.mean_all()
    }
}