hodu_nn 0.1.9

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

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

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

impl Default for NLLLoss {
    fn default() -> Self {
        Self { dim: -1 }
    }
}

impl NLLLoss {
    pub fn new() -> Self {
        Self::default()
    }

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

    pub fn forward(&self, (log_probs, target): (&Tensor, &Tensor)) -> HoduResult<Tensor> {
        // 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, &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()
    }
}