fastnn 0.3.3

A robust deep learning library in Rust with CUDA kernels
Documentation 
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

//! Chaining layers.

use crate::tensor::Tensor;

use super::module::{Module, scoped};
use super::param::{Buffer, Param};

/// Runs its layers in order.
///
/// ```
/// use fastnn::prelude::*;
///
/// let model = Sequential::new()
///     .add(Linear::new(784, 128))
///     .add(ReLU)
///     .add(Linear::new(128, 10));
///
/// assert_eq!(model.forward(&Tensor::randn(&[4, 784])).shape(), &[4, 10]);
/// ```
///
/// Parameter names combine the layer's index with its own name, so the second
/// `Linear` above owns `"2.weight"` and `"2.bias"`.
#[derive(Default)]
pub struct Sequential {
    layers: Vec<Box<dyn Module>>,
}

impl Sequential {
    pub fn new() -> Sequential {
        Sequential::default()
    }

    /// Append a layer, returning `self` so calls chain.
    #[allow(clippy::should_implement_trait)]
    pub fn add(mut self, layer: impl Module + 'static) -> Sequential {
        self.layers.push(Box::new(layer));
        self
    }

    /// Append a layer in place.
    pub fn push(&mut self, layer: impl Module + 'static) {
        self.layers.push(Box::new(layer));
    }

    pub fn len(&self) -> usize {
        self.layers.len()
    }

    pub fn is_empty(&self) -> bool {
        self.layers.is_empty()
    }
}

impl Module for Sequential {
    fn forward(&self, input: &Tensor) -> Tensor {
        self.layers
            .iter()
            .fold(input.clone(), |x, layer| layer.forward(&x))
    }

    fn named_parameters(&self) -> Vec<(String, Param)> {
        self.layers
            .iter()
            .enumerate()
            .flat_map(|(i, layer)| scoped(i, layer.named_parameters()))
            .collect()
    }

    fn named_buffers(&self) -> Vec<(String, Buffer)> {
        self.layers
            .iter()
            .enumerate()
            .flat_map(|(i, layer)| scoped(i, layer.named_buffers()))
            .collect()
    }

    fn set_training(&self, training: bool) {
        for layer in &self.layers {
            layer.set_training(training);
        }
    }
}