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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

//! The trait every layer implements.

use crate::tensor::{Device, Tensor};

use super::param::{Buffer, Param};

/// A piece of a model. Only [`forward`](Module::forward) is required, plus
/// [`named_parameters`](Module::named_parameters) if it has weights.
///
/// Everything takes `&self`, since parameters have interior mutability.
///
/// ```
/// use fastnn::prelude::*;
///
/// struct Residual(Linear);
///
/// impl Module for Residual {
///     fn forward(&self, x: &Tensor) -> Tensor {
///         x.add(&self.0.forward(x).relu())
///     }
///     fn named_parameters(&self) -> Vec<(String, Param)> {
///         self.0.named_parameters()
///     }
/// }
/// ```
pub trait Module: Send + Sync {
    /// Transform an input into an output.
    fn forward(&self, input: &Tensor) -> Tensor;

    /// Learnable tensors, named like `"weight"` or `"2.bias"`.
    fn named_parameters(&self) -> Vec<(String, Param)> {
        Vec::new()
    }

    /// Persistent non-learnable state, such as batch-norm running statistics.
    fn named_buffers(&self) -> Vec<(String, Buffer)> {
        Vec::new()
    }

    /// Switch between training and inference behaviour. Only
    /// [`Dropout`](super::Dropout) and [`BatchNorm2d`](super::BatchNorm2d) act on it.
    fn set_training(&self, _training: bool) {}

    /// Every learnable tensor, ready to hand to an optimizer.
    fn parameters(&self) -> Vec<Param> {
        self.named_parameters()
            .into_iter()
            .map(|(_, p)| p)
            .collect()
    }

    /// Total number of learnable scalars.
    fn num_parameters(&self) -> usize {
        self.parameters().iter().map(|p| p.numel()).sum()
    }

    /// Move every parameter and buffer to `device`.
    fn to_device(&self, device: Device) {
        for (_, p) in self.named_parameters() {
            p.to_device(device);
        }
        for (_, b) in self.named_buffers() {
            b.to_device(device);
        }
    }

    /// Freeze every parameter.
    fn freeze(&self) {
        for (_, p) in self.named_parameters() {
            p.freeze();
        }
    }

    /// Unfreeze every parameter.
    fn unfreeze(&self) {
        for (_, p) in self.named_parameters() {
            p.unfreeze();
        }
    }

    /// The parameters an optimizer will actually update.
    fn trainable_parameters(&self) -> Vec<Param> {
        self.parameters()
            .into_iter()
            .filter(Param::is_trainable)
            .collect()
    }

    /// Discard all accumulated gradients.
    fn zero_grad(&self) {
        for (_, p) in self.named_parameters() {
            p.zero_grad();
        }
    }

    /// Training mode.
    fn train(&self) {
        self.set_training(true);
    }

    /// Inference mode.
    fn eval(&self) {
        self.set_training(false);
    }
}

/// Prefix each name with `scope`.
///
/// ```
/// # use fastnn::prelude::*;
/// # use fastnn::nn::module::scoped;
/// # struct Net { encoder: Linear }
/// # impl Net {
/// fn named_parameters(&self) -> Vec<(String, Param)> {
///     scoped("encoder", self.encoder.named_parameters())  // -> "encoder.weight", ...
/// }
/// # }
/// ```
pub fn scoped<T>(scope: impl std::fmt::Display, items: Vec<(String, T)>) -> Vec<(String, T)> {
    items
        .into_iter()
        .map(|(name, item)| (format!("{scope}.{name}"), item))
        .collect()
}