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

//! Fully connected layer.

use crate::tensor::Tensor;

use super::module::Module;
use super::param::Param;

/// `y = x · Wᵀ + b` over the last dimension, any number of leading dimensions.
///
/// ```
/// use fastnn::prelude::*;
///
/// let layer = Linear::new(4, 2);
/// let y = layer.forward(&Tensor::randn(&[8, 4]));
/// assert_eq!(y.shape(), &[8, 2]);
/// ```
pub struct Linear {
    pub weight: Param,
    pub bias: Option<Param>,
    in_features: usize,
    out_features: usize,
}

impl Linear {
    /// A layer with bias, Kaiming-initialised.
    pub fn new(in_features: usize, out_features: usize) -> Linear {
        Linear {
            weight: Param::new(Tensor::kaiming_uniform(
                &[out_features, in_features],
                in_features,
            )),
            bias: Some(Param::new(uniform_bias(out_features, in_features))),
            in_features,
            out_features,
        }
    }

    /// A layer without bias.
    pub fn no_bias(in_features: usize, out_features: usize) -> Linear {
        Linear {
            weight: Param::new(Tensor::kaiming_uniform(
                &[out_features, in_features],
                in_features,
            )),
            bias: None,
            in_features,
            out_features,
        }
    }

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

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

impl Module for Linear {
    fn forward(&self, input: &Tensor) -> Tensor {
        assert_eq!(
            input.last_dim(),
            self.in_features,
            "Linear expects {} input features, got {:?}",
            self.in_features,
            input.shape()
        );

        // matmul_nt takes the weight as stored, [out, in], so no transposed copy
        let flat = input.reshape(&[-1, self.in_features as i64]);
        let mut out = flat.matmul_nt(&self.weight.tensor());

        if let Some(bias) = &self.bias {
            out = out.add(&bias.tensor().reshape(&[1, self.out_features as i64]));
        }

        let mut shape: Vec<i64> = input.shape()[..input.ndim() - 1]
            .iter()
            .map(|&d| d as i64)
            .collect();
        shape.push(self.out_features as i64);
        out.reshape(&shape)
    }

    fn named_parameters(&self) -> Vec<(String, Param)> {
        let mut params = vec![("weight".into(), self.weight.clone())];
        if let Some(bias) = &self.bias {
            params.push(("bias".into(), bias.clone()));
        }
        params
    }
}

/// Uniform over `±1/√fan_in`, the bias init shared with the conv layers.
pub(super) fn uniform_bias(out_features: usize, fan_in: usize) -> Tensor {
    let bound = 1.0 / (fan_in as f32).sqrt();
    Tensor::uniform(&[out_features], -bound, bound)
}