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
148
149
150

//! Derivatives of the fused norms.

use crate::autograd::Backward;
use crate::cuda::kernels;
use crate::tensor::Tensor;
use crate::tensor::storage::Storage;

/// LayerNorm on the host. Returns gradients for `(input, gamma, beta)`.
pub struct LayerNormBackward {
    pub normalized: Tensor,
    pub gamma: Vec<f32>,
    pub inv_std: Vec<f32>,
    pub cols: usize,
}

impl Backward for LayerNormBackward {
    fn backward(&self, grad: &Tensor) -> Vec<Tensor> {
        let g = grad.to_vec();
        let unit = self.normalized.to_vec();
        let cols = self.cols;
        let width = cols as f32;
        let rows = g.len() / cols;

        let mut d_input = vec![0.0f32; g.len()];
        let mut d_gamma = vec![0.0f32; cols];
        let mut d_beta = vec![0.0f32; cols];

        for r in 0..rows {
            let base = r * cols;

            // the mean and variance both move with every input in the row
            let mut mean_shift = 0.0f32;
            let mut scale_shift = 0.0f32;
            for c in 0..cols {
                let scaled = self.gamma[c] * g[base + c];
                mean_shift += scaled;
                scale_shift += scaled * unit[base + c];
                d_gamma[c] += g[base + c] * unit[base + c];
                d_beta[c] += g[base + c];
            }

            for c in 0..cols {
                let scaled = self.gamma[c] * g[base + c];
                d_input[base + c] = self.inv_std[r]
                    * (scaled - mean_shift / width - unit[base + c] * scale_shift / width);
            }
        }

        let device = grad.device();
        vec![
            Tensor::from_vec(d_input, grad.shape()).to(device),
            Tensor::from_vec(d_gamma, &[cols]).to(device),
            Tensor::from_vec(d_beta, &[cols]).to(device),
        ]
    }
    fn name(&self) -> &'static str {
        "LayerNorm"
    }
}

/// LayerNorm on the GPU, one kernel for all three gradients.
pub struct LayerNormCudaBackward {
    pub input: Tensor,
    pub gamma: Tensor,
    pub mean: Tensor,
    pub inv_std: Tensor,
    pub cols: usize,
}

impl Backward for LayerNormCudaBackward {
    fn backward(&self, grad: &Tensor) -> Vec<Tensor> {
        let rows = grad.numel() / self.cols;
        let buffers = (
            grad.storage().as_cuda(),
            self.input.storage().as_cuda(),
            self.gamma.storage().as_cuda(),
            self.mean.storage().as_cuda(),
            self.inv_std.storage().as_cuda(),
        );
        let (Some(g), Some(x), Some(gamma), Some(mean), Some(inv_std)) = buffers else {
            panic!("LayerNormCudaBackward reached with a non-CUDA gradient");
        };

        let (d_input, d_gamma, d_beta) =
            kernels::layer_norm_backward(g, x, gamma, mean, inv_std, rows, self.cols)
                .expect("cuda layer_norm backward");

        let device = grad.device();
        vec![
            Tensor::raw(Storage::Cuda(d_input), grad.shape().to_vec(), device),
            Tensor::raw(Storage::Cuda(d_gamma), vec![self.cols], device),
            Tensor::raw(Storage::Cuda(d_beta), vec![self.cols], device),
        ]
    }
    fn name(&self) -> &'static str {
        "LayerNormCuda"
    }
}

/// BatchNorm2d. Like LayerNorm but the group is a channel across the batch.
pub struct BatchNorm2dBackward {
    pub normalized: Tensor,
    pub gamma: Vec<f32>,
    pub inv_std: Vec<f32>,
}

impl Backward for BatchNorm2dBackward {
    fn backward(&self, grad: &Tensor) -> Vec<Tensor> {
        let (n, c, h, w) = (grad.dim(0), grad.dim(1), grad.dim(2), grad.dim(3));
        let plane = h * w;
        let count = (n * plane) as f32;

        let g = grad.to_vec();
        let unit = self.normalized.to_vec();
        let mut d_input = vec![0.0f32; g.len()];
        let mut d_gamma = vec![0.0f32; c];
        let mut d_beta = vec![0.0f32; c];

        for channel in 0..c {
            for image in 0..n {
                for i in 0..plane {
                    let index = (image * c + channel) * plane + i;
                    d_gamma[channel] += g[index] * unit[index];
                    d_beta[channel] += g[index];
                }
            }

            let scale = self.gamma[channel] * self.inv_std[channel];
            for image in 0..n {
                for i in 0..plane {
                    let index = (image * c + channel) * plane + i;
                    d_input[index] = scale
                        * (g[index]
                            - d_beta[channel] / count
                            - unit[index] * d_gamma[channel] / count);
                }
            }
        }

        let device = grad.device();
        vec![
            Tensor::from_vec(d_input, grad.shape()).to(device),
            Tensor::from_vec(d_gamma, &[c]).to(device),
            Tensor::from_vec(d_beta, &[c]).to(device),
        ]
    }
    fn name(&self) -> &'static str {
        "BatchNorm2d"
    }
}