flodl 0.5.2

floDl — a flow-graph deep learning framework built on libtorch
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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

use crate::autograd::{self, Variable};
use crate::tensor::{Device, Result};

use super::init;
use super::parameter::Parameter;
use super::Module;

/// Fully connected layer: `y = x @ W^T + b`.
///
/// Weight shape: `[out_features, in_features]`.
/// Bias shape: `[out_features]` (optional).
///
/// Input shape: `[batch, in_features]`.
/// Output shape: `[batch, out_features]`.
///
/// ```ignore
/// let layer = Linear::new(4, 2)?;
/// let x = Variable::new(Tensor::randn(&[8, 4], opts)?, false);
/// let y = layer.forward(&x)?;
/// assert_eq!(y.shape(), vec![8, 2]);
/// ```
pub struct Linear {
    pub weight: Parameter,
    pub bias: Option<Parameter>,
}

impl Linear {
    /// Create a linear layer on CPU with bias.
    pub fn new(in_features: i64, out_features: i64) -> Result<Self> {
        Self::on_device(in_features, out_features, Device::CPU)
    }

    /// Create a linear layer on a specific device with bias.
    pub fn on_device(in_features: i64, out_features: i64, device: Device) -> Result<Self> {
        let w = init::kaiming_uniform(&[out_features, in_features], in_features, 5.0_f64.sqrt(), device)?;
        let b = init::uniform_bias(in_features, &[out_features], device)?;
        Ok(Linear {
            weight: Parameter::new(w, "weight"),
            bias: Some(Parameter::new(b, "bias")),
        })
    }

    /// Create a linear layer without bias on CPU.
    pub fn no_bias(in_features: i64, out_features: i64) -> Result<Self> {
        Self::no_bias_on_device(in_features, out_features, Device::CPU)
    }

    /// Create a linear layer without bias on a specific device.
    pub fn no_bias_on_device(in_features: i64, out_features: i64, device: Device) -> Result<Self> {
        let w = init::kaiming_uniform(&[out_features, in_features], in_features, 5.0_f64.sqrt(), device)?;
        Ok(Linear {
            weight: Parameter::new(w, "weight"),
            bias: None,
        })
    }
}

impl Module for Linear {
    fn name(&self) -> &str { "linear" }

    fn forward(&self, input: &Variable) -> Result<Variable> {
        autograd::linear(
            input,
            &self.weight.variable,
            self.bias.as_ref().map(|b| &b.variable),
        )
    }

    fn parameters(&self) -> Vec<Parameter> {
        let mut params = vec![self.weight.clone()];
        if let Some(ref b) = self.bias {
            params.push(b.clone());
        }
        params
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tensor::{Tensor, test_device, test_opts};

    #[test]
    fn test_linear_forward_shape() {
        let dev = test_device();
        let layer = Linear::on_device(4, 2, dev).unwrap();
        let x = Variable::new(Tensor::randn(&[8, 4], test_opts()).unwrap(), false);
        let y = layer.forward(&x).unwrap();
        assert_eq!(y.shape(), vec![8, 2]);
    }

    #[test]
    fn test_linear_parameters_with_bias() {
        let layer = Linear::on_device(4, 2, test_device()).unwrap();
        let params = layer.parameters();
        assert_eq!(params.len(), 2);
        assert_eq!(params[0].variable.shape(), vec![2, 4]); // weight
        assert_eq!(params[1].variable.shape(), vec![2]);     // bias
    }

    #[test]
    fn test_linear_no_bias() {
        let layer = Linear::no_bias_on_device(4, 2, test_device()).unwrap();
        let params = layer.parameters();
        assert_eq!(params.len(), 1);
        assert!(layer.bias.is_none());

        let x = Variable::new(Tensor::randn(&[3, 4], test_opts()).unwrap(), false);
        let y = layer.forward(&x).unwrap();
        assert_eq!(y.shape(), vec![3, 2]);
    }

    #[test]
    fn test_linear_gradient_flow() {
        let dev = test_device();
        let layer = Linear::on_device(3, 2, dev).unwrap();
        let x = Variable::new(Tensor::randn(&[4, 3], test_opts()).unwrap(), false);
        let y = layer.forward(&x).unwrap();
        let loss = y.sum().unwrap();
        loss.backward().unwrap();

        let params = layer.parameters();
        assert!(params[0].variable.grad().is_some(), "weight should have gradient");
        assert!(params[1].variable.grad().is_some(), "bias should have gradient");
    }

    #[test]
    fn test_linear_on_device() {
        let dev = test_device();
        let layer = Linear::on_device(4, 2, dev).unwrap();
        assert_eq!(layer.weight.variable.device(), dev);
        if let Some(ref b) = layer.bias {
            assert_eq!(b.variable.device(), dev);
        }
    }

    #[test]
    fn test_linear_name() {
        let layer = Linear::new(4, 2).unwrap();
        assert_eq!(layer.name(), "linear");
    }
}