hodu_nn/losses/
huber.rs

1use crate::compat::*;
2use crate::module::Module;
3use hodu_core::{error::HoduResult, scalar::Scalar, tensor::Tensor};
4
5#[derive(Module, Clone)]
6#[module(inputs = 2)]
7pub struct HuberLoss {
8    delta: Scalar,
9}
10
11impl HuberLoss {
12    pub fn new(delta: impl Into<Scalar>) -> Self {
13        Self { delta: delta.into() }
14    }
15
16    pub fn forward(&self, (pred, target): (&Tensor, &Tensor)) -> HoduResult<Tensor> {
17        let diff = pred.sub(target)?;
18        let abs_diff = diff.abs()?;
19
20        let quadratic_loss = diff
21            .pow_scalar(Scalar::from_f32(2.0, diff.dtype()))?
22            .div_scalar(Scalar::from_f32(2.0, diff.dtype()))?;
23        let linear_loss = abs_diff
24            .mul_scalar(self.delta)?
25            .sub_scalar(self.delta.powi(2) / Scalar::from_f32(2.0, self.delta.dtype()))?;
26
27        let mask = abs_diff.le_scalar(self.delta)?;
28
29        let loss = mask
30            .mul(&quadratic_loss)?
31            .add(&mask.logical_not()?.mul(&linear_loss)?)?;
32
33        let pred_shape = pred.shape();
34        let batch_size = Scalar::from_f32(pred_shape[0] as f32, diff.dtype());
35        let mean = loss.sum_all()?.div_scalar(batch_size)?;
36
37        Ok(mean)
38    }
39}
hodu_nn/losses/huber.rs

hodu_nn/losses/
huber.rs
