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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

//! Shared handles to a model's tensors. Cloning a [`Param`] gives another handle
//! to the same weight and gradient, which is how an optimizer updates a model
//! it does not borrow. [`Buffer`] is the same without a gradient.

use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, RwLock};

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

/// A learnable tensor plus its gradient.
#[derive(Clone)]
pub struct Param {
    value: Arc<RwLock<Tensor>>,
    grad: GradSlot,
    /// Shared with every clone, so freezing through one handle freezes them all.
    trainable: Arc<AtomicBool>,
}

impl Param {
    /// Wrap `value` as a learnable parameter.
    pub fn new(value: Tensor) -> Param {
        Param {
            value: Arc::new(RwLock::new(value.detach())),
            grad: GradSlot::new(),
            trainable: Arc::new(AtomicBool::new(true)),
        }
    }

    /// The parameter as a graph leaf for a forward pass. Detached while frozen,
    /// so backward never computes a gradient for it.
    pub fn tensor(&self) -> Tensor {
        let value = self.value.read().unwrap();
        if !self.is_trainable() {
            return value.detach();
        }
        Tensor::leaf(
            value.storage_arc(),
            value.shape().to_vec(),
            value.device(),
            self.grad.clone(),
        )
    }

    /// Stop training this parameter.
    pub fn freeze(&self) {
        self.trainable.store(false, Ordering::Relaxed);
        self.grad.clear();
    }

    /// Undo [`freeze`](Param::freeze).
    pub fn unfreeze(&self) {
        self.trainable.store(true, Ordering::Relaxed);
    }

    /// Whether optimizers should update this parameter.
    pub fn is_trainable(&self) -> bool {
        self.trainable.load(Ordering::Relaxed)
    }

    /// The current value, detached from the graph.
    pub fn value(&self) -> Tensor {
        self.value.read().unwrap().detach()
    }

    /// Replace the value for every handle. Shape must not change.
    pub fn set_value(&self, value: Tensor) {
        assert_eq!(
            value.shape(),
            self.value.read().unwrap().shape(),
            "set_value: shape must not change"
        );
        *self.value.write().unwrap() = value.detach();
    }

    /// The accumulated gradient, if `backward()` has reached this parameter.
    pub fn grad(&self) -> Option<Tensor> {
        self.grad.get()
    }

    /// Discard the accumulated gradient.
    pub fn zero_grad(&self) {
        self.grad.clear();
    }

    /// Replace the gradient, for clipping and custom optimizers.
    pub fn set_grad(&self, grad: Tensor) {
        self.grad.clear();
        self.grad.accumulate(grad.detach());
    }

    /// Move the value to `device`, dropping the gradient.
    pub fn to_device(&self, device: Device) {
        let mut value = self.value.write().unwrap();
        if value.device() != device {
            *value = value.to(device).detach();
            self.grad.clear();
        }
    }

    pub fn shape(&self) -> Vec<usize> {
        self.value.read().unwrap().shape().to_vec()
    }

    pub fn numel(&self) -> usize {
        self.value.read().unwrap().numel()
    }

    pub fn device(&self) -> Device {
        self.value.read().unwrap().device()
    }
}

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

/// State that is checkpointed but not learned.
#[derive(Clone)]
pub struct Buffer {
    value: Arc<RwLock<Tensor>>,
}

impl Buffer {
    pub fn new(value: Tensor) -> Buffer {
        Buffer {
            value: Arc::new(RwLock::new(value.detach())),
        }
    }

    pub fn value(&self) -> Tensor {
        self.value.read().unwrap().detach()
    }

    pub fn set_value(&self, value: Tensor) {
        *self.value.write().unwrap() = value.detach();
    }

    pub fn to_device(&self, device: Device) {
        let mut value = self.value.write().unwrap();
        if value.device() != device {
            *value = value.to(device).detach();
        }
    }
}