nevermind_neu/layers/
euclidean_loss_layer.rs

1use std::collections::HashMap;
2use std::ops::{Deref, DerefMut};
3
4use ndarray::{Zip, indices};
5
6use log::{debug, info};
7
8use crate::layers::*;
9use crate::cpu_params::*;
10use crate::util::*;
11
12#[derive(Clone)]
13pub struct EuclideanLossLayer<T: Fn(f32) -> f32 + Clone, TD: Fn(f32) -> f32 + Clone> {
14    pub size: usize,
15    pub lr_params: CpuParams,
16    pub l2_regul: f32,
17    pub l1_regul: f32,
18    pub activation: Activation<T, TD>,
19}
20
21impl<T, TD> AbstractLayer for EuclideanLossLayer<T, TD>
22where
23    T: Fn(f32) -> f32 + Sync + Clone + 'static,
24    TD: Fn(f32) -> f32 + Sync + Clone + 'static,
25{
26    fn forward(&mut self, input: ParamsBlob) -> LayerForwardResult {
27        let inp_m = input[0].get_2d_buf_t(TypeBuffer::Output);
28        let inp_m = inp_m.borrow();
29        let inp_m = inp_m.deref();
30
31        let out_m = self.lr_params.get_2d_buf_t(TypeBuffer::Output);
32        let mut out_m = out_m.borrow_mut();
33        let out_m = out_m.deref_mut();
34
35        let ws = self.lr_params.get_2d_buf_t(TypeBuffer::Weights);
36        let ws = ws.borrow();
37        let ws = ws.deref();
38
39        let bias_out = self.lr_params.get_1d_buf_t(TypeBuffer::Bias);
40        let bias_out = bias_out.borrow();
41        let bias_out = bias_out.deref();
42
43        Zip::from(inp_m.rows())
44            .and(out_m.rows_mut())
45            .par_for_each(|inp_r, out_b| {
46                // for each batch
47                let mul_res = ws.dot(&inp_r);
48
49                // for each neuron
50                Zip::from(out_b)
51                    .and(&mul_res)
52                    .and(bias_out)
53                    .for_each(|out_el, in_row, bias_el| {
54                        // for each "neuron"
55                        *out_el = (self.activation.func)(*in_row + bias_el);
56                    });
57            });
58
59        debug!("[ok] ErrorLayer forward()");
60
61        Ok(vec![self.lr_params.clone()])
62    }
63
64    fn backward_output(
65        &mut self,
66        prev_input: ParamsBlob,
67        expected_vec: Array2D,
68    ) -> LayerBackwardResult {
69        let prev_input = prev_input[0].get_2d_buf_t(TypeBuffer::Output);
70        let prev_input = prev_input.borrow();
71        let prev_input = prev_input.deref();
72
73        let self_neu_grad = self.lr_params.get_2d_buf_t(TypeBuffer::NeuGrad);
74        let mut self_neu_grad = self_neu_grad.borrow_mut();
75        let self_neu_grad = self_neu_grad.deref_mut();
76
77        let self_output = self.lr_params.get_2d_buf_t(TypeBuffer::Output);
78        let self_output = self_output.borrow();
79        let self_output = self_output.deref();
80
81        let self_bias_grad = self
82            .lr_params
83            .get_1d_buf_t(TypeBuffer::BiasGrad);
84        let mut self_bias_grad = self_bias_grad.borrow_mut();
85        let self_bias_grad = self_bias_grad.deref_mut();
86
87        // for each batch
88        Zip::from(self_neu_grad.rows_mut())
89            .and(expected_vec.rows())
90            .and(self_output.rows()) 
91            .par_for_each(|err_val_r, expected_r, output_r| {
92                Zip::from(err_val_r).and(expected_r).and(output_r).for_each(
93                    |err_val, expected, output| {
94                        *err_val = (expected - output) * (self.activation.func_deriv)(*output);
95                    },
96                );
97            });
98
99        let ws_grad = self
100            .lr_params
101            .get_2d_buf_t(TypeBuffer::WeightsGrad);
102        let mut ws_grad = ws_grad.borrow_mut();
103        let ws_grad = ws_grad.deref_mut();
104
105        let ws = self
106            .lr_params
107            .get_2d_buf_t(TypeBuffer::Weights);
108        let ws = ws.borrow();
109        let ws = ws.deref();
110
111        let self_bias = self
112            .lr_params
113            .get_1d_buf_t(TypeBuffer::Bias);
114        let mut self_bias = self_bias.borrow_mut();
115        let self_bias = self_bias.deref_mut();
116
117        // calc grad for weights
118        let ws_idxs = indices(ws_grad.dim());
119        Zip::from(ws_grad)
120            .and(ws)
121            .and(ws_idxs)
122            .par_for_each(|val_ws_grad, val_ws, ws_idx| {
123                let self_neu_idx = ws_idx.0;
124                let prev_neu_idx = ws_idx.1;
125
126                let mut avg = 0.0;
127
128                Zip::from(prev_input.column(prev_neu_idx))
129                    .and(self_neu_grad.column(self_neu_idx))
130                    .for_each(|prev_val, err_val| {
131                        avg += prev_val * err_val;
132                    });
133
134                avg = avg / prev_input.column(prev_neu_idx).len() as f32;
135
136                let mut l2_penalty = 0.0;
137                if self.l2_regul != 0.0 {
138                    l2_penalty = self.l2_regul * val_ws;
139                }
140
141                let mut l1_penalty = 0.0;
142                if self.l1_regul == 0.0 {
143                    l1_penalty = self.l1_regul * sign(*val_ws);
144                }
145
146                *val_ws_grad = avg - l2_penalty - l1_penalty;
147            });
148
149        Zip::from(self_neu_grad.columns()).and(self_bias_grad).and(self_bias).for_each(
150            |err_vals, bias_grad, bias| {
151                let grad = err_vals.mean().unwrap();
152
153                let mut l2_penalty = 0.0;
154                if self.l2_regul != 0.0 {
155                    l2_penalty = self.l2_regul * *bias;
156                }
157
158                let mut l1_penalty = 0.0;
159                if self.l1_regul == 0.0 {
160                    l1_penalty = self.l1_regul * sign(*bias);
161                }
162
163                *bias_grad = grad - l2_penalty - l1_penalty;
164            }
165        );
166
167        debug!("[ok] ErrorLayer backward()");
168
169        Ok(vec![self.lr_params.clone()])
170    }
171
172    fn layer_type(&self) -> &str {
173        "EuclideanLossLayer"
174    }
175
176    fn cpu_params(&self) -> Option<CpuParams> {
177        Some(self.lr_params.clone())
178    }
179
180    fn set_cpu_params(&mut self, lp: CpuParams) {
181        self.lr_params = lp;
182    }
183
184    /// Carefull this method overwrites weights and all other params
185    fn set_input_shape(&mut self, sh: &[usize]) {
186        self.lr_params = CpuParams::new_with_bias(self.size, sh[0]);
187    }
188
189    fn size(&self) -> usize {
190        self.size
191    }
192
193    fn copy_layer(&self) -> Box<dyn AbstractLayer> {
194        let mut copy_l = EuclideanLossLayer::new(self.size, self.activation.clone());
195        copy_l.set_cpu_params(self.lr_params.copy());
196        Box::new(copy_l)
197    }
198
199    fn clone_layer(&self) -> Box<dyn AbstractLayer> {
200        Box::new(self.clone())
201    }
202}
203
204impl<T, TD> EuclideanLossLayer<T, TD>
205where
206    T: Fn(f32) -> f32 + Clone,
207    TD: Fn(f32) -> f32 + Clone,
208{
209    pub fn new(size: usize, activation: Activation<T, TD>) -> Self {
210        Self {
211            size,
212            lr_params: CpuParams::empty(),
213            activation,
214            l1_regul: 0.0,
215            l2_regul: 0.0,
216        }
217    }
218
219    pub fn new_box(size: usize, activation: Activation<T, TD>) -> Box<Self> {
220        Box::new(EuclideanLossLayer::new(size, activation))
221    }
222
223    pub fn l2_regularization(mut self, coef: f32) -> Self {
224        self.l2_regul = coef;
225        self
226    }
227
228    pub fn l1_regularization(mut self, coef: f32) -> Self {
229        self.l1_regul = coef;
230        self
231    }
232}
233
234impl<T, TD> WithParams for EuclideanLossLayer<T, TD>
235where
236    T: Fn(f32) -> f32 + Sync + Clone + 'static,
237    TD: Fn(f32) -> f32 + Sync + Clone + 'static,
238{
239    fn cfg(&self) -> HashMap<String, Variant> {
240        let mut cfg: HashMap<String, Variant> = HashMap::new();
241
242        cfg.insert("size".to_owned(), Variant::Int(self.size as i32));
243        // cfg.insert("prev_size".to_owned(), Variant::Int(self.prev_size as i32));
244
245        cfg.insert(
246            "activation".to_owned(),
247            Variant::String(self.activation.name.clone()),
248        );
249
250        cfg.insert("l2_regul".to_owned(), Variant::Float(self.l2_regul));
251        cfg.insert("l1_regul".to_owned(), Variant::Float(self.l1_regul));
252
253        cfg
254    }
255
256    fn set_cfg(&mut self, cfg: &HashMap<String, Variant>) {
257        let mut size: usize = 0;
258
259        if let Some(var_size) = cfg.get("size") {
260            if let Variant::Int(var_size) = var_size {
261                size = *var_size as usize;
262            }
263        }
264
265        if size > 0 {
266            self.size = size;
267            self.lr_params = CpuParams::empty();
268        }
269
270        if let Some(l1_regul) = cfg.get("l1_regul") {
271            if let Variant::Float(l1_regul) = l1_regul {
272                self.l1_regul = *l1_regul;
273            }
274        }
275
276        if let Some(l2_regul) = cfg.get("l2_regul") {
277            if let Variant::Float(l2_regul) = l2_regul {
278                self.l2_regul = *l2_regul;
279            }
280        }
281    }
282}