style: cargo fmt --workspace (whitespace/wrapping only, no semantic change)
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
This commit is contained in:
osobh
@@ -5,16 +5,16 @@
|
||||
//!
|
||||
//! Key features:
|
||||
//! - Cross-correlation matrix computation between twin representations
|
||||
//! - Redundancy reduction loss (making correlation matrix close to identity)
|
||||
//! - Redundancy reduction loss (making correlation matrix close to identity)
|
||||
//! - Projector network for embedding transformation
|
||||
//! - Support for different backbone architectures
|
||||
//! - Batch normalization in projector
|
||||
//! - Symmetric loss computation
|
||||
|
||||
use super::byol::Backbone; // Reuse the Backbone trait
|
||||
use crate::prelude::*;
|
||||
use super::byol::Backbone; // Reuse the Backbone trait
|
||||
use std::sync::Arc;
|
||||
use parking_lot::RwLock;
|
||||
use std::sync::Arc;
|
||||
|
||||
/// Barlow Twins configuration parameters
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -34,11 +34,11 @@ pub struct BarlowTwinsConfig {
|
||||
impl Default for BarlowTwinsConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
backbone_dim: 2048, // Will be overridden based on backbone
|
||||
backbone_dim: 2048, // Will be overridden based on backbone
|
||||
projector_dims: vec![8192, 8192, 8192],
|
||||
lambda_coeff: 0.005,
|
||||
batch_norm: true,
|
||||
scale_loss: 1.0 / 32.0, // For numerical stability
|
||||
scale_loss: 1.0 / 32.0, // For numerical stability
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -115,26 +115,26 @@ impl ProjectorNetwork {
|
||||
) -> Result<Self> {
|
||||
let mut layers = Vec::new();
|
||||
let mut batch_norms = Vec::new();
|
||||
|
||||
|
||||
let mut current_dim = input_dim;
|
||||
|
||||
|
||||
for (i, &output_dim) in layer_dims.iter().enumerate() {
|
||||
// Linear layer
|
||||
let weight = Tensor::randn(&[current_dim, output_dim], device)?;
|
||||
let bias = Tensor::zeros(&[output_dim], device)?;
|
||||
|
||||
|
||||
layers.push(LinearLayer {
|
||||
weight: Arc::new(RwLock::new(weight)),
|
||||
bias: Arc::new(RwLock::new(bias)),
|
||||
});
|
||||
|
||||
|
||||
// Batch normalization (except for the last layer)
|
||||
if use_batch_norm && i < layer_dims.len() - 1 {
|
||||
let bn_weight = Tensor::ones(&[output_dim], device)?;
|
||||
let bn_bias = Tensor::zeros(&[output_dim], device)?;
|
||||
let running_mean = Tensor::zeros(&[output_dim], device)?;
|
||||
let running_var = Tensor::ones(&[output_dim], device)?;
|
||||
|
||||
|
||||
batch_norms.push(Some(BatchNorm {
|
||||
weight: Arc::new(RwLock::new(bn_weight)),
|
||||
bias: Arc::new(RwLock::new(bn_bias)),
|
||||
@@ -146,10 +146,10 @@ impl ProjectorNetwork {
|
||||
} else {
|
||||
batch_norms.push(None);
|
||||
}
|
||||
|
||||
|
||||
current_dim = output_dim;
|
||||
}
|
||||
|
||||
|
||||
Ok(Self {
|
||||
layers,
|
||||
batch_norms,
|
||||
@@ -163,24 +163,24 @@ impl ProjectorNetwork {
|
||||
/// Forward pass through projector network
|
||||
pub fn forward(&self, input: &Tensor) -> Result<Tensor> {
|
||||
let mut x = input.clone();
|
||||
|
||||
|
||||
for (i, layer) in self.layers.iter().enumerate() {
|
||||
// Linear transformation
|
||||
let weight = layer.weight.read();
|
||||
let bias = layer.bias.read();
|
||||
x = x.matmul(&*weight)?.add(&*bias)?;
|
||||
|
||||
|
||||
// Batch normalization (if enabled and not last layer)
|
||||
if let Some(ref batch_norm) = self.batch_norms[i] {
|
||||
x = self.apply_batch_norm(&x, batch_norm)?;
|
||||
}
|
||||
|
||||
|
||||
// ReLU activation (except for the last layer)
|
||||
if i < self.layers.len() - 1 {
|
||||
x = x.relu()?;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Ok(x)
|
||||
}
|
||||
|
||||
@@ -190,16 +190,16 @@ impl ProjectorNetwork {
|
||||
let bias = batch_norm.bias.read();
|
||||
let running_mean = batch_norm.running_mean.read();
|
||||
let running_var = batch_norm.running_var.read();
|
||||
|
||||
|
||||
// Normalize: (x - mean) / sqrt(var + eps)
|
||||
let eps_tensor = Tensor::full(running_var.dims(), batch_norm.eps, &self.device)?;
|
||||
let var_eps = running_var.add(&eps_tensor)?;
|
||||
let std = var_eps.sqrt()?;
|
||||
|
||||
|
||||
let normalized = x.sub(&*running_mean)?.div(&std)?;
|
||||
let scaled = normalized.mul(&*weight)?;
|
||||
let output = scaled.add(&*bias)?;
|
||||
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
@@ -226,28 +226,28 @@ impl ProjectorNetwork {
|
||||
|
||||
/// Normalize features to have zero mean and unit variance per feature dimension
|
||||
pub fn normalize_features(features: &Tensor) -> Result<Tensor> {
|
||||
let mean = features.mean(&[0i32], false)?; // Mean across batch dimension
|
||||
let mean = features.mean(&[0i32], false)?; // Mean across batch dimension
|
||||
let centered = features.sub(&mean)?;
|
||||
|
||||
let variance = centered.pow_scalar(2.0)?.mean(&[0i32], false)?;
|
||||
let eps = Tensor::full(variance.dims(), 1e-8, variance.device())?;
|
||||
let std = variance.add(&eps)?.sqrt()?;
|
||||
|
||||
|
||||
Ok(centered.div(&std)?)
|
||||
}
|
||||
|
||||
/// Compute cross-correlation matrix between two normalized embeddings
|
||||
pub fn compute_cross_correlation_matrix(y1: &Tensor, y2: &Tensor) -> Result<Tensor> {
|
||||
let batch_size = y1.shape()[0] as f32;
|
||||
|
||||
|
||||
// Normalize features
|
||||
let y1_norm = normalize_features(y1)?;
|
||||
let y2_norm = normalize_features(y2)?;
|
||||
|
||||
|
||||
// Compute cross-correlation: C[i,j] = sum(y1_norm[i] * y2_norm[j]) / batch_size
|
||||
let y1_t = y1_norm.transpose(0, 1)?; // [feature_dim, batch_size]
|
||||
let cross_corr = y1_t.matmul(&y2_norm)?; // [feature_dim, feature_dim]
|
||||
|
||||
let y1_t = y1_norm.transpose(0, 1)?; // [feature_dim, batch_size]
|
||||
let cross_corr = y1_t.matmul(&y2_norm)?; // [feature_dim, feature_dim]
|
||||
|
||||
let batch_size_tensor = Tensor::full(&[], batch_size, cross_corr.device())?;
|
||||
Ok(cross_corr.div(&batch_size_tensor)?)
|
||||
}
|
||||
@@ -256,15 +256,15 @@ pub fn compute_cross_correlation_matrix(y1: &Tensor, y2: &Tensor) -> Result<Tens
|
||||
pub fn extract_diagonal(matrix: &Tensor) -> Result<Tensor> {
|
||||
let shape = matrix.shape();
|
||||
let size = shape[0];
|
||||
|
||||
|
||||
let mut diag_values = Vec::with_capacity(size);
|
||||
let matrix_data = matrix.to_vec()?;
|
||||
|
||||
|
||||
for i in 0..size {
|
||||
let idx = i * size + i; // Diagonal index in flattened matrix
|
||||
let idx = i * size + i; // Diagonal index in flattened matrix
|
||||
diag_values.push(matrix_data[idx]);
|
||||
}
|
||||
|
||||
|
||||
Ok(Tensor::from_vec(diag_values, &[size], matrix.device())?)
|
||||
}
|
||||
|
||||
@@ -291,37 +291,37 @@ pub fn compute_barlow_twins_loss(
|
||||
// Compute cross-correlation matrix
|
||||
let cross_corr = compute_cross_correlation_matrix(y1, y2)?;
|
||||
let feature_dim = cross_corr.shape()[0];
|
||||
|
||||
|
||||
// Create identity matrix
|
||||
let identity = Tensor::eye(feature_dim, cross_corr.device())?;
|
||||
|
||||
|
||||
// Invariance loss: sum((1 - C[i,i])^2) - diagonal should be 1
|
||||
let diag_diff = identity.sub(&cross_corr)?;
|
||||
let diag_squared = diag_diff.pow_scalar(2.0)?;
|
||||
|
||||
|
||||
// Extract diagonal elements for invariance loss
|
||||
let diagonal = extract_diagonal(&cross_corr)?;
|
||||
let ones = Tensor::ones(diagonal.dims(), diagonal.device())?;
|
||||
let diag_loss_vec = ones.sub(&diagonal)?.pow_scalar(2.0)?;
|
||||
let invariance_loss_tensor = diag_loss_vec.sum(None)?;
|
||||
let invariance_loss = invariance_loss_tensor.to_vec()?[0];
|
||||
|
||||
|
||||
// Redundancy reduction loss: sum(C[i,j]^2 for i≠j) - off-diagonal should be 0
|
||||
let cross_corr_squared = cross_corr.pow_scalar(2.0)?;
|
||||
let total_squared = cross_corr_squared.sum(None)?;
|
||||
let diag_squared_sum = diag_squared.sum(None)?;
|
||||
let redundancy_loss_tensor = total_squared.sub(&diag_squared_sum)?;
|
||||
let redundancy_loss = redundancy_loss_tensor.to_vec()?[0];
|
||||
|
||||
|
||||
// Total loss = invariance_loss + lambda * redundancy_loss
|
||||
let lambda_tensor = Tensor::full(&[], lambda_coeff, cross_corr.device())?;
|
||||
let weighted_redundancy = redundancy_loss_tensor.mul(&lambda_tensor)?;
|
||||
let total_loss = invariance_loss_tensor.add(&weighted_redundancy)?;
|
||||
|
||||
|
||||
// Apply scaling
|
||||
let scale_tensor = Tensor::full(&[], scale_loss, cross_corr.device())?;
|
||||
let scaled_loss = total_loss.mul(&scale_tensor)?;
|
||||
|
||||
|
||||
Ok(BarlowTwinsLossResult {
|
||||
loss: scaled_loss,
|
||||
invariance_loss,
|
||||
@@ -360,18 +360,18 @@ impl BarlowTwinsTrainer {
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
let backbone_dim = backbone.output_dim();
|
||||
|
||||
|
||||
// Update config with actual backbone dimension
|
||||
let mut config = config;
|
||||
config.backbone_dim = backbone_dim;
|
||||
|
||||
|
||||
let projector = ProjectorNetwork::new(
|
||||
backbone_dim,
|
||||
config.projector_dims.clone(),
|
||||
config.batch_norm,
|
||||
device,
|
||||
)?;
|
||||
|
||||
|
||||
Ok(Self {
|
||||
backbone,
|
||||
projector,
|
||||
@@ -382,28 +382,35 @@ impl BarlowTwinsTrainer {
|
||||
}
|
||||
|
||||
/// Perform one training step with two augmented views
|
||||
pub fn train_step(&mut self, images: &Tensor, _seed: Option<u64>) -> Result<BarlowTwinsTrainingResult> {
|
||||
pub fn train_step(
|
||||
&mut self,
|
||||
images: &Tensor,
|
||||
_seed: Option<u64>,
|
||||
) -> Result<BarlowTwinsTrainingResult> {
|
||||
if !self.training {
|
||||
return Err(TransformerError::InvalidInput("Trainer must be in training mode".to_string()));
|
||||
return Err(TransformerError::InvalidInput(
|
||||
"Trainer must be in training mode".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
|
||||
// For now, create two simple "augmented" views by adding noise
|
||||
// In practice, this would use proper augmentation pipeline
|
||||
let noise1 = Tensor::randn(images.dims(), &self.device)?.mul_scalar(0.01)?;
|
||||
let noise2 = Tensor::randn(images.dims(), &self.device)?.mul_scalar(0.01)?;
|
||||
let view1 = images.add(&noise1)?;
|
||||
let view2 = images.add(&noise2)?;
|
||||
|
||||
|
||||
// Forward pass through backbone and projector
|
||||
let z1 = self.backbone.forward(&view1)?;
|
||||
let z2 = self.backbone.forward(&view2)?;
|
||||
|
||||
|
||||
let y1 = self.projector.forward(&z1)?;
|
||||
let y2 = self.projector.forward(&z2)?;
|
||||
|
||||
|
||||
// Compute Barlow Twins loss
|
||||
let loss_result = compute_barlow_twins_loss(&y1, &y2, self.config.lambda_coeff, self.config.scale_loss)?;
|
||||
|
||||
let loss_result =
|
||||
compute_barlow_twins_loss(&y1, &y2, self.config.lambda_coeff, self.config.scale_loss)?;
|
||||
|
||||
Ok(BarlowTwinsTrainingResult {
|
||||
loss: loss_result.loss,
|
||||
invariance_loss: loss_result.invariance_loss,
|
||||
@@ -431,4 +438,4 @@ impl BarlowTwinsTrainer {
|
||||
pub fn config(&self) -> &BarlowTwinsConfig {
|
||||
&self.config
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user