boostr 0.1.0

ML framework built on numr - attention, quantization, model architectures
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

//! CPU SGD fused kernel implementations (f32 and f64).

use crate::error::Result;
use numr::runtime::cpu::CpuRuntime;
use numr::tensor::Tensor;

#[allow(clippy::too_many_arguments)]
pub(super) fn fused_sgd_f32(
    param: &Tensor<CpuRuntime>,
    grad: &Tensor<CpuRuntime>,
    momentum_buf: Option<&Tensor<CpuRuntime>>,
    lr: f64,
    momentum: f64,
    dampening: f64,
    wd: f64,
    nesterov: bool,
) -> Result<(Tensor<CpuRuntime>, Tensor<CpuRuntime>)> {
    let n: usize = param.shape().iter().product();
    let p = param.to_vec::<f32>();
    let g = grad.to_vec::<f32>();
    let has_buf = momentum_buf.is_some();
    let buf_data = momentum_buf.map(|b| b.to_vec::<f32>());

    let mut new_p = vec![0.0f32; n];
    let mut new_buf = vec![0.0f32; n];

    let lr_f = lr as f32;
    let mom = momentum as f32;
    let damp = dampening as f32;
    let w = wd as f32;

    for i in 0..n {
        let grad_wd = if w > 0.0 { g[i] + w * p[i] } else { g[i] };

        let buf = if mom > 0.0 {
            if has_buf {
                let prev = buf_data.as_ref().expect("has_buf is true")[i];
                mom * prev + (1.0 - damp) * grad_wd
            } else {
                grad_wd
            }
        } else {
            grad_wd
        };
        new_buf[i] = buf;

        let update = if nesterov && mom > 0.0 {
            grad_wd + mom * buf
        } else if mom > 0.0 {
            buf
        } else {
            grad_wd
        };

        new_p[i] = p[i] - lr_f * update;
    }

    let shape = param.shape();
    let device = param.device();
    Ok((
        Tensor::<CpuRuntime>::from_slice(&new_p, shape, device),
        Tensor::<CpuRuntime>::from_slice(&new_buf, shape, device),
    ))
}

#[allow(clippy::too_many_arguments)]
pub(super) fn fused_sgd_f64(
    param: &Tensor<CpuRuntime>,
    grad: &Tensor<CpuRuntime>,
    momentum_buf: Option<&Tensor<CpuRuntime>>,
    lr: f64,
    momentum: f64,
    dampening: f64,
    wd: f64,
    nesterov: bool,
) -> Result<(Tensor<CpuRuntime>, Tensor<CpuRuntime>)> {
    let n: usize = param.shape().iter().product();
    let p = param.to_vec::<f64>();
    let g = grad.to_vec::<f64>();
    let has_buf = momentum_buf.is_some();
    let buf_data = momentum_buf.map(|b| b.to_vec::<f64>());

    let mut new_p = vec![0.0f64; n];
    let mut new_buf = vec![0.0f64; n];

    for i in 0..n {
        let grad_wd = if wd > 0.0 { g[i] + wd * p[i] } else { g[i] };

        let buf = if momentum > 0.0 {
            if has_buf {
                let prev = buf_data.as_ref().expect("has_buf is true")[i];
                momentum * prev + (1.0 - dampening) * grad_wd
            } else {
                grad_wd
            }
        } else {
            grad_wd
        };
        new_buf[i] = buf;

        let update = if nesterov && momentum > 0.0 {
            grad_wd + momentum * buf
        } else if momentum > 0.0 {
            buf
        } else {
            grad_wd
        };

        new_p[i] = p[i] - lr * update;
    }

    let shape = param.shape();
    let device = param.device();
    Ok((
        Tensor::<CpuRuntime>::from_slice(&new_p, shape, device),
        Tensor::<CpuRuntime>::from_slice(&new_buf, shape, device),
    ))
}