koho 0.1.1

A deep learning model for spectral diffusion over k-cells in arbitrary cell complexes, built on `candle`
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

use candle_core::{Result, Tensor};
use candle_nn::loss::{binary_cross_entropy_with_logit, cross_entropy, mse, nll};

/// Pick your poison.
#[derive(Debug, Clone)]
pub enum LossKind {
    NLL,           // inp is log-probs, target is class indices
    CrossEntropy,  // inp is raw logits, target is class indices
    MSE,           // inp and target are same shape
    BCEWithLogits, // inp is raw logits, target is same shape
}

/// A callable loss function.
#[derive(Debug, Clone)]
pub struct LossFn {
    kind: LossKind,
}

impl LossFn {
    /// Choose which loss to use.
    pub fn new(kind: LossKind) -> Self {
        Self { kind }
    }

    /// Compute the loss scalar for a batch.
    pub fn compute(&self, inp: &Tensor, target: &Tensor) -> Result<Tensor> {
        match self.kind {
            LossKind::NLL => nll(inp, target),
            LossKind::CrossEntropy => cross_entropy(inp, target),
            LossKind::MSE => mse(inp, target),
            LossKind::BCEWithLogits => binary_cross_entropy_with_logit(inp, target),
        }
    }
}