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

//! Fused layer and batch norm. They could be built from small ops but the
//! backward needs the statistics again, and saving them is cheaper than the graph.

use crate::autograd::ops::norm::{BatchNorm2dBackward, LayerNormBackward, LayerNormCudaBackward};
use crate::cuda::kernels;
use crate::tensor::Tensor;
use crate::tensor::storage::Storage;

/// Per-channel statistics of one batch, for updating running estimates.
pub struct BatchStats {
    pub mean: Vec<f32>,
    pub var: Vec<f32>,
}

impl Tensor {
    /// Normalize each row over the last dimension, then scale and shift.
    ///
    /// `gamma` and `beta` must both be `[last_dim]`.
    pub fn layer_norm(&self, gamma: &Tensor, beta: &Tensor, eps: f32) -> Tensor {
        let (rows, cols) = self.as_rows();
        assert_eq!(gamma.numel(), cols, "layer_norm: gamma must be [{cols}]");
        assert_eq!(beta.numel(), cols, "layer_norm: beta must be [{cols}]");

        if let (Storage::Cuda(x), Storage::Cuda(g), Storage::Cuda(b)) =
            (self.storage(), gamma.storage(), beta.storage())
        {
            let (out, mean, inv_std) =
                kernels::layer_norm_forward(x, g, b, rows, cols, eps).expect("cuda layer_norm");
            let out = Tensor::raw(Storage::Cuda(out), self.shape().to_vec(), self.device());
            let mean = Tensor::raw(Storage::Cuda(mean), vec![rows], self.device());
            let inv_std = Tensor::raw(Storage::Cuda(inv_std), vec![rows], self.device());
            return out.with_grad(&[self, gamma, beta], || LayerNormCudaBackward {
                input: self.detach(),
                gamma: gamma.detach(),
                mean,
                inv_std,
                cols,
            });
        }

        let (data, g, b) = (self.to_vec(), gamma.to_vec(), beta.to_vec());
        let mut out = vec![0.0f32; data.len()];
        let mut normalized = vec![0.0f32; data.len()];
        let mut inv_std = vec![0.0f32; rows];

        for r in 0..rows {
            let row = &data[r * cols..(r + 1) * cols];
            let mean = row.iter().sum::<f32>() / cols as f32;
            let var = row.iter().map(|&x| (x - mean).powi(2)).sum::<f32>() / cols as f32;
            inv_std[r] = 1.0 / (var + eps).sqrt();
            for c in 0..cols {
                let unit = (row[c] - mean) * inv_std[r];
                normalized[r * cols + c] = unit;
                out[r * cols + c] = g[c] * unit + b[c];
            }
        }

        let shape = self.shape().to_vec();
        Tensor::from_vec(out, &shape)
            .to(self.device())
            .with_grad(&[self, gamma, beta], || LayerNormBackward {
                normalized: Tensor::from_vec(normalized, &shape),
                gamma: g,
                inv_std,
                cols,
            })
    }

    /// Normalize `[N, C, H, W]` per channel with this batch's statistics. Training
    /// only, eval uses the layer's running statistics instead.
    pub fn batch_norm2d(&self, gamma: &Tensor, beta: &Tensor, eps: f32) -> (Tensor, BatchStats) {
        assert_eq!(
            self.ndim(),
            4,
            "batch_norm2d expects [N, C, H, W], got {:?}",
            self.shape()
        );
        let (n, c, h, w) = (self.dim(0), self.dim(1), self.dim(2), self.dim(3));
        let plane = h * w;
        let count = (n * plane) as f32;

        let (data, g, b) = (self.to_vec(), gamma.to_vec(), beta.to_vec());
        let mut out = vec![0.0f32; data.len()];
        let mut normalized = vec![0.0f32; data.len()];
        let mut stats = BatchStats {
            mean: vec![0.0; c],
            var: vec![0.0; c],
        };
        let mut inv_std = vec![0.0f32; c];

        for channel in 0..c {
            let values = |image: usize, i: usize| data[(image * c + channel) * plane + i];

            let mean = (0..n)
                .flat_map(|i| (0..plane).map(move |p| (i, p)))
                .map(|(i, p)| values(i, p))
                .sum::<f32>()
                / count;
            let var = (0..n)
                .flat_map(|i| (0..plane).map(move |p| (i, p)))
                .map(|(i, p)| (values(i, p) - mean).powi(2))
                .sum::<f32>()
                / count;

            stats.mean[channel] = mean;
            stats.var[channel] = var;
            inv_std[channel] = 1.0 / (var + eps).sqrt();

            for image in 0..n {
                for i in 0..plane {
                    let index = (image * c + channel) * plane + i;
                    let unit = (data[index] - mean) * inv_std[channel];
                    normalized[index] = unit;
                    out[index] = g[channel] * unit + b[channel];
                }
            }
        }

        let shape = self.shape().to_vec();
        let result =
            Tensor::from_vec(out, &shape)
                .to(self.device())
                .with_grad(&[self, gamma, beta], || BatchNorm2dBackward {
                    normalized: Tensor::from_vec(normalized, &shape),
                    gamma: g,
                    inv_std,
                });
        (result, stats)
    }
}