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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162

//! TODO: DOC
//!

use crate::capnp_util::*;
use crate::co::{IBackend, ITensorDesc, SharedTensor};
use crate::juice_capnp::negative_log_likelihood_config as capnp_config;
use crate::layer::*;
use crate::util::{native_backend, ArcLock};

#[derive(Debug, Clone)]
#[allow(missing_copy_implementations)]
/// NegativeLogLikelihood Loss Layer
pub struct NegativeLogLikelihood {
    num_classes: usize,
}

impl NegativeLogLikelihood {
    /// Create a NegativeLogLikelihood layer from a NegativeLogLikelihoodConfig.
    pub fn from_config(config: &NegativeLogLikelihoodConfig) -> NegativeLogLikelihood {
        NegativeLogLikelihood {
            num_classes: config.num_classes,
        }
    }

    fn calculate_outer_num(softmax_axis: usize, input_shape: &[usize]) -> usize {
        input_shape.iter().take(softmax_axis + 1).fold(1, |prod, i| prod * i)
    }

    fn calculate_inner_num(softmax_axis: usize, input_shape: &[usize]) -> usize {
        input_shape.iter().skip(softmax_axis + 1).fold(1, |prod, i| prod * i)
    }

    fn batch_size(input_shape: &[usize]) -> usize {
        match input_shape.len() {
            1 => 1,
            2 => input_shape[0],
            _ => panic!("NegativeLogLikelihood layer only supports 1D/2D inputs"),
        }
    }
}

impl<B: IBackend> ILayer<B> for NegativeLogLikelihood {
    impl_ilayer_loss!();

    fn sync_native(&self) -> bool {
        true
    }

    fn reshape(
        &mut self,
        backend: ::std::rc::Rc<B>,
        input_data: &mut Vec<ArcLock<SharedTensor<f32>>>,
        input_gradient: &mut Vec<ArcLock<SharedTensor<f32>>>,
        weights_data: &mut Vec<ArcLock<SharedTensor<f32>>>,
        weights_gradient: &mut Vec<ArcLock<SharedTensor<f32>>>,
        output_data: &mut Vec<ArcLock<SharedTensor<f32>>>,
        output_gradient: &mut Vec<ArcLock<SharedTensor<f32>>>,
    ) {
        let data = input_data[0].read().unwrap();
        let label = input_data[1].read().unwrap();

        input_gradient[0].write().unwrap().resize(data.desc()).unwrap();
        output_data[0].write().unwrap().resize(label.desc()).unwrap();
    }
}

impl<B: IBackend> ComputeOutput<f32, B> for NegativeLogLikelihood {
    fn compute_output(
        &self,
        backend: &B,
        _weights: &[&SharedTensor<f32>],
        input_data: &[&SharedTensor<f32>],
        output_data: &mut [&mut SharedTensor<f32>],
    ) {
        let probabilities = input_data[0];
        let labels = input_data[1];

        let batch_size = Self::batch_size(labels.desc());

        let native = native_backend();
        let native_labels = labels.read(native.device()).unwrap().as_slice::<f32>();
        let native_probabilities = probabilities.read(native.device()).unwrap().as_slice::<f32>();

        let mut writable_loss = Vec::<f32>::new();
        for &label_value in native_labels {
            let probability_value = native_probabilities[label_value as usize];
            writable_loss.push(-probability_value);
        }

        let mut loss = writable_loss.iter().fold(0f32, |sum, &val| sum + val);
        loss = loss / (batch_size as f32);
        writable_loss = vec![loss];

        crate::util::write_to_memory(output_data[0].write_only(native.device()).unwrap(), &writable_loss);
    }
}

impl<B: IBackend> ComputeInputGradient<f32, B> for NegativeLogLikelihood {
    fn compute_input_gradient(
        &self,
        backend: &B,
        weights_data: &[&SharedTensor<f32>],
        output_data: &[&SharedTensor<f32>],
        output_gradients: &[&SharedTensor<f32>],
        input_data: &[&SharedTensor<f32>],
        input_gradients: &mut [&mut SharedTensor<f32>],
    ) {
        let labels = input_data[1];
        let batch_size = Self::batch_size(input_data[0].desc());
        let num_classes = self.num_classes;

        let native = native_backend();
        let native_labels = labels.read(native.device()).unwrap().as_slice::<f32>();
        let mut writable_gradient = vec![0f32; input_gradients[0].desc().size()];

        for (batch_n, &label_value) in native_labels.iter().enumerate() {
            let index = (num_classes * batch_n) + label_value as usize;
            writable_gradient[index] = -1f32;
        }
        crate::util::write_to_memory(
            input_gradients[0].write_only(native.device()).unwrap(),
            &writable_gradient,
        );
    }
}

impl<B: IBackend> ComputeParametersGradient<f32, B> for NegativeLogLikelihood {}

#[derive(Debug, Clone)]
#[allow(missing_copy_implementations)]
/// Specifies configuration parameters for a NegativeLogLikelihood Layer.
pub struct NegativeLogLikelihoodConfig {
    /// How many different classes can be classified.
    pub num_classes: usize,
}

impl<'a> CapnpWrite<'a> for NegativeLogLikelihoodConfig {
    type Builder = capnp_config::Builder<'a>;

    /// Write the NegativeLogLikelihoodConfig into a capnp message.
    fn write_capnp(&self, builder: &mut Self::Builder) {
        builder.set_num_classes(self.num_classes as u64);
    }
}

impl<'a> CapnpRead<'a> for NegativeLogLikelihoodConfig {
    type Reader = capnp_config::Reader<'a>;

    fn read_capnp(reader: Self::Reader) -> Self {
        let num_classes = reader.get_num_classes() as usize;

        NegativeLogLikelihoodConfig {
            num_classes: num_classes,
        }
    }
}

impl Into<LayerType> for NegativeLogLikelihoodConfig {
    fn into(self) -> LayerType {
        LayerType::NegativeLogLikelihood(self)
    }
}