Skip to main content

ruvector_cnn/layers/
mod.rs

1//! Neural Network Layers
2//!
3//! This module provides standard CNN layers with SIMD-optimized implementations:
4//! - **Conv2d**: 2D convolution with configurable kernel, stride, padding
5//! - **DepthwiseSeparableConv**: MobileNet-style efficient convolutions
6//! - **BatchNorm**: Batch normalization with learned parameters
7//! - **Activations**: ReLU, ReLU6, Swish, HardSwish, Sigmoid
8//! - **Pooling**: GlobalAvgPool, MaxPool2d, AvgPool2d
9//! - **Linear**: Fully connected layer
10
11pub mod activation;
12pub mod batchnorm;
13pub mod conv;
14pub mod linear;
15pub mod pooling;
16
17// Quantized layers (ADR-091 Phase 4)
18pub mod quantized_conv2d;
19pub mod quantized_depthwise;
20pub mod quantized_linear;
21pub mod quantized_pooling;
22pub mod quantized_residual;
23
24pub use activation::{Activation, ActivationType, HardSwish, ReLU, ReLU6, Sigmoid, Swish};
25pub use batchnorm::{BatchNorm, BatchNorm2d};
26pub use conv::{Conv2d, DepthwiseSeparableConv};
27pub use linear::Linear;
28pub use pooling::{AvgPool2d, GlobalAvgPool, GlobalAvgPool2d, MaxPool2d};
29
30// Quantized layer exports
31pub use quantized_conv2d::QuantizedConv2d;
32pub use quantized_depthwise::QuantizedDepthwiseConv2d;
33pub use quantized_linear::QuantizedLinear;
34pub use quantized_pooling::{QuantizedAvgPool2d, QuantizedMaxPool2d};
35pub use quantized_residual::QuantizedResidualAdd;
36
37use crate::{CnnResult, Tensor};
38
39/// Tensor shape for 4D tensors (N, C, H, W).
40#[derive(Clone, Copy, Debug, PartialEq, Eq)]
41pub struct TensorShape {
42    /// Batch size
43    pub n: usize,
44    /// Number of channels
45    pub c: usize,
46    /// Height
47    pub h: usize,
48    /// Width
49    pub w: usize,
50}
51
52impl TensorShape {
53    /// Creates a new tensor shape.
54    pub fn new(n: usize, c: usize, h: usize, w: usize) -> Self {
55        Self { n, c, h, w }
56    }
57
58    /// Returns the total number of elements.
59    pub fn numel(&self) -> usize {
60        self.n * self.c * self.h * self.w
61    }
62
63    /// Returns the spatial size (H * W).
64    pub fn spatial_size(&self) -> usize {
65        self.h * self.w
66    }
67
68    /// Returns the channel size (C * H * W).
69    pub fn channel_size(&self) -> usize {
70        self.c * self.h * self.w
71    }
72}
73
74impl std::fmt::Display for TensorShape {
75    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76        write!(f, "[{}, {}, {}, {}]", self.n, self.c, self.h, self.w)
77    }
78}
79
80/// Computes output size for convolution or pooling.
81pub fn conv_output_size(input: usize, kernel: usize, stride: usize, padding: usize, dilation: usize) -> usize {
82    let effective_kernel = dilation * (kernel - 1) + 1;
83    (input + 2 * padding - effective_kernel) / stride + 1
84}
85
86/// Trait for all neural network layers
87pub trait Layer {
88    /// Perform the forward pass
89    fn forward(&self, input: &Tensor) -> CnnResult<Tensor>;
90
91    /// Get the layer name
92    fn name(&self) -> &'static str;
93
94    /// Get the number of parameters
95    fn num_params(&self) -> usize {
96        0
97    }
98}
99
100// =============================================================================
101// Standalone layer functions (for backward compatibility with old backbone code)
102// =============================================================================
103
104/// 3x3 convolution function (standalone, for backward compatibility)
105///
106/// Input layout: NCHW flattened as [N * C * H * W]
107/// Output layout: NCHW flattened
108pub fn conv2d_3x3(
109    input: &[f32],
110    weights: &[f32],
111    in_channels: usize,
112    out_channels: usize,
113    height: usize,
114    width: usize,
115) -> Vec<f32> {
116    let out_h = height; // Same padding assumed
117    let out_w = width;
118    let mut output = vec![0.0f32; out_h * out_w * out_channels];
119
120    crate::simd::scalar::conv_3x3_scalar(
121        input,
122        weights,
123        &mut output,
124        height,
125        width,
126        in_channels,
127        out_channels,
128        1, // stride
129        1, // padding for same output size
130    );
131
132    output
133}
134
135/// Batch normalization function (standalone, for backward compatibility)
136///
137/// y = gamma * (x - mean) / sqrt(var + epsilon) + beta
138pub fn batch_norm(
139    input: &[f32],
140    gamma: &[f32],
141    beta: &[f32],
142    mean: &[f32],
143    var: &[f32],
144    epsilon: f32,
145) -> Vec<f32> {
146    let mut output = vec![0.0f32; input.len()];
147    let channels = gamma.len();
148
149    crate::simd::batch_norm_simd(
150        input,
151        &mut output,
152        gamma,
153        beta,
154        mean,
155        var,
156        epsilon,
157        channels,
158    );
159
160    output
161}
162
163/// HardSwish activation function (standalone, for backward compatibility)
164///
165/// hard_swish(x) = x * relu6(x + 3) / 6
166pub fn hard_swish(input: &[f32]) -> Vec<f32> {
167    let mut output = vec![0.0f32; input.len()];
168    crate::simd::scalar::hard_swish_scalar(input, &mut output);
169    output
170}
171
172/// ReLU activation function (standalone, for backward compatibility)
173pub fn relu(input: &[f32]) -> Vec<f32> {
174    let mut output = vec![0.0f32; input.len()];
175    crate::simd::relu_simd(input, &mut output);
176    output
177}
178
179/// ReLU6 activation function (standalone, for backward compatibility)
180pub fn relu6(input: &[f32]) -> Vec<f32> {
181    let mut output = vec![0.0f32; input.len()];
182    crate::simd::relu6_simd(input, &mut output);
183    output
184}
185
186/// Global average pooling function (standalone, for backward compatibility)
187///
188/// Assumes NCHW layout, pools over H*W dimensions
189pub fn global_avg_pool(input: &[f32], channels: usize) -> Vec<f32> {
190    let spatial_size = input.len() / channels;
191    let mut output = vec![0.0f32; channels];
192
193    // Sum over spatial dimensions
194    for i in 0..input.len() {
195        let c = i % channels;
196        output[c] += input[i];
197    }
198
199    // Average
200    let inv_spatial = 1.0 / spatial_size as f32;
201    for o in output.iter_mut() {
202        *o *= inv_spatial;
203    }
204
205    output
206}
207
208// Re-export Conv2dBuilder from conv module
209pub use conv::Conv2dBuilder;
210
211// =============================================================================
212// Activation helper methods
213// =============================================================================
214
215impl Activation {
216    /// Creates a ReLU activation.
217    pub fn relu() -> Self {
218        Self::new(ActivationType::ReLU)
219    }
220
221    /// Creates a ReLU6 activation.
222    pub fn relu6() -> Self {
223        Self::new(ActivationType::ReLU6)
224    }
225
226    /// Creates a HardSwish activation.
227    pub fn hard_swish() -> Self {
228        Self::new(ActivationType::HardSwish)
229    }
230
231    /// Creates a HardSigmoid activation (using Sigmoid as approximation).
232    pub fn hard_sigmoid() -> Self {
233        Self::new(ActivationType::Sigmoid)
234    }
235
236    /// Creates an identity (no-op) activation.
237    pub fn identity() -> Self {
238        Self::new(ActivationType::Identity)
239    }
240}