hodu_nn 0.1.2

hodu nn
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

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

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

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

    pub fn forward(&self, (pred, target): (&Tensor, &Tensor)) -> HoduResult<Tensor> {
        let diff = pred.sub(target)?;
        let abs_diff = diff.abs()?;

        let quadratic_loss = diff
            .pow_scalar(Scalar::from_f32(2.0, diff.get_dtype()))?
            .div_scalar(Scalar::from_f32(2.0, diff.get_dtype()))?;
        let linear_loss = abs_diff
            .mul_scalar(self.delta)?
            .sub_scalar(self.delta.powi(2) / Scalar::from_f32(2.0, self.delta.get_dtype()))?;

        let mask = abs_diff.le_scalar(self.delta)?;

        let loss = mask
            .mul(&quadratic_loss)?
            .add(&mask.logical_not()?.mul(&linear_loss)?)?;

        let pred_layout = pred.get_layout();
        let pred_shape = pred_layout.get_shape();
        let batch_size = Scalar::from_f32(pred_shape[0] as f32, diff.get_dtype());
        let mean = loss.sum_all()?.div_scalar(batch_size)?;

        Ok(mean)
    }
}