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

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

/// A named learnable parameter wrapping a Variable with requires_grad=true.
#[derive(Clone)]
pub struct Parameter {
    pub variable: Variable,
    pub name: String,
}

impl Parameter {
    /// Create a named parameter from a tensor (always requires_grad=true).
    pub fn new(data: Tensor, name: &str) -> Self {
        Parameter {
            variable: Variable::new(data, true),
            name: name.to_string(),
        }
    }

    /// Move this parameter to a different device. Returns a new Parameter.
    /// Preserves the current requires_grad state (frozen params stay frozen).
    pub fn to_device(&self, device: Device) -> Result<Parameter> {
        let rg = self.variable.requires_grad();
        let moved = self.variable.data().to_device(device)?;
        Ok(Parameter {
            variable: Variable::new(moved, rg),
            name: self.name.clone(),
        })
    }

    /// Freeze this parameter: disable gradient tracking.
    /// The optimizer will skip frozen parameters (they have no grad).
    pub fn freeze(&self) -> Result<()> {
        self.variable.set_requires_grad(false)
    }

    /// Unfreeze this parameter: re-enable gradient tracking.
    pub fn unfreeze(&self) -> Result<()> {
        self.variable.set_requires_grad(true)
    }

    /// Whether this parameter is frozen (not tracking gradients).
    pub fn is_frozen(&self) -> bool {
        !self.variable.requires_grad()
    }
}

impl std::fmt::Debug for Parameter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Parameter({}, {:?})", self.name, self.variable)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tensor::TensorOptions;

    fn make_param(name: &str) -> Parameter {
        let t = Tensor::randn(&[3, 2], TensorOptions {
            dtype: crate::tensor::DType::Float32,
            device: crate::tensor::test_device(),
        }).unwrap();
        Parameter::new(t, name)
    }

    #[test]
    fn test_freeze_no_gradient() {
        let p = make_param("w");
        assert!(!p.is_frozen());
        assert!(p.variable.requires_grad());

        p.freeze().unwrap();
        assert!(p.is_frozen());
        assert!(!p.variable.requires_grad());

        // Verify no gradient accumulates on a frozen param
        // Use a simple sum so we don't need matmul
        let y = p.variable.data().sum().unwrap();
        let _ = y.backward();
        assert!(p.variable.grad().is_none(), "frozen param should have no gradient");
    }

    #[test]
    fn test_unfreeze_gradient_reappears() {
        let p = make_param("w");
        p.freeze().unwrap();
        assert!(p.is_frozen());

        p.unfreeze().unwrap();
        assert!(!p.is_frozen());
        assert!(p.variable.requires_grad());

        // Forward + backward: unfrozen param should accumulate gradient
        let x = Variable::new(
            Tensor::from_f32(&[1.0, 2.0, 3.0], &[1, 3], crate::tensor::test_device()).unwrap(),
            false,
        );
        let y = x.matmul(&p.variable).unwrap();
        let loss = y.sum().unwrap();
        loss.backward().unwrap();
        assert!(p.variable.grad().is_some(), "unfrozen param should have gradient");
    }

    #[test]
    fn test_optimizer_step_with_frozen_params() {
        use crate::nn::optim::{Adam, Optimizer};

        let p1 = make_param("w1");
        let p2 = make_param("w2");
        p1.freeze().unwrap();

        let mut opt = Adam::new(&[p1.clone(), p2.clone()], 0.01);

        let x = Variable::new(
            Tensor::from_f32(&[1.0, 2.0, 3.0], &[1, 3], crate::tensor::test_device()).unwrap(),
            false,
        );
        // Only p2 participates in grad-tracked computation
        let y = x.matmul(&p2.variable).unwrap();
        let loss = y.sum().unwrap();
        loss.backward().unwrap();

        let p1_before = p1.variable.data().to_f32_vec().unwrap();
        opt.step().unwrap();
        let p1_after = p1.variable.data().to_f32_vec().unwrap();
        // Frozen param unchanged
        assert_eq!(p1_before, p1_after, "frozen param should not change");
    }
}