scirs2_neural/layers/
mod.rs

1//! Neural network layers implementation
2//!
3//! This module provides implementations of various neural network layers
4//! such as dense (fully connected), attention, convolution, pooling, etc.
5//! Layers are the fundamental building blocks of neural networks.
6
7use crate::error::Result;
8use scirs2_core::ndarray::{Array, ScalarOperand};
9use scirs2_core::numeric::{Float, NumAssign};
10use std::fmt::Debug;
11
12/// Base trait for neural network layers
13///
14/// This trait defines the core interface that all neural network layers must implement.
15/// It supports forward propagation, backpropagation, parameter management, and
16/// training/evaluation mode switching.
17pub trait Layer<F: Float + Debug + ScalarOperand + NumAssign>: Send + Sync {
18    /// Forward pass of the layer
19    ///
20    /// Computes the output of the layer given an input tensor.
21    fn forward(
22        &self,
23        input: &Array<F, scirs2_core::ndarray::IxDyn>,
24    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>>;
25
26    /// Backward pass of the layer to compute gradients
27    ///
28    /// Computes gradients with respect to the layer's input, which is needed
29    /// for backpropagation.
30    fn backward(
31        &self,
32        input: &Array<F, scirs2_core::ndarray::IxDyn>,
33        grad_output: &Array<F, scirs2_core::ndarray::IxDyn>,
34    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>>;
35
36    /// Update the layer parameters with the given learning rate
37    fn update(&mut self, learningrate: F) -> Result<()>;
38
39    /// Get the layer as a dyn Any for downcasting
40    fn as_any(&self) -> &dyn std::any::Any;
41
42    /// Get the layer as a mutable dyn Any for downcasting
43    fn as_any_mut(&mut self) -> &mut dyn std::any::Any;
44
45    /// Get the parameters of the layer
46    fn params(&self) -> Vec<Array<F, scirs2_core::ndarray::IxDyn>> {
47        Vec::new()
48    }
49
50    /// Get the gradients of the layer parameters
51    fn gradients(&self) -> Vec<Array<F, scirs2_core::ndarray::IxDyn>> {
52        Vec::new()
53    }
54
55    /// Set the gradients of the layer parameters
56    fn set_gradients(
57        &mut self,
58        _gradients: &[Array<F, scirs2_core::ndarray::IxDyn>],
59    ) -> Result<()> {
60        Ok(())
61    }
62
63    /// Set the parameters of the layer
64    fn set_params(&mut self, _params: &[Array<F, scirs2_core::ndarray::IxDyn>]) -> Result<()> {
65        Ok(())
66    }
67
68    /// Set the layer to training mode (true) or evaluation mode (false)
69    fn set_training(&mut self, _training: bool) {
70        // Default implementation: do nothing
71    }
72
73    /// Get the current training mode
74    fn is_training(&self) -> bool {
75        true // Default implementation: always in training mode
76    }
77
78    /// Get the type of the layer (e.g., "Dense", "Conv2D")
79    fn layer_type(&self) -> &str {
80        "Unknown"
81    }
82
83    /// Get the number of trainable parameters in this layer
84    fn parameter_count(&self) -> usize {
85        0
86    }
87
88    /// Get a detailed description of this layer
89    fn layer_description(&self) -> String {
90        format!("type:{}", self.layer_type())
91    }
92
93    /// Get the input shape if known
94    fn inputshape(&self) -> Option<Vec<usize>> {
95        None
96    }
97
98    /// Get the output shape if known  
99    fn outputshape(&self) -> Option<Vec<usize>> {
100        None
101    }
102
103    /// Get the name of the layer if set
104    fn name(&self) -> Option<&str> {
105        None
106    }
107}
108
109/// Trait for layers with parameters (weights, biases)
110pub trait ParamLayer<F: Float + Debug + ScalarOperand + NumAssign>: Layer<F> {
111    /// Get the parameters of the layer as a vector of arrays
112    fn get_parameters(&self) -> Vec<Array<F, scirs2_core::ndarray::IxDyn>>;
113
114    /// Get the gradients of the parameters
115    fn get_gradients(&self) -> Vec<Array<F, scirs2_core::ndarray::IxDyn>>;
116
117    /// Set the parameters
118    fn set_parameters(&mut self, params: Vec<Array<F, scirs2_core::ndarray::IxDyn>>) -> Result<()>;
119}
120
121/// Information about a layer for visualization purposes
122#[derive(Debug, Clone)]
123pub struct LayerInfo {
124    /// Index of the layer in the sequence
125    pub index: usize,
126    /// Name of the layer
127    pub name: String,
128    /// Type of the layer
129    pub layer_type: String,
130    /// Number of parameters in the layer
131    pub parameter_count: usize,
132    /// Input shape of the layer
133    pub inputshape: Option<Vec<usize>>,
134    /// Output shape of the layer
135    pub outputshape: Option<Vec<usize>>,
136}
137
138/// Sequential container for neural network layers
139///
140/// A Sequential model is a linear stack of layers where data flows through
141/// each layer in order.
142pub struct Sequential<F: Float + Debug + ScalarOperand + NumAssign> {
143    layers: Vec<Box<dyn Layer<F> + Send + Sync>>,
144    training: bool,
145    /// Output of every layer of the last forward pass, so that `backward` can
146    /// hand each layer the input it actually saw
147    layer_outputs: std::sync::RwLock<Vec<Array<F, scirs2_core::ndarray::IxDyn>>>,
148}
149
150impl<F: Float + Debug + ScalarOperand + NumAssign> std::fmt::Debug for Sequential<F> {
151    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
152        f.debug_struct("Sequential")
153            .field("num_layers", &self.layers.len())
154            .field("training", &self.training)
155            .finish()
156    }
157}
158
159impl<F: Float + Debug + ScalarOperand + NumAssign + 'static> Clone for Sequential<F> {
160    fn clone(&self) -> Self {
161        // We can't clone the layers, so we just create an empty Sequential
162        // with the same training flag
163        Self {
164            layers: Vec::new(),
165            training: self.training,
166            layer_outputs: std::sync::RwLock::new(Vec::new()),
167        }
168    }
169}
170
171impl<F: Float + Debug + ScalarOperand + NumAssign> Default for Sequential<F> {
172    fn default() -> Self {
173        Self::new()
174    }
175}
176
177impl<F: Float + Debug + ScalarOperand + NumAssign> Sequential<F> {
178    /// Create a new Sequential container
179    pub fn new() -> Self {
180        Self {
181            layers: Vec::new(),
182            training: true,
183            layer_outputs: std::sync::RwLock::new(Vec::new()),
184        }
185    }
186
187    /// Add a layer to the container
188    pub fn add<L: Layer<F> + Send + Sync + 'static>(&mut self, layer: L) {
189        self.layers.push(Box::new(layer));
190    }
191
192    /// Get the number of layers
193    pub fn len(&self) -> usize {
194        self.layers.len()
195    }
196
197    /// Check if there are no layers
198    pub fn is_empty(&self) -> bool {
199        self.layers.is_empty()
200    }
201
202    /// Get total parameter count across all layers
203    pub fn total_parameters(&self) -> usize {
204        self.layers
205            .iter()
206            .map(|layer| layer.parameter_count())
207            .sum()
208    }
209
210    /// Get layer information for visualization purposes
211    pub fn layer_info(&self) -> Vec<LayerInfo> {
212        self.layers
213            .iter()
214            .enumerate()
215            .map(|(i, layer)| LayerInfo {
216                index: i,
217                name: layer.name().unwrap_or(&format!("Layer_{i}")).to_string(),
218                layer_type: layer.layer_type().to_string(),
219                parameter_count: layer.parameter_count(),
220                inputshape: layer.inputshape(),
221                outputshape: layer.outputshape(),
222            })
223            .collect()
224    }
225}
226
227impl<F: Float + Debug + ScalarOperand + NumAssign + Send + Sync> Layer<F> for Sequential<F> {
228    fn forward(
229        &self,
230        input: &Array<F, scirs2_core::ndarray::IxDyn>,
231    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>> {
232        let mut outputs = Vec::with_capacity(self.layers.len());
233        let mut output = input.clone();
234        for layer in &self.layers {
235            output = layer.forward(&output)?;
236            outputs.push(output.clone());
237        }
238        if let Ok(mut cache) = self.layer_outputs.write() {
239            *cache = outputs;
240        }
241        Ok(output)
242    }
243
244    /// Propagate the gradient through the stack in reverse, giving every layer
245    /// the input it received during the forward pass.
246    fn backward(
247        &self,
248        input: &Array<F, scirs2_core::ndarray::IxDyn>,
249        grad_output: &Array<F, scirs2_core::ndarray::IxDyn>,
250    ) -> Result<Array<F, scirs2_core::ndarray::IxDyn>> {
251        if self.layers.is_empty() {
252            // An empty container is the identity, so is its gradient.
253            return Ok(grad_output.clone());
254        }
255        let outputs = self
256            .layer_outputs
257            .read()
258            .map_err(|_| {
259                crate::error::NeuralError::InferenceError(
260                    "Failed to acquire read lock on layer outputs".to_string(),
261                )
262            })?
263            .clone();
264        if outputs.len() != self.layers.len() {
265            return Err(crate::error::NeuralError::InferenceError(format!(
266                "Cached {} layer outputs for {} layers. Call forward() first.",
267                outputs.len(),
268                self.layers.len()
269            )));
270        }
271
272        let mut grad = grad_output.clone();
273        for (idx, layer) in self.layers.iter().enumerate().rev() {
274            let layer_input = if idx == 0 { input } else { &outputs[idx - 1] };
275            grad = layer.backward(layer_input, &grad)?;
276        }
277        Ok(grad)
278    }
279
280    fn update(&mut self, learningrate: F) -> Result<()> {
281        for layer in &mut self.layers {
282            layer.update(learningrate)?;
283        }
284        Ok(())
285    }
286
287    fn params(&self) -> Vec<Array<F, scirs2_core::ndarray::IxDyn>> {
288        let mut params = Vec::new();
289        for layer in &self.layers {
290            params.extend(layer.params());
291        }
292        params
293    }
294
295    fn gradients(&self) -> Vec<Array<F, scirs2_core::ndarray::IxDyn>> {
296        let mut grads = Vec::new();
297        for layer in &self.layers {
298            grads.extend(layer.gradients());
299        }
300        grads
301    }
302
303    fn set_training(&mut self, training: bool) {
304        self.training = training;
305        for layer in &mut self.layers {
306            layer.set_training(training);
307        }
308    }
309
310    fn is_training(&self) -> bool {
311        self.training
312    }
313
314    fn as_any(&self) -> &dyn std::any::Any {
315        self
316    }
317
318    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
319        self
320    }
321
322    fn layer_type(&self) -> &str {
323        "Sequential"
324    }
325
326    fn parameter_count(&self) -> usize {
327        self.layers
328            .iter()
329            .map(|layer| layer.parameter_count())
330            .sum()
331    }
332}
333
334/// Configuration enum for different types of layers
335#[derive(Debug, Clone)]
336pub enum LayerConfig {
337    /// Dense (fully connected) layer
338    Dense {
339        input_size: usize,
340        output_size: usize,
341        activation: Option<String>,
342    },
343    /// 2D Convolutional layer
344    Conv2D {
345        in_channels: usize,
346        out_channels: usize,
347        kernel_size: (usize, usize),
348    },
349    /// Dropout layer
350    Dropout { rate: f64 },
351}
352
353// Fixed modules
354pub mod conv;
355pub mod dense;
356pub mod dropout;
357pub mod graph_conv;
358pub mod layer_norm_2d;
359pub mod norm_variants;
360pub mod normalization;
361pub mod patch_embed;
362pub mod recurrent;
363
364// Additional layer modules
365mod attention;
366mod embedding;
367mod flash_attention;
368mod flash_attention_v2;
369mod grouped_query_attention;
370mod multi_query_attention;
371mod regularization;
372pub mod rnn_thread_safe;
373
374// Re-export all modules
375pub use attention::{AttentionConfig, AttentionMask, MultiHeadAttention, SelfAttention};
376pub use conv::{AvgPool2D, Conv2D, GlobalAvgPool2D, MaxPool2D};
377pub use dense::Dense;
378pub use dropout::Dropout;
379pub use embedding::{Embedding, EmbeddingConfig, PositionalEmbedding};
380pub use flash_attention::{flash_attention_compute, FlashAttention, FlashAttentionConfig};
381pub use flash_attention_v2::{
382    flash_attention_v2_compute, FlashAttentionV2, FlashAttentionV2Config,
383};
384pub use graph_conv::{
385    GraphActivation, GraphAttentionLayer, GraphConvLayer, GraphSageLayer, SageAggregator,
386};
387pub use grouped_query_attention::{
388    GqaKvCache, GroupedQueryAttention, GroupedQueryAttentionConfig, RotaryPositionEmbedding,
389};
390pub use layer_norm_2d::LayerNorm2D;
391pub use multi_query_attention::{KvCache, MultiQueryAttention, MultiQueryAttentionConfig};
392pub use norm_variants::{GroupNorm, InstanceNorm, RMSNorm, WeightNorm};
393pub use normalization::{BatchNorm, LayerNorm};
394pub use patch_embed::PatchEmbedding;
395pub use recurrent::rnn::{RNNConfig, RecurrentActivation as RecurrentActivationRNN};
396pub use recurrent::{LSTM, RNN};
397pub use regularization::{
398    ActivityRegularization, L1ActivityRegularization, L2ActivityRegularization,
399};
400pub use rnn_thread_safe::{
401    RecurrentActivation as ThreadSafeRecurrentActivation, ThreadSafeBidirectional, ThreadSafeRNN,
402};
scirs2_neural/layers/mod.rs

scirs2_neural/layers/
mod.rs
