open_stomp/nam/
activation.rs

1//! Activation functions, as NeuralAmpModelerCore defines them (`NAM/activations.h`).
2
3use serde_json::Value;
4
5use super::NamError;
6
7#[derive(Clone, Debug, PartialEq)]
8pub enum Activation {
9    Tanh,
10    Hardtanh,
11    Fasttanh,
12    ReLU,
13    LeakyReLU(f32),
14    /// One slope per channel (or a single slope for all).
15    PReLU(Vec<f32>),
16    Sigmoid,
17    SiLU,
18    Hardswish,
19    LeakyHardtanh { min_val: f32, max_val: f32, min_slope: f32, max_slope: f32 },
20    Softsign,
21}
22
23fn fast_tanh(x: f32) -> f32 {
24    let ax = x.abs();
25    let x2 = x * x;
26    (x * (2.455_507_5 + 2.455_507_5 * ax + (0.893_229_85 + 0.821_226_7 * ax) * x2))
27        / (2.445_066_3 + (2.445_066_3 + x2) * (x + 0.814_642_7 * x * ax).abs())
28}
29
30fn sigmoid(x: f32) -> f32 {
31    1.0 / (1.0 + (-x).exp())
32}
33
34impl Activation {
35    /// From a config entry: a name (`"Tanh"`) or an object (`{"type": "LeakyReLU", "negative_slope": 0.01}`).
36    pub fn from_json(v: &Value) -> Result<Activation, NamError> {
37        let (name, obj) = match v {
38            Value::String(s) => (s.as_str(), None),
39            Value::Object(o) => (o.get("type").and_then(Value::as_str).unwrap_or(""), Some(o)),
40            _ => return Err(NamError::Unsupported(format!("activation {v}"))),
41        };
42        let num = |key: &str, default: f32| {
43            obj.and_then(|o| o.get(key)).and_then(Value::as_f64).map_or(default, |x| x as f32)
44        };
45        Ok(match name {
46            "Tanh" => Activation::Tanh,
47            "Hardtanh" => Activation::Hardtanh,
48            "Fasttanh" => Activation::Fasttanh,
49            "ReLU" => Activation::ReLU,
50            "LeakyReLU" => Activation::LeakyReLU(num("negative_slope", 0.01)),
51            "PReLU" => {
52                let slopes = obj
53                    .and_then(|o| o.get("negative_slopes"))
54                    .and_then(Value::as_array)
55                    .map(|a| a.iter().filter_map(Value::as_f64).map(|x| x as f32).collect())
56                    .unwrap_or_else(|| vec![num("negative_slope", 0.01)]);
57                Activation::PReLU(slopes)
58            }
59            "Sigmoid" => Activation::Sigmoid,
60            "SiLU" => Activation::SiLU,
61            "Hardswish" => Activation::Hardswish,
62            "LeakyHardtanh" | "LeakyHardTanh" => Activation::LeakyHardtanh {
63                min_val: num("min_val", -1.0),
64                max_val: num("max_val", 1.0),
65                min_slope: num("min_slope", 0.01),
66                max_slope: num("max_slope", 0.01),
67            },
68            "Softsign" => Activation::Softsign,
69            other => return Err(NamError::Unsupported(format!("activation {other:?}"))),
70        })
71    }
72
73    /// Apply in place to frame-major data with `channels` values per frame.
74    pub fn apply(&self, data: &mut [f32], channels: usize) {
75        match self {
76            Activation::Tanh => data.iter_mut().for_each(|x| *x = x.tanh()),
77            Activation::Hardtanh => data.iter_mut().for_each(|x| *x = x.clamp(-1.0, 1.0)),
78            Activation::Fasttanh => data.iter_mut().for_each(|x| *x = fast_tanh(*x)),
79            Activation::ReLU => data.iter_mut().for_each(|x| *x = if *x > 0.0 { *x } else { 0.0 }),
80            Activation::LeakyReLU(s) => data.iter_mut().for_each(|x| *x = if *x > 0.0 { *x } else { s * *x }),
81            Activation::PReLU(slopes) => {
82                for (i, x) in data.iter_mut().enumerate() {
83                    let s = slopes[if slopes.len() == 1 { 0 } else { i % channels }];
84                    *x = if *x > 0.0 { *x } else { s * *x };
85                }
86            }
87            Activation::Sigmoid => data.iter_mut().for_each(|x| *x = sigmoid(*x)),
88            Activation::SiLU => data.iter_mut().for_each(|x| *x *= sigmoid(*x)),
89            Activation::Hardswish => data.iter_mut().for_each(|x| *x = *x * (*x + 3.0).clamp(0.0, 6.0) / 6.0),
90            Activation::LeakyHardtanh { min_val, max_val, min_slope, max_slope } => data.iter_mut().for_each(|x| {
91                *x = if *x < *min_val {
92                    (*x - min_val) * min_slope + min_val
93                } else if *x > *max_val {
94                    (*x - max_val) * max_slope + max_val
95                } else {
96                    *x
97                }
98            }),
99            Activation::Softsign => data.iter_mut().for_each(|x| *x /= 1.0 + x.abs()),
100        }
101    }
102
103    /// Apply to one value of channel `ch`.
104    pub fn apply_one(&self, x: f32, ch: usize) -> f32 {
105        let mut v = [x];
106        match self {
107            Activation::PReLU(slopes) => {
108                let s = slopes[if slopes.len() == 1 { 0 } else { ch }];
109                if x > 0.0 { x } else { s * x }
110            }
111            _ => {
112                self.apply(&mut v, 1);
113                v[0]
114            }
115        }
116    }
117}
open_stomp/nam/activation.rs

open_stomp/nam/
activation.rs
