minidx_core/layers/
conv1d.rs

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

minidx_core/layers/
conv1d.rs
