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

//! Activations and the softmax family. Softmax runs over the last dimension, so
//! reshape to `[rows, classes]` first for other layouts.

use crate::autograd::ops::activation::{
    GeluBackward, LeakyReluBackward, LogSoftmaxBackward, ReluBackward, SigmoidBackward,
    SiluBackward, SoftmaxBackward, TanhBackward,
};
use crate::cuda::{ffi, kernels};
use crate::tensor::Tensor;
use crate::tensor::storage::Storage;

use super::unary_op;

impl Tensor {
    /// `max(0, x)`.
    pub fn relu(&self) -> Tensor {
        unary_op(self, |x| x.max(0.0), ffi::fastnn_cuda_relu).with_grad(&[self], || ReluBackward {
            input: self.detach(),
        })
    }

    /// `1 / (1 + e^-x)`.
    pub fn sigmoid(&self) -> Tensor {
        let out = unary_op(self, sigmoid_scalar, ffi::fastnn_cuda_sigmoid);
        let saved = out.clone();
        out.with_grad(&[self], || SigmoidBackward { output: saved })
    }

    /// Hyperbolic tangent.
    pub fn tanh(&self) -> Tensor {
        let out = unary_op(self, f32::tanh, ffi::fastnn_cuda_tanh_forward);
        let saved = out.clone();
        out.with_grad(&[self], || TanhBackward { output: saved })
    }

    /// Exact GELU, `x · Φ(x)`, the transformer default.
    pub fn gelu(&self) -> Tensor {
        unary_op(self, gelu_scalar, ffi::fastnn_cuda_gelu).with_grad(&[self], || GeluBackward {
            input: self.detach(),
        })
    }

    /// SiLU / swish, `x · sigmoid(x)`.
    pub fn silu(&self) -> Tensor {
        unary_op(self, |x| x * sigmoid_scalar(x), ffi::fastnn_cuda_silu).with_grad(&[self], || {
            SiluBackward {
                input: self.detach(),
            }
        })
    }

    /// ReLU with a non-zero slope for negative inputs.
    pub fn leaky_relu(&self, slope: f32) -> Tensor {
        let out = match self.storage() {
            Storage::Cpu(data) => Tensor::from_vec(
                super::map(data, move |x| if x > 0.0 { x } else { slope * x }),
                self.shape(),
            ),
            Storage::Cuda(buf) => {
                let out = kernels::leaky_relu(buf, slope, self.numel()).expect("cuda leaky_relu");
                Tensor::raw(Storage::Cuda(out), self.shape().to_vec(), self.device())
            }
        };
        out.with_grad(&[self], || LeakyReluBackward {
            input: self.detach(),
            slope,
        })
    }

    /// Softmax over the last dimension.
    pub fn softmax(&self) -> Tensor {
        let (rows, cols) = self.as_rows();
        let out = match self.storage() {
            Storage::Cpu(data) => {
                let mut result = vec![0.0f32; data.len()];
                for r in 0..rows {
                    let row = &data[r * cols..(r + 1) * cols];
                    let shift = row_max(row);
                    let total: f32 = row.iter().map(|&x| (x - shift).exp()).sum();
                    for c in 0..cols {
                        result[r * cols + c] = (row[c] - shift).exp() / total;
                    }
                }
                Tensor::from_vec(result, self.shape())
            }
            Storage::Cuda(buf) => {
                let out = kernels::softmax(buf, rows, cols).expect("cuda softmax");
                Tensor::raw(Storage::Cuda(out), self.shape().to_vec(), self.device())
            }
        };
        let saved = out.clone();
        out.with_grad(&[self], || SoftmaxBackward { output: saved })
    }

    /// `log(softmax(x))` over the last dimension, without the overflow-prone exp.
    pub fn log_softmax(&self) -> Tensor {
        let (rows, cols) = self.as_rows();
        let out = match self.storage() {
            Storage::Cpu(data) => {
                let mut result = vec![0.0f32; data.len()];
                for r in 0..rows {
                    let row = &data[r * cols..(r + 1) * cols];
                    let shift = row_max(row);
                    let log_sum: f32 = row.iter().map(|&x| (x - shift).exp()).sum::<f32>().ln();
                    for c in 0..cols {
                        result[r * cols + c] = row[c] - shift - log_sum;
                    }
                }
                Tensor::from_vec(result, self.shape())
            }
            Storage::Cuda(buf) => {
                let out = kernels::log_softmax(buf, rows, cols).expect("cuda log_softmax");
                Tensor::raw(Storage::Cuda(out), self.shape().to_vec(), self.device())
            }
        };
        let saved = out.clone();
        out.with_grad(&[self], || LogSoftmaxBackward { output: saved })
    }

    /// The tensor as `(rows, last_dim)`.
    pub(crate) fn as_rows(&self) -> (usize, usize) {
        let cols = self.last_dim();
        (self.numel() / cols, cols)
    }
}

fn sigmoid_scalar(x: f32) -> f32 {
    1.0 / (1.0 + (-x).exp())
}

fn gelu_scalar(x: f32) -> f32 {
    0.5 * x * (1.0 + libm::erff(x * std::f32::consts::FRAC_1_SQRT_2))
}

// subtracted before exp so large logits do not overflow
fn row_max(row: &[f32]) -> f32 {
    row.iter().copied().fold(f32::NEG_INFINITY, f32::max)
}