Skip to main content

tensorlogic_train/loss/
kldivergenceloss_traits.rs

1//! # KLDivergenceLoss - Trait Implementations
2//!
3//! This module contains trait implementations for `KLDivergenceLoss`.
4//!
5//! ## Implemented Traits
6//!
7//! - `Default`
8//! - `Loss`
9
10use crate::{TrainError, TrainResult};
11use scirs2_core::ndarray::{Array, ArrayView, Ix2};
12
13use super::functions::Loss;
14use super::types::KLDivergenceLoss;
15
16impl Default for KLDivergenceLoss {
17    fn default() -> Self {
18        Self { epsilon: 1e-10 }
19    }
20}
21
22impl Loss for KLDivergenceLoss {
23    fn compute(
24        &self,
25        predictions: &ArrayView<f64, Ix2>,
26        targets: &ArrayView<f64, Ix2>,
27    ) -> TrainResult<f64> {
28        if predictions.shape() != targets.shape() {
29            return Err(TrainError::LossError(format!(
30                "Shape mismatch: predictions {:?} vs targets {:?}",
31                predictions.shape(),
32                targets.shape()
33            )));
34        }
35        let mut total_loss = 0.0;
36        for i in 0..predictions.nrows() {
37            for j in 0..predictions.ncols() {
38                let pred = predictions[[i, j]].max(self.epsilon);
39                let target = targets[[i, j]].max(self.epsilon);
40                total_loss += target * (target / pred).ln();
41            }
42        }
43        Ok(total_loss)
44    }
45    fn gradient(
46        &self,
47        predictions: &ArrayView<f64, Ix2>,
48        targets: &ArrayView<f64, Ix2>,
49    ) -> TrainResult<Array<f64, Ix2>> {
50        let mut grad = Array::zeros(predictions.raw_dim());
51        for i in 0..predictions.nrows() {
52            for j in 0..predictions.ncols() {
53                let pred = predictions[[i, j]].max(self.epsilon);
54                let target = targets[[i, j]].max(self.epsilon);
55                grad[[i, j]] = -target / pred;
56            }
57        }
58        Ok(grad)
59    }
60}