tensorlogic_train/loss/
kldivergenceloss_traits.rs1use 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}