mirror of
https://github.com/ruvnet/RuView
synced 2026-08-08 20:11:43 +00:00
feat: vendor midstream and sublinear-time-solver libraries
Add ruvnet/midstream (AIMDS real-time inference) and ruvnet/sublinear-time-solver (sublinear optimization algorithms) as vendored dependencies under vendor/. Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
@@ -0,0 +1,198 @@
|
||||
use ndarray::{Array2, Array1, Axis};
|
||||
use std::f32;
|
||||
|
||||
/// Self-attention mechanism for temporal sequences
|
||||
pub struct TemporalAttention {
|
||||
d_model: usize,
|
||||
n_heads: usize,
|
||||
d_k: usize,
|
||||
// Query, Key, Value projections for each head
|
||||
w_q: Vec<Array2<f32>>,
|
||||
w_k: Vec<Array2<f32>>,
|
||||
w_v: Vec<Array2<f32>>,
|
||||
w_o: Array2<f32>,
|
||||
// Positional encoding
|
||||
pos_encoding: Array2<f32>,
|
||||
}
|
||||
|
||||
impl TemporalAttention {
|
||||
pub fn new(d_model: usize, n_heads: usize, max_seq_len: usize) -> Self {
|
||||
assert_eq!(d_model % n_heads, 0, "d_model must be divisible by n_heads");
|
||||
let d_k = d_model / n_heads;
|
||||
|
||||
use rand::{thread_rng, Rng};
|
||||
let mut rng = thread_rng();
|
||||
let scale = (1.0 / d_k as f32).sqrt();
|
||||
|
||||
// Initialize projection matrices for each head
|
||||
let mut w_q = Vec::new();
|
||||
let mut w_k = Vec::new();
|
||||
let mut w_v = Vec::new();
|
||||
|
||||
for _ in 0..n_heads {
|
||||
w_q.push(Array2::from_shape_fn((d_k, d_model), |_|
|
||||
rng.gen::<f32>() * scale - scale/2.0));
|
||||
w_k.push(Array2::from_shape_fn((d_k, d_model), |_|
|
||||
rng.gen::<f32>() * scale - scale/2.0));
|
||||
w_v.push(Array2::from_shape_fn((d_k, d_model), |_|
|
||||
rng.gen::<f32>() * scale - scale/2.0));
|
||||
}
|
||||
|
||||
let w_o = Array2::from_shape_fn((d_model, d_model), |_|
|
||||
rng.gen::<f32>() * scale - scale/2.0);
|
||||
|
||||
// Create sinusoidal positional encoding
|
||||
let pos_encoding = Self::create_positional_encoding(max_seq_len, d_model);
|
||||
|
||||
Self {
|
||||
d_model,
|
||||
n_heads,
|
||||
d_k,
|
||||
w_q,
|
||||
w_k,
|
||||
w_v,
|
||||
w_o,
|
||||
pos_encoding,
|
||||
}
|
||||
}
|
||||
|
||||
fn create_positional_encoding(max_len: usize, d_model: usize) -> Array2<f32> {
|
||||
let mut encoding = Array2::zeros((max_len, d_model));
|
||||
|
||||
for pos in 0..max_len {
|
||||
for i in 0..d_model/2 {
|
||||
let angle = pos as f32 / (10000.0_f32.powf(2.0 * i as f32 / d_model as f32));
|
||||
encoding[[pos, 2*i]] = angle.sin();
|
||||
encoding[[pos, 2*i + 1]] = angle.cos();
|
||||
}
|
||||
}
|
||||
|
||||
encoding
|
||||
}
|
||||
|
||||
/// Scaled dot-product attention
|
||||
fn attention(&self, q: &Array2<f32>, k: &Array2<f32>, v: &Array2<f32>) -> Array2<f32> {
|
||||
let d_k_sqrt = (self.d_k as f32).sqrt();
|
||||
|
||||
// Compute attention scores: Q @ K^T / sqrt(d_k)
|
||||
let scores = q.dot(&k.t()) / d_k_sqrt;
|
||||
|
||||
// Apply softmax
|
||||
let exp_scores = scores.mapv(|x| x.exp());
|
||||
let sum_exp = exp_scores.sum_axis(Axis(1));
|
||||
let attention_weights = &exp_scores / &sum_exp.insert_axis(Axis(1));
|
||||
|
||||
// Apply attention to values
|
||||
attention_weights.dot(v)
|
||||
}
|
||||
|
||||
/// Multi-head attention forward pass
|
||||
pub fn forward(&self, x: &Array2<f32>) -> Array2<f32> {
|
||||
let seq_len = x.nrows();
|
||||
let batch_d = x.ncols();
|
||||
|
||||
// Add positional encoding
|
||||
let pos_slice = self.pos_encoding.slice(s![..seq_len, ..batch_d]);
|
||||
let x_pos = x + &pos_slice;
|
||||
|
||||
let mut head_outputs = Vec::new();
|
||||
|
||||
// Process each attention head
|
||||
for h in 0..self.n_heads {
|
||||
let q = x_pos.dot(&self.w_q[h].t());
|
||||
let k = x_pos.dot(&self.w_k[h].t());
|
||||
let v = x_pos.dot(&self.w_v[h].t());
|
||||
|
||||
let head_out = self.attention(&q, &k, &v);
|
||||
head_outputs.push(head_out);
|
||||
}
|
||||
|
||||
// Concatenate heads
|
||||
let mut concat = Array2::zeros((seq_len, self.d_model));
|
||||
for (h, head_out) in head_outputs.iter().enumerate() {
|
||||
let start = h * self.d_k;
|
||||
let end = start + self.d_k;
|
||||
concat.slice_mut(s![.., start..end]).assign(head_out);
|
||||
}
|
||||
|
||||
// Final linear projection
|
||||
concat.dot(&self.w_o.t())
|
||||
}
|
||||
|
||||
/// Extract temporal features with attention
|
||||
pub fn extract_features(&self, sequence: &[Vec<f32>]) -> Vec<f32> {
|
||||
let seq_len = sequence.len();
|
||||
let feat_dim = sequence[0].len();
|
||||
|
||||
// Convert to ndarray
|
||||
let mut x = Array2::zeros((seq_len, feat_dim));
|
||||
for (i, features) in sequence.iter().enumerate() {
|
||||
for (j, &val) in features.iter().enumerate() {
|
||||
x[[i, j]] = val;
|
||||
}
|
||||
}
|
||||
|
||||
// Apply attention
|
||||
let attended = self.forward(&x);
|
||||
|
||||
// Global average pooling over time
|
||||
attended.mean_axis(Axis(0))
|
||||
.unwrap()
|
||||
.to_vec()
|
||||
}
|
||||
}
|
||||
|
||||
/// Causal attention for autoregressive prediction
|
||||
pub struct CausalAttention {
|
||||
attention: TemporalAttention,
|
||||
mask: Array2<bool>,
|
||||
}
|
||||
|
||||
impl CausalAttention {
|
||||
pub fn new(d_model: usize, n_heads: usize, max_seq_len: usize) -> Self {
|
||||
let attention = TemporalAttention::new(d_model, n_heads, max_seq_len);
|
||||
|
||||
// Create causal mask (lower triangular)
|
||||
let mut mask = Array2::from_elem((max_seq_len, max_seq_len), false);
|
||||
for i in 0..max_seq_len {
|
||||
for j in 0..=i {
|
||||
mask[[i, j]] = true;
|
||||
}
|
||||
}
|
||||
|
||||
Self { attention, mask }
|
||||
}
|
||||
|
||||
/// Apply causal masking to attention scores
|
||||
pub fn forward_causal(&self, x: &Array2<f32>) -> Array2<f32> {
|
||||
let seq_len = x.nrows();
|
||||
|
||||
// Apply standard attention
|
||||
let output = self.attention.forward(x);
|
||||
|
||||
// Apply causal mask (in practice, this would be done inside attention computation)
|
||||
let mask_slice = self.mask.slice(s![..seq_len, ..seq_len]);
|
||||
|
||||
// Masked output
|
||||
output.masked_fill(&mask_slice.mapv(|b| !b), 0.0)
|
||||
}
|
||||
}
|
||||
|
||||
// Simplified implementation for masked_fill
|
||||
trait MaskedFill {
|
||||
fn masked_fill(&self, mask: &Array2<bool>, value: f32) -> Self;
|
||||
}
|
||||
|
||||
impl MaskedFill for Array2<f32> {
|
||||
fn masked_fill(&self, mask: &Array2<bool>, value: f32) -> Self {
|
||||
let mut result = self.clone();
|
||||
for ((i, j), &m) in mask.indexed_iter() {
|
||||
if !m {
|
||||
result[[i, j]] = value;
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
}
|
||||
|
||||
use ndarray::s;
|
||||
@@ -0,0 +1,17 @@
|
||||
use crate::data::{Sample, to_class};
|
||||
|
||||
/// Naive baseline: next = last_in_window; class from that
|
||||
pub struct Baseline;
|
||||
|
||||
impl Baseline {
|
||||
pub fn predict_reg(samples: &[Sample], window: usize) -> Vec<f32> {
|
||||
samples.iter().map(|s| s.x[window-1]).collect()
|
||||
}
|
||||
|
||||
pub fn predict_cls(samples: &[Sample], window: usize) -> Vec<usize> {
|
||||
samples.iter().map(|s| {
|
||||
let yhat = s.x[window-1];
|
||||
to_class(yhat)
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,62 @@
|
||||
use rand::{rngs::StdRng, SeedableRng, Rng};
|
||||
use rand_distr::{Normal, Distribution};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Sample {
|
||||
pub x: Vec<f32>, // windowed features
|
||||
pub y: f32, // next value or class
|
||||
}
|
||||
|
||||
pub struct Dataset {
|
||||
pub train: Vec<Sample>,
|
||||
pub val: Vec<Sample>,
|
||||
pub test: Vec<Sample>,
|
||||
}
|
||||
|
||||
/// Synthetic temporal process with regime shifts and delays.
|
||||
/// Tasks:
|
||||
/// 1) next_value regression
|
||||
/// 2) future_bucket classification (coarse future state)
|
||||
pub fn make_synthetic(window: usize, n: usize, seed: u64) -> Dataset {
|
||||
let mut rng = StdRng::seed_from_u64(seed);
|
||||
let noise = Normal::new(0.0, 0.3).unwrap();
|
||||
|
||||
// latent regime that flips with low prob
|
||||
let mut regime = 0.0f32;
|
||||
let mut series: Vec<f32> = Vec::with_capacity(n + window + 10);
|
||||
let mut last = 0.0f32;
|
||||
|
||||
for t in 0..(n + window + 10) {
|
||||
if rng.gen::<f32>() < 0.02 { regime = if regime == 0.0 { 1.0 } else { 0.0 }; }
|
||||
let drift = if regime == 0.0 { 0.02 } else { -0.015 };
|
||||
let val = 0.8 * last + drift + noise.sample(&mut rng) as f32;
|
||||
series.push(val);
|
||||
last = val;
|
||||
// occasional delayed impulse
|
||||
if t % 37 == 0 { last += 0.9; }
|
||||
}
|
||||
|
||||
let mut rows = Vec::new();
|
||||
for i in 0..n {
|
||||
let w = &series[i..i+window];
|
||||
let y = series[i+window];
|
||||
let mut x = w.to_vec();
|
||||
// add simple time features
|
||||
x.push(((i as f32) % 24.0) / 24.0);
|
||||
x.push(regime);
|
||||
rows.push(Sample { x, y });
|
||||
}
|
||||
|
||||
let split1 = (0.7 * rows.len() as f32) as usize;
|
||||
let split2 = (0.85 * rows.len() as f32) as usize;
|
||||
Dataset {
|
||||
train: rows[..split1].to_vec(),
|
||||
val: rows[split1..split2].to_vec(),
|
||||
test: rows[split2..].to_vec(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_class(y: f32) -> usize {
|
||||
// 3-bucket classification for coarse future state
|
||||
if y < -0.25 { 0 } else if y > 0.25 { 2 } else { 1 }
|
||||
}
|
||||
@@ -0,0 +1,300 @@
|
||||
use crate::mlp::Mlp;
|
||||
use crate::mlp_optimized::OptimizedMlp;
|
||||
use crate::mlp_ultra::UltraMlp;
|
||||
use crate::mlp_classifier::ClassifierMlp;
|
||||
use rayon::prelude::*;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Ensemble model combining multiple predictors with weighted voting
|
||||
pub struct EnsembleModel {
|
||||
models: Vec<ModelType>,
|
||||
weights: Vec<f32>,
|
||||
use_adaptive_weights: bool,
|
||||
}
|
||||
|
||||
enum ModelType {
|
||||
Simple(Mlp),
|
||||
Optimized(OptimizedMlp),
|
||||
Ultra(UltraMlp),
|
||||
Classifier(ClassifierMlp),
|
||||
}
|
||||
|
||||
impl EnsembleModel {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
models: Vec::new(),
|
||||
weights: Vec::new(),
|
||||
use_adaptive_weights: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_model_simple(&mut self, input: usize, hidden: usize, output: usize) {
|
||||
self.models.push(ModelType::Simple(Mlp::new(input, hidden, output)));
|
||||
self.weights.push(1.0);
|
||||
}
|
||||
|
||||
pub fn add_model_optimized(&mut self, input: usize, hidden: usize, output: usize) {
|
||||
self.models.push(ModelType::Optimized(OptimizedMlp::new(input, hidden, output)));
|
||||
self.weights.push(1.0);
|
||||
}
|
||||
|
||||
pub fn add_model_ultra(&mut self, input: usize, hidden: usize, output: usize) {
|
||||
self.models.push(ModelType::Ultra(UltraMlp::new(input, hidden, output)));
|
||||
self.weights.push(1.0);
|
||||
}
|
||||
|
||||
pub fn add_model_classifier(&mut self, input: usize, output: usize) {
|
||||
self.models.push(ModelType::Classifier(ClassifierMlp::new(input, output)));
|
||||
self.weights.push(1.0);
|
||||
}
|
||||
|
||||
/// Train all models in parallel with different random initializations
|
||||
pub fn train_ensemble(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>,
|
||||
epochs: usize, lr: f32, val_x: &Vec<Vec<f32>>, val_y: &Vec<usize>) {
|
||||
use rand::seq::SliceRandom;
|
||||
use rand::thread_rng;
|
||||
|
||||
// Train each model with different data shuffling for diversity
|
||||
let model_count = self.models.len();
|
||||
|
||||
// Parallel training with thread-safe access
|
||||
let x_arc = Arc::new(x.clone());
|
||||
let y_arc = Arc::new(y.clone());
|
||||
|
||||
// Sequential training (models contain mutable state)
|
||||
for (i, model) in self.models.iter_mut().enumerate() {
|
||||
let mut indices: Vec<usize> = (0..x.len()).collect();
|
||||
let mut rng = thread_rng();
|
||||
use rand::Rng;
|
||||
indices.shuffle(&mut rng);
|
||||
|
||||
// Bootstrap sampling for diversity
|
||||
let bootstrap_size = (x.len() as f32 * 0.8) as usize;
|
||||
let bootstrap_indices: Vec<usize> = (0..bootstrap_size)
|
||||
.map(|_| indices[rng.gen::<usize>() % indices.len()])
|
||||
.collect();
|
||||
|
||||
let x_bootstrap: Vec<Vec<f32>> = bootstrap_indices.iter()
|
||||
.map(|&i| x[i].clone())
|
||||
.collect();
|
||||
let y_bootstrap: Vec<f32> = bootstrap_indices.iter()
|
||||
.map(|&i| y[i])
|
||||
.collect();
|
||||
|
||||
// Train based on model type
|
||||
match model {
|
||||
ModelType::Simple(ref mut m) => {
|
||||
m.train_regression(&x_bootstrap, &y_bootstrap, epochs, lr);
|
||||
}
|
||||
ModelType::Optimized(ref mut m) => {
|
||||
m.train_batch(&x_bootstrap, &y_bootstrap, epochs, lr, 32);
|
||||
}
|
||||
ModelType::Ultra(ref mut m) => {
|
||||
m.train_batch_parallel(&x_bootstrap, &y_bootstrap, epochs, lr, 32);
|
||||
}
|
||||
ModelType::Classifier(ref mut m) => {
|
||||
m.train_classification(&x_bootstrap, &y_bootstrap, epochs, 32);
|
||||
}
|
||||
}
|
||||
|
||||
println!("Trained model {}/{}", i + 1, model_count);
|
||||
}
|
||||
|
||||
// Update weights based on validation performance
|
||||
if self.use_adaptive_weights && !val_x.is_empty() {
|
||||
self.update_weights(val_x, val_y);
|
||||
}
|
||||
}
|
||||
|
||||
/// Update ensemble weights based on validation accuracy
|
||||
fn update_weights(&mut self, val_x: &Vec<Vec<f32>>, val_y: &Vec<usize>) {
|
||||
let accuracies: Vec<f32> = self.models.iter_mut().map(|model| {
|
||||
let predictions = match model {
|
||||
ModelType::Simple(ref m) => m.predict_cls3(val_x),
|
||||
ModelType::Optimized(ref m) => m.predict_cls3(val_x),
|
||||
ModelType::Ultra(ref mut m) => m.predict_cls3_parallel(val_x),
|
||||
ModelType::Classifier(ref mut m) => m.predict_cls3(val_x),
|
||||
};
|
||||
|
||||
let correct = predictions.iter().zip(val_y.iter())
|
||||
.filter(|(p, y)| p == y)
|
||||
.count();
|
||||
correct as f32 / val_y.len() as f32
|
||||
}).collect();
|
||||
|
||||
// Convert accuracies to weights (squared for emphasis)
|
||||
let total_acc: f32 = accuracies.iter().map(|&a| a * a).sum();
|
||||
if total_acc > 0.0 {
|
||||
self.weights = accuracies.iter()
|
||||
.map(|&a| (a * a) / total_acc)
|
||||
.collect();
|
||||
}
|
||||
|
||||
println!("Ensemble weights: {:?}", self.weights);
|
||||
}
|
||||
|
||||
/// Predict using weighted voting
|
||||
pub fn predict_ensemble(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
let predictions: Vec<Vec<usize>> = self.models.iter_mut().map(|model| {
|
||||
match model {
|
||||
ModelType::Simple(ref m) => m.predict_cls3(x),
|
||||
ModelType::Optimized(ref m) => m.predict_cls3(x),
|
||||
ModelType::Ultra(ref mut m) => m.predict_cls3_parallel(x),
|
||||
ModelType::Classifier(ref mut m) => m.predict_cls3(x),
|
||||
}
|
||||
}).collect();
|
||||
|
||||
// Weighted voting for each sample
|
||||
(0..x.len()).map(|i| {
|
||||
let mut votes = vec![0.0; 3]; // 3 classes
|
||||
|
||||
for (model_idx, model_preds) in predictions.iter().enumerate() {
|
||||
let pred = model_preds[i];
|
||||
if pred < 3 {
|
||||
votes[pred] += self.weights[model_idx];
|
||||
}
|
||||
}
|
||||
|
||||
// Return class with highest weighted votes
|
||||
votes.iter()
|
||||
.enumerate()
|
||||
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
|
||||
.map(|(idx, _)| idx)
|
||||
.unwrap_or(1) // Default to middle class
|
||||
}).collect()
|
||||
}
|
||||
|
||||
/// Fast ensemble prediction with majority voting (no weights)
|
||||
pub fn predict_fast(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
let predictions: Vec<Vec<usize>> = self.models.par_iter_mut().map(|model| {
|
||||
match model {
|
||||
ModelType::Simple(ref m) => m.predict_cls3(x),
|
||||
ModelType::Optimized(ref m) => m.predict_cls3(x),
|
||||
ModelType::Ultra(ref mut m) => m.predict_cls3_parallel(x),
|
||||
ModelType::Classifier(ref mut m) => m.predict_cls3(x),
|
||||
}
|
||||
}).collect();
|
||||
|
||||
// Simple majority voting
|
||||
(0..x.len()).into_par_iter().map(|i| {
|
||||
let mut votes = [0u32; 3];
|
||||
for model_preds in &predictions {
|
||||
if model_preds[i] < 3 {
|
||||
votes[model_preds[i]] += 1;
|
||||
}
|
||||
}
|
||||
|
||||
votes.iter()
|
||||
.enumerate()
|
||||
.max_by_key(|&(_, &v)| v)
|
||||
.map(|(idx, _)| idx)
|
||||
.unwrap_or(1)
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Boosted ensemble using AdaBoost-style weighting
|
||||
pub struct BoostedEnsemble {
|
||||
weak_learners: Vec<Mlp>,
|
||||
alphas: Vec<f32>,
|
||||
input_dim: usize,
|
||||
hidden_dim: usize,
|
||||
}
|
||||
|
||||
impl BoostedEnsemble {
|
||||
pub fn new(input: usize, hidden: usize) -> Self {
|
||||
Self {
|
||||
weak_learners: Vec::new(),
|
||||
alphas: Vec::new(),
|
||||
input_dim: input,
|
||||
hidden_dim: hidden,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn train_boosted(&mut self, x: &Vec<Vec<f32>>, y: &Vec<usize>,
|
||||
n_estimators: usize, epochs: usize) {
|
||||
let n_samples = x.len();
|
||||
let mut weights = vec![1.0 / n_samples as f32; n_samples];
|
||||
|
||||
for t in 0..n_estimators {
|
||||
// Train weak learner on weighted data
|
||||
let mut learner = Mlp::new(self.input_dim, self.hidden_dim, 3);
|
||||
|
||||
// Convert classes to continuous for regression
|
||||
let y_cont: Vec<f32> = y.iter().map(|&c| c as f32 - 1.0).collect();
|
||||
|
||||
// Weight samples by resampling
|
||||
let mut weighted_x = Vec::new();
|
||||
let mut weighted_y = Vec::new();
|
||||
|
||||
use rand::thread_rng;
|
||||
use rand::distributions::{Distribution, WeightedIndex};
|
||||
let mut rng = thread_rng();
|
||||
let dist = WeightedIndex::new(&weights).unwrap();
|
||||
|
||||
for _ in 0..n_samples {
|
||||
let idx = dist.sample(&mut rng);
|
||||
weighted_x.push(x[idx].clone());
|
||||
weighted_y.push(y_cont[idx]);
|
||||
}
|
||||
|
||||
learner.train_regression(&weighted_x, &weighted_y, epochs, 0.01);
|
||||
|
||||
// Calculate error
|
||||
let predictions = learner.predict_cls3(x);
|
||||
let mut error = 0.0;
|
||||
for i in 0..n_samples {
|
||||
if predictions[i] != y[i] {
|
||||
error += weights[i];
|
||||
}
|
||||
}
|
||||
|
||||
// Avoid division by zero
|
||||
if error >= 0.5 {
|
||||
break; // Stop if not better than random
|
||||
}
|
||||
|
||||
// Calculate alpha
|
||||
let alpha = 0.5 * ((1.0 - error) / error.max(1e-10)).ln();
|
||||
|
||||
// Update weights
|
||||
for i in 0..n_samples {
|
||||
if predictions[i] != y[i] {
|
||||
weights[i] *= (alpha.exp());
|
||||
} else {
|
||||
weights[i] *= ((-alpha).exp());
|
||||
}
|
||||
}
|
||||
|
||||
// Normalize weights
|
||||
let sum: f32 = weights.iter().sum();
|
||||
weights.iter_mut().for_each(|w| *w /= sum);
|
||||
|
||||
self.weak_learners.push(learner);
|
||||
self.alphas.push(alpha);
|
||||
|
||||
println!("Boosting round {}/{}, error: {:.4}", t + 1, n_estimators, error);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn predict_boosted(&self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
let n_classes = 3;
|
||||
|
||||
x.iter().map(|xi| {
|
||||
let mut class_scores = vec![0.0; n_classes];
|
||||
|
||||
for (learner, &alpha) in self.weak_learners.iter().zip(&self.alphas) {
|
||||
let pred = learner.predict_cls3(&[xi.clone()])[0];
|
||||
if pred < n_classes {
|
||||
class_scores[pred] += alpha;
|
||||
}
|
||||
}
|
||||
|
||||
class_scores.iter()
|
||||
.enumerate()
|
||||
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
|
||||
.map(|(idx, _)| idx)
|
||||
.unwrap_or(1)
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,181 @@
|
||||
use ndarray::{Array1, Array2};
|
||||
use rand::{thread_rng, Rng, distributions::Uniform};
|
||||
use std::f32::consts::PI;
|
||||
|
||||
/// Random Fourier Features for kernel approximation
|
||||
/// Maps input to high-dimensional space where linear methods work well
|
||||
pub struct FourierFeatures {
|
||||
// Random projection parameters
|
||||
omega: Array2<f32>, // Random frequencies (D x d)
|
||||
b: Array1<f32>, // Random phase shifts
|
||||
input_dim: usize,
|
||||
feature_dim: usize,
|
||||
sigma: f32, // RBF kernel bandwidth
|
||||
|
||||
// Linear model on top
|
||||
weights: Array1<f32>,
|
||||
bias: f32,
|
||||
}
|
||||
|
||||
impl FourierFeatures {
|
||||
pub fn new(input_dim: usize, feature_dim: usize, sigma: f32) -> Self {
|
||||
let mut rng = thread_rng();
|
||||
let normal = Uniform::new(0.0, 2.0 * PI);
|
||||
|
||||
// Sample random frequencies from Gaussian (for RBF kernel)
|
||||
let scale = 1.0 / sigma;
|
||||
let omega = Array2::from_shape_fn((feature_dim, input_dim), |_|
|
||||
rng.gen::<f32>() * scale);
|
||||
|
||||
// Random phase shifts
|
||||
let b = Array1::from_shape_fn(feature_dim, |_| rng.sample(normal));
|
||||
|
||||
Self {
|
||||
omega,
|
||||
b,
|
||||
input_dim,
|
||||
feature_dim,
|
||||
sigma,
|
||||
weights: Array1::zeros(feature_dim),
|
||||
bias: 0.0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Transform input using random Fourier features
|
||||
pub fn transform(&self, x: &[f32]) -> Array1<f32> {
|
||||
let x_arr = Array1::from_vec(x.to_vec());
|
||||
let projections = self.omega.dot(&x_arr) + &self.b;
|
||||
|
||||
// Apply cosine to get Fourier features
|
||||
let scale = (2.0 / self.feature_dim as f32).sqrt();
|
||||
projections.mapv(|p| (p.cos() * scale))
|
||||
}
|
||||
|
||||
/// Batch transform
|
||||
pub fn transform_batch(&self, x: &[Vec<f32>]) -> Array2<f32> {
|
||||
let n_samples = x.len();
|
||||
let mut features = Array2::zeros((n_samples, self.feature_dim));
|
||||
|
||||
for (i, xi) in x.iter().enumerate() {
|
||||
let feat = self.transform(xi);
|
||||
features.row_mut(i).assign(&feat);
|
||||
}
|
||||
|
||||
features
|
||||
}
|
||||
|
||||
/// Train using closed-form ridge regression
|
||||
pub fn train(&mut self, x: &[Vec<f32>], y: &[f32], lambda: f32) {
|
||||
let features = self.transform_batch(x);
|
||||
let n = features.nrows();
|
||||
|
||||
// Closed-form solution: w = (X^T X + λI)^{-1} X^T y
|
||||
let xtx = features.t().dot(&features);
|
||||
let xty = features.t().dot(&Array1::from_vec(y.to_vec()));
|
||||
|
||||
// Add regularization
|
||||
let mut reg_xtx = xtx + Array2::<f32>::eye(self.feature_dim) * lambda * n as f32;
|
||||
|
||||
// Solve (simplified - production would use LAPACK)
|
||||
self.weights = self.solve_regularized(®_xtx, &xty);
|
||||
|
||||
// Compute bias
|
||||
let predictions = features.dot(&self.weights);
|
||||
let mean_pred = predictions.mean().unwrap();
|
||||
let mean_y = y.iter().sum::<f32>() / y.len() as f32;
|
||||
self.bias = mean_y - mean_pred;
|
||||
}
|
||||
|
||||
fn solve_regularized(&self, a: &Array2<f32>, b: &Array1<f32>) -> Array1<f32> {
|
||||
// Simplified solver using gradient descent
|
||||
let mut x = Array1::zeros(self.feature_dim);
|
||||
let lr = 0.01;
|
||||
|
||||
for _ in 0..100 {
|
||||
let grad = a.dot(&x) - b;
|
||||
x = x - &grad * lr;
|
||||
}
|
||||
|
||||
x
|
||||
}
|
||||
|
||||
pub fn predict(&self, x: &[Vec<f32>]) -> Vec<f32> {
|
||||
x.iter().map(|xi| {
|
||||
let features = self.transform(xi);
|
||||
self.weights.dot(&features) + self.bias
|
||||
}).collect()
|
||||
}
|
||||
|
||||
pub fn predict_class(&self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
self.predict(x).iter().map(|&y| {
|
||||
if y < -0.25 { 0 }
|
||||
else if y > 0.25 { 2 }
|
||||
else { 1 }
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Adaptive Fourier Features with frequency learning
|
||||
pub struct AdaptiveFourierFeatures {
|
||||
base: FourierFeatures,
|
||||
frequency_lr: f32,
|
||||
adapt_frequencies: bool,
|
||||
}
|
||||
|
||||
impl AdaptiveFourierFeatures {
|
||||
pub fn new(input_dim: usize, feature_dim: usize, sigma: f32) -> Self {
|
||||
Self {
|
||||
base: FourierFeatures::new(input_dim, feature_dim, sigma),
|
||||
frequency_lr: 0.001,
|
||||
adapt_frequencies: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Train with frequency adaptation
|
||||
pub fn train_adaptive(&mut self, x: &[Vec<f32>], y: &[f32], epochs: usize) {
|
||||
for epoch in 0..epochs {
|
||||
// Standard training
|
||||
self.base.train(x, y, 0.01);
|
||||
|
||||
if self.adapt_frequencies && epoch % 10 == 0 {
|
||||
// Adapt frequencies based on gradient
|
||||
self.adapt_frequencies_step(x, y);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn adapt_frequencies_step(&mut self, x: &[Vec<f32>], y: &[f32]) {
|
||||
// Compute gradient w.r.t frequencies
|
||||
let mut grad_omega = Array2::zeros((self.base.feature_dim, self.base.input_dim));
|
||||
|
||||
for (xi, &yi) in x.iter().zip(y.iter()) {
|
||||
let x_arr = Array1::from_vec(xi.clone());
|
||||
let projections = self.base.omega.dot(&x_arr) + &self.base.b;
|
||||
let features = projections.mapv(|p| p.cos());
|
||||
|
||||
let pred = self.base.weights.dot(&features) + self.base.bias;
|
||||
let error = pred - yi;
|
||||
|
||||
// Gradient through cosine
|
||||
for j in 0..self.base.feature_dim {
|
||||
let grad_cos = -projections[j].sin();
|
||||
let grad_j = error * self.base.weights[j] * grad_cos;
|
||||
|
||||
for k in 0..self.base.input_dim {
|
||||
grad_omega[[j, k]] += grad_j * xi[k];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Update frequencies
|
||||
self.base.omega = &self.base.omega - &grad_omega * self.frequency_lr;
|
||||
}
|
||||
|
||||
pub fn predict(&self, x: &[Vec<f32>]) -> Vec<f32> {
|
||||
self.base.predict(x)
|
||||
}
|
||||
|
||||
pub fn predict_class(&self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
self.base.predict_class(x)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
// lib.rs - Expose modules for library use
|
||||
|
||||
pub mod data;
|
||||
pub mod metrics;
|
||||
pub mod baseline;
|
||||
pub mod mlp;
|
||||
pub mod mlp_optimized;
|
||||
pub mod mlp_ultra;
|
||||
pub mod mlp_classifier;
|
||||
pub mod ensemble;
|
||||
pub mod attention;
|
||||
pub mod reservoir;
|
||||
pub mod fourier;
|
||||
pub mod sparse;
|
||||
pub mod mlp_avx512;
|
||||
pub mod quantization;
|
||||
pub mod mlp_quantized;
|
||||
|
||||
#[cfg(feature = "ruv-fann")]
|
||||
pub mod ruv_fann_impl;
|
||||
|
||||
#[cfg(not(feature = "ruv-fann"))]
|
||||
pub mod ruv_fann_adapter;
|
||||
@@ -0,0 +1,298 @@
|
||||
mod data;
|
||||
mod metrics;
|
||||
mod baseline;
|
||||
mod mlp;
|
||||
mod mlp_optimized;
|
||||
mod mlp_ultra;
|
||||
mod mlp_classifier;
|
||||
mod ensemble;
|
||||
mod attention;
|
||||
mod reservoir;
|
||||
mod fourier;
|
||||
mod sparse;
|
||||
mod mlp_avx512;
|
||||
mod quantization;
|
||||
mod mlp_quantized;
|
||||
mod ruv_fann_adapter;
|
||||
mod ruv_fann_impl;
|
||||
|
||||
use clap::{Parser, ValueEnum};
|
||||
use data::{make_synthetic, to_class};
|
||||
use metrics::{mse, acc};
|
||||
use baseline::Baseline;
|
||||
use mlp::Mlp;
|
||||
use mlp_optimized::OptimizedMlp;
|
||||
use mlp_ultra::UltraMlp;
|
||||
use mlp_classifier::ClassifierMlp;
|
||||
use ensemble::{EnsembleModel, BoostedEnsemble};
|
||||
use reservoir::{ReservoirComputer, QuantumReservoir};
|
||||
use fourier::{FourierFeatures, AdaptiveFourierFeatures};
|
||||
use sparse::{SparseNetwork, LotteryTicketNetwork};
|
||||
use mlp_avx512::DynamicAvx512Mlp;
|
||||
use mlp_quantized::QuantizedMlpBackend;
|
||||
#[cfg(feature = "ruv-fann")]
|
||||
use ruv_fann_impl::RuvFannModel;
|
||||
#[cfg(not(feature = "ruv-fann"))]
|
||||
use ruv_fann_adapter::ruv_fann_backend::RuvFannModel;
|
||||
|
||||
#[derive(Copy, Clone, ValueEnum)]
|
||||
enum Backend { Mlp, MlpOpt, MlpUltra, MlpAvx512, MlpQuantized, MlpClassifier, Ensemble, Boosted, Reservoir, QuantumReservoir, Fourier, AdaptiveFourier, Sparse, LotteryTicket, RuvFann, Baseline }
|
||||
|
||||
#[derive(Parser)]
|
||||
struct Args {
|
||||
#[arg(long, default_value_t=32)]
|
||||
window: usize,
|
||||
#[arg(long, default_value_t=5000)]
|
||||
n: usize,
|
||||
#[arg(long, default_value_t=42)]
|
||||
seed: u64,
|
||||
#[arg(long, value_enum, default_value_t=Backend::Mlp)]
|
||||
backend: Backend,
|
||||
#[arg(long, default_value_t=64)]
|
||||
hidden: usize,
|
||||
#[arg(long, default_value_t=8)]
|
||||
epochs: usize,
|
||||
#[arg(long, default_value_t=0.01)]
|
||||
lr: f32,
|
||||
#[arg(long, default_value_t=false)]
|
||||
classify: bool
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let args = Args::parse();
|
||||
let ds = make_synthetic(args.window, args.n, args.seed);
|
||||
|
||||
// prepare tensors
|
||||
let to_xy = |v: &Vec<data::Sample>| {
|
||||
let x: Vec<Vec<f32>> = v.iter().map(|s| s.x.clone()).collect();
|
||||
let y_reg: Vec<f32> = v.iter().map(|s| s.y).collect();
|
||||
let y_cls: Vec<usize> = v.iter().map(|s| to_class(s.y)).collect();
|
||||
(x, y_reg, y_cls)
|
||||
};
|
||||
let (xtr, ytr, ytrc) = to_xy(&ds.train);
|
||||
let (xva, yva, yvac) = to_xy(&ds.val);
|
||||
let (xte, yte, ytec) = to_xy(&ds.test);
|
||||
|
||||
match args.backend {
|
||||
Backend::Baseline => {
|
||||
let yhat = Baseline::predict_reg(&ds.test, args.window);
|
||||
let yhatc = Baseline::predict_cls(&ds.test, args.window);
|
||||
println!("baseline_mse_test={:.6}", mse(&yte, &yhat));
|
||||
println!("baseline_acc_test={:.4}", acc(&ytec, &yhatc));
|
||||
}
|
||||
Backend::Mlp => {
|
||||
let mut model = if args.classify { Mlp::new(xtr[0].len(), args.hidden, 3) }
|
||||
else { Mlp::new(xtr[0].len(), args.hidden, 1) };
|
||||
if args.classify {
|
||||
// train via regression to continuous y, then map to buckets
|
||||
model.train_regression(&xtr, &ytr, args.epochs, args.lr);
|
||||
let yhat = model.predict_cls3(&xte);
|
||||
println!("mlp_acc_test={:.4}", acc(&ytec, &yhat));
|
||||
} else {
|
||||
model.train_regression(&xtr, &ytr, args.epochs, args.lr);
|
||||
let yhat = model.predict_reg(&xte);
|
||||
println!("mlp_mse_val={:.6}", mse(&yva, &model.predict_reg(&xva)));
|
||||
println!("mlp_mse_test={:.6}", mse(&yte, &yhat));
|
||||
}
|
||||
}
|
||||
Backend::MlpOpt => {
|
||||
let mut model = if args.classify { OptimizedMlp::new(xtr[0].len(), args.hidden, 3) }
|
||||
else { OptimizedMlp::new(xtr[0].len(), args.hidden, 1) };
|
||||
if args.classify {
|
||||
model.train_batch(&xtr, &ytr, args.epochs, args.lr, 32);
|
||||
let yhat = model.predict_cls3(&xte);
|
||||
println!("mlp_opt_acc_test={:.4}", acc(&ytec, &yhat));
|
||||
} else {
|
||||
model.train_batch(&xtr, &ytr, args.epochs, args.lr, 32);
|
||||
let yhat = model.predict_reg(&xte);
|
||||
println!("mlp_opt_mse_val={:.6}", mse(&yva, &model.predict_reg(&xva)));
|
||||
println!("mlp_opt_mse_test={:.6}", mse(&yte, &yhat));
|
||||
}
|
||||
}
|
||||
Backend::MlpUltra => {
|
||||
let mut model = if args.classify { UltraMlp::new(xtr[0].len(), args.hidden, 3) }
|
||||
else { UltraMlp::new(xtr[0].len(), args.hidden, 1) };
|
||||
if args.classify {
|
||||
model.train_batch_parallel(&xtr, &ytr, args.epochs, args.lr, 32);
|
||||
let yhat = model.predict_cls3_parallel(&xte);
|
||||
println!("mlp_ultra_acc_test={:.4}", acc(&ytec, &yhat));
|
||||
} else {
|
||||
model.train_batch_parallel(&xtr, &ytr, args.epochs, args.lr, 32);
|
||||
let yhat = model.predict_parallel(&xte);
|
||||
println!("mlp_ultra_mse_val={:.6}", mse(&yva, &model.predict_parallel(&xva)));
|
||||
println!("mlp_ultra_mse_test={:.6}", mse(&yte, &yhat));
|
||||
}
|
||||
}
|
||||
Backend::MlpAvx512 => {
|
||||
let model = DynamicAvx512Mlp::new(xtr[0].len(), args.hidden, if args.classify { 3 } else { 1 });
|
||||
if args.classify {
|
||||
// Note: AVX512 model doesn't have train method yet, using predict only
|
||||
let yhat = model.predict_class(&xte);
|
||||
println!("mlp_avx512_acc_test={:.4} (inference only)", acc(&ytec, &yhat));
|
||||
} else {
|
||||
let yhat = model.predict(&xte);
|
||||
println!("mlp_avx512_mse_test={:.6} (inference only)", mse(&yte, &yhat));
|
||||
}
|
||||
}
|
||||
Backend::MlpQuantized => {
|
||||
let mut model = QuantizedMlpBackend::new(xtr[0].len(), args.hidden, if args.classify { 3 } else { 1 });
|
||||
|
||||
// Train in FP32
|
||||
model.train(&xtr, &ytr, args.epochs, args.lr);
|
||||
|
||||
// Benchmark INT8 vs FP32
|
||||
println!("\n=== INT8 Quantization Performance ===");
|
||||
model.benchmark_inference(&xte[..100.min(xte.len())], 100);
|
||||
|
||||
if args.classify {
|
||||
let yhat = model.predict_class(&xte);
|
||||
println!("\nmlp_quantized_acc_test={:.4}", acc(&ytec, &yhat));
|
||||
} else {
|
||||
// Compare FP32 vs INT8 accuracy
|
||||
let yhat_fp32 = model.predict_fp32(&xte);
|
||||
let yhat_int8 = model.predict(&xte);
|
||||
|
||||
println!("\nmlp_quantized_mse_test_fp32={:.6}", mse(&yte, &yhat_fp32));
|
||||
println!("mlp_quantized_mse_test_int8={:.6}", mse(&yte, &yhat_int8));
|
||||
|
||||
// Calculate accuracy loss
|
||||
let mse_fp32 = mse(&yte, &yhat_fp32);
|
||||
let mse_int8 = mse(&yte, &yhat_int8);
|
||||
let accuracy_loss = ((mse_int8 - mse_fp32).abs() / mse_fp32) * 100.0;
|
||||
println!("Quantization accuracy loss: {:.2}%", accuracy_loss);
|
||||
}
|
||||
|
||||
let (orig, quant, ratio) = model.get_compression_stats();
|
||||
println!("\nModel compression: {} -> {} bytes ({:.2}x)", orig, quant, ratio);
|
||||
}
|
||||
Backend::MlpClassifier => {
|
||||
let mut model = ClassifierMlp::new(xtr[0].len(), 3);
|
||||
if args.classify {
|
||||
model.train_classification(&xtr, &ytr, args.epochs, 32);
|
||||
let yhat = model.predict_cls3(&xte);
|
||||
let yhat_val = model.predict_cls3(&xva);
|
||||
println!("mlp_classifier_acc_val={:.4}", acc(&yvac, &yhat_val));
|
||||
println!("mlp_classifier_acc_test={:.4}", acc(&ytec, &yhat));
|
||||
} else {
|
||||
println!("MlpClassifier is for classification only, use --classify flag");
|
||||
}
|
||||
}
|
||||
Backend::Ensemble => {
|
||||
if args.classify {
|
||||
let mut ensemble = EnsembleModel::new();
|
||||
let input_dim = xtr[0].len();
|
||||
|
||||
// Add diverse models
|
||||
ensemble.add_model_simple(input_dim, args.hidden, 3);
|
||||
ensemble.add_model_optimized(input_dim, args.hidden * 2, 3);
|
||||
ensemble.add_model_ultra(input_dim, args.hidden, 3);
|
||||
ensemble.add_model_classifier(input_dim, 3);
|
||||
|
||||
ensemble.train_ensemble(&xtr, &ytr, args.epochs, args.lr, &xva, &yvac);
|
||||
let yhat = ensemble.predict_ensemble(&xte);
|
||||
println!("ensemble_acc_test={:.4}", acc(&ytec, &yhat));
|
||||
} else {
|
||||
println!("Ensemble is for classification only, use --classify flag");
|
||||
}
|
||||
}
|
||||
Backend::Boosted => {
|
||||
if args.classify {
|
||||
let mut boosted = BoostedEnsemble::new(xtr[0].len(), args.hidden);
|
||||
boosted.train_boosted(&xtr, &ytrc, 10, args.epochs / 2);
|
||||
let yhat = boosted.predict_boosted(&xte);
|
||||
println!("boosted_acc_test={:.4}", acc(&ytec, &yhat));
|
||||
} else {
|
||||
println!("Boosted is for classification only, use --classify flag");
|
||||
}
|
||||
}
|
||||
Backend::Reservoir => {
|
||||
let mut model = ReservoirComputer::new(xtr[0].len(), 100, 1);
|
||||
if args.classify {
|
||||
model.train_ridge(&xtr, &ytr, 0.001);
|
||||
let yhat = model.predict_class(&xte);
|
||||
println!("reservoir_acc_test={:.4}", acc(&ytec, &yhat));
|
||||
} else {
|
||||
model.train_ridge(&xtr, &ytr, 0.001);
|
||||
let yhat = model.predict(&xte);
|
||||
println!("reservoir_mse_test={:.6}", mse(&yte, &yhat));
|
||||
}
|
||||
}
|
||||
Backend::QuantumReservoir => {
|
||||
let mut model = QuantumReservoir::new(xtr[0].len(), 100, 1);
|
||||
if args.classify {
|
||||
model.train(&xtr, &ytr);
|
||||
let yhat = model.predict_quantum(&xte);
|
||||
println!("quantum_reservoir_acc_test={:.4}", acc(&ytec, &yhat));
|
||||
} else {
|
||||
model.train(&xtr, &ytr);
|
||||
// For regression, use classical prediction
|
||||
let yhat = model.classical_reservoir.predict(&xte);
|
||||
println!("quantum_reservoir_mse_test={:.6}", mse(&yte, &yhat));
|
||||
}
|
||||
}
|
||||
Backend::Fourier => {
|
||||
let mut model = FourierFeatures::new(xtr[0].len(), 500, 1.0);
|
||||
if args.classify {
|
||||
model.train(&xtr, &ytr, 0.001);
|
||||
let yhat = model.predict_class(&xte);
|
||||
println!("fourier_acc_test={:.4}", acc(&ytec, &yhat));
|
||||
} else {
|
||||
model.train(&xtr, &ytr, 0.001);
|
||||
let yhat = model.predict(&xte);
|
||||
println!("fourier_mse_test={:.6}", mse(&yte, &yhat));
|
||||
}
|
||||
}
|
||||
Backend::AdaptiveFourier => {
|
||||
let mut model = AdaptiveFourierFeatures::new(xtr[0].len(), 500, 1.0);
|
||||
if args.classify {
|
||||
model.train_adaptive(&xtr, &ytr, 50);
|
||||
let yhat = model.predict_class(&xte);
|
||||
println!("adaptive_fourier_acc_test={:.4}", acc(&ytec, &yhat));
|
||||
} else {
|
||||
model.train_adaptive(&xtr, &ytr, 50);
|
||||
let yhat = model.predict(&xte);
|
||||
println!("adaptive_fourier_mse_test={:.6}", mse(&yte, &yhat));
|
||||
}
|
||||
}
|
||||
Backend::Sparse => {
|
||||
let mut model = SparseNetwork::new(xtr[0].len(), args.hidden, if args.classify { 3 } else { 1 }, 0.1);
|
||||
if args.classify {
|
||||
model.train(&xtr, &ytr, args.epochs * 10, 0.01);
|
||||
let yhat = model.predict_class(&xte);
|
||||
let (active, pruned, sparsity) = model.get_sparsity_stats();
|
||||
println!("sparse_acc_test={:.4} (active={}, pruned={}, sparsity={:.1}%)",
|
||||
acc(&ytec, &yhat), active, pruned, sparsity * 100.0);
|
||||
} else {
|
||||
model.train(&xtr, &ytr, args.epochs * 10, 0.01);
|
||||
let yhat = model.predict(&xte);
|
||||
let (active, pruned, sparsity) = model.get_sparsity_stats();
|
||||
println!("sparse_mse_test={:.6} (active={}, pruned={}, sparsity={:.1}%)",
|
||||
mse(&yte, &yhat), active, pruned, sparsity * 100.0);
|
||||
}
|
||||
}
|
||||
Backend::LotteryTicket => {
|
||||
let mut model = LotteryTicketNetwork::new(xtr[0].len(), args.hidden, if args.classify { 3 } else { 1 });
|
||||
if args.classify {
|
||||
model.find_winning_ticket(&xtr, &ytr, 0.2, 5);
|
||||
let yhat = model.predict_class(&xte);
|
||||
println!("lottery_ticket_acc_test={:.4}", acc(&ytec, &yhat));
|
||||
} else {
|
||||
model.find_winning_ticket(&xtr, &ytr, 0.2, 5);
|
||||
let yhat = model.predict(&xte);
|
||||
println!("lottery_ticket_mse_test={:.6}", mse(&yte, &yhat));
|
||||
}
|
||||
}
|
||||
Backend::RuvFann => {
|
||||
let mut model = if args.classify { RuvFannModel::new(xtr[0].len(), args.hidden, 3) }
|
||||
else { RuvFannModel::new(xtr[0].len(), args.hidden, 1) };
|
||||
model.train_regression(&xtr, &ytr, args.epochs, args.lr);
|
||||
if args.classify {
|
||||
let yhat = model.predict_cls3(&xte);
|
||||
println!("ruv_fann_acc_test={:.4}", acc(&ytec, &yhat));
|
||||
} else {
|
||||
let yhat = model.predict_reg(&xte);
|
||||
println!("ruv_fann_mse_test={:.6}", mse(&yte, &yhat));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
pub fn mse(y: &[f32], yhat: &[f32]) -> f32 {
|
||||
y.iter().zip(yhat).map(|(a,b)| (a-b)*(a-b)).sum::<f32>() / y.len() as f32
|
||||
}
|
||||
|
||||
pub fn acc(y: &[usize], yhat: &[usize]) -> f32 {
|
||||
let c = y.iter().zip(yhat).filter(|(a,b)| a==b).count();
|
||||
c as f32 / y.len() as f32
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
use ndarray::{Array2, Array1};
|
||||
use rand::{thread_rng, Rng};
|
||||
|
||||
pub struct Mlp {
|
||||
w1: Array2<f32>,
|
||||
b1: Array1<f32>,
|
||||
w2: Array2<f32>,
|
||||
b2: Array1<f32>,
|
||||
output_dim: usize,
|
||||
}
|
||||
|
||||
fn relu(x: &mut Array1<f32>) {
|
||||
x.iter_mut().for_each(|v| if *v < 0.0 { *v = 0.0 });
|
||||
}
|
||||
|
||||
impl Mlp {
|
||||
pub fn new(input: usize, hidden: usize, output: usize) -> Self {
|
||||
let mut rng = thread_rng();
|
||||
let scale = (2.0 / input as f32).sqrt();
|
||||
|
||||
let w1 = Array2::from_shape_fn((hidden, input), |_| rng.gen::<f32>() * scale - scale/2.0);
|
||||
let b1 = Array1::zeros(hidden);
|
||||
let w2 = Array2::from_shape_fn((output, hidden), |_| rng.gen::<f32>() * scale - scale/2.0);
|
||||
let b2 = Array1::zeros(output);
|
||||
|
||||
Self { w1, b1, w2, b2, output_dim: output }
|
||||
}
|
||||
|
||||
fn forward(&self, x: &Array1<f32>) -> Array1<f32> {
|
||||
let mut h = self.w1.dot(x) + &self.b1;
|
||||
relu(&mut h);
|
||||
self.w2.dot(&h) + &self.b2
|
||||
}
|
||||
|
||||
pub fn train_regression(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>, epochs: usize, lr: f32) {
|
||||
// Simplified SGD without full backprop for now
|
||||
for _ in 0..epochs {
|
||||
for (xi, yi) in x.iter().zip(y) {
|
||||
let x_arr = Array1::from_vec(xi.clone());
|
||||
let output = self.forward(&x_arr);
|
||||
|
||||
// Only train if we have 1D output for regression
|
||||
if self.output_dim == 1 {
|
||||
let err = output[0] - yi;
|
||||
|
||||
// Numerical gradient approximation for simplicity
|
||||
let _eps = 0.0001;
|
||||
|
||||
// Update w2 and b2
|
||||
let mut h = self.w1.dot(&x_arr) + &self.b1;
|
||||
relu(&mut h);
|
||||
|
||||
for i in 0..self.w2.nrows() {
|
||||
for j in 0..self.w2.ncols() {
|
||||
self.w2[[i, j]] -= lr * err * h[j] * 0.1;
|
||||
}
|
||||
self.b2[i] -= lr * err * 0.1;
|
||||
}
|
||||
|
||||
// Update w1 and b1 with smaller learning rate
|
||||
for i in 0..self.w1.nrows() {
|
||||
for j in 0..self.w1.ncols() {
|
||||
self.w1[[i, j]] -= lr * err * x_arr[j] * 0.01;
|
||||
}
|
||||
self.b1[i] -= lr * err * 0.01;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn predict_reg(&self, x: &[Vec<f32>]) -> Vec<f32> {
|
||||
x.iter().map(|xi| {
|
||||
let out = self.forward(&Array1::from_vec(xi.clone()));
|
||||
if self.output_dim == 1 {
|
||||
out[0]
|
||||
} else {
|
||||
// For multi-output, return first element as regression
|
||||
out[0]
|
||||
}
|
||||
}).collect()
|
||||
}
|
||||
|
||||
pub fn predict_cls3(&self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
x.iter().map(|xi| {
|
||||
let out = self.forward(&Array1::from_vec(xi.clone()));
|
||||
|
||||
if self.output_dim >= 3 {
|
||||
// Find argmax for 3-class classification
|
||||
let mut best = 0;
|
||||
let mut best_val = out[0];
|
||||
for i in 1..3 {
|
||||
if out[i] > best_val {
|
||||
best_val = out[i];
|
||||
best = i;
|
||||
}
|
||||
}
|
||||
best
|
||||
} else {
|
||||
// Fall back to threshold-based classification
|
||||
let val = out[0];
|
||||
if val < -0.25 { 0 }
|
||||
else if val > 0.25 { 2 }
|
||||
else { 1 }
|
||||
}
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,341 @@
|
||||
use std::arch::x86_64::*;
|
||||
use std::alloc::{alloc, dealloc, Layout};
|
||||
use std::ptr;
|
||||
|
||||
/// Ultra-optimized CPU-only MLP with AVX-512 for lowest possible latency
|
||||
/// Key optimizations:
|
||||
/// - AVX-512: Process 16 floats at once (vs 8 with AVX2)
|
||||
/// - Cache-aligned memory: Prevent false sharing
|
||||
/// - Prefetching: Hide memory latency
|
||||
/// - Loop unrolling: Reduce branch overhead
|
||||
/// - Compile-time dimensions: Enable better optimization
|
||||
pub struct UltraLowLatencyMlp<const INPUT: usize, const HIDDEN: usize, const OUTPUT: usize> {
|
||||
// Cache-aligned weight storage (64-byte aligned for AVX-512)
|
||||
w1: *mut f32, // HIDDEN x INPUT (row-major for sequential access)
|
||||
b1: *mut f32, // HIDDEN
|
||||
w2: *mut f32, // OUTPUT x HIDDEN
|
||||
b2: *mut f32, // OUTPUT
|
||||
|
||||
// Pre-allocated buffers for zero-copy operation
|
||||
hidden_buf: *mut f32, // HIDDEN (aligned)
|
||||
}
|
||||
|
||||
impl<const I: usize, const H: usize, const O: usize> UltraLowLatencyMlp<I, H, O> {
|
||||
pub fn new() -> Self {
|
||||
unsafe {
|
||||
// Allocate cache-aligned memory (64 bytes for AVX-512)
|
||||
let align = 64;
|
||||
|
||||
let w1 = Self::alloc_aligned(H * I, align);
|
||||
let b1 = Self::alloc_aligned(H, align);
|
||||
let w2 = Self::alloc_aligned(O * H, align);
|
||||
let b2 = Self::alloc_aligned(O, align);
|
||||
let hidden_buf = Self::alloc_aligned(H, align);
|
||||
|
||||
// Initialize with small random weights
|
||||
Self::init_weights(w1, H * I);
|
||||
Self::init_weights(w2, O * H);
|
||||
ptr::write_bytes(b1, 0, H);
|
||||
ptr::write_bytes(b2, 0, O);
|
||||
|
||||
Self { w1, b1, w2, b2, hidden_buf }
|
||||
}
|
||||
}
|
||||
|
||||
#[inline(always)]
|
||||
unsafe fn alloc_aligned(size: usize, align: usize) -> *mut f32 {
|
||||
let layout = Layout::from_size_align(size * 4, align).unwrap();
|
||||
alloc(layout) as *mut f32
|
||||
}
|
||||
|
||||
unsafe fn init_weights(ptr: *mut f32, size: usize) {
|
||||
let scale = (2.0 / size as f32).sqrt();
|
||||
for i in 0..size {
|
||||
*ptr.add(i) = (rand::random::<f32>() - 0.5) * scale;
|
||||
}
|
||||
}
|
||||
|
||||
/// Ultra-fast forward pass with AVX-512
|
||||
/// Latency target: <100ns for typical sizes
|
||||
#[target_feature(enable = "avx512f")]
|
||||
#[inline]
|
||||
pub unsafe fn forward_avx512(&self, input: &[f32; I], output: &mut [f32; O]) {
|
||||
// Layer 1: Input -> Hidden with AVX-512 (16 floats at once)
|
||||
self.matmul_avx512(input.as_ptr(), self.w1, self.b1, self.hidden_buf, H, I);
|
||||
|
||||
// ReLU activation (vectorized)
|
||||
self.relu_avx512(self.hidden_buf, H);
|
||||
|
||||
// Layer 2: Hidden -> Output
|
||||
self.matmul_avx512(self.hidden_buf, self.w2, self.b2, output.as_mut_ptr(), O, H);
|
||||
}
|
||||
|
||||
/// AVX-512 matrix multiply with prefetching
|
||||
#[target_feature(enable = "avx512f")]
|
||||
unsafe fn matmul_avx512(&self, x: *const f32, w: *const f32, b: *const f32,
|
||||
out: *mut f32, rows: usize, cols: usize) {
|
||||
const SIMD_WIDTH: usize = 16; // AVX-512 processes 16 floats
|
||||
|
||||
// Process each output neuron
|
||||
for i in 0..rows {
|
||||
let row_offset = i * cols;
|
||||
|
||||
// Prefetch next cache line
|
||||
_mm_prefetch(w.add(row_offset + 64) as *const i8, _MM_HINT_T0);
|
||||
|
||||
// Initialize accumulator with bias
|
||||
let mut sum = _mm512_set1_ps(*b.add(i));
|
||||
|
||||
// Process 16 elements at a time
|
||||
let chunks = cols / SIMD_WIDTH;
|
||||
let mut j = 0;
|
||||
|
||||
// Unroll by 4 for better pipelining
|
||||
while j + 3 < chunks {
|
||||
let w0 = _mm512_load_ps(w.add(row_offset + j * SIMD_WIDTH));
|
||||
let x0 = _mm512_load_ps(x.add(j * SIMD_WIDTH));
|
||||
sum = _mm512_fmadd_ps(w0, x0, sum);
|
||||
|
||||
let w1 = _mm512_load_ps(w.add(row_offset + (j + 1) * SIMD_WIDTH));
|
||||
let x1 = _mm512_load_ps(x.add((j + 1) * SIMD_WIDTH));
|
||||
sum = _mm512_fmadd_ps(w1, x1, sum);
|
||||
|
||||
let w2 = _mm512_load_ps(w.add(row_offset + (j + 2) * SIMD_WIDTH));
|
||||
let x2 = _mm512_load_ps(x.add((j + 2) * SIMD_WIDTH));
|
||||
sum = _mm512_fmadd_ps(w2, x2, sum);
|
||||
|
||||
let w3 = _mm512_load_ps(w.add(row_offset + (j + 3) * SIMD_WIDTH));
|
||||
let x3 = _mm512_load_ps(x.add((j + 3) * SIMD_WIDTH));
|
||||
sum = _mm512_fmadd_ps(w3, x3, sum);
|
||||
|
||||
j += 4;
|
||||
}
|
||||
|
||||
// Handle remaining chunks
|
||||
while j < chunks {
|
||||
let wv = _mm512_load_ps(w.add(row_offset + j * SIMD_WIDTH));
|
||||
let xv = _mm512_load_ps(x.add(j * SIMD_WIDTH));
|
||||
sum = _mm512_fmadd_ps(wv, xv, sum);
|
||||
j += 1;
|
||||
}
|
||||
|
||||
// Horizontal sum (reduce 16 -> 1)
|
||||
let result = _mm512_reduce_add_ps(sum);
|
||||
|
||||
// Handle remaining elements (non-SIMD)
|
||||
let mut scalar_sum = result;
|
||||
for k in (chunks * SIMD_WIDTH)..cols {
|
||||
scalar_sum += *w.add(row_offset + k) * *x.add(k);
|
||||
}
|
||||
|
||||
*out.add(i) = scalar_sum;
|
||||
}
|
||||
}
|
||||
|
||||
/// Vectorized ReLU with AVX-512
|
||||
#[target_feature(enable = "avx512f")]
|
||||
unsafe fn relu_avx512(&self, data: *mut f32, size: usize) {
|
||||
let zero = _mm512_setzero_ps();
|
||||
let chunks = size / 16;
|
||||
|
||||
for i in 0..chunks {
|
||||
let val = _mm512_load_ps(data.add(i * 16));
|
||||
let relu = _mm512_max_ps(val, zero);
|
||||
_mm512_store_ps(data.add(i * 16), relu);
|
||||
}
|
||||
|
||||
// Handle remainder
|
||||
for i in (chunks * 16)..size {
|
||||
let val = *data.add(i);
|
||||
*data.add(i) = val.max(0.0);
|
||||
}
|
||||
}
|
||||
|
||||
/// Fallback for non-AVX512 systems (still optimized)
|
||||
pub fn forward_fallback(&self, input: &[f32; I], output: &mut [f32; O]) {
|
||||
unsafe {
|
||||
// Use AVX2 if available, otherwise scalar
|
||||
#[cfg(target_feature = "avx2")]
|
||||
{
|
||||
self.forward_avx512(input, output);
|
||||
}
|
||||
#[cfg(not(target_feature = "avx2"))]
|
||||
{
|
||||
// Scalar fallback implementation
|
||||
unsafe {
|
||||
for i in 0..H {
|
||||
let mut sum = *self.b1.add(i);
|
||||
for j in 0..I {
|
||||
sum += input[j] * *self.w1.add(i * I + j);
|
||||
}
|
||||
*self.hidden_buf.add(i) = sum.max(0.0);
|
||||
}
|
||||
for i in 0..O {
|
||||
let mut sum = *self.b2.add(i);
|
||||
for j in 0..H {
|
||||
sum += *self.hidden_buf.add(j) * *self.w2.add(i * H + j);
|
||||
}
|
||||
output[i] = sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Train with low-latency SGD (no allocations in hot path)
|
||||
pub fn train_fast(&mut self, x: &[f32; I], y: f32, lr: f32) {
|
||||
unsafe {
|
||||
let mut output = [0.0; O];
|
||||
|
||||
// Forward pass
|
||||
self.forward_avx512(x, &mut output);
|
||||
|
||||
// Compute error
|
||||
let error = output[0] - y;
|
||||
|
||||
// Backward pass (simplified, no allocation)
|
||||
// Update output weights
|
||||
for i in 0..H {
|
||||
let grad = error * (*self.hidden_buf.add(i));
|
||||
*self.w2.add(i) -= lr * grad;
|
||||
}
|
||||
*self.b2 -= lr * error;
|
||||
|
||||
// Backprop to hidden
|
||||
for i in 0..H {
|
||||
if *self.hidden_buf.add(i) > 0.0 { // ReLU gradient
|
||||
let hidden_error = error * (*self.w2.add(i));
|
||||
|
||||
// Update hidden weights
|
||||
for j in 0..I {
|
||||
*self.w1.add(i * I + j) -= lr * hidden_error * x[j];
|
||||
}
|
||||
*self.b1.add(i) -= lr * hidden_error;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Batch prediction with minimal latency
|
||||
#[inline]
|
||||
pub fn predict_batch(&self, inputs: &[[f32; I]], outputs: &mut [[f32; O]]) {
|
||||
// Sequential processing for thread safety
|
||||
unsafe {
|
||||
for (x, y) in inputs.iter().zip(outputs.iter_mut()) {
|
||||
self.forward_avx512(x, y);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<const I: usize, const H: usize, const O: usize> Drop for UltraLowLatencyMlp<I, H, O> {
|
||||
fn drop(&mut self) {
|
||||
unsafe {
|
||||
let align = 64;
|
||||
dealloc(self.w1 as *mut u8, Layout::from_size_align(H * I * 4, align).unwrap());
|
||||
dealloc(self.b1 as *mut u8, Layout::from_size_align(H * 4, align).unwrap());
|
||||
dealloc(self.w2 as *mut u8, Layout::from_size_align(O * H * 4, align).unwrap());
|
||||
dealloc(self.b2 as *mut u8, Layout::from_size_align(O * 4, align).unwrap());
|
||||
dealloc(self.hidden_buf as *mut u8, Layout::from_size_align(H * 4, align).unwrap());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper for dynamic dimensions
|
||||
pub struct DynamicAvx512Mlp {
|
||||
weights_flat: Vec<f32>,
|
||||
dims: (usize, usize, usize),
|
||||
}
|
||||
|
||||
impl DynamicAvx512Mlp {
|
||||
pub fn new(input: usize, hidden: usize, output: usize) -> Self {
|
||||
let total_params = (input * hidden) + hidden + (hidden * output) + output;
|
||||
let mut weights_flat = Vec::with_capacity(total_params);
|
||||
|
||||
// Initialize
|
||||
let scale = (2.0 / input as f32).sqrt();
|
||||
for _ in 0..total_params {
|
||||
weights_flat.push((rand::random::<f32>() - 0.5) * scale);
|
||||
}
|
||||
|
||||
Self { weights_flat, dims: (input, hidden, output) }
|
||||
}
|
||||
|
||||
/// Predict with dynamic dimensions (still uses SIMD where possible)
|
||||
pub fn predict(&self, x: &[Vec<f32>]) -> Vec<f32> {
|
||||
let (input_dim, hidden_dim, _) = self.dims;
|
||||
|
||||
x.iter().map(|xi| {
|
||||
let mut hidden = vec![0.0f32; hidden_dim];
|
||||
|
||||
// Layer 1: Use AVX2 if available
|
||||
#[cfg(target_feature = "avx2")]
|
||||
unsafe {
|
||||
self.matmul_avx2_dynamic(&xi, &self.weights_flat[0..input_dim * hidden_dim],
|
||||
&mut hidden, hidden_dim, input_dim);
|
||||
}
|
||||
|
||||
#[cfg(not(target_feature = "avx2"))]
|
||||
{
|
||||
// Scalar fallback
|
||||
for i in 0..hidden_dim {
|
||||
let mut sum = self.weights_flat[input_dim * hidden_dim + i]; // bias
|
||||
for j in 0..input_dim {
|
||||
sum += xi[j] * self.weights_flat[i * input_dim + j];
|
||||
}
|
||||
hidden[i] = sum.max(0.0);
|
||||
}
|
||||
}
|
||||
|
||||
// Layer 2 (simplified for single output)
|
||||
let w2_start = input_dim * hidden_dim + hidden_dim;
|
||||
let mut output = self.weights_flat[w2_start + hidden_dim]; // bias
|
||||
|
||||
for i in 0..hidden_dim {
|
||||
output += hidden[i] * self.weights_flat[w2_start + i];
|
||||
}
|
||||
|
||||
output
|
||||
}).collect()
|
||||
}
|
||||
|
||||
#[cfg(target_feature = "avx2")]
|
||||
#[target_feature(enable = "avx2")]
|
||||
unsafe fn matmul_avx2_dynamic(&self, x: &[f32], w: &[f32], out: &mut [f32],
|
||||
rows: usize, cols: usize) {
|
||||
const SIMD_WIDTH: usize = 8;
|
||||
|
||||
for i in 0..rows {
|
||||
let row_offset = i * cols;
|
||||
let mut sum = _mm256_setzero_ps();
|
||||
|
||||
// Process 8 elements at a time
|
||||
let chunks = cols / SIMD_WIDTH;
|
||||
for j in 0..chunks {
|
||||
let idx = row_offset + j * SIMD_WIDTH;
|
||||
let wv = _mm256_loadu_ps(&w[idx]);
|
||||
let xv = _mm256_loadu_ps(&x[j * SIMD_WIDTH]);
|
||||
sum = _mm256_fmadd_ps(wv, xv, sum);
|
||||
}
|
||||
|
||||
// Horizontal sum
|
||||
let sum_array: [f32; 8] = std::mem::transmute(sum);
|
||||
let mut result: f32 = sum_array.iter().sum();
|
||||
|
||||
// Handle remainder
|
||||
for j in (chunks * SIMD_WIDTH)..cols {
|
||||
result += w[row_offset + j] * x[j];
|
||||
}
|
||||
|
||||
out[i] = result.max(0.0); // ReLU
|
||||
}
|
||||
}
|
||||
|
||||
pub fn predict_class(&self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
self.predict(x).iter().map(|&y| {
|
||||
if y < -0.25 { 0 }
|
||||
else if y > 0.25 { 2 }
|
||||
else { 1 }
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,325 @@
|
||||
use ndarray::{Array2, Array1, Axis};
|
||||
use rand::{thread_rng, Rng};
|
||||
|
||||
/// Specialized classifier with proper softmax, dropout, and batch norm
|
||||
pub struct ClassifierMlp {
|
||||
// Weights
|
||||
w1: Array2<f32>,
|
||||
b1: Array1<f32>,
|
||||
w2: Array2<f32>,
|
||||
b2: Array1<f32>,
|
||||
w3: Array2<f32>, // Extra layer for better classification
|
||||
b3: Array1<f32>,
|
||||
|
||||
// Batch normalization parameters
|
||||
bn1_gamma: Array1<f32>,
|
||||
bn1_beta: Array1<f32>,
|
||||
bn1_running_mean: Array1<f32>,
|
||||
bn1_running_var: Array1<f32>,
|
||||
|
||||
bn2_gamma: Array1<f32>,
|
||||
bn2_beta: Array1<f32>,
|
||||
bn2_running_mean: Array1<f32>,
|
||||
bn2_running_var: Array1<f32>,
|
||||
|
||||
// Adam optimizer state
|
||||
mw1: Array2<f32>, mb1: Array1<f32>,
|
||||
mw2: Array2<f32>, mb2: Array1<f32>,
|
||||
mw3: Array2<f32>, mb3: Array1<f32>,
|
||||
sw1: Array2<f32>, sb1: Array1<f32>,
|
||||
sw2: Array2<f32>, sb2: Array1<f32>,
|
||||
sw3: Array2<f32>, sb3: Array1<f32>,
|
||||
t: f32,
|
||||
|
||||
// Architecture
|
||||
input_dim: usize,
|
||||
hidden1_dim: usize,
|
||||
hidden2_dim: usize,
|
||||
output_dim: usize,
|
||||
|
||||
// Training settings
|
||||
dropout_rate: f32,
|
||||
is_training: bool,
|
||||
lr_schedule: LRSchedule,
|
||||
}
|
||||
|
||||
pub enum LRSchedule {
|
||||
Constant(f32),
|
||||
Cosine { initial: f32, min: f32, period: usize },
|
||||
StepDecay { initial: f32, decay: f32, step_size: usize },
|
||||
}
|
||||
|
||||
impl ClassifierMlp {
|
||||
pub fn new(input: usize, output: usize) -> Self {
|
||||
let mut rng = thread_rng();
|
||||
|
||||
// Deeper network for better classification
|
||||
let hidden1 = 128;
|
||||
let hidden2 = 64;
|
||||
|
||||
// Xavier initialization
|
||||
let scale1 = (2.0 / input as f32).sqrt();
|
||||
let scale2 = (2.0 / hidden1 as f32).sqrt();
|
||||
let scale3 = (2.0 / hidden2 as f32).sqrt();
|
||||
|
||||
let w1 = Array2::from_shape_fn((hidden1, input), |_|
|
||||
rng.gen::<f32>() * scale1 - scale1/2.0);
|
||||
let b1 = Array1::zeros(hidden1);
|
||||
|
||||
let w2 = Array2::from_shape_fn((hidden2, hidden1), |_|
|
||||
rng.gen::<f32>() * scale2 - scale2/2.0);
|
||||
let b2 = Array1::zeros(hidden2);
|
||||
|
||||
let w3 = Array2::from_shape_fn((output, hidden2), |_|
|
||||
rng.gen::<f32>() * scale3 - scale3/2.0);
|
||||
let b3 = Array1::zeros(output);
|
||||
|
||||
// Batch norm parameters
|
||||
let bn1_gamma = Array1::ones(hidden1);
|
||||
let bn1_beta = Array1::zeros(hidden1);
|
||||
let bn1_running_mean = Array1::zeros(hidden1);
|
||||
let bn1_running_var = Array1::ones(hidden1);
|
||||
|
||||
let bn2_gamma = Array1::ones(hidden2);
|
||||
let bn2_beta = Array1::zeros(hidden2);
|
||||
let bn2_running_mean = Array1::zeros(hidden2);
|
||||
let bn2_running_var = Array1::ones(hidden2);
|
||||
|
||||
// Adam states
|
||||
let mw1 = Array2::zeros((hidden1, input));
|
||||
let mb1 = Array1::zeros(hidden1);
|
||||
let mw2 = Array2::zeros((hidden2, hidden1));
|
||||
let mb2 = Array1::zeros(hidden2);
|
||||
let mw3 = Array2::zeros((output, hidden2));
|
||||
let mb3 = Array1::zeros(output);
|
||||
|
||||
let sw1 = Array2::zeros((hidden1, input));
|
||||
let sb1 = Array1::zeros(hidden1);
|
||||
let sw2 = Array2::zeros((hidden2, hidden1));
|
||||
let sb2 = Array1::zeros(hidden2);
|
||||
let sw3 = Array2::zeros((output, hidden2));
|
||||
let sb3 = Array1::zeros(output);
|
||||
|
||||
Self {
|
||||
w1, b1, w2, b2, w3, b3,
|
||||
bn1_gamma, bn1_beta, bn1_running_mean, bn1_running_var,
|
||||
bn2_gamma, bn2_beta, bn2_running_mean, bn2_running_var,
|
||||
mw1, mb1, mw2, mb2, mw3, mb3,
|
||||
sw1, sb1, sw2, sb2, sw3, sb3,
|
||||
t: 0.0,
|
||||
input_dim: input,
|
||||
hidden1_dim: hidden1,
|
||||
hidden2_dim: hidden2,
|
||||
output_dim: output,
|
||||
dropout_rate: 0.3,
|
||||
is_training: true,
|
||||
lr_schedule: LRSchedule::Cosine {
|
||||
initial: 0.001,
|
||||
min: 0.00001,
|
||||
period: 1000
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn batch_norm(&self, x: &Array1<f32>, gamma: &Array1<f32>, beta: &Array1<f32>,
|
||||
running_mean: &Array1<f32>, running_var: &Array1<f32>) -> Array1<f32> {
|
||||
if self.is_training {
|
||||
let mean = x.mean().unwrap();
|
||||
let var = x.var(0.0);
|
||||
let x_norm = (x - mean) / (var + 1e-5).sqrt();
|
||||
gamma * &x_norm + beta
|
||||
} else {
|
||||
let x_norm = (x - running_mean) / (running_var + 1e-5).mapv(f32::sqrt);
|
||||
gamma * &x_norm + beta
|
||||
}
|
||||
}
|
||||
|
||||
fn dropout(&self, x: &Array1<f32>) -> Array1<f32> {
|
||||
if self.is_training && self.dropout_rate > 0.0 {
|
||||
let mut rng = thread_rng();
|
||||
x.mapv(|v| if rng.gen::<f32>() > self.dropout_rate {
|
||||
v / (1.0 - self.dropout_rate)
|
||||
} else {
|
||||
0.0
|
||||
})
|
||||
} else {
|
||||
x.clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn leaky_relu(x: &Array1<f32>) -> Array1<f32> {
|
||||
x.mapv(|v| if v > 0.0 { v } else { 0.01 * v })
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Array1<f32>) -> (Array1<f32>, Array1<f32>, Array1<f32>) {
|
||||
// Layer 1: Linear + BatchNorm + LeakyReLU + Dropout
|
||||
let z1 = self.w1.dot(x) + &self.b1;
|
||||
let bn1 = self.batch_norm(&z1, &self.bn1_gamma, &self.bn1_beta,
|
||||
&self.bn1_running_mean, &self.bn1_running_var);
|
||||
let h1 = Self::leaky_relu(&bn1);
|
||||
let h1_drop = self.dropout(&h1);
|
||||
|
||||
// Layer 2: Linear + BatchNorm + LeakyReLU + Dropout
|
||||
let z2 = self.w2.dot(&h1_drop) + &self.b2;
|
||||
let bn2 = self.batch_norm(&z2, &self.bn2_gamma, &self.bn2_beta,
|
||||
&self.bn2_running_mean, &self.bn2_running_var);
|
||||
let h2 = Self::leaky_relu(&bn2);
|
||||
let h2_drop = self.dropout(&h2);
|
||||
|
||||
// Output layer: Linear (no activation, will apply softmax in loss)
|
||||
let logits = self.w3.dot(&h2_drop) + &self.b3;
|
||||
|
||||
(logits, h1_drop, h2_drop)
|
||||
}
|
||||
|
||||
fn softmax(logits: &Array1<f32>) -> Array1<f32> {
|
||||
let max = logits.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
||||
let exp_vals = logits.mapv(|v| (v - max).exp());
|
||||
let sum = exp_vals.sum();
|
||||
exp_vals / sum
|
||||
}
|
||||
|
||||
fn get_lr(&self, epoch: usize) -> f32 {
|
||||
match &self.lr_schedule {
|
||||
LRSchedule::Constant(lr) => *lr,
|
||||
LRSchedule::Cosine { initial, min, period } => {
|
||||
let progress = (epoch % period) as f32 / *period as f32;
|
||||
min + (initial - min) * (1.0 + (std::f32::consts::PI * progress).cos()) / 2.0
|
||||
}
|
||||
LRSchedule::StepDecay { initial, decay, step_size } => {
|
||||
initial * decay.powi((epoch / step_size) as i32)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn backward(&mut self, x: &Array1<f32>, y_class: usize, epoch: usize) {
|
||||
let (logits, h1, h2) = self.forward(x);
|
||||
let probs = Self::softmax(&logits);
|
||||
|
||||
// Cross-entropy gradient
|
||||
let mut grad_logits = probs;
|
||||
if y_class < self.output_dim {
|
||||
grad_logits[y_class] -= 1.0;
|
||||
}
|
||||
|
||||
// Gradient w.r.t W3, b3
|
||||
let grad_w3 = grad_logits.clone().insert_axis(Axis(1)) * h2.clone().insert_axis(Axis(0));
|
||||
let grad_b3 = grad_logits.clone();
|
||||
|
||||
// Backprop to h2
|
||||
let grad_h2 = self.w3.t().dot(&grad_logits);
|
||||
let grad_h2_relu = grad_h2 * h2.mapv(|v| if v > 0.0 { 1.0 } else { 0.01 });
|
||||
|
||||
// Gradient w.r.t W2, b2
|
||||
let grad_w2 = grad_h2_relu.clone().insert_axis(Axis(1)) * h1.clone().insert_axis(Axis(0));
|
||||
let grad_b2 = grad_h2_relu.clone();
|
||||
|
||||
// Backprop to h1
|
||||
let grad_h1 = self.w2.t().dot(&grad_h2_relu);
|
||||
let grad_h1_relu = grad_h1 * h1.mapv(|v| if v > 0.0 { 1.0 } else { 0.01 });
|
||||
|
||||
// Gradient w.r.t W1, b1
|
||||
let grad_w1 = grad_h1_relu.clone().insert_axis(Axis(1)) * x.clone().insert_axis(Axis(0));
|
||||
let grad_b1 = grad_h1_relu;
|
||||
|
||||
// Adam update with scheduled learning rate
|
||||
let lr = self.get_lr(epoch);
|
||||
self.adam_update(
|
||||
grad_w1, grad_b1,
|
||||
grad_w2.into_shape(self.w2.dim()).unwrap(), grad_b2,
|
||||
grad_w3.into_shape(self.w3.dim()).unwrap(), grad_b3,
|
||||
lr
|
||||
);
|
||||
}
|
||||
|
||||
fn adam_update(&mut self, grad_w1: Array2<f32>, grad_b1: Array1<f32>,
|
||||
grad_w2: Array2<f32>, grad_b2: Array1<f32>,
|
||||
grad_w3: Array2<f32>, grad_b3: Array1<f32>, lr: f32) {
|
||||
let beta1 = 0.9;
|
||||
let beta2 = 0.999;
|
||||
let epsilon = 1e-8;
|
||||
|
||||
self.t += 1.0;
|
||||
|
||||
// Update first moments
|
||||
self.mw1 = &self.mw1 * beta1 + &grad_w1 * (1.0 - beta1);
|
||||
self.mb1 = &self.mb1 * beta1 + &grad_b1 * (1.0 - beta1);
|
||||
self.mw2 = &self.mw2 * beta1 + &grad_w2 * (1.0 - beta1);
|
||||
self.mb2 = &self.mb2 * beta1 + &grad_b2 * (1.0 - beta1);
|
||||
self.mw3 = &self.mw3 * beta1 + &grad_w3 * (1.0 - beta1);
|
||||
self.mb3 = &self.mb3 * beta1 + &grad_b3 * (1.0 - beta1);
|
||||
|
||||
// Update second moments
|
||||
self.sw1 = &self.sw1 * beta2 + grad_w1.mapv(|x| x * x) * (1.0 - beta2);
|
||||
self.sb1 = &self.sb1 * beta2 + grad_b1.mapv(|x| x * x) * (1.0 - beta2);
|
||||
self.sw2 = &self.sw2 * beta2 + grad_w2.mapv(|x| x * x) * (1.0 - beta2);
|
||||
self.sb2 = &self.sb2 * beta2 + grad_b2.mapv(|x| x * x) * (1.0 - beta2);
|
||||
self.sw3 = &self.sw3 * beta2 + grad_w3.mapv(|x| x * x) * (1.0 - beta2);
|
||||
self.sb3 = &self.sb3 * beta2 + grad_b3.mapv(|x| x * x) * (1.0 - beta2);
|
||||
|
||||
// Bias correction
|
||||
let bias1 = 1.0 - beta1.powf(self.t);
|
||||
let bias2 = 1.0 - beta2.powf(self.t);
|
||||
|
||||
// Update weights
|
||||
self.w1 = &self.w1 - lr * &self.mw1 / bias1 / ((&self.sw1 / bias2).mapv(f32::sqrt) + epsilon);
|
||||
self.b1 = &self.b1 - lr * &self.mb1 / bias1 / ((&self.sb1 / bias2).mapv(f32::sqrt) + epsilon);
|
||||
self.w2 = &self.w2 - lr * &self.mw2 / bias1 / ((&self.sw2 / bias2).mapv(f32::sqrt) + epsilon);
|
||||
self.b2 = &self.b2 - lr * &self.mb2 / bias1 / ((&self.sb2 / bias2).mapv(f32::sqrt) + epsilon);
|
||||
self.w3 = &self.w3 - lr * &self.mw3 / bias1 / ((&self.sw3 / bias2).mapv(f32::sqrt) + epsilon);
|
||||
self.b3 = &self.b3 - lr * &self.mb3 / bias1 / ((&self.sb3 / bias2).mapv(f32::sqrt) + epsilon);
|
||||
}
|
||||
|
||||
pub fn train_classification(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>,
|
||||
epochs: usize, batch_size: usize) {
|
||||
self.is_training = true;
|
||||
|
||||
for epoch in 0..epochs {
|
||||
let mut indices: Vec<usize> = (0..x.len()).collect();
|
||||
let mut rng = thread_rng();
|
||||
use rand::seq::SliceRandom;
|
||||
indices.shuffle(&mut rng);
|
||||
|
||||
for batch_start in (0..x.len()).step_by(batch_size) {
|
||||
let batch_end = (batch_start + batch_size).min(x.len());
|
||||
|
||||
for i in batch_start..batch_end {
|
||||
let idx = indices[i];
|
||||
let x_arr = Array1::from_vec(x[idx].clone());
|
||||
|
||||
// Convert continuous y to class
|
||||
let y_class = if y[idx] < -0.25 { 0 }
|
||||
else if y[idx] > 0.25 { 2 }
|
||||
else { 1 };
|
||||
|
||||
self.backward(&x_arr, y_class, epoch);
|
||||
}
|
||||
|
||||
// Update batch norm running stats
|
||||
self.bn1_running_mean = &self.bn1_running_mean * 0.9 + &self.bn1_beta * 0.1;
|
||||
self.bn1_running_var = &self.bn1_running_var * 0.9 + &self.bn1_gamma * 0.1;
|
||||
self.bn2_running_mean = &self.bn2_running_mean * 0.9 + &self.bn2_beta * 0.1;
|
||||
self.bn2_running_var = &self.bn2_running_var * 0.9 + &self.bn2_gamma * 0.1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn predict_cls3(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
self.is_training = false;
|
||||
|
||||
x.iter().map(|xi| {
|
||||
let x_arr = Array1::from_vec(xi.clone());
|
||||
let (logits, _, _) = self.forward(&x_arr);
|
||||
let probs = Self::softmax(&logits);
|
||||
|
||||
// Return argmax
|
||||
let mut best = 0;
|
||||
for i in 1..self.output_dim.min(3) {
|
||||
if probs[i] > probs[best] {
|
||||
best = i;
|
||||
}
|
||||
}
|
||||
best
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,246 @@
|
||||
use ndarray::{Array2, Array1, Axis};
|
||||
use rand::{thread_rng, Rng};
|
||||
use rayon::prelude::*;
|
||||
|
||||
pub struct OptimizedMlp {
|
||||
w1: Array2<f32>,
|
||||
b1: Array1<f32>,
|
||||
w2: Array2<f32>,
|
||||
b2: Array1<f32>,
|
||||
// Momentum terms
|
||||
vw1: Array2<f32>,
|
||||
vb1: Array1<f32>,
|
||||
vw2: Array2<f32>,
|
||||
vb2: Array1<f32>,
|
||||
// Adam optimizer state
|
||||
mw1: Array2<f32>,
|
||||
mb1: Array1<f32>,
|
||||
mw2: Array2<f32>,
|
||||
mb2: Array1<f32>,
|
||||
sw1: Array2<f32>,
|
||||
sb1: Array1<f32>,
|
||||
sw2: Array2<f32>,
|
||||
sb2: Array1<f32>,
|
||||
t: f32,
|
||||
output_dim: usize,
|
||||
use_adam: bool,
|
||||
}
|
||||
|
||||
impl OptimizedMlp {
|
||||
pub fn new(input: usize, hidden: usize, output: usize) -> Self {
|
||||
let mut rng = thread_rng();
|
||||
|
||||
// He initialization for ReLU
|
||||
let scale1 = (2.0 / input as f32).sqrt();
|
||||
let scale2 = (2.0 / hidden as f32).sqrt();
|
||||
|
||||
let w1 = Array2::from_shape_fn((hidden, input), |_|
|
||||
rng.gen::<f32>() * scale1 - scale1/2.0);
|
||||
let b1 = Array1::zeros(hidden);
|
||||
let w2 = Array2::from_shape_fn((output, hidden), |_|
|
||||
rng.gen::<f32>() * scale2 - scale2/2.0);
|
||||
let b2 = Array1::zeros(output);
|
||||
|
||||
// Initialize momentum and Adam states
|
||||
let vw1 = Array2::zeros((hidden, input));
|
||||
let vb1 = Array1::zeros(hidden);
|
||||
let vw2 = Array2::zeros((output, hidden));
|
||||
let vb2 = Array1::zeros(output);
|
||||
|
||||
let mw1 = Array2::zeros((hidden, input));
|
||||
let mb1 = Array1::zeros(hidden);
|
||||
let mw2 = Array2::zeros((output, hidden));
|
||||
let mb2 = Array1::zeros(output);
|
||||
|
||||
let sw1 = Array2::zeros((hidden, input));
|
||||
let sb1 = Array1::zeros(hidden);
|
||||
let sw2 = Array2::zeros((output, hidden));
|
||||
let sb2 = Array1::zeros(output);
|
||||
|
||||
Self {
|
||||
w1, b1, w2, b2,
|
||||
vw1, vb1, vw2, vb2,
|
||||
mw1, mb1, mw2, mb2,
|
||||
sw1, sb1, sw2, sb2,
|
||||
t: 0.0,
|
||||
output_dim: output,
|
||||
use_adam: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn forward(&self, x: &Array1<f32>) -> (Array1<f32>, Array1<f32>) {
|
||||
let z1 = self.w1.dot(x) + &self.b1;
|
||||
let h = z1.mapv(|v| v.max(0.0)); // ReLU
|
||||
let output = self.w2.dot(&h) + &self.b2;
|
||||
(output, h)
|
||||
}
|
||||
|
||||
pub fn backward(&mut self, x: &Array1<f32>, y_true: f32, lr: f32) {
|
||||
let (output, h) = self.forward(x);
|
||||
|
||||
// Compute loss gradient
|
||||
let grad_output = if self.output_dim == 1 {
|
||||
Array1::from_elem(1, output[0] - y_true)
|
||||
} else {
|
||||
// Softmax + cross-entropy gradient for classification
|
||||
let exp_out = output.mapv(|v| v.exp());
|
||||
let sum_exp = exp_out.sum();
|
||||
let softmax = &exp_out / sum_exp;
|
||||
|
||||
// Create one-hot target
|
||||
let class = if y_true < -0.25 { 0 }
|
||||
else if y_true > 0.25 { 2 }
|
||||
else { 1 };
|
||||
let mut target = Array1::zeros(self.output_dim);
|
||||
if class < self.output_dim {
|
||||
target[class] = 1.0;
|
||||
}
|
||||
|
||||
softmax - target
|
||||
};
|
||||
|
||||
// Gradient w.r.t w2 and b2
|
||||
let grad_w2 = grad_output.clone().insert_axis(Axis(1)) * h.clone().insert_axis(Axis(0));
|
||||
let grad_b2 = grad_output.clone();
|
||||
|
||||
// Backprop through hidden layer
|
||||
let grad_h = self.w2.t().dot(&grad_output);
|
||||
// Gradient of ReLU: 1 if input > 0, else 0
|
||||
let z1 = self.w1.dot(x) + &self.b1;
|
||||
let grad_z1 = grad_h * z1.mapv(|v| if v > 0.0 { 1.0 } else { 0.0 });
|
||||
|
||||
// Gradient w.r.t w1 and b1
|
||||
let grad_w1 = grad_z1.clone().insert_axis(Axis(1)) * x.clone().insert_axis(Axis(0));
|
||||
let grad_b1 = grad_z1;
|
||||
|
||||
// Update weights using Adam or momentum
|
||||
if self.use_adam {
|
||||
self.adam_update(grad_w1, grad_b1, grad_w2.into_shape(self.w2.dim()).unwrap(), grad_b2, lr);
|
||||
} else {
|
||||
self.momentum_update(grad_w1, grad_b1, grad_w2.into_shape(self.w2.dim()).unwrap(), grad_b2, lr);
|
||||
}
|
||||
}
|
||||
|
||||
fn adam_update(&mut self, grad_w1: Array2<f32>, grad_b1: Array1<f32>,
|
||||
grad_w2: Array2<f32>, grad_b2: Array1<f32>, lr: f32) {
|
||||
let beta1 = 0.9;
|
||||
let beta2 = 0.999;
|
||||
let epsilon = 1e-8;
|
||||
|
||||
self.t += 1.0;
|
||||
|
||||
// Update biased first moment estimate
|
||||
self.mw1 = &self.mw1 * beta1 + &grad_w1 * (1.0 - beta1);
|
||||
self.mb1 = &self.mb1 * beta1 + &grad_b1 * (1.0 - beta1);
|
||||
self.mw2 = &self.mw2 * beta1 + &grad_w2 * (1.0 - beta1);
|
||||
self.mb2 = &self.mb2 * beta1 + &grad_b2 * (1.0 - beta1);
|
||||
|
||||
// Update biased second raw moment estimate
|
||||
self.sw1 = &self.sw1 * beta2 + grad_w1.mapv(|x| x * x) * (1.0 - beta2);
|
||||
self.sb1 = &self.sb1 * beta2 + grad_b1.mapv(|x| x * x) * (1.0 - beta2);
|
||||
self.sw2 = &self.sw2 * beta2 + grad_w2.mapv(|x| x * x) * (1.0 - beta2);
|
||||
self.sb2 = &self.sb2 * beta2 + grad_b2.mapv(|x| x * x) * (1.0 - beta2);
|
||||
|
||||
// Compute bias-corrected moments
|
||||
let bias_correction1 = 1.0 - beta1.powf(self.t);
|
||||
let bias_correction2 = 1.0 - beta2.powf(self.t);
|
||||
|
||||
// Update weights
|
||||
self.w1 = &self.w1 - lr * &self.mw1 / bias_correction1 / ((&self.sw1 / bias_correction2).mapv(f32::sqrt) + epsilon);
|
||||
self.b1 = &self.b1 - lr * &self.mb1 / bias_correction1 / ((&self.sb1 / bias_correction2).mapv(f32::sqrt) + epsilon);
|
||||
self.w2 = &self.w2 - lr * &self.mw2 / bias_correction1 / ((&self.sw2 / bias_correction2).mapv(f32::sqrt) + epsilon);
|
||||
self.b2 = &self.b2 - lr * &self.mb2 / bias_correction1 / ((&self.sb2 / bias_correction2).mapv(f32::sqrt) + epsilon);
|
||||
}
|
||||
|
||||
fn momentum_update(&mut self, grad_w1: Array2<f32>, grad_b1: Array1<f32>,
|
||||
grad_w2: Array2<f32>, grad_b2: Array1<f32>, lr: f32) {
|
||||
let momentum = 0.9;
|
||||
|
||||
// Update velocity
|
||||
self.vw1 = &self.vw1 * momentum - &grad_w1 * lr;
|
||||
self.vb1 = &self.vb1 * momentum - &grad_b1 * lr;
|
||||
self.vw2 = &self.vw2 * momentum - &grad_w2 * lr;
|
||||
self.vb2 = &self.vb2 * momentum - &grad_b2 * lr;
|
||||
|
||||
// Update weights
|
||||
self.w1 = &self.w1 + &self.vw1;
|
||||
self.b1 = &self.b1 + &self.vb1;
|
||||
self.w2 = &self.w2 + &self.vw2;
|
||||
self.b2 = &self.b2 + &self.vb2;
|
||||
}
|
||||
|
||||
pub fn train_regression(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>, epochs: usize, lr: f32) {
|
||||
for _ in 0..epochs {
|
||||
// Shuffle indices for SGD
|
||||
let mut indices: Vec<usize> = (0..x.len()).collect();
|
||||
let mut rng = thread_rng();
|
||||
use rand::seq::SliceRandom;
|
||||
indices.shuffle(&mut rng);
|
||||
|
||||
for &i in &indices {
|
||||
let x_arr = Array1::from_vec(x[i].clone());
|
||||
self.backward(&x_arr, y[i], lr);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn train_batch(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>, epochs: usize,
|
||||
lr: f32, batch_size: usize) {
|
||||
for _ in 0..epochs {
|
||||
let mut indices: Vec<usize> = (0..x.len()).collect();
|
||||
let mut rng = thread_rng();
|
||||
use rand::seq::SliceRandom;
|
||||
indices.shuffle(&mut rng);
|
||||
|
||||
for batch_start in (0..x.len()).step_by(batch_size) {
|
||||
let batch_end = (batch_start + batch_size).min(x.len());
|
||||
|
||||
// Accumulate gradients for batch
|
||||
for i in batch_start..batch_end {
|
||||
let idx = indices[i];
|
||||
let x_arr = Array1::from_vec(x[idx].clone());
|
||||
self.backward(&x_arr, y[idx], lr / batch_size as f32);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn predict_reg(&self, x: &[Vec<f32>]) -> Vec<f32> {
|
||||
x.par_iter().map(|xi| {
|
||||
let (out, _) = self.forward(&Array1::from_vec(xi.clone()));
|
||||
if self.output_dim == 1 {
|
||||
out[0]
|
||||
} else {
|
||||
out[0]
|
||||
}
|
||||
}).collect()
|
||||
}
|
||||
|
||||
pub fn predict_cls3(&self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
x.par_iter().map(|xi| {
|
||||
let (out, _) = self.forward(&Array1::from_vec(xi.clone()));
|
||||
|
||||
if self.output_dim >= 3 {
|
||||
// Softmax + argmax for classification
|
||||
let exp_out = out.mapv(|v| v.exp());
|
||||
let sum_exp = exp_out.sum();
|
||||
let probs = exp_out / sum_exp;
|
||||
|
||||
let mut best = 0;
|
||||
let mut best_val = probs[0];
|
||||
for i in 1..3.min(probs.len()) {
|
||||
if probs[i] > best_val {
|
||||
best_val = probs[i];
|
||||
best = i;
|
||||
}
|
||||
}
|
||||
best
|
||||
} else {
|
||||
let val = out[0];
|
||||
if val < -0.25 { 0 }
|
||||
else if val > 0.25 { 2 }
|
||||
else { 1 }
|
||||
}
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,267 @@
|
||||
use crate::quantization::{QuantizedMlp, QuantizedWeights};
|
||||
use rand::Rng;
|
||||
|
||||
/// Quantized MLP wrapper for temporal-compare
|
||||
/// Provides 4x model size reduction with minimal accuracy loss
|
||||
pub struct QuantizedMlpBackend {
|
||||
quantized: Option<QuantizedMlp>,
|
||||
|
||||
// Training happens in FP32, quantize after
|
||||
weights1: Vec<f32>,
|
||||
bias1: Vec<f32>,
|
||||
weights2: Vec<f32>,
|
||||
bias2: Vec<f32>,
|
||||
|
||||
input_dim: usize,
|
||||
hidden_dim: usize,
|
||||
output_dim: usize,
|
||||
|
||||
// Track compression stats
|
||||
original_size: usize,
|
||||
quantized_size: usize,
|
||||
}
|
||||
|
||||
impl QuantizedMlpBackend {
|
||||
pub fn new(input_dim: usize, hidden_dim: usize, output_dim: usize) -> Self {
|
||||
let mut rng = rand::thread_rng();
|
||||
|
||||
// Xavier initialization
|
||||
let scale1 = (2.0 / input_dim as f32).sqrt();
|
||||
let scale2 = (2.0 / hidden_dim as f32).sqrt();
|
||||
|
||||
let weights1: Vec<f32> = (0..hidden_dim * input_dim)
|
||||
.map(|_| rng.gen_range(-scale1..scale1))
|
||||
.collect();
|
||||
|
||||
let weights2: Vec<f32> = (0..output_dim * hidden_dim)
|
||||
.map(|_| rng.gen_range(-scale2..scale2))
|
||||
.collect();
|
||||
|
||||
let original_size = (weights1.len() + weights2.len() + hidden_dim + output_dim) * 4;
|
||||
|
||||
Self {
|
||||
quantized: None,
|
||||
weights1,
|
||||
bias1: vec![0.0; hidden_dim],
|
||||
weights2,
|
||||
bias2: vec![0.0; output_dim],
|
||||
input_dim,
|
||||
hidden_dim,
|
||||
output_dim,
|
||||
original_size,
|
||||
quantized_size: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Train in FP32 for best accuracy
|
||||
pub fn train(&mut self, x: &[Vec<f32>], y: &[f32], epochs: usize, lr: f32) {
|
||||
for epoch in 0..epochs {
|
||||
let mut total_loss = 0.0;
|
||||
|
||||
for (xi, &yi) in x.iter().zip(y.iter()) {
|
||||
// Forward pass (FP32)
|
||||
let mut hidden = vec![0.0f32; self.hidden_dim];
|
||||
|
||||
// Layer 1
|
||||
for i in 0..self.hidden_dim {
|
||||
let mut sum = self.bias1[i];
|
||||
for j in 0..self.input_dim {
|
||||
sum += self.weights1[i * self.input_dim + j] * xi[j];
|
||||
}
|
||||
hidden[i] = sum.max(0.0); // ReLU
|
||||
}
|
||||
|
||||
// Layer 2
|
||||
let mut output = self.bias2[0];
|
||||
for i in 0..self.hidden_dim {
|
||||
output += self.weights2[i] * hidden[i];
|
||||
}
|
||||
|
||||
// Loss (MSE)
|
||||
let error = output - yi;
|
||||
total_loss += error * error;
|
||||
|
||||
// Backward pass
|
||||
// Output layer gradients
|
||||
for i in 0..self.hidden_dim {
|
||||
self.weights2[i] -= lr * error * hidden[i];
|
||||
}
|
||||
self.bias2[0] -= lr * error;
|
||||
|
||||
// Hidden layer gradients
|
||||
for i in 0..self.hidden_dim {
|
||||
if hidden[i] > 0.0 {
|
||||
let grad = error * self.weights2[i];
|
||||
|
||||
for j in 0..self.input_dim {
|
||||
self.weights1[i * self.input_dim + j] -= lr * grad * xi[j];
|
||||
}
|
||||
self.bias1[i] -= lr * grad;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if epoch % 10 == 0 {
|
||||
println!("Epoch {}: Loss = {:.6}", epoch, total_loss / x.len() as f32);
|
||||
}
|
||||
}
|
||||
|
||||
// Quantize after training
|
||||
self.quantize();
|
||||
}
|
||||
|
||||
/// Quantize the trained FP32 model to INT8
|
||||
pub fn quantize(&mut self) {
|
||||
let qmlp = QuantizedMlp::from_float_mlp(
|
||||
&self.weights1,
|
||||
&self.bias1,
|
||||
&self.weights2,
|
||||
&self.bias2,
|
||||
self.input_dim,
|
||||
self.hidden_dim,
|
||||
self.output_dim
|
||||
);
|
||||
|
||||
self.quantized_size = qmlp.model_size();
|
||||
self.quantized = Some(qmlp);
|
||||
}
|
||||
|
||||
/// Predict using quantized weights (fast)
|
||||
pub fn predict(&self, x: &[Vec<f32>]) -> Vec<f32> {
|
||||
match &self.quantized {
|
||||
Some(qmlp) => {
|
||||
x.iter().map(|xi| {
|
||||
let mut output = vec![0.0f32; self.output_dim];
|
||||
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
{
|
||||
qmlp.forward_avx2(xi, &mut output);
|
||||
}
|
||||
#[cfg(not(target_arch = "x86_64"))]
|
||||
{
|
||||
qmlp.forward(xi, &mut output);
|
||||
}
|
||||
|
||||
output[0]
|
||||
}).collect()
|
||||
}
|
||||
None => {
|
||||
// Fallback to FP32 if not quantized
|
||||
self.predict_fp32(x)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Predict using FP32 weights (for comparison)
|
||||
pub fn predict_fp32(&self, x: &[Vec<f32>]) -> Vec<f32> {
|
||||
x.iter().map(|xi| {
|
||||
let mut hidden = vec![0.0f32; self.hidden_dim];
|
||||
|
||||
// Layer 1
|
||||
for i in 0..self.hidden_dim {
|
||||
let mut sum = self.bias1[i];
|
||||
for j in 0..self.input_dim {
|
||||
sum += self.weights1[i * self.input_dim + j] * xi[j];
|
||||
}
|
||||
hidden[i] = sum.max(0.0);
|
||||
}
|
||||
|
||||
// Layer 2
|
||||
let mut output = self.bias2[0];
|
||||
for i in 0..self.hidden_dim {
|
||||
output += self.weights2[i] * hidden[i];
|
||||
}
|
||||
|
||||
output
|
||||
}).collect()
|
||||
}
|
||||
|
||||
/// Classification prediction
|
||||
pub fn predict_class(&self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
self.predict(x).iter().map(|&y| {
|
||||
if y < -0.25 { 0 }
|
||||
else if y > 0.25 { 2 }
|
||||
else { 1 }
|
||||
}).collect()
|
||||
}
|
||||
|
||||
/// Get compression statistics
|
||||
pub fn get_compression_stats(&self) -> (usize, usize, f32) {
|
||||
let ratio = if self.quantized_size > 0 {
|
||||
self.original_size as f32 / self.quantized_size as f32
|
||||
} else {
|
||||
1.0
|
||||
};
|
||||
|
||||
(self.original_size, self.quantized_size, ratio)
|
||||
}
|
||||
|
||||
/// Compare FP32 vs INT8 performance
|
||||
pub fn benchmark_inference(&self, x: &[Vec<f32>], iterations: usize) {
|
||||
use std::time::Instant;
|
||||
|
||||
// Warm up
|
||||
let _ = self.predict_fp32(&x[..1.min(x.len())]);
|
||||
let _ = self.predict(&x[..1.min(x.len())]);
|
||||
|
||||
// Benchmark FP32
|
||||
let start = Instant::now();
|
||||
for _ in 0..iterations {
|
||||
let _ = self.predict_fp32(x);
|
||||
}
|
||||
let fp32_time = start.elapsed();
|
||||
|
||||
// Benchmark INT8
|
||||
let start = Instant::now();
|
||||
for _ in 0..iterations {
|
||||
let _ = self.predict(x);
|
||||
}
|
||||
let int8_time = start.elapsed();
|
||||
|
||||
let speedup = fp32_time.as_secs_f32() / int8_time.as_secs_f32();
|
||||
|
||||
println!("\n=== Quantization Benchmark ===");
|
||||
println!("FP32 time: {:.3}s", fp32_time.as_secs_f32());
|
||||
println!("INT8 time: {:.3}s", int8_time.as_secs_f32());
|
||||
println!("Speedup: {:.2}x", speedup);
|
||||
|
||||
let (orig, quant, ratio) = self.get_compression_stats();
|
||||
println!("\n=== Model Size ===");
|
||||
println!("Original: {} bytes", orig);
|
||||
println!("Quantized: {} bytes", quant);
|
||||
println!("Compression: {:.2}x", ratio);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_quantized_mlp() {
|
||||
let mut model = QuantizedMlpBackend::new(32, 64, 1);
|
||||
|
||||
// Create dummy data
|
||||
let x: Vec<Vec<f32>> = (0..100)
|
||||
.map(|_| (0..32).map(|_| rand::random()).collect())
|
||||
.collect();
|
||||
let y: Vec<f32> = (0..100).map(|_| rand::random()).collect();
|
||||
|
||||
// Train and quantize
|
||||
model.train(&x, &y, 10, 0.01);
|
||||
|
||||
// Check predictions work
|
||||
let pred_fp32 = model.predict_fp32(&x[..10]);
|
||||
let pred_int8 = model.predict(&x[..10]);
|
||||
|
||||
// Should be similar but not identical
|
||||
for (p32, p8) in pred_fp32.iter().zip(&pred_int8) {
|
||||
let diff = (p32 - p8).abs();
|
||||
assert!(diff < 0.1, "Quantization error too large: {}", diff);
|
||||
}
|
||||
|
||||
// Check compression
|
||||
let (_, _, ratio) = model.get_compression_stats();
|
||||
assert!(ratio > 3.0, "Compression ratio too low: {}", ratio);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,314 @@
|
||||
use std::arch::x86_64::*;
|
||||
use ndarray::Axis;
|
||||
use rand::{thread_rng, Rng};
|
||||
use rayon::prelude::*;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Ultra-optimized MLP with SIMD, cache blocking, and parallel processing
|
||||
pub struct UltraMlp {
|
||||
// Weights stored in row-major for cache efficiency
|
||||
w1_flat: Vec<f32>,
|
||||
b1: Vec<f32>,
|
||||
w2_flat: Vec<f32>,
|
||||
b2: Vec<f32>,
|
||||
|
||||
// Dimensions
|
||||
input_dim: usize,
|
||||
hidden_dim: usize,
|
||||
output_dim: usize,
|
||||
|
||||
// Momentum buffers (flat)
|
||||
vw1: Vec<f32>,
|
||||
vb1: Vec<f32>,
|
||||
vw2: Vec<f32>,
|
||||
vb2: Vec<f32>,
|
||||
|
||||
// Pre-allocated buffers for forward pass
|
||||
hidden_buffer: Vec<f32>,
|
||||
output_buffer: Vec<f32>,
|
||||
}
|
||||
|
||||
impl UltraMlp {
|
||||
pub fn new(input: usize, hidden: usize, output: usize) -> Self {
|
||||
let mut rng = thread_rng();
|
||||
|
||||
// He initialization
|
||||
let scale1 = (2.0 / input as f32).sqrt();
|
||||
let scale2 = (2.0 / hidden as f32).sqrt();
|
||||
|
||||
// Initialize weights flat for SIMD
|
||||
let w1_flat: Vec<f32> = (0..hidden*input)
|
||||
.map(|_| rng.gen::<f32>() * scale1 - scale1/2.0)
|
||||
.collect();
|
||||
let w2_flat: Vec<f32> = (0..output*hidden)
|
||||
.map(|_| rng.gen::<f32>() * scale2 - scale2/2.0)
|
||||
.collect();
|
||||
|
||||
Self {
|
||||
w1_flat,
|
||||
b1: vec![0.0; hidden],
|
||||
w2_flat,
|
||||
b2: vec![0.0; output],
|
||||
input_dim: input,
|
||||
hidden_dim: hidden,
|
||||
output_dim: output,
|
||||
vw1: vec![0.0; hidden * input],
|
||||
vb1: vec![0.0; hidden],
|
||||
vw2: vec![0.0; output * hidden],
|
||||
vb2: vec![0.0; output],
|
||||
hidden_buffer: vec![0.0; hidden],
|
||||
output_buffer: vec![0.0; output],
|
||||
}
|
||||
}
|
||||
|
||||
/// SIMD-accelerated matrix-vector multiplication
|
||||
#[target_feature(enable = "avx2")]
|
||||
#[inline]
|
||||
unsafe fn simd_matmul(weights: &[f32], input: &[f32], output: &mut [f32],
|
||||
rows: usize, cols: usize) {
|
||||
const SIMD_WIDTH: usize = 8;
|
||||
|
||||
for i in 0..rows {
|
||||
let row_offset = i * cols;
|
||||
let mut sum = _mm256_setzero_ps();
|
||||
|
||||
// Process 8 elements at a time with AVX2
|
||||
let chunks = cols / SIMD_WIDTH;
|
||||
for j in 0..chunks {
|
||||
let idx = j * SIMD_WIDTH;
|
||||
let w = _mm256_loadu_ps(&weights[row_offset + idx]);
|
||||
let x = _mm256_loadu_ps(&input[idx]);
|
||||
sum = _mm256_fmadd_ps(w, x, sum);
|
||||
}
|
||||
|
||||
// Horizontal sum
|
||||
let sum_array = std::mem::transmute::<__m256, [f32; 8]>(sum);
|
||||
let mut result = sum_array.iter().sum::<f32>();
|
||||
|
||||
// Handle remaining elements
|
||||
for j in (chunks * SIMD_WIDTH)..cols {
|
||||
result += weights[row_offset + j] * input[j];
|
||||
}
|
||||
|
||||
output[i] = result;
|
||||
}
|
||||
}
|
||||
|
||||
/// Vectorized ReLU activation
|
||||
#[target_feature(enable = "avx2")]
|
||||
#[inline]
|
||||
unsafe fn simd_relu(data: &mut [f32]) {
|
||||
const SIMD_WIDTH: usize = 8;
|
||||
let zero = _mm256_setzero_ps();
|
||||
|
||||
let chunks = data.len() / SIMD_WIDTH;
|
||||
for i in 0..chunks {
|
||||
let idx = i * SIMD_WIDTH;
|
||||
let val = _mm256_loadu_ps(&data[idx]);
|
||||
let relu = _mm256_max_ps(val, zero);
|
||||
_mm256_storeu_ps(&mut data[idx], relu);
|
||||
}
|
||||
|
||||
// Handle remaining
|
||||
for i in (chunks * SIMD_WIDTH)..data.len() {
|
||||
data[i] = data[i].max(0.0);
|
||||
}
|
||||
}
|
||||
|
||||
/// Ultra-fast forward pass
|
||||
pub fn forward_fast(&mut self, x: &[f32]) -> &[f32] {
|
||||
unsafe {
|
||||
// Layer 1: input -> hidden
|
||||
Self::simd_matmul(&self.w1_flat, x, &mut self.hidden_buffer,
|
||||
self.hidden_dim, self.input_dim);
|
||||
|
||||
// Add bias with SIMD
|
||||
for i in 0..self.hidden_dim {
|
||||
self.hidden_buffer[i] += self.b1[i];
|
||||
}
|
||||
|
||||
// ReLU activation
|
||||
Self::simd_relu(&mut self.hidden_buffer);
|
||||
|
||||
// Layer 2: hidden -> output
|
||||
Self::simd_matmul(&self.w2_flat, &self.hidden_buffer, &mut self.output_buffer,
|
||||
self.output_dim, self.hidden_dim);
|
||||
|
||||
// Add bias
|
||||
for i in 0..self.output_dim {
|
||||
self.output_buffer[i] += self.b2[i];
|
||||
}
|
||||
}
|
||||
|
||||
&self.output_buffer
|
||||
}
|
||||
|
||||
/// Optimized backpropagation with momentum
|
||||
pub fn backward_fast(&mut self, x: &[f32], y_true: f32, lr: f32) {
|
||||
// Clone output to avoid borrow issues
|
||||
let output_copy = self.forward_fast(x).to_vec();
|
||||
let momentum = 0.9;
|
||||
|
||||
// Output gradient
|
||||
let grad_out = if self.output_dim == 1 {
|
||||
vec![output_copy[0] - y_true]
|
||||
} else {
|
||||
// Softmax gradient for classification
|
||||
let max = output_copy.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
||||
let exp_sum = output_copy.iter().map(|&v| (v - max).exp()).sum::<f32>();
|
||||
let mut grad = output_copy.iter().map(|&v| (v - max).exp() / exp_sum).collect::<Vec<_>>();
|
||||
|
||||
let class = if y_true < -0.25 { 0 } else if y_true > 0.25 { 2 } else { 1 };
|
||||
if class < grad.len() {
|
||||
grad[class] -= 1.0;
|
||||
}
|
||||
grad
|
||||
};
|
||||
|
||||
// Gradient w.r.t W2 and b2 (vectorized)
|
||||
for i in 0..self.output_dim {
|
||||
let grad_i = grad_out[i];
|
||||
|
||||
// Update bias with momentum
|
||||
self.vb2[i] = momentum * self.vb2[i] - lr * grad_i;
|
||||
self.b2[i] += self.vb2[i];
|
||||
|
||||
// Update weights with momentum (vectorized)
|
||||
let w2_offset = i * self.hidden_dim;
|
||||
for j in 0..self.hidden_dim {
|
||||
let idx = w2_offset + j;
|
||||
let grad_w = grad_i * self.hidden_buffer[j];
|
||||
self.vw2[idx] = momentum * self.vw2[idx] - lr * grad_w;
|
||||
self.w2_flat[idx] += self.vw2[idx];
|
||||
}
|
||||
}
|
||||
|
||||
// Gradient backprop to hidden layer
|
||||
let mut grad_hidden = vec![0.0; self.hidden_dim];
|
||||
for i in 0..self.hidden_dim {
|
||||
for j in 0..self.output_dim {
|
||||
grad_hidden[i] += self.w2_flat[j * self.hidden_dim + i] * grad_out[j];
|
||||
}
|
||||
// ReLU gradient
|
||||
if self.hidden_buffer[i] <= 0.0 {
|
||||
grad_hidden[i] = 0.0;
|
||||
}
|
||||
}
|
||||
|
||||
// Update W1 and b1
|
||||
for i in 0..self.hidden_dim {
|
||||
let grad_i = grad_hidden[i];
|
||||
|
||||
// Update bias
|
||||
self.vb1[i] = momentum * self.vb1[i] - lr * grad_i;
|
||||
self.b1[i] += self.vb1[i];
|
||||
|
||||
// Update weights
|
||||
let w1_offset = i * self.input_dim;
|
||||
for j in 0..self.input_dim {
|
||||
let idx = w1_offset + j;
|
||||
let grad_w = grad_i * x[j];
|
||||
self.vw1[idx] = momentum * self.vw1[idx] - lr * grad_w;
|
||||
self.w1_flat[idx] += self.vw1[idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parallel batch training with cache-friendly access
|
||||
pub fn train_batch_parallel(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>,
|
||||
epochs: usize, lr: f32, batch_size: usize) {
|
||||
use rand::seq::SliceRandom;
|
||||
|
||||
for _ in 0..epochs {
|
||||
let mut indices: Vec<usize> = (0..x.len()).collect();
|
||||
let mut rng = thread_rng();
|
||||
indices.shuffle(&mut rng);
|
||||
|
||||
// Process in mini-batches
|
||||
for batch_start in (0..x.len()).step_by(batch_size) {
|
||||
let batch_end = (batch_start + batch_size).min(x.len());
|
||||
|
||||
// Sequential within batch for weight updates
|
||||
for i in batch_start..batch_end {
|
||||
let idx = indices[i];
|
||||
self.backward_fast(&x[idx], y[idx], lr / batch_size as f32);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Parallel prediction for multiple samples
|
||||
pub fn predict_parallel(&mut self, x: &[Vec<f32>]) -> Vec<f32> {
|
||||
// Create thread-local copies for parallel execution
|
||||
let w1 = Arc::new(self.w1_flat.clone());
|
||||
let b1 = Arc::new(self.b1.clone());
|
||||
let w2 = Arc::new(self.w2_flat.clone());
|
||||
let b2 = Arc::new(self.b2.clone());
|
||||
let hidden_dim = self.hidden_dim;
|
||||
let input_dim = self.input_dim;
|
||||
let output_dim = self.output_dim;
|
||||
|
||||
x.par_iter().map(|xi| {
|
||||
let mut hidden = vec![0.0; hidden_dim];
|
||||
let mut output = vec![0.0; output_dim];
|
||||
|
||||
unsafe {
|
||||
// Forward pass with local buffers
|
||||
Self::simd_matmul(&w1, xi, &mut hidden, hidden_dim, input_dim);
|
||||
for i in 0..hidden_dim {
|
||||
hidden[i] += b1[i];
|
||||
}
|
||||
Self::simd_relu(&mut hidden);
|
||||
|
||||
Self::simd_matmul(&w2, &hidden, &mut output, output_dim, hidden_dim);
|
||||
for i in 0..output_dim {
|
||||
output[i] += b2[i];
|
||||
}
|
||||
}
|
||||
|
||||
if output_dim == 1 { output[0] } else { output[0] }
|
||||
}).collect()
|
||||
}
|
||||
|
||||
pub fn predict_cls3_parallel(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
let w1 = Arc::new(self.w1_flat.clone());
|
||||
let b1 = Arc::new(self.b1.clone());
|
||||
let w2 = Arc::new(self.w2_flat.clone());
|
||||
let b2 = Arc::new(self.b2.clone());
|
||||
let hidden_dim = self.hidden_dim;
|
||||
let input_dim = self.input_dim;
|
||||
let output_dim = self.output_dim;
|
||||
|
||||
x.par_iter().map(|xi| {
|
||||
let mut hidden = vec![0.0; hidden_dim];
|
||||
let mut output = vec![0.0; output_dim];
|
||||
|
||||
unsafe {
|
||||
Self::simd_matmul(&w1, xi, &mut hidden, hidden_dim, input_dim);
|
||||
for i in 0..hidden_dim {
|
||||
hidden[i] += b1[i];
|
||||
}
|
||||
Self::simd_relu(&mut hidden);
|
||||
|
||||
Self::simd_matmul(&w2, &hidden, &mut output, output_dim, hidden_dim);
|
||||
for i in 0..output_dim {
|
||||
output[i] += b2[i];
|
||||
}
|
||||
}
|
||||
|
||||
if output_dim >= 3 {
|
||||
// Argmax
|
||||
let mut best = 0;
|
||||
for i in 1..3.min(output.len()) {
|
||||
if output[i] > output[best] {
|
||||
best = i;
|
||||
}
|
||||
}
|
||||
best
|
||||
} else {
|
||||
let val = output[0];
|
||||
if val < -0.25 { 0 } else if val > 0.25 { 2 } else { 1 }
|
||||
}
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,323 @@
|
||||
use std::ops::{Add, Mul};
|
||||
|
||||
/// INT8 Quantization for 4x size reduction and 2x faster inference
|
||||
/// Symmetric quantization: maps [-max, max] to [-127, 127]
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Int8Quantizer {
|
||||
pub scale: f32,
|
||||
pub zero_point: i8,
|
||||
}
|
||||
|
||||
impl Int8Quantizer {
|
||||
/// Create quantizer from float weights
|
||||
pub fn from_weights(weights: &[f32]) -> Self {
|
||||
// Find min/max for dynamic range
|
||||
let (min, max) = weights.iter().fold(
|
||||
(f32::INFINITY, f32::NEG_INFINITY),
|
||||
|(min, max), &w| (min.min(w), max.max(w))
|
||||
);
|
||||
|
||||
// Symmetric quantization for better accuracy
|
||||
let abs_max = min.abs().max(max.abs());
|
||||
let scale = abs_max / 127.0;
|
||||
|
||||
Self {
|
||||
scale,
|
||||
zero_point: 0, // Symmetric around zero
|
||||
}
|
||||
}
|
||||
|
||||
/// Quantize float to INT8
|
||||
#[inline]
|
||||
pub fn quantize(&self, value: f32) -> i8 {
|
||||
let scaled = value / self.scale;
|
||||
scaled.round().clamp(-127.0, 127.0) as i8
|
||||
}
|
||||
|
||||
/// Dequantize INT8 back to float
|
||||
#[inline]
|
||||
pub fn dequantize(&self, value: i8) -> f32 {
|
||||
value as f32 * self.scale
|
||||
}
|
||||
|
||||
/// Quantize entire array
|
||||
pub fn quantize_array(&self, values: &[f32]) -> Vec<i8> {
|
||||
values.iter().map(|&v| self.quantize(v)).collect()
|
||||
}
|
||||
|
||||
/// Dequantize entire array
|
||||
pub fn dequantize_array(&self, values: &[i8]) -> Vec<f32> {
|
||||
values.iter().map(|&v| self.dequantize(v)).collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Quantized weight storage for neural networks
|
||||
pub struct QuantizedWeights {
|
||||
pub weights: Vec<i8>,
|
||||
pub quantizer: Int8Quantizer,
|
||||
pub shape: (usize, usize),
|
||||
}
|
||||
|
||||
impl QuantizedWeights {
|
||||
pub fn from_float_matrix(weights: &[f32], rows: usize, cols: usize) -> Self {
|
||||
let quantizer = Int8Quantizer::from_weights(weights);
|
||||
let quantized = quantizer.quantize_array(weights);
|
||||
|
||||
Self {
|
||||
weights: quantized,
|
||||
quantizer,
|
||||
shape: (rows, cols),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get dequantized weight at position
|
||||
#[inline]
|
||||
pub fn get(&self, row: usize, col: usize) -> f32 {
|
||||
let idx = row * self.shape.1 + col;
|
||||
self.quantizer.dequantize(self.weights[idx])
|
||||
}
|
||||
|
||||
/// Quantized matrix multiply with on-the-fly dequantization
|
||||
/// Still faster than FP32 due to better cache utilization
|
||||
pub fn matmul_quantized(&self, input: &[f32], output: &mut [f32]) {
|
||||
let (rows, cols) = self.shape;
|
||||
|
||||
for i in 0..rows {
|
||||
let mut sum = 0.0f32;
|
||||
let row_offset = i * cols;
|
||||
|
||||
// Process in chunks for better cache performance
|
||||
for j in 0..cols {
|
||||
let w_int8 = self.weights[row_offset + j];
|
||||
// Delay dequantization to minimize float operations
|
||||
sum += (w_int8 as f32) * input[j];
|
||||
}
|
||||
|
||||
// Apply scale once at the end
|
||||
output[i] = sum * self.quantizer.scale;
|
||||
}
|
||||
}
|
||||
|
||||
/// SIMD-accelerated quantized matmul (AVX2)
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
#[target_feature(enable = "avx2")]
|
||||
pub unsafe fn matmul_quantized_avx2(&self, input: &[f32], output: &mut [f32]) {
|
||||
use std::arch::x86_64::*;
|
||||
|
||||
let (rows, cols) = self.shape;
|
||||
let scale = _mm256_set1_ps(self.quantizer.scale);
|
||||
|
||||
for i in 0..rows {
|
||||
let mut sum = _mm256_setzero_ps();
|
||||
let row_offset = i * cols;
|
||||
|
||||
// Process 8 elements at once
|
||||
let chunks = cols / 8;
|
||||
for j in 0..chunks {
|
||||
let idx = j * 8;
|
||||
|
||||
// Load 8 INT8 weights and convert to float
|
||||
let w_ptr = self.weights.as_ptr().add(row_offset + idx);
|
||||
let w_i8 = _mm_loadl_epi64(w_ptr as *const __m128i);
|
||||
let w_i32 = _mm256_cvtepi8_epi32(w_i8);
|
||||
let w_f32 = _mm256_cvtepi32_ps(w_i32);
|
||||
|
||||
// Load 8 input floats
|
||||
let x = _mm256_loadu_ps(&input[idx]);
|
||||
|
||||
// Multiply and accumulate
|
||||
sum = _mm256_fmadd_ps(w_f32, x, sum);
|
||||
}
|
||||
|
||||
// Horizontal sum
|
||||
let sum_array: [f32; 8] = std::mem::transmute(sum);
|
||||
let mut result = sum_array.iter().sum::<f32>();
|
||||
|
||||
// Handle remainder
|
||||
for j in (chunks * 8)..cols {
|
||||
result += (self.weights[row_offset + j] as f32) * input[j];
|
||||
}
|
||||
|
||||
// Apply scale
|
||||
output[i] = result * self.quantizer.scale;
|
||||
}
|
||||
}
|
||||
|
||||
/// Get memory size in bytes
|
||||
pub fn memory_size(&self) -> usize {
|
||||
self.weights.len() // INT8 = 1 byte each
|
||||
}
|
||||
}
|
||||
|
||||
/// Quantized MLP with INT8 weights
|
||||
pub struct QuantizedMlp {
|
||||
pub w1: QuantizedWeights,
|
||||
pub b1: Vec<f32>, // Keep biases as FP32 (small overhead)
|
||||
pub w2: QuantizedWeights,
|
||||
pub b2: Vec<f32>,
|
||||
hidden_dim: usize,
|
||||
}
|
||||
|
||||
impl QuantizedMlp {
|
||||
/// Create from existing float MLP
|
||||
pub fn from_float_mlp(w1: &[f32], b1: &[f32], w2: &[f32], b2: &[f32],
|
||||
input_dim: usize, hidden_dim: usize, output_dim: usize) -> Self {
|
||||
Self {
|
||||
w1: QuantizedWeights::from_float_matrix(w1, hidden_dim, input_dim),
|
||||
b1: b1.to_vec(),
|
||||
w2: QuantizedWeights::from_float_matrix(w2, output_dim, hidden_dim),
|
||||
b2: b2.to_vec(),
|
||||
hidden_dim,
|
||||
}
|
||||
}
|
||||
|
||||
/// Forward pass with INT8 weights
|
||||
pub fn forward(&self, input: &[f32], output: &mut [f32]) {
|
||||
let mut hidden = vec![0.0f32; self.hidden_dim];
|
||||
|
||||
// Layer 1: Input -> Hidden
|
||||
self.w1.matmul_quantized(input, &mut hidden);
|
||||
|
||||
// Add bias and ReLU
|
||||
for (h, &b) in hidden.iter_mut().zip(&self.b1) {
|
||||
*h = (*h + b).max(0.0);
|
||||
}
|
||||
|
||||
// Layer 2: Hidden -> Output
|
||||
self.w2.matmul_quantized(&hidden, output);
|
||||
|
||||
// Add bias
|
||||
for (o, &b) in output.iter_mut().zip(&self.b2) {
|
||||
*o += b;
|
||||
}
|
||||
}
|
||||
|
||||
/// Forward pass with AVX2 acceleration
|
||||
#[cfg(target_arch = "x86_64")]
|
||||
pub fn forward_avx2(&self, input: &[f32], output: &mut [f32]) {
|
||||
unsafe {
|
||||
let mut hidden = vec![0.0f32; self.hidden_dim];
|
||||
|
||||
// Layer 1 with SIMD
|
||||
self.w1.matmul_quantized_avx2(input, &mut hidden);
|
||||
|
||||
// Vectorized bias + ReLU
|
||||
use std::arch::x86_64::*;
|
||||
let zero = _mm256_setzero_ps();
|
||||
|
||||
for i in (0..self.hidden_dim).step_by(8) {
|
||||
if i + 8 <= self.hidden_dim {
|
||||
let h = _mm256_loadu_ps(&hidden[i]);
|
||||
let b = _mm256_loadu_ps(&self.b1[i]);
|
||||
let sum = _mm256_add_ps(h, b);
|
||||
let relu = _mm256_max_ps(sum, zero);
|
||||
_mm256_storeu_ps(&mut hidden[i], relu);
|
||||
} else {
|
||||
// Handle remainder
|
||||
for j in i..self.hidden_dim {
|
||||
hidden[j] = (hidden[j] + self.b1[j]).max(0.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Layer 2 with SIMD
|
||||
self.w2.matmul_quantized_avx2(&hidden, output);
|
||||
|
||||
// Add output bias
|
||||
for (o, &b) in output.iter_mut().zip(&self.b2) {
|
||||
*o += b;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get total model size in bytes
|
||||
pub fn model_size(&self) -> usize {
|
||||
self.w1.memory_size() +
|
||||
self.w2.memory_size() +
|
||||
(self.b1.len() + self.b2.len()) * 4 // FP32 biases
|
||||
}
|
||||
|
||||
/// Compression ratio vs FP32
|
||||
pub fn compression_ratio(&self, original_params: usize) -> f32 {
|
||||
let original_bytes = original_params * 4; // FP32
|
||||
let quantized_bytes = self.model_size();
|
||||
original_bytes as f32 / quantized_bytes as f32
|
||||
}
|
||||
}
|
||||
|
||||
/// Quantization-aware training for better INT8 accuracy
|
||||
pub struct QuantizationAwareTraining {
|
||||
pub fake_quantize: bool,
|
||||
pub num_bits: u8,
|
||||
}
|
||||
|
||||
impl QuantizationAwareTraining {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
fake_quantize: true,
|
||||
num_bits: 8,
|
||||
}
|
||||
}
|
||||
|
||||
/// Fake quantization during training to simulate INT8 effects
|
||||
pub fn fake_quantize_weights(&self, weights: &mut [f32]) {
|
||||
if !self.fake_quantize {
|
||||
return;
|
||||
}
|
||||
|
||||
let quantizer = Int8Quantizer::from_weights(weights);
|
||||
|
||||
for w in weights.iter_mut() {
|
||||
// Quantize and immediately dequantize to simulate INT8 effects
|
||||
let quantized = quantizer.quantize(*w);
|
||||
*w = quantizer.dequantize(quantized);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_quantization_round_trip() {
|
||||
let weights = vec![-1.5, -0.5, 0.0, 0.5, 1.5];
|
||||
let quantizer = Int8Quantizer::from_weights(&weights);
|
||||
|
||||
for &w in &weights {
|
||||
let q = quantizer.quantize(w);
|
||||
let dq = quantizer.dequantize(q);
|
||||
|
||||
// Should be close but not exact due to quantization
|
||||
assert!((w - dq).abs() < 0.02);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_compression_ratio() {
|
||||
let input_dim = 32;
|
||||
let hidden_dim = 64;
|
||||
let output_dim = 3;
|
||||
|
||||
let total_params = (input_dim * hidden_dim) + hidden_dim +
|
||||
(hidden_dim * output_dim) + output_dim;
|
||||
|
||||
let w1 = vec![0.1; input_dim * hidden_dim];
|
||||
let b1 = vec![0.0; hidden_dim];
|
||||
let w2 = vec![0.1; hidden_dim * output_dim];
|
||||
let b2 = vec![0.0; output_dim];
|
||||
|
||||
let qmlp = QuantizedMlp::from_float_mlp(
|
||||
&w1, &b1, &w2, &b2,
|
||||
input_dim, hidden_dim, output_dim
|
||||
);
|
||||
|
||||
let ratio = qmlp.compression_ratio(total_params);
|
||||
|
||||
// Should achieve ~3.5-4x compression (weights are INT8, biases stay FP32)
|
||||
assert!(ratio > 3.0 && ratio < 4.5);
|
||||
println!("Compression ratio: {:.2}x", ratio);
|
||||
println!("Original size: {} bytes", total_params * 4);
|
||||
println!("Quantized size: {} bytes", qmlp.model_size());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,311 @@
|
||||
use ndarray::{Array2, Array1, ArrayView1};
|
||||
use rand::{thread_rng, Rng, distributions::Uniform};
|
||||
use std::f32::consts::PI;
|
||||
|
||||
/// Echo State Network (Reservoir Computing) - Ultra-efficient for time series
|
||||
/// Key insight: Random reservoir + simple readout often beats complex architectures
|
||||
pub struct ReservoirComputer {
|
||||
// Reservoir parameters
|
||||
reservoir_size: usize,
|
||||
spectral_radius: f32,
|
||||
sparsity: f32,
|
||||
leak_rate: f32,
|
||||
|
||||
// Reservoir matrices
|
||||
w_in: Array2<f32>, // Input weights (random, fixed)
|
||||
w_res: Array2<f32>, // Reservoir weights (random, sparse, fixed)
|
||||
w_out: Array2<f32>, // Output weights (trained)
|
||||
|
||||
// State
|
||||
state: Array1<f32>,
|
||||
|
||||
// Memory optimization: ring buffer for states
|
||||
state_history: Vec<Array1<f32>>,
|
||||
history_index: usize,
|
||||
max_history: usize,
|
||||
}
|
||||
|
||||
impl ReservoirComputer {
|
||||
pub fn new(input_dim: usize, reservoir_size: usize, output_dim: usize) -> Self {
|
||||
let mut rng = thread_rng();
|
||||
let uniform = Uniform::new(-1.0, 1.0);
|
||||
|
||||
// Input weights: random projection
|
||||
let w_in = Array2::from_shape_fn((reservoir_size, input_dim), |_|
|
||||
rng.sample(uniform));
|
||||
|
||||
// Reservoir weights: sparse random matrix
|
||||
let sparsity = 0.9; // 90% zeros for efficiency
|
||||
let mut w_res = Array2::zeros((reservoir_size, reservoir_size));
|
||||
|
||||
// Create sparse reservoir with specific spectral radius
|
||||
for i in 0..reservoir_size {
|
||||
for j in 0..reservoir_size {
|
||||
if rng.gen::<f32>() > sparsity {
|
||||
w_res[[i, j]] = rng.sample(uniform);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Scale to desired spectral radius (0.9 for edge of chaos)
|
||||
let spectral_radius = 0.9;
|
||||
w_res = Self::scale_spectral_radius(w_res, spectral_radius);
|
||||
|
||||
// Output weights will be learned
|
||||
let w_out = Array2::zeros((output_dim, reservoir_size + input_dim));
|
||||
|
||||
Self {
|
||||
reservoir_size,
|
||||
spectral_radius,
|
||||
sparsity,
|
||||
leak_rate: 0.3,
|
||||
w_in,
|
||||
w_res,
|
||||
w_out,
|
||||
state: Array1::zeros(reservoir_size),
|
||||
state_history: Vec::new(),
|
||||
history_index: 0,
|
||||
max_history: 1000,
|
||||
}
|
||||
}
|
||||
|
||||
/// Scale matrix to desired spectral radius
|
||||
fn scale_spectral_radius(mut matrix: Array2<f32>, target_radius: f32) -> Array2<f32> {
|
||||
// Approximate largest eigenvalue with power iteration
|
||||
let n = matrix.nrows();
|
||||
let mut v = Array1::from_shape_fn(n, |_| rand::random::<f32>());
|
||||
|
||||
for _ in 0..20 {
|
||||
v = matrix.dot(&v);
|
||||
let norm = v.dot(&v).sqrt();
|
||||
if norm > 0.0 {
|
||||
v = v / norm;
|
||||
}
|
||||
}
|
||||
|
||||
let eigenvalue = v.dot(&matrix.dot(&v)) / v.dot(&v);
|
||||
if eigenvalue.abs() > 0.0 {
|
||||
matrix = matrix * (target_radius / eigenvalue.abs());
|
||||
}
|
||||
|
||||
matrix
|
||||
}
|
||||
|
||||
/// Update reservoir state with new input
|
||||
pub fn update_state(&mut self, input: &Array1<f32>) -> Array1<f32> {
|
||||
// Reservoir dynamics: state = (1-α)*state + α*tanh(W_in*input + W_res*state)
|
||||
let input_contribution = self.w_in.dot(input);
|
||||
let recurrent_contribution = self.w_res.dot(&self.state);
|
||||
|
||||
let new_state = &input_contribution + &recurrent_contribution;
|
||||
let activated = new_state.mapv(|x| x.tanh());
|
||||
|
||||
// Leaky integration
|
||||
self.state = &self.state * (1.0 - self.leak_rate) + &activated * self.leak_rate;
|
||||
|
||||
// Store in history (ring buffer)
|
||||
if self.state_history.len() < self.max_history {
|
||||
self.state_history.push(self.state.clone());
|
||||
} else {
|
||||
self.state_history[self.history_index] = self.state.clone();
|
||||
}
|
||||
self.history_index = (self.history_index + 1) % self.max_history;
|
||||
|
||||
self.state.clone()
|
||||
}
|
||||
|
||||
/// Collect states for training (washout period to remove initial transients)
|
||||
pub fn collect_states(&mut self, inputs: &[Vec<f32>], washout: usize)
|
||||
-> (Array2<f32>, Vec<Array1<f32>>) {
|
||||
|
||||
let n_samples = inputs.len();
|
||||
let mut all_states = Vec::new();
|
||||
|
||||
// Reset reservoir
|
||||
self.state = Array1::zeros(self.reservoir_size);
|
||||
|
||||
for (i, input) in inputs.iter().enumerate() {
|
||||
let input_arr = Array1::from_vec(input.clone());
|
||||
let state = self.update_state(&input_arr);
|
||||
|
||||
if i >= washout {
|
||||
// Concatenate state with input (for direct connections)
|
||||
let mut extended = Vec::from(state.as_slice().unwrap());
|
||||
extended.extend_from_slice(input);
|
||||
all_states.push(Array1::from_vec(extended));
|
||||
}
|
||||
}
|
||||
|
||||
// Convert to matrix for ridge regression
|
||||
let n_train = all_states.len();
|
||||
let state_dim = self.reservoir_size + inputs[0].len();
|
||||
let mut state_matrix = Array2::zeros((n_train, state_dim));
|
||||
|
||||
for (i, state) in all_states.iter().enumerate() {
|
||||
state_matrix.row_mut(i).assign(state);
|
||||
}
|
||||
|
||||
(state_matrix, all_states)
|
||||
}
|
||||
|
||||
/// Train output weights using ridge regression
|
||||
pub fn train_ridge(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>,
|
||||
regularization: f32) {
|
||||
let washout = 100.min(x.len() / 10);
|
||||
let (states, _) = self.collect_states(x, washout);
|
||||
|
||||
// Prepare targets (skip washout)
|
||||
let targets = &y[washout..];
|
||||
|
||||
// Ridge regression: W_out = Y * X^T * (X * X^T + λI)^-1
|
||||
let lambda = regularization;
|
||||
let n = states.nrows();
|
||||
let d = states.ncols();
|
||||
|
||||
// X * X^T
|
||||
let gram = states.t().dot(&states);
|
||||
|
||||
// Add regularization
|
||||
let mut reg_gram = gram + Array2::<f32>::eye(d) * lambda;
|
||||
|
||||
// Solve using pseudo-inverse (in practice, use better solver)
|
||||
let y_mat = Array2::from_shape_vec((1, targets.len()), targets.to_vec())
|
||||
.expect("Shape mismatch");
|
||||
|
||||
// Simple solution (production would use LAPACK)
|
||||
// w_out should be (output_dim x extended_dim) = (1 x (reservoir_size + input_dim))
|
||||
// y_mat is (1 x n), states is (n x d), so compute Y * pinv(X^T)
|
||||
let pinv = Self::simple_pinv(®_gram); // (d x d)
|
||||
let temp = states.dot(&pinv); // (n x d)
|
||||
self.w_out = y_mat.dot(&temp); // (1 x d)
|
||||
}
|
||||
|
||||
/// Simple pseudo-inverse (production code would use LAPACK)
|
||||
fn simple_pinv(matrix: &Array2<f32>) -> Array2<f32> {
|
||||
// Simplified: just add strong regularization for stability
|
||||
let n = matrix.nrows();
|
||||
let reg = Array2::<f32>::eye(n) * 0.01;
|
||||
let stabilized = matrix + ®
|
||||
|
||||
// Return stabilized inverse approximation
|
||||
// In production, use proper SVD or QR decomposition
|
||||
stabilized.mapv(|x| 1.0 / (x + 0.001))
|
||||
}
|
||||
|
||||
/// Predict using trained reservoir
|
||||
pub fn predict(&mut self, x: &[Vec<f32>]) -> Vec<f32> {
|
||||
// Reset state for prediction
|
||||
self.state = Array1::zeros(self.reservoir_size);
|
||||
|
||||
let mut predictions = Vec::new();
|
||||
|
||||
for input in x {
|
||||
let input_arr = Array1::from_vec(input.clone());
|
||||
let state = self.update_state(&input_arr);
|
||||
|
||||
// Extended state (reservoir + input)
|
||||
let mut extended = Vec::from(state.as_slice().unwrap());
|
||||
extended.extend_from_slice(input);
|
||||
let extended_arr = Array1::from_vec(extended);
|
||||
|
||||
// Linear readout
|
||||
let output = self.w_out.dot(&extended_arr);
|
||||
predictions.push(output[0]);
|
||||
}
|
||||
|
||||
predictions
|
||||
}
|
||||
|
||||
/// Classify using reservoir
|
||||
pub fn predict_class(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
let outputs = self.predict(x);
|
||||
|
||||
outputs.iter().map(|&val| {
|
||||
if val < -0.25 { 0 }
|
||||
else if val > 0.25 { 2 }
|
||||
else { 1 }
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Quantum Reservoir Computing - Exploit quantum superposition
|
||||
pub struct QuantumReservoir {
|
||||
pub classical_reservoir: ReservoirComputer,
|
||||
quantum_layer_size: usize,
|
||||
phase_matrix: Array2<f32>,
|
||||
entanglement_strength: f32,
|
||||
}
|
||||
|
||||
impl QuantumReservoir {
|
||||
pub fn new(input_dim: usize, reservoir_size: usize, output_dim: usize) -> Self {
|
||||
let classical_reservoir = ReservoirComputer::new(input_dim, reservoir_size, output_dim);
|
||||
let quantum_layer_size = 16; // Small quantum layer
|
||||
|
||||
// Random phase matrix for quantum interference
|
||||
let mut rng = thread_rng();
|
||||
let phase_matrix = Array2::from_shape_fn((quantum_layer_size, quantum_layer_size),
|
||||
|_| rng.gen::<f32>() * 2.0 * PI);
|
||||
|
||||
Self {
|
||||
classical_reservoir,
|
||||
quantum_layer_size,
|
||||
phase_matrix,
|
||||
entanglement_strength: 0.5,
|
||||
}
|
||||
}
|
||||
|
||||
/// Simulate quantum interference patterns
|
||||
fn quantum_transform(&self, state: &Array1<f32>) -> Array1<f32> {
|
||||
let n = self.quantum_layer_size.min(state.len());
|
||||
let mut quantum_state = Array1::zeros(n);
|
||||
|
||||
// Take first n elements
|
||||
for i in 0..n {
|
||||
quantum_state[i] = state[i % state.len()];
|
||||
}
|
||||
|
||||
// Apply phase rotations (simulated quantum gates)
|
||||
let mut result = Array1::<f32>::zeros(n);
|
||||
for i in 0..n {
|
||||
for j in 0..n {
|
||||
let phase = self.phase_matrix[[i, j]];
|
||||
let amplitude = quantum_state[j] * phase.cos()
|
||||
+ quantum_state[(j + 1) % n] * phase.sin();
|
||||
result[i] += amplitude * self.entanglement_strength;
|
||||
}
|
||||
}
|
||||
|
||||
// Normalize (quantum measurement)
|
||||
let norm = result.dot(&result).sqrt();
|
||||
if norm > 0.0 {
|
||||
result = result / norm;
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
pub fn train(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>) {
|
||||
// First pass through classical reservoir
|
||||
self.classical_reservoir.train_ridge(x, y, 0.001);
|
||||
|
||||
// Enhance with quantum layer (in practice, would train quantum parameters)
|
||||
}
|
||||
|
||||
pub fn predict_quantum(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
// Classical prediction
|
||||
let classical_pred = self.classical_reservoir.predict_class(x);
|
||||
|
||||
// Quantum enhancement (simplified)
|
||||
classical_pred.iter().enumerate().map(|(i, &pred)| {
|
||||
// Add quantum fluctuation based on position
|
||||
let quantum_factor = (i as f32 * 0.1).sin();
|
||||
if quantum_factor > 0.3 && pred < 2 {
|
||||
pred + 1
|
||||
} else if quantum_factor < -0.3 && pred > 0 {
|
||||
pred - 1
|
||||
} else {
|
||||
pred
|
||||
}
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
#[cfg(feature = "ruv-fann")]
|
||||
pub mod ruv_fann_backend {
|
||||
// Implement this against ruv-fann's API.
|
||||
// Expected surface:
|
||||
// - new(input, hidden, output) -> Self
|
||||
// - train_regression(x, y, epochs, lr)
|
||||
// - predict_reg(x) -> Vec<f32>
|
||||
// - predict_cls3(x) -> Vec<usize>
|
||||
// Keep identical signatures to mlp.rs so the runner can swap backends.
|
||||
// Example skeleton:
|
||||
pub struct RuvFannModel { /* fields from ruv-fann */ }
|
||||
|
||||
impl RuvFannModel {
|
||||
pub fn new(_input: usize, _hidden: usize, _output: usize) -> Self { Self { } }
|
||||
pub fn train_regression(&mut self, _x: &Vec<Vec<f32>>, _y: &Vec<f32>, _epochs: usize, _lr: f32) { /* ... */ }
|
||||
pub fn predict_reg(&self, _x: &[Vec<f32>]) -> Vec<f32> { vec![] }
|
||||
pub fn predict_cls3(&self, _x: &[Vec<f32>]) -> Vec<usize> { vec![] }
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "ruv-fann"))]
|
||||
pub mod ruv_fann_backend {
|
||||
// No-op stub so the crate compiles without the feature.
|
||||
pub struct RuvFannModel;
|
||||
impl RuvFannModel {
|
||||
pub fn new(_: usize, _: usize, _: usize) -> Self { Self }
|
||||
pub fn train_regression(&mut self, _: &Vec<Vec<f32>>, _: &Vec<f32>, _: usize, _: f32) {}
|
||||
pub fn predict_reg(&self, _: &[Vec<f32>]) -> Vec<f32> { Vec::new() }
|
||||
pub fn predict_cls3(&self, _: &[Vec<f32>]) -> Vec<usize> { Vec::new() }
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
#[cfg(feature = "ruv-fann")]
|
||||
use ruv_fann_dep::{Network, Activation};
|
||||
|
||||
#[cfg(feature = "ruv-fann")]
|
||||
pub struct RuvFannModel {
|
||||
network: Network,
|
||||
input_dim: usize,
|
||||
output_dim: usize,
|
||||
}
|
||||
|
||||
#[cfg(feature = "ruv-fann")]
|
||||
impl RuvFannModel {
|
||||
pub fn new(input: usize, hidden: usize, output: usize) -> Self {
|
||||
// Create a 3-layer network: input -> hidden -> output
|
||||
let network = Network::new(&[input, hidden, output])
|
||||
.with_activation(Activation::Sigmoid)
|
||||
.with_learning_rate(0.7);
|
||||
|
||||
Self {
|
||||
network,
|
||||
input_dim: input,
|
||||
output_dim: output,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn train_regression(&mut self, x: &Vec<Vec<f32>>, y: &Vec<f32>, epochs: usize, lr: f32) {
|
||||
// Get mutable access to network
|
||||
let network = Arc::get_mut(&mut self.network)
|
||||
.expect("Cannot get mutable reference to network");
|
||||
|
||||
network.set_learning_rate(lr);
|
||||
|
||||
// Prepare training data in FANN format
|
||||
let mut inputs: Vec<Vec<f32>> = Vec::new();
|
||||
let mut outputs: Vec<Vec<f32>> = Vec::new();
|
||||
|
||||
for (xi, yi) in x.iter().zip(y) {
|
||||
inputs.push(xi.clone());
|
||||
if self.output_dim == 1 {
|
||||
outputs.push(vec![*yi]);
|
||||
} else {
|
||||
// For classification, create one-hot encoding
|
||||
let class = if *yi < -0.25 { 0 }
|
||||
else if *yi > 0.25 { 2 }
|
||||
else { 1 };
|
||||
let mut one_hot = vec![0.0; self.output_dim];
|
||||
if class < self.output_dim {
|
||||
one_hot[class] = 1.0;
|
||||
}
|
||||
outputs.push(one_hot);
|
||||
}
|
||||
}
|
||||
|
||||
// Train using batch training
|
||||
for _ in 0..epochs {
|
||||
for (input, output) in inputs.iter().zip(&outputs) {
|
||||
network.train(input, output);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn predict_reg(&self, x: &[Vec<f32>]) -> Vec<f32> {
|
||||
x.iter().map(|xi| {
|
||||
let output = self.network.run(xi).expect("Failed to run network");
|
||||
if self.output_dim == 1 {
|
||||
output[0]
|
||||
} else {
|
||||
// Return first output for regression
|
||||
output[0]
|
||||
}
|
||||
}).collect()
|
||||
}
|
||||
|
||||
pub fn predict_cls3(&self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
x.iter().map(|xi| {
|
||||
let output = self.network.run(xi).expect("Failed to run network");
|
||||
|
||||
if self.output_dim >= 3 {
|
||||
// Find argmax for classification
|
||||
let mut best = 0;
|
||||
let mut best_val = output[0];
|
||||
for i in 1..3.min(output.len()) {
|
||||
if output[i] > best_val {
|
||||
best_val = output[i];
|
||||
best = i;
|
||||
}
|
||||
}
|
||||
best
|
||||
} else {
|
||||
// Threshold-based classification for single output
|
||||
let val = output[0];
|
||||
if val < -0.25 { 0 }
|
||||
else if val > 0.25 { 2 }
|
||||
else { 1 }
|
||||
}
|
||||
}).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "ruv-fann"))]
|
||||
pub struct RuvFannModel;
|
||||
|
||||
#[cfg(not(feature = "ruv-fann"))]
|
||||
impl RuvFannModel {
|
||||
pub fn new(_: usize, _: usize, _: usize) -> Self { Self }
|
||||
pub fn train_regression(&mut self, _: &Vec<Vec<f32>>, _: &Vec<f32>, _: usize, _: f32) {}
|
||||
pub fn predict_reg(&self, _: &[Vec<f32>]) -> Vec<f32> { Vec::new() }
|
||||
pub fn predict_cls3(&self, _: &[Vec<f32>]) -> Vec<usize> { Vec::new() }
|
||||
}
|
||||
@@ -0,0 +1,397 @@
|
||||
use ndarray::{Array1, Array2};
|
||||
use rand::{thread_rng, Rng};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Sparse Neural Network with dynamic sparsity
|
||||
/// Only keeps top-k% connections active
|
||||
pub struct SparseNetwork {
|
||||
// Sparse weight storage
|
||||
w1_indices: Vec<(usize, usize)>, // Active connections
|
||||
w1_values: Vec<f32>, // Weight values
|
||||
w2_indices: Vec<(usize, usize)>,
|
||||
w2_values: Vec<f32>,
|
||||
|
||||
// Biases (dense)
|
||||
b1: Array1<f32>,
|
||||
b2: Array1<f32>,
|
||||
|
||||
// Dimensions
|
||||
input_dim: usize,
|
||||
hidden_dim: usize,
|
||||
output_dim: usize,
|
||||
sparsity: f32, // Fraction of weights to keep
|
||||
|
||||
// Hidden activations
|
||||
hidden: Array1<f32>,
|
||||
|
||||
// Pruning statistics
|
||||
pruning_threshold: f32,
|
||||
pruned_count: usize,
|
||||
}
|
||||
|
||||
impl SparseNetwork {
|
||||
pub fn new(input: usize, hidden: usize, output: usize, sparsity: f32) -> Self {
|
||||
let mut rng = thread_rng();
|
||||
|
||||
// Initialize with random sparse connections
|
||||
let n_connections1 = ((input * hidden) as f32 * sparsity) as usize;
|
||||
let n_connections2 = ((hidden * output) as f32 * sparsity) as usize;
|
||||
|
||||
let mut w1_indices = Vec::new();
|
||||
let mut w1_values = Vec::new();
|
||||
|
||||
// Random sparse initialization for layer 1
|
||||
let scale1 = (2.0 / input as f32).sqrt();
|
||||
let mut used1 = std::collections::HashSet::new();
|
||||
|
||||
while w1_indices.len() < n_connections1 {
|
||||
let i = rng.gen_range(0..hidden);
|
||||
let j = rng.gen_range(0..input);
|
||||
if used1.insert((i, j)) {
|
||||
w1_indices.push((i, j));
|
||||
w1_values.push(rng.gen::<f32>() * scale1 - scale1/2.0);
|
||||
}
|
||||
}
|
||||
|
||||
// Layer 2
|
||||
let mut w2_indices = Vec::new();
|
||||
let mut w2_values = Vec::new();
|
||||
let scale2 = (2.0 / hidden as f32).sqrt();
|
||||
let mut used2 = std::collections::HashSet::new();
|
||||
|
||||
while w2_indices.len() < n_connections2 {
|
||||
let i = rng.gen_range(0..output);
|
||||
let j = rng.gen_range(0..hidden);
|
||||
if used2.insert((i, j)) {
|
||||
w2_indices.push((i, j));
|
||||
w2_values.push(rng.gen::<f32>() * scale2 - scale2/2.0);
|
||||
}
|
||||
}
|
||||
|
||||
Self {
|
||||
w1_indices,
|
||||
w1_values,
|
||||
w2_indices,
|
||||
w2_values,
|
||||
b1: Array1::zeros(hidden),
|
||||
b2: Array1::zeros(output),
|
||||
input_dim: input,
|
||||
hidden_dim: hidden,
|
||||
output_dim: output,
|
||||
sparsity,
|
||||
hidden: Array1::zeros(hidden),
|
||||
pruning_threshold: 0.01,
|
||||
pruned_count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Sparse forward pass
|
||||
pub fn forward(&mut self, x: &[f32]) -> Vec<f32> {
|
||||
// Reset hidden
|
||||
self.hidden.fill(0.0);
|
||||
|
||||
// Sparse matrix-vector multiplication for layer 1
|
||||
for ((&(i, j), &val)) in self.w1_indices.iter().zip(&self.w1_values) {
|
||||
if j < x.len() {
|
||||
self.hidden[i] += x[j] * val;
|
||||
}
|
||||
}
|
||||
|
||||
// Add bias and apply ReLU
|
||||
self.hidden = &self.hidden + &self.b1;
|
||||
self.hidden.mapv_inplace(|x| x.max(0.0));
|
||||
|
||||
// Layer 2
|
||||
let mut output = vec![0.0; self.output_dim];
|
||||
for ((&(i, j), &val)) in self.w2_indices.iter().zip(&self.w2_values) {
|
||||
output[i] += self.hidden[j] * val;
|
||||
}
|
||||
|
||||
// Add bias
|
||||
for i in 0..self.output_dim {
|
||||
output[i] += self.b2[i];
|
||||
}
|
||||
|
||||
output
|
||||
}
|
||||
|
||||
/// Dynamic pruning - remove small weights
|
||||
pub fn prune_weights(&mut self, threshold: f32) {
|
||||
// Prune layer 1
|
||||
let mut new_indices1 = Vec::new();
|
||||
let mut new_values1 = Vec::new();
|
||||
|
||||
for (idx, val) in self.w1_indices.iter().zip(&self.w1_values) {
|
||||
if val.abs() > threshold {
|
||||
new_indices1.push(*idx);
|
||||
new_values1.push(*val);
|
||||
}
|
||||
}
|
||||
|
||||
let pruned1 = self.w1_indices.len() - new_indices1.len();
|
||||
self.w1_indices = new_indices1;
|
||||
self.w1_values = new_values1;
|
||||
|
||||
// Prune layer 2
|
||||
let mut new_indices2 = Vec::new();
|
||||
let mut new_values2 = Vec::new();
|
||||
|
||||
for (idx, val) in self.w2_indices.iter().zip(&self.w2_values) {
|
||||
if val.abs() > threshold {
|
||||
new_indices2.push(*idx);
|
||||
new_values2.push(*val);
|
||||
}
|
||||
}
|
||||
|
||||
let pruned2 = self.w2_indices.len() - new_indices2.len();
|
||||
self.w2_indices = new_indices2;
|
||||
self.w2_values = new_values2;
|
||||
|
||||
self.pruned_count += pruned1 + pruned2;
|
||||
self.pruning_threshold = threshold;
|
||||
}
|
||||
|
||||
/// Regrow connections (lottery ticket hypothesis)
|
||||
pub fn regrow_connections(&mut self, n_regrow: usize) {
|
||||
let mut rng = thread_rng();
|
||||
|
||||
// Regrow layer 1
|
||||
let mut used1: std::collections::HashSet<_> = self.w1_indices.iter().cloned().collect();
|
||||
let scale1 = (2.0 / self.input_dim as f32).sqrt();
|
||||
|
||||
for _ in 0..n_regrow/2 {
|
||||
let i = rng.gen_range(0..self.hidden_dim);
|
||||
let j = rng.gen_range(0..self.input_dim);
|
||||
if used1.insert((i, j)) {
|
||||
self.w1_indices.push((i, j));
|
||||
self.w1_values.push(rng.gen::<f32>() * scale1 - scale1/2.0);
|
||||
}
|
||||
}
|
||||
|
||||
// Regrow layer 2
|
||||
let mut used2: std::collections::HashSet<_> = self.w2_indices.iter().cloned().collect();
|
||||
let scale2 = (2.0 / self.hidden_dim as f32).sqrt();
|
||||
|
||||
for _ in 0..n_regrow/2 {
|
||||
let i = rng.gen_range(0..self.output_dim);
|
||||
let j = rng.gen_range(0..self.hidden_dim);
|
||||
if used2.insert((i, j)) {
|
||||
self.w2_indices.push((i, j));
|
||||
self.w2_values.push(rng.gen::<f32>() * scale2 - scale2/2.0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Train with sparse gradient updates
|
||||
pub fn train(&mut self, x: &[Vec<f32>], y: &[f32], epochs: usize, lr: f32) {
|
||||
for epoch in 0..epochs {
|
||||
let mut total_loss = 0.0;
|
||||
|
||||
for (xi, &yi) in x.iter().zip(y.iter()) {
|
||||
let output = self.forward(xi);
|
||||
let pred = if self.output_dim == 1 { output[0] } else { output[0] };
|
||||
let error = pred - yi;
|
||||
total_loss += error * error;
|
||||
|
||||
// Sparse backpropagation
|
||||
self.sparse_backward(xi, error, lr);
|
||||
}
|
||||
|
||||
// Dynamic sparsity - prune and regrow periodically
|
||||
if epoch > 0 && epoch % 10 == 0 {
|
||||
let old_count = self.w1_indices.len() + self.w2_indices.len();
|
||||
self.prune_weights(self.pruning_threshold);
|
||||
let pruned = old_count - (self.w1_indices.len() + self.w2_indices.len());
|
||||
if pruned > 0 {
|
||||
self.regrow_connections(pruned);
|
||||
}
|
||||
}
|
||||
|
||||
if epoch % 100 == 0 {
|
||||
println!("Sparse epoch {}: loss={:.6}, active_weights={}",
|
||||
epoch, total_loss / x.len() as f32,
|
||||
self.w1_indices.len() + self.w2_indices.len());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn sparse_backward(&mut self, x: &[f32], error: f32, lr: f32) {
|
||||
// Output gradient
|
||||
let grad_out = error;
|
||||
|
||||
// Update layer 2 weights (sparse)
|
||||
let mut grad_hidden = Array1::zeros(self.hidden_dim);
|
||||
|
||||
for i in 0..self.w2_indices.len() {
|
||||
let (out_idx, hid_idx) = self.w2_indices[i];
|
||||
if out_idx == 0 { // For single output
|
||||
// Update weight
|
||||
self.w2_values[i] -= lr * grad_out * self.hidden[hid_idx];
|
||||
// Accumulate gradient for hidden layer
|
||||
grad_hidden[hid_idx] += self.w2_values[i] * grad_out;
|
||||
}
|
||||
}
|
||||
|
||||
// Update bias
|
||||
self.b2[0] -= lr * grad_out;
|
||||
|
||||
// Apply ReLU gradient
|
||||
grad_hidden.mapv_inplace(|g| if self.hidden[0] > 0.0 { g } else { 0.0 });
|
||||
|
||||
// Update layer 1 weights (sparse)
|
||||
for i in 0..self.w1_indices.len() {
|
||||
let (hid_idx, in_idx) = self.w1_indices[i];
|
||||
if in_idx < x.len() {
|
||||
self.w1_values[i] -= lr * grad_hidden[hid_idx] * x[in_idx];
|
||||
}
|
||||
}
|
||||
|
||||
// Update hidden bias
|
||||
self.b1 = &self.b1 - &grad_hidden * lr;
|
||||
}
|
||||
|
||||
pub fn predict(&mut self, x: &[Vec<f32>]) -> Vec<f32> {
|
||||
x.iter().map(|xi| {
|
||||
let output = self.forward(xi);
|
||||
if self.output_dim == 1 { output[0] } else { output[0] }
|
||||
}).collect()
|
||||
}
|
||||
|
||||
pub fn predict_class(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
self.predict(x).iter().map(|&y| {
|
||||
if y < -0.25 { 0 }
|
||||
else if y > 0.25 { 2 }
|
||||
else { 1 }
|
||||
}).collect()
|
||||
}
|
||||
|
||||
pub fn get_sparsity_stats(&self) -> (usize, usize, f32) {
|
||||
let active = self.w1_indices.len() + self.w2_indices.len();
|
||||
let total = self.input_dim * self.hidden_dim + self.hidden_dim * self.output_dim;
|
||||
let sparsity = active as f32 / total as f32;
|
||||
(active, self.pruned_count, sparsity)
|
||||
}
|
||||
}
|
||||
|
||||
/// Lottery Ticket Network - finds winning sparse subnetworks
|
||||
pub struct LotteryTicketNetwork {
|
||||
base_network: SparseNetwork,
|
||||
initial_weights1: HashMap<(usize, usize), f32>,
|
||||
initial_weights2: HashMap<(usize, usize), f32>,
|
||||
winning_mask1: HashMap<(usize, usize), bool>,
|
||||
winning_mask2: HashMap<(usize, usize), bool>,
|
||||
iteration: usize,
|
||||
}
|
||||
|
||||
impl LotteryTicketNetwork {
|
||||
pub fn new(input: usize, hidden: usize, output: usize) -> Self {
|
||||
let base = SparseNetwork::new(input, hidden, output, 1.0); // Start dense
|
||||
|
||||
// Store initial weights
|
||||
let mut initial_weights1 = HashMap::new();
|
||||
for (&idx, &val) in base.w1_indices.iter().zip(&base.w1_values) {
|
||||
initial_weights1.insert(idx, val);
|
||||
}
|
||||
|
||||
let mut initial_weights2 = HashMap::new();
|
||||
for (&idx, &val) in base.w2_indices.iter().zip(&base.w2_values) {
|
||||
initial_weights2.insert(idx, val);
|
||||
}
|
||||
|
||||
Self {
|
||||
base_network: base,
|
||||
initial_weights1,
|
||||
initial_weights2,
|
||||
winning_mask1: HashMap::new(),
|
||||
winning_mask2: HashMap::new(),
|
||||
iteration: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Iterative magnitude pruning to find winning tickets
|
||||
pub fn find_winning_ticket(&mut self, x: &[Vec<f32>], y: &[f32],
|
||||
prune_rate: f32, iterations: usize) {
|
||||
for iter in 0..iterations {
|
||||
println!("Lottery iteration {}/{}", iter + 1, iterations);
|
||||
|
||||
// Reset to initial weights with current mask
|
||||
self.reset_to_initial();
|
||||
|
||||
// Train
|
||||
self.base_network.train(x, y, 100, 0.01);
|
||||
|
||||
// Prune based on magnitude
|
||||
self.magnitude_prune(prune_rate);
|
||||
|
||||
self.iteration += 1;
|
||||
}
|
||||
}
|
||||
|
||||
fn reset_to_initial(&mut self) {
|
||||
// Reset to initial weights, keeping only winning connections
|
||||
let mut new_indices1 = Vec::new();
|
||||
let mut new_values1 = Vec::new();
|
||||
|
||||
for (&idx, &init_val) in &self.initial_weights1 {
|
||||
if self.iteration == 0 || self.winning_mask1.get(&idx) == Some(&true) {
|
||||
new_indices1.push(idx);
|
||||
new_values1.push(init_val);
|
||||
}
|
||||
}
|
||||
|
||||
self.base_network.w1_indices = new_indices1;
|
||||
self.base_network.w1_values = new_values1;
|
||||
|
||||
// Layer 2
|
||||
let mut new_indices2 = Vec::new();
|
||||
let mut new_values2 = Vec::new();
|
||||
|
||||
for (&idx, &init_val) in &self.initial_weights2 {
|
||||
if self.iteration == 0 || self.winning_mask2.get(&idx) == Some(&true) {
|
||||
new_indices2.push(idx);
|
||||
new_values2.push(init_val);
|
||||
}
|
||||
}
|
||||
|
||||
self.base_network.w2_indices = new_indices2;
|
||||
self.base_network.w2_values = new_values2;
|
||||
}
|
||||
|
||||
fn magnitude_prune(&mut self, prune_rate: f32) {
|
||||
// Collect all weight magnitudes
|
||||
let mut magnitudes: Vec<f32> = self.base_network.w1_values.iter()
|
||||
.chain(&self.base_network.w2_values)
|
||||
.map(|v| v.abs())
|
||||
.collect();
|
||||
|
||||
magnitudes.sort_by(|a, b| a.partial_cmp(b).unwrap());
|
||||
|
||||
let cutoff_idx = (magnitudes.len() as f32 * prune_rate) as usize;
|
||||
let threshold = if cutoff_idx < magnitudes.len() {
|
||||
magnitudes[cutoff_idx]
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
// Update masks
|
||||
self.winning_mask1.clear();
|
||||
for (&idx, &val) in self.base_network.w1_indices.iter()
|
||||
.zip(&self.base_network.w1_values) {
|
||||
self.winning_mask1.insert(idx, val.abs() > threshold);
|
||||
}
|
||||
|
||||
self.winning_mask2.clear();
|
||||
for (&idx, &val) in self.base_network.w2_indices.iter()
|
||||
.zip(&self.base_network.w2_values) {
|
||||
self.winning_mask2.insert(idx, val.abs() > threshold);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn predict(&mut self, x: &[Vec<f32>]) -> Vec<f32> {
|
||||
self.base_network.predict(x)
|
||||
}
|
||||
|
||||
pub fn predict_class(&mut self, x: &[Vec<f32>]) -> Vec<usize> {
|
||||
self.base_network.predict_class(x)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user