flodl 0.5.0

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

use crate::autograd::{Variable, bilinear};
use crate::tensor::{Result, Device, Tensor, DType, TensorOptions};

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

/// Bilinear transformation: `y = x1^T A x2 + b`.
///
/// Weight shape: `[out_features, in1_features, in2_features]`.
/// Bias shape: `[out_features]` (optional).
///
/// Equivalent to PyTorch's `nn.Bilinear`.
pub struct Bilinear {
    pub weight: Parameter,
    pub bias: Option<Parameter>,
}

impl Bilinear {
    /// Create a Bilinear layer.
    pub fn new(
        in1_features: i64, in2_features: i64, out_features: i64, with_bias: bool,
    ) -> Result<Self> {
        Self::build(in1_features, in2_features, out_features, with_bias, Device::CPU)
    }

    /// Create on a specific device.
    pub fn on_device(
        in1_features: i64, in2_features: i64, out_features: i64, device: Device,
    ) -> Result<Self> {
        Self::build(in1_features, in2_features, out_features, true, device)
    }

    fn build(
        in1_features: i64, in2_features: i64, out_features: i64,
        with_bias: bool, device: Device,
    ) -> Result<Self> {
        let opts = TensorOptions { dtype: DType::Float32, device };
        let bound = 1.0 / (in1_features as f64).sqrt();
        let w = Tensor::rand(&[out_features, in1_features, in2_features], opts)?
            .mul_scalar(2.0 * bound)?
            .add_scalar(-bound)?;
        let weight = Parameter::new(w, "weight");

        let bias = if with_bias {
            let b = Tensor::rand(&[out_features], opts)?
                .mul_scalar(2.0 * bound)?
                .add_scalar(-bound)?;
            Some(Parameter::new(b, "bias"))
        } else {
            None
        };

        Ok(Self { weight, bias })
    }
}

impl Module for Bilinear {
    fn name(&self) -> &str { "bilinear" }

    fn forward(&self, _input: &Variable) -> Result<Variable> {
        Err(crate::tensor::TensorError::new(
            "Bilinear requires two inputs; use forward_bilinear(x1, x2) instead",
        ))
    }

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

impl Bilinear {
    /// Forward pass with two inputs.
    pub fn forward_bilinear(&self, input1: &Variable, input2: &Variable) -> Result<Variable> {
        bilinear(
            input1, input2,
            &self.weight.variable,
            self.bias.as_ref().map(|b| &b.variable),
        )
    }
}

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

    #[test]
    fn test_bilinear_forward() {
        let layer = Bilinear::on_device(3, 4, 5, test_device()).unwrap();
        let x1 = Variable::new(Tensor::randn(&[2, 3], test_opts()).unwrap(), false);
        let x2 = Variable::new(Tensor::randn(&[2, 4], test_opts()).unwrap(), false);
        let y = layer.forward_bilinear(&x1, &x2).unwrap();
        assert_eq!(y.shape(), vec![2, 5]);
    }

    #[test]
    fn test_bilinear_params() {
        let layer = Bilinear::on_device(3, 4, 5, test_device()).unwrap();
        assert_eq!(layer.parameters().len(), 2); // weight + bias
    }

    #[test]
    fn test_bilinear_gradient() {
        let layer = Bilinear::on_device(3, 4, 5, test_device()).unwrap();
        let x1 = Variable::new(Tensor::randn(&[2, 3], test_opts()).unwrap(), true);
        let x2 = Variable::new(Tensor::randn(&[2, 4], test_opts()).unwrap(), true);
        let y = layer.forward_bilinear(&x1, &x2).unwrap().sum().unwrap();
        y.backward().unwrap();
        assert!(x1.grad().is_some());
        assert!(x2.grad().is_some());
        assert!(layer.weight.variable.grad().is_some());
    }

    #[test]
    fn test_bilinear_no_bias() {
        let layer = Bilinear::build(3, 4, 5, false, test_device()).unwrap();
        assert_eq!(layer.parameters().len(), 1); // weight only
        assert!(layer.bias.is_none());
        let x1 = Variable::new(Tensor::randn(&[2, 3], test_opts()).unwrap(), false);
        let x2 = Variable::new(Tensor::randn(&[2, 4], test_opts()).unwrap(), false);
        let y = layer.forward_bilinear(&x1, &x2).unwrap();
        assert_eq!(y.shape(), vec![2, 5]);
    }

    #[test]
    fn test_bilinear_forward_rejects_single_input() {
        let layer = Bilinear::on_device(3, 4, 5, test_device()).unwrap();
        let x = Variable::new(Tensor::randn(&[2, 3], test_opts()).unwrap(), false);
        assert!(layer.forward(&x).is_err());
    }
}