hodu_nn 0.1.4

hodu nn
Documentation 
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
40
41
42
43
44
45
46
47
48

use crate::compat::*;
use crate::module::Module;
use hodu_core::{error::HoduResult, tensor::Tensor};

#[derive(Module, Clone)]
#[module(inputs = 2)]
pub struct CrossEntropyLoss {
    dim: i32,
}

impl CrossEntropyLoss {
    pub fn new() -> Self {
        Self { dim: -1 }
    }

    pub fn with_dim(dim: i32) -> Self {
        Self { dim }
    }

    pub fn forward(&self, (logits, target): (&Tensor, &Tensor)) -> HoduResult<Tensor> {
        // Apply log_softmax to logits
        let log_probs = logits.log_softmax(self.dim)?;

        // Normalize dim to positive index
        let rank = log_probs.get_layout().get_shape().len() as i32;
        let gather_dim = if self.dim < 0 {
            rank + self.dim
        } else {
            self.dim
        };

        // Gather the log probabilities for the correct classes
        // target shape: [...] -> unsqueeze to [..., 1] at gather_dim
        let target_unsqueezed = target.unsqueeze(-1)?;

        // Gather along class dimension
        let gathered = log_probs.gather(gather_dim as i64, &target_unsqueezed)?;

        // Remove the extra dimension
        let gathered_squeezed = gathered.squeeze(Some(-1))?;

        // Negative log likelihood
        let nll = gathered_squeezed.neg()?;

        // Mean over batch
        nll.mean_all()
    }
}