burn_nn/activation/
activation_wrapper.rs

1use burn_core as burn;
2
3use crate::activation::{
4    Gelu, HardSigmoid, HardSigmoidConfig, LeakyRelu, LeakyReluConfig, PRelu, PReluConfig, Relu,
5    Sigmoid, SwiGlu, SwiGluConfig, Tanh,
6};
7use burn::config::Config;
8use burn::module::Module;
9use burn::tensor::Tensor;
10use burn::tensor::backend::Backend;
11
12/// [`Activation`] Configuration.
13#[derive(Config, Debug)]
14#[non_exhaustive]
15pub enum ActivationConfig {
16    /// [`Gelu`] activation layer.
17    Gelu,
18
19    /// [`PRelu`] activation layer.
20    PRelu(PReluConfig),
21
22    /// [`Relu`] activation layer.
23    Relu,
24
25    /// [`LeakyRelu`] activation layer.
26    LeakyRelu(LeakyReluConfig),
27
28    /// [`SwiGlu`] activation layer.
29    SwiGlu(SwiGluConfig),
30
31    /// [`Sigmoid`] activation layer.
32    Sigmoid,
33
34    /// [`Tanh`] activation layer.
35    Tanh,
36
37    /// [`HardSigmoid`] activation layer.
38    HardSigmoid(HardSigmoidConfig),
39}
40
41impl From<PReluConfig> for ActivationConfig {
42    fn from(config: PReluConfig) -> Self {
43        Self::PRelu(config)
44    }
45}
46
47impl From<LeakyReluConfig> for ActivationConfig {
48    fn from(config: LeakyReluConfig) -> Self {
49        Self::LeakyRelu(config)
50    }
51}
52
53impl From<SwiGluConfig> for ActivationConfig {
54    fn from(config: SwiGluConfig) -> Self {
55        Self::SwiGlu(config)
56    }
57}
58
59impl From<HardSigmoidConfig> for ActivationConfig {
60    fn from(config: HardSigmoidConfig) -> Self {
61        Self::HardSigmoid(config)
62    }
63}
64
65impl ActivationConfig {
66    /// Initialize a wrapped activation layer.
67    pub fn init<B: Backend>(&self, device: &B::Device) -> Activation<B> {
68        match self {
69            ActivationConfig::Relu => Relu.into(),
70            ActivationConfig::LeakyRelu(conf) => conf.init().into(),
71            ActivationConfig::Gelu => Gelu.into(),
72            ActivationConfig::PRelu(conf) => conf.init(device).into(),
73            ActivationConfig::SwiGlu(conf) => conf.init(device).into(),
74            ActivationConfig::HardSigmoid(conf) => conf.init().into(),
75            ActivationConfig::Sigmoid => Sigmoid.into(),
76            ActivationConfig::Tanh => Tanh.into(),
77        }
78    }
79}
80
81/// Activation Layer Wrapper.
82///
83/// Provides support for many in-built `burn::nn` activations.
84#[derive(Module, Debug)]
85#[non_exhaustive]
86#[allow(clippy::large_enum_variant)]
87pub enum Activation<B: Backend> {
88    /// [`Gelu`] activation layer.
89    Gelu(Gelu),
90
91    /// [`PRelu`] activation layer.
92    PRelu(PRelu<B>),
93
94    /// [`Relu`] activation layer.
95    Relu(Relu),
96
97    /// [`LeakyRelu`] activation layer.
98    LeakyRelu(LeakyRelu),
99
100    /// [`SwiGlu`] activation layer.
101    SwiGlu(SwiGlu<B>),
102
103    /// [`Sigmoid`] activation layer.
104    Sigmoid(Sigmoid),
105
106    /// [`Tanh`] activation layer.
107    Tanh(Tanh),
108
109    /// [`HardSigmoid`] activation layer.
110    HardSigmoid(HardSigmoid),
111}
112
113impl<B: Backend> From<Gelu> for Activation<B> {
114    fn from(layer: Gelu) -> Self {
115        Self::Gelu(layer)
116    }
117}
118
119impl<B: Backend> From<PRelu<B>> for Activation<B> {
120    fn from(layer: PRelu<B>) -> Self {
121        Self::PRelu(layer)
122    }
123}
124
125impl<B: Backend> From<Relu> for Activation<B> {
126    fn from(layer: Relu) -> Self {
127        Self::Relu(layer)
128    }
129}
130
131impl<B: Backend> From<LeakyRelu> for Activation<B> {
132    fn from(layer: LeakyRelu) -> Self {
133        Self::LeakyRelu(layer)
134    }
135}
136
137impl<B: Backend> From<SwiGlu<B>> for Activation<B> {
138    fn from(layer: SwiGlu<B>) -> Self {
139        Self::SwiGlu(layer)
140    }
141}
142
143impl<B: Backend> From<Sigmoid> for Activation<B> {
144    fn from(layer: Sigmoid) -> Self {
145        Self::Sigmoid(layer)
146    }
147}
148
149impl<B: Backend> From<Tanh> for Activation<B> {
150    fn from(layer: Tanh) -> Self {
151        Self::Tanh(layer)
152    }
153}
154
155impl<B: Backend> From<HardSigmoid> for Activation<B> {
156    fn from(layer: HardSigmoid) -> Self {
157        Self::HardSigmoid(layer)
158    }
159}
160
161impl<B: Backend> Activation<B> {
162    /// Forward pass.
163    pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
164        match self {
165            Activation::Relu(layer) => layer.forward(input),
166            Activation::LeakyRelu(layer) => layer.forward(input),
167            Activation::Gelu(layer) => layer.forward(input),
168            Activation::PRelu(layer) => layer.forward(input),
169            Activation::SwiGlu(layer) => layer.forward(input),
170            Activation::HardSigmoid(layer) => layer.forward(input),
171            Activation::Sigmoid(layer) => layer.forward(input),
172            Activation::Tanh(layer) => layer.forward(input),
173        }
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use crate::TestBackend;
181    use burn::module::Module;
182
183    fn make_input<B: Backend>(device: &B::Device) -> Tensor<B, 2> {
184        Tensor::from_data([[-1.0, -0.5, 0.0], [1.0, 0.5, 0.0]], device)
185    }
186
187    fn expect_tensor<B: Backend, const D: usize>(actual: Tensor<B, D>, expected: Tensor<B, D>) {
188        actual.to_data().assert_eq(&expected.to_data(), true);
189    }
190
191    fn check_stateless_config_output<B: Backend, const D: usize>(
192        config: ActivationConfig,
193        input: Tensor<B, D>,
194        expected: Tensor<B, D>,
195        device: &B::Device,
196    ) {
197        let act = config.init(device);
198        let output = act.forward(input);
199        expect_tensor(output, expected);
200    }
201
202    #[test]
203    fn test_gelu() {
204        let device = Default::default();
205        let input = make_input::<TestBackend>(&device);
206
207        let expected = Gelu.forward(input.clone());
208
209        check_stateless_config_output(ActivationConfig::Gelu, input, expected, &device)
210    }
211
212    #[test]
213    fn test_prelu() {
214        let device = Default::default();
215        let input = make_input::<TestBackend>(&device);
216
217        let inner_config = PReluConfig::new();
218        let expected = inner_config.init(&device).forward(input.clone());
219
220        check_stateless_config_output(inner_config.into(), input, expected, &device)
221    }
222
223    #[test]
224    fn test_relu() {
225        let device = Default::default();
226        let input = make_input::<TestBackend>(&device);
227
228        let expected = Relu.forward(input.clone());
229
230        check_stateless_config_output(ActivationConfig::Relu, input, expected, &device)
231    }
232
233    #[test]
234    fn test_leaky_relu() {
235        let device = Default::default();
236        let input = make_input::<TestBackend>(&device);
237
238        let inner_config = LeakyReluConfig::new();
239        let expected = inner_config.init().forward(input.clone());
240
241        check_stateless_config_output(inner_config.into(), input, expected, &device)
242    }
243
244    #[test]
245    fn test_swi_glu() {
246        let device = Default::default();
247        let input = make_input::<TestBackend>(&device);
248
249        let d_input = input.shape().dims[1];
250        let d_output = 2 * d_input;
251
252        let inner_config = SwiGluConfig::new(d_input, d_output);
253        let mut reference: SwiGlu<TestBackend> = inner_config.init(&device);
254
255        let config: ActivationConfig = inner_config.into();
256        let layer = config.init(&device);
257
258        match &layer {
259            Activation::SwiGlu(inner) => {
260                // Clone the initialized weights.
261                let state = inner.clone().into_record();
262                reference = reference.load_record(state);
263            }
264            _ => unreachable!(),
265        };
266
267        expect_tensor(
268            layer.forward(input.clone()),
269            reference.forward(input.clone()),
270        )
271    }
272
273    #[test]
274    fn test_sigmoid() {
275        let device = Default::default();
276        let input = make_input::<TestBackend>(&device);
277
278        let expected = Sigmoid.forward(input.clone());
279
280        check_stateless_config_output(ActivationConfig::Sigmoid, input, expected, &device)
281    }
282
283    #[test]
284    fn test_tanh() {
285        let device = Default::default();
286        let input = make_input::<TestBackend>(&device);
287
288        let expected = Tanh.forward(input.clone());
289
290        check_stateless_config_output(ActivationConfig::Tanh, input, expected, &device)
291    }
292
293    #[test]
294    fn test_hard_sigmoid() {
295        let device = Default::default();
296        let input = make_input::<TestBackend>(&device);
297
298        let inner_config = HardSigmoidConfig::new();
299        let expected = inner_config.init().forward(input.clone());
300
301        check_stateless_config_output(inner_config.into(), input, expected, &device)
302    }
303}
burn_nn/activation/activation_wrapper.rs

burn_nn/activation/
activation_wrapper.rs
