minidx_core/layers/
linear.rs

1use crate::matmul::MatMulImpl;
2use crate::{Dtype, Shape};
3
4/// A fully-connected layer with a fixed number of inputs and outputs. No bias.
5#[derive(Clone, Debug)]
6pub struct Dense<E: Dtype + MatMulImpl, const I: usize, const O: usize> {
7    pub(crate) weights: [[E; I]; O],
8}
9
10impl<E: Dtype + MatMulImpl, const I: usize, const O: usize> Default for Dense<E, I, O> {
11    fn default() -> Self {
12        Dense {
13            weights: [[E::default(); I]; O],
14        }
15    }
16}
17
18impl<E: Dtype + MatMulImpl, const I: usize, const O: usize> Dense<E, I, O> {
19    #[inline]
20    fn forward(&self, input: &[E; I]) -> [E; O] {
21        let mut out: [E; O] = [E::default(); O];
22
23        let (m, k) = (1, I);
24        let n = O;
25        E::matmul(
26            (m, k, n),
27            true, // I think this needs to be false?
28            input.as_ptr(),
29            Shape::strides(&(1, I)),
30            self.weights.as_ptr() as *const E,
31            Shape::strides(&(I, O)),
32            out.as_mut_ptr(),
33            Shape::strides(&(1, O)),
34        );
35
36        out
37    }
38
39    #[inline]
40    fn gradients_wrt_input(&self, output_gradients: &[E; O]) -> [E; I] {
41        let mut out: [E; I] = [E::default(); I];
42
43        let (m, k) = (1, I);
44        let n = O;
45        let strides = (m, n).strides();
46        E::matmul(
47            (m, n, k),
48            true,
49            output_gradients.as_ptr(),
50            strides,
51            self.weights.as_ptr() as *const E,
52            Shape::strides(&(O, I)),
53            out.as_mut_ptr(),
54            Shape::strides(&(1, I)),
55        );
56        out
57    }
58
59    #[inline]
60    fn gradients_wrt_weights(&self, input: &[E; I], output_gradients: &[E; O]) -> [[E; I]; O] {
61        let mut out: [[E; I]; O] = [[E::default(); I]; O];
62
63        let (m, k) = (1, I);
64        let n = O;
65        let strides = (m, n).strides();
66        E::matmul(
67            (k, m, n),
68            false, // Just flipped this to false and it still works?
69            input.as_ptr(),
70            Shape::strides(&(I, 1)),
71            output_gradients.as_ptr(),
72            strides,
73            out.as_mut_ptr() as *mut E,
74            Shape::strides(&(I, O)),
75        );
76        out
77    }
78}
79
80impl<E: Dtype + MatMulImpl, const I: usize, const O: usize> crate::BaseModule for Dense<E, I, O> {}
81
82impl<E: Dtype + MatMulImpl, const I: usize, const O: usize> crate::Module<[E; I]>
83    for Dense<E, I, O>
84{
85    type Output = [E; O];
86
87    fn forward(&self, x: &[E; I]) -> Result<Self::Output, crate::Error> {
88        Ok(Dense::forward(self, &x))
89    }
90}
91
92impl<E: Dtype + MatMulImpl, const I: usize, const O: usize> crate::RevModule<[E; I]>
93    for Dense<E, I, O>
94{
95    type SelfGrads = [[E; I]; O];
96
97    fn reverse(&self, inputs: &[E; I], grads_wrt_output: &[E; O]) -> ([E; I], Self::SelfGrads) {
98        (
99            Dense::gradients_wrt_input(self, grads_wrt_output),
100            Dense::gradients_wrt_weights(self, inputs, grads_wrt_output),
101        )
102    }
103
104    fn apply(
105        &mut self,
106        applyer: &mut impl crate::optimizers::GradApplyer,
107        updates: Self::SelfGrads,
108    ) -> Result<(), crate::Error> {
109        applyer.apply(updates, &mut self.weights)
110    }
111}
112
113impl<E: Dtype + MatMulImpl, const I: usize, const O: usize> crate::ResetParams for Dense<E, I, O> {
114    fn rand_params<RNG: rand::Rng>(
115        &mut self,
116        rng: &mut RNG,
117        scale: f32,
118    ) -> Result<(), crate::Error> {
119        // Xavier/Glorot Initialization: initial values from a distribution with
120        // zero mean and a variance of 2 / (inp + outp).
121        // Can use either normal or uniform distribution, we use normal for now.
122        let normal = rand_distr::Normal::new(0.0, 2.0 / (I as f32 + O as f32).sqrt()).unwrap();
123
124        self.weights.iter_mut().for_each(|r| {
125            r.iter_mut().for_each(|w| {
126                let s: f32 = rng.sample::<f32, _>(normal) * scale;
127                *w = E::from_f32(s).unwrap();
128            })
129        });
130        Ok(())
131    }
132}
133
134impl<E: Dtype + MatMulImpl, const I: usize, const O: usize> crate::LoadableModule
135    for Dense<E, I, O>
136{
137    fn save(
138        &self,
139        path: String,
140        dict: &mut std::collections::HashMap<String, Vec<f64>>,
141    ) -> Result<(), crate::LoadSaveError> {
142        dict.insert(
143            path,
144            self.weights
145                .iter()
146                .flatten()
147                .map(|f| f.to_f64().unwrap())
148                .collect(),
149        );
150        Ok(())
151    }
152
153    fn load(
154        &mut self,
155        path: String,
156        dict: &std::collections::HashMap<String, Vec<f64>>,
157    ) -> Result<(), crate::LoadSaveError> {
158        let params = dict.get(&path).ok_or(crate::LoadSaveError {
159            path: path.clone(),
160            err: "Parameters missing".into(),
161        })?;
162        if params.len() != I * O {
163            return Err(crate::LoadSaveError {
164                path,
165                err: format!(
166                    "Parameters have wrong size: got {}, want {}",
167                    params.len(),
168                    I * O
169                )
170                .into(),
171            });
172        }
173
174        for (a, b) in self.weights.iter_mut().flatten().zip(params.into_iter()) {
175            *a = E::from_f64(*b).unwrap();
176        }
177        Ok(())
178    }
179}
180
181impl<E: Dtype + MatMulImpl, const I: usize, const O: usize> crate::VisualizableUnit
182    for Dense<E, I, O>
183{
184    const KIND: &'static str = "dense";
185    type Params = [[E; I]; O];
186    fn params(&self) -> &Self::Params {
187        &self.weights
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194
195    #[test]
196    fn test_zero_weights() {
197        let layer = Dense::<f32, 1, 2>::default();
198        assert_eq!(layer.forward(&[1.0]), [0.0, 0.0],);
199    }
200
201    #[test]
202    fn test_1_2() {
203        let layer = Dense::<f32, 1, 2> {
204            weights: [[0.5], [1.0]],
205        };
206        assert_eq!(layer.forward(&[1.0]), [0.5, 1.0],);
207        assert_eq!(layer.gradients_wrt_input(&[2.0, 0.0]), [1.0],);
208
209        assert_eq!(
210            layer.gradients_wrt_weights(&[1.0], &[2.0, 0.0]),
211            [[2.0], [0.0]],
212        );
213        assert_eq!(
214            layer.gradients_wrt_weights(&[1.0], &[0.0, 2.0]),
215            [[0.0], [2.0]],
216        );
217        assert_eq!(
218            layer.gradients_wrt_weights(&[1.0], &[1.0, 2.0]),
219            [[1.0], [2.0]],
220        );
221    }
222
223    #[test]
224    fn test_2_1() {
225        let layer = Dense::<f32, 2, 1> {
226            weights: [[0.05, 0.5]],
227        };
228        assert_eq!(layer.forward(&[1.0, 2.0]), [1.05],);
229        assert_eq!(layer.gradients_wrt_input(&[2.0]), [0.1, 1.0],);
230    }
231
232    #[test]
233    fn test_2_2() {
234        let layer = Dense::<f32, 2, 2> {
235            weights: [[0.1, 0.4], [0.5, 0.2]],
236        };
237        assert_eq!(layer.forward(&[1.0, 2.0]), [1.1, 0.8],);
238        assert_eq!(layer.gradients_wrt_input(&[0.0, 1.0]), [0.5, 0.2],);
239    }
240
241    #[test]
242    fn grad_wrt_input() {
243        let layer = Dense::<f32, 2, 2> {
244            weights: [[0.0, 1.0], [1.0, 0.0]],
245        };
246        assert_eq!(layer.forward(&[1.0, 0.0]), [0.0, 1.0],);
247        assert_eq!(layer.gradients_wrt_input(&[1.0, -1.0]), [-1.0, 1.0],);
248    }
249}
minidx_core/layers/linear.rs

minidx_core/layers/
linear.rs
