add HuberLoss and add Loss Trait

candle-nn/src/loss.rs

@@ -2,6 +2,53 @@
 //!
 use candle::{Result, Tensor};
 
+pub trait Loss: Sized {
+    type Config: Sized;
+    fn new(config: Self::Config) -> Result<Self>;
+    fn forward(&mut self, inp: &Tensor, target: &Tensor) -> Result<Tensor>;
+}
+
+/// HuberLoss
+///
+/// A robust loss function that combines `MAE` and `MSE` losses:
+/// - When the absolute element-wise error is less than `delta`, it uses a squared term (MSE loss).
+/// - When the absolute element-wise error is greater than or equal to `delta`, it uses a linear term (MAE loss scaled by `delta`).
+/// # Formula
+///
+/// HuberLoss =
+/// 0.5(x_n - y_n)^2, & |x_n - y_n| < delta
+/// delta(|x_n - y_n| - 0.5delta), & |x_n - y_n| >= delta
+///
+pub struct HuberLoss {
+    delta: f64,
+}
+impl Default for HuberLoss {
+    fn default() -> Self {
+        Self { delta: 1.0 }
+    }
+}
+impl Loss for HuberLoss {
+    type Config = f64;
+    fn new(delta: f64) -> Result<Self> {
+        Ok(Self { delta: delta })
+    }
+    fn forward(&mut self, inp: &Tensor, target: &Tensor) -> Result<Tensor> {
+        if inp.dims() != target.dims() {
+            candle::bail!(
+                "input and target must have the same shape, got inp: {:?}, target: {:?}",
+                inp.dims(),
+                target.dims()
+            );
+        }
+        let diff = (inp - target)?;
+        let abs_diff = diff.abs()?;
+        let mask = abs_diff.le(self.delta)?;
+        let squared_loss = ((&diff * &diff)? * 0.5)?;
+        let linear_loss = ((abs_diff * self.delta)? - 0.5 * self.delta.powi(2))?;
+        let loss = mask.where_cond(&squared_loss, &linear_loss)?;
+        loss.mean_all()
+    }
+}
 /// The negative log likelihood loss.
 ///
 /// Arguments

candle-nn/tests/loss.rs

@@ -6,7 +6,7 @@ extern crate accelerate_src;
 
 use candle::test_utils::to_vec0_round;
 use candle::{Device, Result, Tensor};
-
+use candle_nn::loss::Loss;
 /* Equivalent python code:
 import torch
 import torch.nn.functional as F
@@ -86,3 +86,52 @@ fn binary_cross_entropy_with_logit() -> Result<()> {
     assert_eq!(to_vec0_round(&loss, 4)?, 0.8224);
     Ok(())
 }
+
+/* Equivalent python code:
+import torch
+import torch.nn.functional as F
+
+inp = torch.Tensor([[ 2.3611, -0.8813, -0.5006, -0.2178],
+        [ 0.0419,  0.0763, -1.0457, -1.6692],
+        [-1.0494,  0.8111,  1.5723,  1.2315],
+        [ 1.3081,  0.6641,  1.1802, -0.2547],
+        [ 0.5292,  0.7636,  0.3692, -0.8318]])
+
+target = torch.Tensor([[0., 1., 0., 0.],
+        [0., 1., 0., 0.],
+        [0., 0., 0., 1.],
+        [1., 0., 0., 0.],
+        [0., 0., 1., 0.]])
+
+print(F.huber_loss(inp, target))
+print(F.huber_loss(inp,target,delta=0.88))
+*/
+#[test]
+fn huber_loss() -> Result<()> {
+    let cpu = Device::Cpu;
+    let inp = [
+        [2.3611f32, -0.8813, -0.5006, -0.2178],
+        [0.0419, 0.0763, -1.0457, -1.6692],
+        [-1.0494, 0.8111, 1.5723, 1.2315],
+        [1.3081, 0.6641, 1.1802, -0.2547],
+        [0.5292, 0.7636, 0.3692, -0.8318],
+    ];
+
+    let target = [
+        [0.0f32, 1., 0., 0.],
+        [0., 1., 0., 0.],
+        [0., 0., 0., 1.],
+        [1., 0., 0., 0.],
+        [0., 0., 1., 0.],
+    ];
+
+    let inp = Tensor::new(&inp, &cpu)?;
+    let target = Tensor::new(&target, &cpu)?;
+    let mut loss_func = candle_nn::loss::HuberLoss::default();
+    let loss = loss_func.forward(&inp, &target)?;
+    assert_eq!(to_vec0_round(&loss, 4)?, 0.4734);
+    let mut loss_func = candle_nn::loss::HuberLoss::new(0.88)?;
+    let loss = loss_func.forward(&inp, &target)?;
+    assert_eq!(to_vec0_round(&loss, 4)?, 0.4483);
+    Ok(())
+}