minidx_core/
modules.rs

1use crate::optimizers::GradApplyer;
2use crate::{Error, Gradients};
3use std::collections::HashMap;
4
5/// An error loading or saving parameters.
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct LoadSaveError {
8    pub path: String,
9    pub err: String,
10}
11
12/// A unit of computation that consumes `Input` and produces [Module::Output].
13pub trait Module<X> {
14    /// The type that this unit produces given `Input`.
15    type Output;
16
17    fn forward(&self, x: &X) -> Result<Self::Output, Error>;
18}
19
20/// A unit of computation which can do backpropagation without knowledge of any
21/// additional state.
22///
23/// This trait is to be implemented by individual network layers, but not compositions of them.
24pub trait RevModule<X>: Module<X> {
25    /// The type describing gradients with respect to the modules' own parameters.
26    type SelfGrads: Gradients;
27
28    /// Returns the gradients with respect to the input, and the
29    /// gradients with respect to any internal parameters.
30    fn reverse(
31        &self,
32        inputs: &X,
33        grads_wrt_output: &<Self as Module<X>>::Output,
34    ) -> (X, Self::SelfGrads);
35
36    /// Applies a gradient update step: adding product of the provided gradients and
37    /// the scalar to the parameters.
38    fn apply(
39        &mut self,
40        applyer: &mut impl GradApplyer,
41        updates: Self::SelfGrads,
42    ) -> Result<(), Error>;
43}
44
45/// Some sequential computation that consumes `Input` and produces [Module::Output],
46/// but also produces artifacts describing the execution that can later be used
47/// during backprop.
48///
49/// Neural-network layers implement this trait.
50///
51/// A supertrait of [Module].
52pub trait TracedModule<X>: Module<X> {
53    /// The type that this unit produces to describe intermediate state.
54    ///
55    /// This is typically the input to the module.
56    type Trace;
57
58    /// Same as [Module::forward], except intermediate computations that are needed
59    /// for backprop are returned.
60    fn traced_forward(&self, x: X) -> Result<(<Self as Module<X>>::Output, Self::Trace), Error>;
61}
62
63impl<Input, M: Module<Input> + BaseModule> TracedModule<Input> for M {
64    type Trace = Input;
65
66    fn traced_forward(
67        &self,
68        x: Input,
69    ) -> Result<(<Self as crate::Module<Input>>::Output, Self::Trace), Error> {
70        Ok((Module::forward(self, &x)?, (x)))
71    }
72}
73
74/// Some sequential computation which can perform backprop on itself given trace state
75/// and the gradients of its outputs: computing parameter updates for itself.
76///
77/// Trainable layers implement this trait.
78///
79/// Relies on behavior from the [TracedModule] trait.
80pub trait BackpropModule<X>: TracedModule<X> {
81    /// Type describing movement in the modules' own parameters in response to backpropagation.
82    type SelfGrads;
83
84    /// Computes gradients for this layer/module, given tracing state from forward
85    /// execution, and the gradients of the output.
86    ///
87    /// Returns the gradients with respect to the input to the layer/module, as well as
88    /// gradients with respect to parameters (needed to call [`update()`](BackpropModule::update)).
89    fn backprop(
90        &self,
91        trace: &<Self as TracedModule<X>>::Trace,
92        grads_wrt_output: <Self as Module<X>>::Output,
93    ) -> (X, Self::SelfGrads);
94
95    /// Applies a gradient update step, given (Self::SelfGrads) and a [`GradApplyer`](GradApplyer).
96    ///
97    /// While `updates` describes the change in parameters, `applyer` is used to change the parameters
98    /// of this layer, which may include applying any kind of clipping or regularization.
99    fn update(
100        &mut self,
101        applyer: &mut impl GradApplyer,
102        updates: Self::SelfGrads,
103    ) -> Result<(), Error>;
104
105    /// Returns a [`GradApplyer`](GradApplyer) object needed to train using SGD + momentum.
106    fn new_momentum(
107        &self,
108        params: crate::optimizers::TrainParams,
109        momentum_coefficient: f32,
110    ) -> crate::optimizers::Momentum<Self::SelfGrads>
111    where
112        <Self as BackpropModule<X>>::SelfGrads: Gradients,
113    {
114        crate::optimizers::Momentum::new(params, momentum_coefficient)
115    }
116
117    /// Returns a [`GradApplyer`](GradApplyer) object needed to train using rmsprop.
118    fn new_rmsprop(
119        &self,
120        params: crate::optimizers::TrainParams,
121        beta: f32,
122    ) -> crate::optimizers::RMSProp<Self::SelfGrads>
123    where
124        <Self as BackpropModule<X>>::SelfGrads: Gradients,
125        <Self::SelfGrads as Gradients>::Concrete: crate::Float,
126    {
127        crate::optimizers::RMSProp::new(params, beta)
128    }
129
130    /// Returns a [`GradApplyer`](GradApplyer) object needed to train using rmsprop + momentum.
131    fn new_rmsprop_with_momentum(
132        &self,
133        params: crate::optimizers::TrainParams,
134        momentum_coefficient: f32,
135        beta: f32,
136    ) -> crate::optimizers::RMSProp<Self::SelfGrads>
137    where
138        <Self as BackpropModule<X>>::SelfGrads: Gradients,
139        <Self::SelfGrads as Gradients>::Concrete: crate::Float,
140    {
141        crate::optimizers::RMSProp::new_with_momentum(params, momentum_coefficient, beta)
142    }
143}
144
145impl<Input, M: TracedModule<Input, Trace = Input> + RevModule<Input> + BaseModule>
146    BackpropModule<Input> for M
147{
148    type SelfGrads = <M as RevModule<Input>>::SelfGrads;
149
150    fn backprop(
151        &self,
152        trace: &<M as TracedModule<Input>>::Trace,
153        grads_wrt_output: <M as Module<Input>>::Output,
154    ) -> (Input, Self::SelfGrads) {
155        M::reverse(self, trace, &grads_wrt_output)
156    }
157
158    fn update(
159        &mut self,
160        applyer: &mut impl GradApplyer,
161        updates: Self::SelfGrads,
162    ) -> Result<(), Error> {
163        M::apply(self, applyer, updates)
164    }
165}
166
167/// A module who's parameters can be loaded or saved.
168pub trait LoadableModule {
169    /// Saves the parameters to the given dictionary.
170    ///
171    /// FIXME: We should be storing parameters as their base type, not f64.
172    fn save(&self, path: String, dict: &mut HashMap<String, Vec<f64>>)
173        -> Result<(), LoadSaveError>;
174
175    /// Loads the parameters from the given dictionary.
176    ///
177    /// FIXME: We should be storing parameters as their base type, not f64.
178    fn load(&mut self, path: String, dict: &HashMap<String, Vec<f64>>)
179        -> Result<(), LoadSaveError>;
180}
181
182/// Marker trait for low-level layers which are composable modules.
183///
184/// Set on layers which are native to minidx: needed to get around
185/// trait conflicts between impls on M and (M,)
186pub(crate) trait BaseModule {}
187
188/// Something that can have its learnable parameters reset.
189pub trait ResetParams {
190    /// Sensibly initializes the learnable parameters of some module.
191    ///
192    /// This typically corresponds to Xavier/Glorot initialization, where
193    /// parameters are sampled from a normal distribution of a given magnitude `scale`,
194    /// divided by the width of the layer.
195    ///
196    /// Scale is typically `1.0`, but smaller values can help if you encounter stability issues
197    /// early in training.
198    fn rand_params<RNG: rand::Rng>(&mut self, rng: &mut RNG, scale: f32) -> Result<(), Error>;
199}
200
201macro_rules! fwd_tuple_impls {
202    ([$($name:ident),+] [$($idx:tt),*], $last:ident, [$($rev_tail:ident),*], [$($trace_name:ident),*]) => {
203        impl<
204            Input,
205            $last:
206            $(crate::Module::<$rev_tail ::Output>, $rev_tail: )*
207            crate::Module<Input>
208        > crate::Module<Input> for ($($name,)+) {
209            type Output = $last ::Output;
210
211            /// Calls forward sequentially on each module in the tuple.
212            fn forward(&self, x: &Input) -> Result<Self::Output, Error> {
213                let x = self.0.forward(x)?;
214                $(let x = self.$idx.forward(&x)?;)*
215                Ok(x)
216            }
217        }
218
219        impl<Input,
220             $last:
221             $(TracedModule::<$rev_tail ::Output>, $rev_tail: )*
222             TracedModule<Input>
223            > TracedModule<Input> for ($($name,)+)
224        {
225            type Trace = ($($name::Trace,)+);
226
227            fn traced_forward(
228                &self,
229                x: Input,
230            ) -> Result<(<Self as Module<Input>>::Output, Self::Trace), Error> {
231                let (x, m1t) = self.0.traced_forward(x)?;
232                $(let (x, $trace_name) = self.$idx.traced_forward(x)?;)*
233                Ok((x, (m1t, $($trace_name,)*)))
234            }
235        }
236
237        impl<
238            $last:
239            $(crate::ResetParams, $rev_tail: )*
240            crate::ResetParams
241        > crate::ResetParams for ($($name,)+) {
242            fn rand_params<RNG: rand::Rng>(&mut self, rng: &mut RNG, scale: f32) -> Result<(), Error> {
243                self.0.rand_params(rng, scale)?;
244                $(self.$idx.rand_params(rng, scale)?;)*
245                Ok(())
246            }
247        }
248
249        impl<
250            $last:
251            $(crate::LoadableModule, $rev_tail: )*
252            crate::LoadableModule
253        > crate::LoadableModule for ($($name,)+) {
254            fn save(&self, path: String, dict: &mut HashMap<String, Vec<f64>>) -> Result<(), LoadSaveError> {
255                self.0.save(path.clone() + ".0", dict)?;
256                $(self.$idx.save(format!("{}.{}", path, $idx), dict)?;)*
257                Ok(())
258            }
259
260            fn load(&mut self, path: String, dict: &HashMap<String, Vec<f64>>) -> Result<(), LoadSaveError> {
261                self.0.load(path.clone() + ".0", dict)?;
262                $(self.$idx.load(format!("{}.{}", path, $idx), dict)?;)*
263                Ok(())
264            }
265        }
266    };
267}
268
269fwd_tuple_impls!([M1][], M1, [], []);
270fwd_tuple_impls!([M1, M2][1], M2, [M1], [m2t]);
271fwd_tuple_impls!([M1, M2, M3] [1, 2], M3, [M2, M1], [m2t, m3t]);
272fwd_tuple_impls!([M1, M2, M3, M4] [1, 2, 3], M4, [M3, M2, M1], [m2t, m3t, m4t]);
273fwd_tuple_impls!([M1, M2, M3, M4, M5] [1, 2, 3, 4], M5, [M4, M3, M2, M1], [m2t, m3t, m4t, m5t]);
274fwd_tuple_impls!([M1, M2, M3, M4, M5, M6] [1, 2, 3, 4, 5], M6, [M5, M4, M3, M2, M1], [m2t, m3t, m4t, m5t, m6t]);
275
276macro_rules! backwd_tuple_impls {
277    ([$($all:ident),+] [$($idx:tt),*] [$($rev_idx:tt),*], $first:ident, [$($rev_grads:ident),+], [$($fwd_grads:ident),+], [$(($mod_for:ident, $mod_from:ident)),*]) => {
278        impl<Input,
279             $first: BackpropModule<Input>,
280             $($mod_for: BackpropModule::<$mod_from ::Output>,)*
281            > BackpropModule<Input>
282            for ($($all,)+)
283        {
284            type SelfGrads = (
285                <$first as BackpropModule<Input>>::SelfGrads,
286                $(<$mod_for as BackpropModule::<$mod_from ::Output>>::SelfGrads,)*
287            );
288
289            fn backprop(
290                &self,
291                trace: &<Self as TracedModule<Input>>::Trace,
292                next_grads: <Self as Module<Input>>::Output,
293            ) -> (Input, Self::SelfGrads) {
294                $(let (next_grads, $rev_grads) = self.$rev_idx.backprop(&trace.$rev_idx, next_grads);)+
295                (next_grads, ($($fwd_grads,)+))
296            }
297
298            fn update(&mut self, applyer: &mut impl GradApplyer, updates: Self::SelfGrads) -> Result<(), Error> {
299                $(self.$idx.update(applyer, updates.$idx)?;)+
300                Ok(())
301            }
302
303        }
304    };
305}
306
307backwd_tuple_impls!([M1][0][0], M1, [u1], [u1], []);
308backwd_tuple_impls!([M1, M2][0, 1][1, 0], M1, [u2, u1], [u1, u2], [(M2, M1)]);
309backwd_tuple_impls!([M1, M2, M3][0, 1, 2][2, 1, 0], M1, [u3, u2, u1], [u1, u2, u3], [(M2, M1), (M3, M2)]);
310backwd_tuple_impls!([M1, M2, M3, M4][0, 1, 2, 3][3, 2, 1, 0], M1, [u4, u3, u2, u1], [u1, u2, u3, u4], [(M2, M1), (M3, M2), (M4, M3)]);
311backwd_tuple_impls!([M1, M2, M3, M4, M5][0, 1, 2, 3, 4][4, 3, 2, 1, 0], M1, [u5, u4, u3, u2, u1], [u1, u2, u3, u4, u5], [(M2, M1), (M3, M2), (M4, M3), (M5, M4)]);
312backwd_tuple_impls!([M1, M2, M3, M4, M5, M6][0, 1, 2, 3, 4, 5][5, 4, 3, 2, 1, 0], M1, [u6, u5, u4, u3, u2, u1], [u1, u2, u3, u4, u5, u6], [(M2, M1), (M3, M2), (M4, M3), (M5, M4), (M6, M5)]);
313
314#[cfg(test)]
315mod tests {
316    use super::*;
317    use crate::layers;
318    use rand::rngs::SmallRng;
319    use rand::SeedableRng;
320
321    #[test]
322    fn test_module_forward() {
323        let mut network = (
324            layers::Dense::<f32, 2, 3>::default(),
325            layers::Activation::Relu,
326            layers::Bias1d::<f32, 3>::default(),
327        );
328        network.0.weights[0][0] = 2.5;
329        network.0.weights[0][1] = 0.5;
330        network.2.bias.grad_iter_mut().for_each(|x| *x += 1.0);
331
332        let output = network.forward(&[1.0, 2.0]);
333        assert_eq!(output, Ok([3.5, 1.5, 1.0]));
334
335        // Also try the TracedModule path
336        let (output2, _) = network.traced_forward([1.0, 2.0]).unwrap();
337        assert_eq!(output.unwrap(), output2);
338    }
339
340    #[test]
341    fn test_module_backward() {
342        let network = (
343            layers::Dense::<f32, 2, 3>::default(),
344            layers::Activation::Relu,
345            layers::Bias1d::<f32, 3>::default(),
346        );
347
348        let (_, trace) = network.traced_forward([1.0, 2.0]).unwrap();
349        let (grad_wrt_input, _gradient_updates) = network.backprop(&trace, [0.0, 0.0, 0.0]);
350        assert_eq!(grad_wrt_input, [0.0, 0.0]);
351
352        let network = (
353            layers::Dense::<f32, 2, 3>::default(),
354            layers::Activation::Relu,
355            layers::Bias1d::<f32, 3>::default(),
356            layers::Bias1d::<f32, 3>::default(),
357            layers::Bias1d::<f32, 3>::default(),
358        );
359        let (_, trace) = network.traced_forward([1.0, 2.0]).unwrap();
360        let (grad_wrt_input, _gradient_updates) = network.backprop(&trace, [0.0, 0.0, 0.0]);
361        assert_eq!(grad_wrt_input, [0.0, 0.0]);
362    }
363
364    #[test]
365    fn test_nested_module() {
366        let network = (
367            (
368                layers::Dense::<f32, 2, 3>::default(),
369                layers::Bias1d::<f32, 3>::default(),
370            ),
371            layers::Activation::Relu,
372        );
373
374        let (_, trace) = network.traced_forward([1.0, 2.0]).unwrap();
375        let (grad_wrt_input, _gradient_updates) = network.backprop(&trace, [5.0, 1.0, 1.0]);
376        assert_eq!(grad_wrt_input, [0.0, 0.0]);
377    }
378
379    #[test]
380    fn test_reset_params() {
381        let mut network = (
382            layers::Dense::<f32, 2, 3>::default(),
383            layers::Bias1d::<f32, 3>::default(),
384            layers::Activation::<f32>::Relu,
385        );
386        let mut rng = SmallRng::seed_from_u64(1);
387        network.rand_params(&mut rng, 1.0).unwrap();
388
389        for x in network.0.weights.iter().flatten() {
390            assert!(*x != 0.0);
391        }
392        for x in network.1.bias.grad_iter() {
393            assert!(*x != 0.0);
394        }
395    }
396
397    #[test]
398    fn test_dense_backprop() {
399        let network = layers::Dense::<f32, 2, 3>::default();
400
401        let (_, trace) = network.traced_forward([1.0, 2.0]).unwrap();
402        let (grad_wrt_input, gradient_updates) = network.backprop(&trace, [0.0, 0.0, 0.0]);
403        assert_eq!(grad_wrt_input, [0.0, 0.0]);
404        assert_eq!(gradient_updates, [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]);
405
406        let (_, gradient_updates) = network.backprop(&trace, [1.0, 0.0, 0.0]);
407        assert_eq!(gradient_updates, [[1.0, 0.0], [0.0, 2.0], [0.0, 0.0]]); // TODO: Not sure if this is right?
408    }
409
410    #[test]
411    fn test_bias1d_backprop() {
412        let network = layers::Bias1d::<f32, 2>::default();
413
414        let (_, trace) = network.traced_forward([1.0, 2.0]).unwrap();
415        let (grad_wrt_input, gradient_updates) = network.backprop(&trace, [0.0, 0.0]);
416        assert_eq!(grad_wrt_input, [0.0, 0.0]);
417        assert_eq!(gradient_updates.raw_grads(), [0.0, 0.0]);
418
419        let (_, gradient_updates) = network.backprop(&trace, [2.0, 7.0]);
420        assert_eq!(gradient_updates.raw_grads(), [2.0, 7.0]);
421    }
422
423    #[test]
424    fn test_activation_backprop() {
425        let network = layers::Activation::Relu;
426
427        let (_, trace) = network.traced_forward([1.0, 2.0]).unwrap();
428        let (grad_wrt_input, gradient_updates) = network.backprop(&trace, [0.0, 0.0]);
429        assert_eq!(grad_wrt_input, [0.0, 0.0]);
430        assert_eq!(gradient_updates, ());
431    }
432}
minidx_core/modules.rs

minidx_core/
modules.rs
