diff --git a/.gitignore b/.gitignore
index c2cd311d..09aaf915 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,13 @@
+# Archive and draft directories
+_arxiv/
+drafts/
+
+# Backup files
+*.bak
+*.old
+*_old.*
+*_backup.*
+
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
diff --git a/README.md b/README.md
index a37dfd52..63e07062 100644
--- a/README.md
+++ b/README.md
@@ -2,23 +2,18 @@
Neural network alignment analysis and intelligent pruning framework.
----
-
## Overview
This framework provides tools for analyzing neural networks through information-theoretic metrics and performing redundancy-aware pruning and quantization.
-**Key capabilities:**
-- Alignment metrics (Rayleigh Quotient, class-conditioned RQ, mutual information)
-- Information-theoretic analysis (redundancy, synergy, PID)
-- Pruning strategies (16 methods including magnitude, gradient-based, redundancy-aware)
-- Quantization (INT8, INT4, mixed-precision with alignment-guided bit allocation)
-- Architecture support (MLPs, CNNs, Transformers, LLMs)
-- Data loading (vision and text datasets)
-- Evaluation (classification accuracy, perplexity for language models)
-- Visualization (plots and reports)
-
----
+Key capabilities:
+- Alignment metrics for analyzing neural network weight-input relationships
+- Information-theoretic analysis tools
+- Pruning strategies with multiple scoring methods
+- Quantization with alignment-guided precision selection
+- Architecture support for MLPs, CNNs, Transformers, and LLMs
+- Data loading for vision and text datasets
+- Evaluation and visualization tools
## Installation
@@ -30,79 +25,40 @@ conda activate alignment
pip install -e .
```
-Details: [docs/installation.md](docs/installation.md)
-
----
+See [docs/installation.md](docs/installation.md) for details.
## Usage
-### Command Line
+Run experiments using configuration files:
```bash
python scripts/run_experiment.py --config configs/examples/mnist_basic.yaml
-python scripts/run_experiment.py --config configs/examples/resnet_pruning.yaml
-python scripts/run_experiment.py --config configs/examples/llama3_scoring.yaml
-python scripts/run_experiment.py --config configs/examples/llama3_quantization.yaml
-```
-
-### Python API
-
-```python
-from alignment import ModelWrapper, get_metric
-
-wrapper = ModelWrapper(model)
-rq = get_metric('rayleigh_quotient')
-
-outputs, acts = wrapper.forward_with_activations(inputs)
-weights = wrapper.get_layer_weights()
-
-scores = rq.compute(acts['layer_input'], weights['layer'])
```
-### Quantization
-
-```python
-from alignment.quantization import quantize_model, find_optimal_bit_allocation
+The framework supports:
+- Training networks from scratch or loading pre-trained models
+- Computing alignment and information-theoretic scores
+- Applying pruning with different strategies and distributions
+- Quantization with various precision settings
-# INT8 quantization
-results = quantize_model(model, precision='int8')
+Example configurations are available in `configs/examples/`:
+- `mnist_basic.yaml` - Basic alignment analysis on MNIST
+- `resnet_pruning.yaml` - Pruning ResNet on CIFAR-10
+- `llama3_scoring.yaml` - Computing scores for LLaMA models
+- `llama3_pruning.yaml` - Pruning transformer networks
-# Mixed precision using alignment scores
-layer_importance = {layer: rq_scores[layer].mean() for layer in layers}
-bit_allocation = find_optimal_bit_allocation(model, layer_importance, target_avg_bits=6.0)
-```
-
----
+See `configs/template.yaml` for all available configuration options.
## Documentation
-- [Installation](docs/installation.md) - Setup and dependencies
-- [Usage Guide](docs/usage.md) - Running experiments
-- [User Guide](docs/user_guide.md) - Complete documentation
-- [API Reference](docs/api_reference.md) - API details
+- [Usage Guide](docs/usage.md) - Running experiments with configs
+- [User Guide](docs/user_guide.md) - Detailed framework documentation
+- [API Reference](docs/api_reference.md) - API documentation
- [Quick Reference](docs/quick_reference.md) - Code examples
-- [Changelog](docs/changelog.md) - Version history
-
-Full documentation: [docs/README.md](docs/README.md)
-
----
-
-## Configuration
-Experiments are configured via YAML files. See `configs/template.yaml` for all parameters.
+## Examples
-Available examples in `configs/examples/`:
-- mnist_basic.yaml
-- resnet_pruning.yaml
-- llama3_scoring.yaml
-- llama3_pruning.yaml
-- llama3_quantization.yaml
-
----
-
-## Code Examples
-
-Python examples in `examples/` directory:
+Python scripts demonstrating framework capabilities:
```bash
python examples/07_mnist_intelligent_pruning.py
@@ -110,7 +66,7 @@ python examples/08_llama_ffn_pruning.py
python examples/09_attention_neuron_vs_head_pruning.py
```
----
+
## Testing
@@ -118,8 +74,6 @@ python examples/09_attention_neuron_vs_head_pruning.py
pytest tests/
```
----
-
## License
See LICENSE file.
diff --git a/_arxiv/_archive/alignment_experiments.py.bak b/_arxiv/_archive/alignment_experiments.py.bak
deleted file mode 100644
index 0c510349..00000000
--- a/_arxiv/_archive/alignment_experiments.py.bak
+++ /dev/null
@@ -1,1288 +0,0 @@
-"""
-Alignment experiment implementations.
-
-This module contains experiment classes for neural network alignment studies,
-focusing on alignment metrics, dropout impacts, and training analysis.
-"""
-
-import logging
-import os
-import sys
-import argparse
-from typing import Dict, List, Tuple, Optional, Any, Union
-import copy
-import pickle
-import datetime
-import json
-import yaml
-
-import numpy as np
-import torch
-import torch.nn as nn
-from tqdm import tqdm
-
-from alignment.config import ExperimentConfig
-from alignment.experiments.experiment import Experiment
-from alignment.metrics import AlignmentMetric, get_metric
-from alignment.models.registry import create_model
-from alignment.dropout import progressive_dropout, eigenvector_dropout
-from alignment.training import train_model, evaluate_model
-from alignment.utils.core import setup_logging
-from alignment.utils.plotting import plot_dropout_results, plot_experiment_summary
-from alignment.datasets import get_dataset, load_dataset
-
-logger = logging.getLogger(__name__)
-
-
-class AlignmentExperiment(Experiment):
- """
- Experiment class for studying neural network alignment properties.
-
- This class implements experiments that assess alignment between layers
- in neural networks, with support for different dropout strategies,
- multiple metrics, and visualization.
- """
-
- def __init__(self, config: Dict) -> None:
- """Initialize the experiment with the given config.
-
- Args:
- config: Experiment configuration.
- """
- super().__init__(config)
-
- # Make sure dataset_config has required fields to prevent loading errors
- if isinstance(self.config.dataset, dict):
- dataset_config = self.config.dataset
-
- # Ensure consistent naming between dataset.name and dataset.dataset_name
- if "name" in dataset_config and "dataset_name" not in dataset_config:
- dataset_config["dataset_name"] = dataset_config["name"]
- elif "dataset_name" in dataset_config and "name" not in dataset_config:
- dataset_config["name"] = dataset_config["dataset_name"]
- elif "name" not in dataset_config and "dataset_name" not in dataset_config:
- if isinstance(self.config.dataset, str):
- # If config.dataset is a string, use it as the dataset name
- dataset_config["name"] = self.config.dataset
- dataset_config["dataset_name"] = self.config.dataset
- logger.info(f"Using dataset name from string: {self.config.dataset}")
- else:
- # Default to CIFAR10 if no dataset name provided
- logger.warning("No dataset name provided in configuration. Using default: cifar10")
- dataset_config["name"] = "cifar10"
- dataset_config["dataset_name"] = "cifar10"
-
- logger.info(f"Using dataset: {dataset_config.get('name', 'unknown')}")
-
- # Set up experiment-specific paths
- self.figure_path = None
- self.weights_path = None
- self.setup_paths()
-
- logger.debug(f"Initialized alignment experiment")
-
- # Add a sanity check
- # Ensure checkpoint parameters exist in the config
- if not hasattr(self.config, "checkpoint"):
- self.config.checkpoint = {}
- logger.warning("Checkpoint configuration not found, using defaults")
-
- self.metric = get_metric(config.alignment.metric)
-
- def get_basename(self) -> str:
- """
- Get the base name for the experiment.
-
- Returns:
- Base name string
- """
- return f"alignment_{self.config.model.model_name}_{self.config.dataset.dataset_name}"
-
- def prepare_path(self) -> List[str]:
- """
- Prepare the experiment path components.
-
- Returns:
- List of path components
- """
- return [
- "alignment",
- self.config.model.model_name,
- self.config.dataset.dataset_name,
- f"metric_{self.config.alignment.metric}"
- ]
-
- def create_networks(self) -> List[nn.Module]:
- """
- Create multiple neural networks for the experiment, each with different initialization.
-
- Following the alignment_v2 approach, this creates multiple independent networks
- rather than copies of a single network.
-
- Returns:
- List containing multiple independently initialized networks
- """
- # Get the number of network replicates from config
- num_replicates = self.config.training.replicates if hasattr(self.config.training, "replicates") else 5
-
- # Create multiple models based on configuration, each with different initialization
- networks = []
- for i in range(num_replicates):
- # Set a different seed for each network to ensure different initializations
- if hasattr(self.config, 'seed') and self.config.seed is not None:
- # Use the base seed plus the replicate index to get different but reproducible initializations
- torch.manual_seed(self.config.seed + i)
- torch.cuda.manual_seed_all(self.config.seed + i)
- np.random.seed(self.config.seed + i)
-
- # Create a new model
- model = create_model(self.config.model)
- model.to(self.device)
- networks.append(model)
-
- logger.info(f"Created {len(networks)} models with independent initializations: {self.config.model.model_name}")
-
- return networks
-
- def main(self) -> Tuple[Dict, List[nn.Module]]:
- """
- Main experiment execution method.
-
- Returns:
- Tuple of (results dictionary, list of networks)
- """
- # Additional initialization
- self.setup_paths()
- logger.info(f"Set up paths. Results will be saved to {self.results_path}")
-
- # Create network models for experiment
- networks = self.create_networks()
-
- # Run the experiment based on the specified type
- experiment_type = getattr(self.config, 'experiment_type', 'alignment_analysis')
-
- if experiment_type == "alignment_analysis" or experiment_type == "alignment":
- # Run alignment analysis and store results
- results = self.run_alignment_analysis(networks)
-
- elif experiment_type == "progressive_dropout":
- # Run progressive dropout experiment for all networks
- results = self._run_progressive_dropout(networks)
-
- # Generate and log visualization
- self._plot_dropout_results(
- results,
- self.figure_path,
- title_prefix=f"{getattr(self.config, 'experiment_name', 'Progressive Dropout')}"
- )
-
- # Save results for later reference
- with open(os.path.join(self.results_path, "progressive_dropout_results.pkl"), "wb") as f:
- pickle.dump(results, f)
-
- elif experiment_type == "eigenvector_dropout":
- # Run eigenvector dropout experiment for one network
- results = self._run_eigenvector_dropout(networks[0])
-
- # Generate and log visualization (this has a different format)
- self._generate_dropout_plots(
- results,
- "Eigenvector Dropout",
- "eigenvector_dropout"
- )
-
- # Save results for later reference
- with open(os.path.join(self.results_path, "eigenvector_dropout_results.pkl"), "wb") as f:
- pickle.dump(results, f)
-
- elif experiment_type == "training_alignment":
- # Run training alignment experiment
- # This is handled separately because it needs different structure
- results = self._run_training_alignment()
-
- # Generate and log visualization
- self._generate_training_alignment_plots(results)
-
- # Save results for later reference
- with open(os.path.join(self.results_path, "training_alignment_results.pkl"), "wb") as f:
- pickle.dump(results, f)
-
- else:
- raise ValueError(f"Unsupported experiment type: {experiment_type}")
-
- logger.info(f"Completed {getattr(self.config, 'experiment_name', experiment_type)} experiment")
-
- return results, networks
-
- def _run_progressive_dropout(self, networks: List[nn.Module]) -> Dict:
- """
- Run progressive dropout experiment with multiple networks.
-
- Uses the progressive_dropout function from alignment.dropout module
- instead of reimplementing the logic directly.
-
- Args:
- networks: List of networks to run the experiment on.
-
- Returns:
- Dictionary with results that match the expected format for plotting.
- """
- # Get dropout fractions from config
- dropin_min = self.config.alignment.dropout_min
- dropin_max = self.config.alignment.dropout_max
- num_dropout_fractions = self.config.alignment.dropout_steps
- dropout_fractions = np.linspace(dropin_min, dropin_max, num_dropout_fractions).tolist()
-
- # Get pruning mode and dropout mode from config
- pruning_mode = getattr(self.config.extra, "dropout_pruning_mode", "global_joint")
- dropout_mode = getattr(self.config.extra, "dropout_mode", "scaled")
-
- # Print exactly what mode is being used
- logger.info(f"Running progressive dropout with pruning_mode={pruning_mode}, dropout_mode={dropout_mode}")
-
- # Debug metric type
- logger.info(f"Using metric type: {type(self.metric)}")
-
- # Prepare dataset
- try:
- batch_size = getattr(self.config.dataset, "batch_size", 128)
- dataset = load_dataset(self.config.dataset, batch_size=batch_size)
- except Exception as e:
- logger.error(f"Error loading dataset: {str(e)}")
- return {"error": f"Dataset configuration error: {str(e)}", "dropout_fractions": dropout_fractions}
-
- # IMPORTANT: Train all networks at once before applying dropout
- logger.info(f"Training {len(networks)} networks before applying progressive dropout")
-
- # Track training history for plotting
- training_history = {
- 'train_loss': [],
- 'train_acc': [],
- 'test_loss': [],
- 'test_acc': []
- }
-
- # Get training parameters from config
- num_epochs = getattr(self.config.training, "epochs", 5)
- learning_rate = getattr(self.config.training, "learning_rate", 0.001)
-
- # Set up optimizer for each network
- optimizers = []
- for network in networks:
- network.to(self.device)
- optimizers.append(torch.optim.Adam(network.parameters(), lr=learning_rate))
-
- # Train all networks for each epoch - with progress bars
- from tqdm import tqdm
- epoch_pbar = tqdm(range(num_epochs), desc="Training epochs", position=0)
- for epoch in epoch_pbar:
- # Initialize epoch stats
- epoch_train_loss = 0.0
- epoch_train_acc = 0.0
- epoch_test_loss = 0.0
- epoch_test_acc = 0.0
-
- # Training phase
- net_pbar = tqdm(enumerate(zip(networks, optimizers)),
- desc=f"Epoch {epoch+1}/{num_epochs} networks",
- total=len(networks),
- position=1,
- leave=False)
-
- for network_idx, (network, optimizer) in net_pbar:
- network.train()
- running_loss = 0.0
- correct = 0
- total = 0
-
- # Train on each batch
- batch_pbar = tqdm(dataset.train_loader,
- desc=f"Network {network_idx+1}/{len(networks)}",
- position=2,
- leave=False)
-
- for inputs, targets in batch_pbar:
- inputs, targets = inputs.to(self.device), targets.to(self.device)
-
- # Zero the parameter gradients
- optimizer.zero_grad()
-
- # Forward pass
- outputs = network(inputs)
- loss = torch.nn.functional.cross_entropy(outputs, targets)
-
- # Backward pass and optimize
- loss.backward()
- optimizer.step()
-
- # Track statistics
- running_loss += loss.item()
- _, predicted = outputs.max(1)
- total += targets.size(0)
- correct += predicted.eq(targets).sum().item()
-
- # Update batch progress bar
- if total > 0:
- batch_pbar.set_postfix({
- 'loss': running_loss / len(batch_pbar),
- 'acc': 100.0 * correct / total
- })
-
- # Calculate network training statistics
- if total > 0:
- network_train_loss = running_loss / len(dataset.train_loader)
- network_train_acc = 100.0 * correct / total
-
- # Accumulate for epoch average
- epoch_train_loss += network_train_loss
- epoch_train_acc += network_train_acc
-
- # Update network progress bar
- net_pbar.set_postfix({
- 'train_loss': network_train_loss,
- 'train_acc': network_train_acc
- })
-
- # Evaluation phase
- network.eval()
- test_correct = 0
- test_total = 0
- test_loss_sum = 0.0
-
- # Evaluation progress bar
- eval_pbar = tqdm(dataset.test_loader,
- desc=f"Evaluating network {network_idx+1}",
- position=2,
- leave=False)
-
- with torch.no_grad():
- for inputs, targets in eval_pbar:
- inputs, targets = inputs.to(self.device), targets.to(self.device)
- outputs = network(inputs)
-
- # Calculate loss
- loss = torch.nn.functional.cross_entropy(outputs, targets, reduction='sum')
- test_loss_sum += loss.item()
-
- # Calculate accuracy
- _, predicted = outputs.max(1)
- test_total += targets.size(0)
- test_correct += predicted.eq(targets).sum().item()
-
- # Update eval progress bar
- if test_total > 0:
- eval_pbar.set_postfix({
- 'test_loss': test_loss_sum / test_total,
- 'test_acc': 100.0 * test_correct / test_total
- })
-
- # Calculate network testing statistics
- if test_total > 0:
- network_test_loss = test_loss_sum / test_total
- network_test_acc = 100.0 * test_correct / test_total
-
- # Accumulate for epoch average
- epoch_test_loss += network_test_loss
- epoch_test_acc += network_test_acc
-
- # Calculate and store epoch averages
- epoch_train_loss /= len(networks)
- epoch_train_acc /= len(networks)
- epoch_test_loss /= len(networks)
- epoch_test_acc /= len(networks)
-
- training_history['train_loss'].append(epoch_train_loss)
- training_history['train_acc'].append(epoch_train_acc)
- training_history['test_loss'].append(epoch_test_loss)
- training_history['test_acc'].append(epoch_test_acc)
-
- # Update epoch progress bar
- epoch_pbar.set_postfix({
- 'train_loss': epoch_train_loss,
- 'train_acc': f"{epoch_train_acc:.2f}%",
- 'test_loss': epoch_test_loss,
- 'test_acc': f"{epoch_test_acc:.2f}%"
- })
-
- # Log progress
- logger.info(f"Epoch {epoch+1}/{num_epochs}: "
- f"Train Loss={epoch_train_loss:.4f}, Train Acc={epoch_train_acc:.2f}%, "
- f"Test Loss={epoch_test_loss:.4f}, Test Acc={epoch_test_acc:.2f}%")
-
- logger.info(f"Completed training {len(networks)} networks. Proceeding to progressive dropout...")
-
- # Initialize results structure with correct format for plotting
- final_results = {
- "dropout_fractions": dropout_fractions,
- "accuracies": {"high_rq": [], "low_rq": [], "random": []},
- "losses": {"high_rq": [], "low_rq": [], "random": []},
- "stds": {"high_rq": [], "low_rq": [], "random": []},
- "training_history": training_history
- }
-
- # Call the progressive_dropout function with our trained networks
- from alignment.dropout import progressive_dropout
-
- try:
- # Define the strategies to run
- strategies = ["high_rq", "low_rq", "random"]
-
- # Run each strategy separately with progress bars
- strategy_pbar = tqdm(strategies, desc="Running pruning strategies", position=0)
- for strategy in strategy_pbar:
- strategy_pbar.set_description(f"Strategy: {strategy}")
-
- # Run progressive dropout with this strategy
- network_accuracies, network_losses = progressive_dropout(
- networks,
- dataset,
- dropout_fractions,
- self.metric,
- self.device,
- pruning_mode=pruning_mode,
- dropout_mode=dropout_mode,
- strategy=strategy
- )
-
- # Calculate averages and standard deviations across networks for this strategy
- fraction_pbar = tqdm(range(len(dropout_fractions)),
- desc="Processing dropout fractions",
- position=1,
- leave=False)
-
- for fraction_idx in fraction_pbar:
- fraction_pbar.set_description(f"Fraction: {dropout_fractions[fraction_idx]:.2f}")
-
- # Collect results for this fraction across all networks
- fraction_accs = []
- fraction_losses = []
-
- for net_idx in network_accuracies:
- if fraction_idx < len(network_accuracies[net_idx]):
- fraction_accs.append(network_accuracies[net_idx][fraction_idx])
- if fraction_idx < len(network_losses[net_idx]):
- fraction_losses.append(network_losses[net_idx][fraction_idx])
-
- # Calculate statistics
- if fraction_accs:
- mean_acc = np.mean(fraction_accs)
- std_acc = np.std(fraction_accs)
- mean_loss = np.mean(fraction_losses) if fraction_losses else 0.0
-
- # Add to appropriate strategy in results
- if fraction_idx >= len(final_results["accuracies"][strategy]):
- final_results["accuracies"][strategy].append(mean_acc)
- final_results["stds"][strategy].append(std_acc)
- final_results["losses"][strategy].append(mean_loss)
- else:
- # Update existing value (shouldn't happen with our current code)
- final_results["accuracies"][strategy][fraction_idx] = mean_acc
- final_results["stds"][strategy][fraction_idx] = std_acc
- final_results["losses"][strategy][fraction_idx] = mean_loss
-
- # Update progress bar
- fraction_pbar.set_postfix({
- 'acc': f"{mean_acc:.2f}%",
- 'std': f"{std_acc:.2f}%"
- })
-
- # Log results for verification
- logger.info(f"Final results:")
- for strategy in strategies:
- logger.info(f" {strategy}: {final_results['accuracies'][strategy]}")
-
- except Exception as e:
- logger.error(f"Error running progressive dropout: {str(e)}")
- import traceback
- logger.error(traceback.format_exc())
- return {
- "error": f"Error running progressive dropout: {str(e)}",
- "dropout_fractions": dropout_fractions,
- "training_history": training_history
- }
-
- # Generate dropout plots
- try:
- logger.info("Generating dropout plots with error bars")
- from alignment.utils.plotting import plot_dropout_results
-
- saved_plots = plot_dropout_results(
- final_results,
- save_dir=self.figure_path,
- title_prefix="Progressive Dropout",
- pruning_mode=pruning_mode,
- dropout_mode=dropout_mode
- )
-
- if saved_plots:
- logger.info(f"Generated {len(saved_plots)} plots: {saved_plots}")
- final_results["plot_files"] = saved_plots
- else:
- logger.warning("No plots were generated. Check your plot_dropout_results function.")
-
- except Exception as e:
- logger.error(f"Error generating dropout plots: {str(e)}")
- import traceback
- logger.error(traceback.format_exc())
-
- return final_results
-
- def _run_eigenvector_dropout(self, net: nn.Module) -> Dict:
- """
- Run eigenvector dropout experiments.
-
- Uses the eigenvector_dropout function from alignment.dropout module
- instead of reimplementing the logic directly.
-
- Args:
- net: Neural network to analyze
-
- Returns:
- Dictionary of eigenvector dropout results
- """
- # Get dropout fractions from config
- dropout_fractions = np.linspace(
- self.config.alignment.dropout_min,
- self.config.alignment.dropout_max,
- self.config.alignment.dropout_steps
- ).tolist()
-
- # Initialize result structure
- results = {
- "dropout_fractions": dropout_fractions,
- "accuracies": {"eigenvector": []},
- "losses": {"eigenvector": []},
- "alignment_values": {"eigenvector": []}
- }
-
- # Prepare dataset config
- dataset_config = self.config.dataset
-
- # Process each dropout fraction
- for dropout_fraction in tqdm(dropout_fractions, desc="Eigenvector Dropout"):
- try:
- # Call the eigenvector_dropout function directly
- accuracy, alignment_values = eigenvector_dropout(
- net,
- dataset_config,
- dropout_fraction=dropout_fraction,
- metric=self.metric,
- device=self.device,
- dropout_mode=getattr(self.config.extra, "dropout_mode", "scaled"),
- dropout_pruning_mode=getattr(self.config.extra, "dropout_pruning_mode", "global_joint")
- )
-
- # Store results
- results["accuracies"]["eigenvector"].append(accuracy)
- results["losses"]["eigenvector"].append(100.0 - accuracy)
- results["alignment_values"]["eigenvector"].append(alignment_values)
-
- except Exception as e:
- logger.error(f"Error in eigenvector dropout at fraction {dropout_fraction}: {str(e)}")
- # Add placeholder values to maintain result structure
- results["accuracies"]["eigenvector"].append(0.0)
- results["losses"]["eigenvector"].append(100.0)
- results["alignment_values"]["eigenvector"].append(None)
-
- return results
-
- def plot(self, results: Dict) -> Dict[str, List[str]]:
- """
- Plot experiment results using enhanced visualization methods.
-
- Args:
- results: Dictionary containing experiment results
-
- Returns:
- Dictionary of plot paths by type
- """
- logger.info("Generating experiment visualizations")
-
- # Check for results to plot
- if not results:
- logger.warning("No results to plot")
- return {}
-
- plot_paths = {}
-
- # Generate experiment summary
- try:
- logger.info("Creating experiment summary visualization")
- summary_path = self._generate_summary_plots(results)
- if summary_path:
- plot_paths["summary"] = summary_path
- logger.info("Successfully created experiment summary visualization")
- except Exception as e:
- logger.error(f"Error creating summary visualization: {str(e)}", exc_info=True)
-
- # Generate progressive dropout plots
- if "progressive_dropout" in results:
- try:
- logger.info("Generating progressive dropout visualizations")
- prog_paths = self._plot_dropout_results(
- results["progressive_dropout"],
- self.figure_path,
- "Progressive Dropout"
- )
- if prog_paths:
- plot_paths["progressive_dropout"] = prog_paths
- logger.info("Successfully generated progressive dropout visualizations")
- except Exception as e:
- logger.error(f"Error generating progressive dropout visualizations: {str(e)}", exc_info=True)
-
- # Generate eigenvector dropout plots
- if "eigenvector_dropout" in results:
- try:
- logger.info("Generating eigenvector dropout visualizations")
- eig_paths = self._plot_dropout_results(
- results["eigenvector_dropout"],
- self.figure_path,
- "Eigenvector Dropout"
- )
- if eig_paths:
- plot_paths["eigenvector_dropout"] = eig_paths
- logger.info("Successfully generated eigenvector dropout visualizations")
- except Exception as e:
- logger.error(f"Error generating eigenvector dropout visualizations: {str(e)}", exc_info=True)
-
- # Generate combined visualizations if both types of results exist
- if "progressive_dropout" in results and "eigenvector_dropout" in results:
- try:
- logger.info("Generating combined dropout comparison visualization")
- comparison_path = self._generate_dropout_comparison(
- results["progressive_dropout"],
- results["eigenvector_dropout"]
- )
- if comparison_path:
- plot_paths["comparison"] = comparison_path
- logger.info("Successfully generated combined dropout comparison")
- except Exception as e:
- logger.error(f"Error generating dropout comparison: {str(e)}", exc_info=True)
-
- logger.info(f"Completed experiment visualization generation: {len(plot_paths)} plot types")
- return plot_paths
-
- def _plot_dropout_results(self, results, figure_path=None, title_prefix="Dropout"):
- """
- Plot dropout experiment results using the plotting module.
-
- Args:
- results: Results dictionary from progressive_dropout
- figure_path: Directory to save plots to
- title_prefix: Prefix for plot titles
-
- Returns:
- List of saved plot files
- """
- # Get pruning mode and dropout mode
- pruning_mode = getattr(self.config.extra, "dropout_pruning_mode", "global_joint")
- dropout_mode = getattr(self.config.extra, "dropout_mode", "scaled")
-
- # Ensure figure path exists
- if figure_path is None:
- figure_path = self.figure_path
-
- # Add debug logs to understand the results
- logger.info(f"Plotting dropout results with pruning_mode={pruning_mode}, dropout_mode={dropout_mode}")
- logger.info(f"Results keys: {list(results.keys())}")
- logger.info(f"Figure path: {figure_path}")
-
- if "accuracies" in results:
- logger.info(f"Accuracies keys: {list(results['accuracies'].keys())}")
- for key in results["accuracies"]:
- logger.info(f"Strategy {key} has {len(results['accuracies'][key])} data points")
- if results["accuracies"][key]:
- logger.info(f"First few values: {results['accuracies'][key][:3]}")
- else:
- logger.info(f"No data for {key}")
-
- # Generate plots using our custom function
- try:
- from alignment.utils.plotting import custom_plot_dropout
-
- # Make a deep copy of results to prevent modification
- import copy
- plot_results = copy.deepcopy(results)
-
- # Ensure we have data for at least one strategy (preferably high_rq)
- if not plot_results.get("accuracies", {}).get("high_rq", []):
- # Try to use data from any strategy that has data
- for strategy in ["low_rq", "random", "eigenvector"]:
- if plot_results.get("accuracies", {}).get(strategy, []):
- # Use this data for high_rq
- logger.info(f"Using {strategy} data for high_rq")
- plot_results["accuracies"]["high_rq"] = plot_results["accuracies"][strategy]
-
- # Also copy to losses if needed
- if strategy in plot_results.get("losses", {}):
- plot_results["losses"]["high_rq"] = plot_results["losses"][strategy]
- break
-
- # Try the numeric indices if we still don't have data
- if not plot_results.get("accuracies", {}).get("high_rq", []) and hasattr(self, "results") and hasattr(self.results, "network_accuracies"):
- # Look for numeric indices
- for idx in range(3): # Try indices 0, 1, 2
- if idx in self.results.network_accuracies and self.results.network_accuracies[idx]:
- strategy_name = {0: "high_rq", 1: "low_rq", 2: "random"}.get(idx, "high_rq")
- logger.info(f"Using numeric index {idx} data for {strategy_name}")
-
- # Make sure the strategy key exists
- if "accuracies" not in plot_results:
- plot_results["accuracies"] = {}
-
- # Copy the data
- plot_results["accuracies"][strategy_name] = self.results.network_accuracies[idx]
-
- # Print final plot data
- logger.info(f"Passing the following to custom_plot_dropout:")
- logger.info(f"- dropout_fractions: {len(plot_results.get('dropout_fractions', []))} points")
- logger.info(f"- high_rq data: {len(plot_results.get('accuracies', {}).get('high_rq', []))} points")
-
- # Call the custom plotting function
- saved_figures = custom_plot_dropout(
- plot_results,
- figure_path,
- pruning_mode=pruning_mode,
- dropout_mode=dropout_mode,
- title_prefix=title_prefix
- )
-
- # Check if any plots were generated
- if saved_figures:
- logger.info(f"Generated {len(saved_figures)} plots: {saved_figures}")
-
- # Log plots to wandb if configured
- if hasattr(self.config.checkpointing, "use_wandb") and self.config.checkpointing.use_wandb:
- try:
- from alignment.utils.plotting import log_plots_to_wandb
- log_plots_to_wandb(saved_figures)
- logger.info(f"Logged {len(saved_figures)} dropout plots to wandb")
- except Exception as e:
- logger.warning(f"Failed to log dropout plots to wandb: {str(e)}")
- else:
- logger.warning("No plots were generated by custom_plot_dropout")
-
- return saved_figures
-
- except Exception as e:
- logger.error(f"Error generating dropout plots: {str(e)}")
- import traceback
- logger.error(traceback.format_exc())
- return []
-
- def _generate_dropout_plots(self, results, title, filename=None):
- """
- Generate enhanced dropout plots.
-
- Args:
- results: Dropout experiment results
- title: Plot title
- filename: Base filename (optional)
-
- Returns:
- List of saved plot files
- """
- # Call the _plot_dropout_results function with our figure path and title
- return self._plot_dropout_results(results, self.figure_path, title)
-
- def _generate_summary_plots(self, results):
- """
- Generate comprehensive summary plots for the experiment using the plotting module.
-
- Args:
- results: Experiment results dictionary
-
- Returns:
- Path to the saved summary plot
- """
- try:
- from alignment.utils.plotting import plot_experiment_summary
-
- # Generate the summary plot
- experiment_name = getattr(self.config, "experiment_name", self.get_basename())
- filepath = plot_experiment_summary(
- results,
- self.figure_path,
- experiment_name=experiment_name
- )
-
- # Log to wandb if configured
- if filepath and hasattr(self.config.checkpointing, "use_wandb") and self.config.checkpointing.use_wandb:
- try:
- import wandb
- if wandb.run is not None:
- wandb.log({"experiment_summary": wandb.Image(filepath)})
- logger.info("Logged experiment summary to wandb")
- except Exception as e:
- logger.warning(f"Failed to log experiment summary to wandb: {str(e)}")
-
- return filepath
-
- except Exception as e:
- logger.error(f"Error creating summary visualization: {str(e)}")
- return None
-
- def _generate_dropout_comparison(self, progressive_results, eigenvector_results):
- """
- Generate comparison plots between progressive and eigenvector dropout using the plotting module.
-
- Args:
- progressive_results: Progressive dropout results
- eigenvector_results: Eigenvector dropout results
-
- Returns:
- Path to the saved comparison plot
- """
- try:
- from alignment.utils.plotting import plot_dropout_comparison
-
- # Generate the comparison plot
- experiment_name = getattr(self.config, "experiment_name", "Experiment")
- filepath = plot_dropout_comparison(
- progressive_results,
- eigenvector_results,
- self.figure_path,
- title=f"Dropout Comparison: {experiment_name}"
- )
-
- # Log to wandb if configured
- if filepath and hasattr(self.config.checkpointing, "use_wandb") and self.config.checkpointing.use_wandb:
- try:
- import wandb
- if wandb.run is not None:
- wandb.log({"dropout_comparison": wandb.Image(filepath)})
- logger.info("Logged dropout comparison to wandb")
- except Exception as e:
- logger.warning(f"Failed to log dropout comparison to wandb: {str(e)}")
-
- return filepath
-
- except Exception as e:
- logger.error(f"Error creating dropout comparison: {str(e)}")
- return None
-
- def _initialize_wandb(self) -> Optional[Any]:
- """
- Initialize wandb for experiment tracking.
-
- Returns:
- wandb module if initialized successfully, None otherwise
- """
- if not getattr(self.config.checkpointing, "use_wandb", False):
- return None
-
- try:
- import wandb
-
- # Get wandb configuration
- entity = getattr(self.config, "wandb_entity", None)
- project = getattr(self.config, "wandb_project", "neural_alignment")
-
- # Generate timestamp if it doesn't exist
- if not hasattr(self, 'timestamp'):
- self.timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
-
- # Create a unique run name
- if hasattr(self.config, "experiment_name") and self.config.experiment_name:
- run_name = f"{self.config.experiment_name}"
- else:
- run_name = f"{self.config.model.model_name}_{self.config.dataset.dataset_name}_{self.timestamp}"
-
- # Convert the config to a dictionary for wandb
- config_dict = self.config.to_dict() if hasattr(self.config, "to_dict") else vars(self.config)
-
- # Initialize wandb run
- wandb.init(
- project=project,
- entity=entity,
- name=run_name,
- config=config_dict,
- reinit=True # Allow reinitializing in the same process
- )
-
- # Set up wandb to watch models if available
- if hasattr(self, "networks") and self.networks:
- # Watch the first network
- wandb.watch(self.networks[0], log="all", log_freq=10)
-
- logger.info(f"Initialized Weights & Biases tracking for run '{run_name}'")
- return wandb
- except Exception as e:
- logger.warning(f"Failed to initialize wandb: {str(e)}")
- return None
-
- def _log_to_wandb(self, data: Dict) -> bool:
- """
- Log data to wandb.
-
- Args:
- data: Data to log to wandb (dict with keys corresponding to metric names)
-
- Returns:
- True if logging was successful, False otherwise
- """
- if not getattr(self.config.checkpointing, "use_wandb", False):
- return False
-
- try:
- import wandb
- if wandb.run is None:
- wandb_module = self._initialize_wandb()
- if wandb_module is None:
- return False
-
- # Log scalar metrics
- metrics_to_log = {}
- for k, v in data.items():
- # Skip non-scalar values
- if isinstance(v, (int, float)):
- metrics_to_log[k] = v
- elif isinstance(v, (list, tuple)) and len(v) == 1:
- # Single-element list or tuple with a scalar
- if isinstance(v[0], (int, float)):
- metrics_to_log[k] = v[0]
-
- # Log the metrics to wandb
- if metrics_to_log:
- wandb.log(metrics_to_log)
-
- # Check if data contains progressive dropout results that we should plot
- if "progressive_dropout" in data and "plot_files" in data and getattr(self.config.extra, "log_images", True):
- # Find all plot files and log them to wandb
- for plot_file in data["plot_files"]:
- if os.path.exists(plot_file):
- plot_name = os.path.basename(plot_file).replace(".png", "")
- wandb.log({plot_name: wandb.Image(plot_file)})
-
- # If 'accuracies' exists in the progressive_dropout results, log them as a line plot
- if "accuracies" in data["progressive_dropout"] and "dropout_fractions" in data["progressive_dropout"]:
- # Create a table to log
- accuracies = np.array(data["progressive_dropout"]["accuracies"])
- dropout_fractions = data["progressive_dropout"]["dropout_fractions"]
-
- # Calculate mean and std
- mean_accuracies = np.mean(accuracies, axis=0)
- std_accuracies = np.std(accuracies, axis=0)
-
- # Prepare data table for the custom plot
- table_data = [[x, y, y_std] for x, y, y_std in zip(dropout_fractions, mean_accuracies, std_accuracies)]
-
- # Log the data as a custom plot
- wandb.log({
- "dropout_curve": wandb.Table(
- data=table_data,
- columns=["dropout_fraction", "mean_accuracy", "std_accuracy"]
- )
- })
-
- # Create interactive plot
- dropout_data = []
- for i, frac in enumerate(dropout_fractions):
- dropout_data.append([frac, mean_accuracies[i], mean_accuracies[i]-std_accuracies[i], mean_accuracies[i]+std_accuracies[i]])
-
- wandb.log({
- "dropout_plot_interactive": wandb.plot.line_series(
- xs=[x[0] for x in dropout_data],
- ys=[[x[1] for x in dropout_data], # mean
- [x[2] for x in dropout_data], # mean-std
- [x[3] for x in dropout_data]], # mean+std
- keys=["Mean Accuracy", "Lower Bound", "Upper Bound"],
- title=f"Progressive Dropout ({self.config.extra.dropout_pruning_mode})",
- xname="Dropout Fraction",
- yname="Accuracy"
- )
- })
-
- return True
- except Exception as e:
- logger.warning(f"Failed to log to wandb: {str(e)}")
- return False
-
- def run(self) -> Tuple[Dict, List[nn.Module]]:
- """
- Run the experiment with custom logging and wandb setup.
-
- Returns:
- Tuple containing (results, networks)
- """
- # Setup paths for experiment
- self.setup_paths()
- self.prepare_path()
-
- # Set up logging and suppress warnings
- set_logging_level(logging.INFO)
-
- # Set random seed if configured
- if hasattr(self.config, 'seed') and self.config.seed is not None:
- torch.manual_seed(self.config.seed)
- torch.cuda.manual_seed_all(self.config.seed)
- np.random.seed(self.config.seed)
- logger.info(f"Set random seed to {self.config.seed}")
-
- # Initialize wandb explicitly before starting experiment
- if hasattr(self.config, 'checkpointing') and self.config.checkpointing.use_wandb:
- self._initialize_wandb()
-
- # Run the main experiment
- results, networks = self.main()
-
- # Store results as an attribute for saving
- self.results = results
- self.networks = networks
-
- # Ensure all plots are saved and logged
- self.plot(results)
-
- # Save experiment state
- self.save()
-
- # Close wandb run if it exists
- try:
- import wandb
- if wandb.run is not None:
- wandb.finish()
- logger.info("Closed wandb run")
- except Exception as e:
- logger.warning(f"Error closing wandb run: {str(e)}")
-
- return results, networks
-
- def run_alignment_analysis(self, networks: List[nn.Module]) -> Dict:
- """
- Run alignment analysis experiment.
-
- This experiment analyzes alignment metrics across networks and can run
- both progressive dropout and eigenvector dropout experiments based on
- configuration.
-
- Args:
- networks: List of neural networks to analyze
-
- Returns:
- Results dictionary
- """
- logger.info("Running alignment analysis experiment")
-
- # Prepare the results structure
- results = {
- "progressive_dropout": {},
- "eigenvector_dropout": {},
- "config": self.config,
- }
-
- # Progressive dropout experiment
- if self.config.alignment.run_progressive:
- logger.info(f"Running progressive dropout experiment (mode: {getattr(self.config.extra, 'dropout_pruning_mode', 'global_joint')})")
- try:
- prog_results = self._run_progressive_dropout(networks)
- results["progressive_dropout"] = prog_results
-
- # Add debug logging to inspect results structure
- if prog_results:
- logger.info(f"Progressive dropout results keys: {list(prog_results.keys())}")
- if 'accuracies' in prog_results:
- logger.info(f"Progressive dropout accuracies keys: {list(prog_results['accuracies'].keys())}")
-
- # Check if there was an error in the dataset configuration
- if "error" in prog_results:
- logger.error(f"Progressive dropout experiment failed: {prog_results['error']}")
- else:
- # Generate plots right after getting results
- try:
- logger.info("Generating progressive dropout visualizations")
- # Add debug print before and after plotting
- figure_path = self._plot_dropout_results(
- results["progressive_dropout"],
- self.figure_path,
- title_prefix="Progressive Dropout"
- )
- logger.info(f"Progressive dropout plot generated at {figure_path}")
- except Exception as e:
- logger.error(f"Error generating progressive dropout visualizations: {str(e)}", exc_info=True)
- except Exception as e:
- logger.error(f"Error running progressive dropout experiment: {str(e)}")
- results["progressive_dropout"] = {"error": str(e)}
-
- # Eigenvector dropout experiment
- if self.config.alignment.run_eigenvector:
- logger.info("Running eigenvector dropout experiment")
- try:
- # Note: For eigenvector dropout, we use just the first network
- eig_results = self._run_eigenvector_dropout(networks[0])
- results["eigenvector_dropout"] = eig_results
-
- # Add debug logging
- if eig_results:
- logger.info(f"Eigenvector dropout results keys: {list(eig_results.keys())}")
- if 'accuracies' in eig_results:
- logger.info(f"Eigenvector dropout accuracies keys: {list(eig_results['accuracies'].keys())}")
-
- # Generate plots right after getting results
- try:
- logger.info("Generating eigenvector dropout visualizations")
- # Use the specialized method for eigenvector dropout
- self._generate_dropout_plots(
- results["eigenvector_dropout"],
- "Eigenvector Dropout",
- "eigenvector_dropout"
- )
- logger.info("Successfully generated eigenvector dropout visualizations")
- except Exception as e:
- logger.error(f"Error generating eigenvector dropout visualizations: {str(e)}", exc_info=True)
- except Exception as e:
- logger.error(f"Error running eigenvector dropout experiment: {str(e)}")
- results["eigenvector_dropout"] = {"error": str(e)}
-
- # Generate and log summary plots for all experiments
- try:
- logger.info("Generating experiment summary")
- self._generate_summary_plots(results)
- logger.info("Successfully generated experiment summary")
- except Exception as e:
- logger.error(f"Error generating experiment summary: {str(e)}", exc_info=True)
-
- return results
-
- def setup_paths(self):
- """
- Set up paths for experiment outputs.
-
- Creates necessary directories for figures, weights, and results.
- """
- # Create timestamp subdirectory if needed
- if getattr(self.config, "use_timestamp", True):
- timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
- base_dir = os.path.join("results", f"{self.get_basename()}_{timestamp}")
- else:
- base_dir = os.path.join("results", self.get_basename())
-
- # Create main result directory
- os.makedirs(base_dir, exist_ok=True)
- self.results_path = base_dir
-
- # Create figure directory
- self.figure_path = os.path.join(base_dir, "figures")
- os.makedirs(self.figure_path, exist_ok=True)
-
- # Create weights directory for saved models
- self.weights_path = os.path.join(base_dir, "weights")
- os.makedirs(self.weights_path, exist_ok=True)
-
- # Set the device to use (default to CUDA if available)
- if hasattr(self.config, "device") and self.config.device:
- self.device = torch.device(self.config.device)
- else:
- self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
- logger.info(f"Set up paths: results={self.results_path}, figures={self.figure_path}, device={self.device}")
-
- def save(self):
- """
- Save experiment state including results and configuration.
-
- This method saves:
- 1. Experiment results as pickle and JSON
- 2. Configuration in YAML format
- 3. Any model checkpoints if available
- """
- try:
- # Create a results directory if it doesn't exist
- os.makedirs(self.results_path, exist_ok=True)
-
- # Save experiment results if available
- if hasattr(self, 'results') and self.results:
- results_file = os.path.join(self.results_path, "results.pkl")
- with open(results_file, "wb") as f:
- pickle.dump(self.results, f)
- logger.info(f"Saved experiment results to {results_file}")
-
- # Also try to save as JSON for human readability
- try:
- # Convert tensors and other non-serializable objects
- def clean_for_json(obj):
- if isinstance(obj, (torch.Tensor, np.ndarray)):
- return obj.tolist() if hasattr(obj, 'tolist') else str(obj)
- elif isinstance(obj, dict):
- return {k: clean_for_json(v) for k, v in obj.items()}
- elif isinstance(obj, list):
- return [clean_for_json(item) for item in obj]
- elif isinstance(obj, tuple):
- return tuple(clean_for_json(item) for item in obj)
- else:
- # Return string representation for other types
- return str(obj) if not isinstance(obj, (str, int, float, bool, type(None))) else obj
-
- json_results = clean_for_json(self.results)
- json_file = os.path.join(self.results_path, "results.json")
- with open(json_file, "w") as f:
- json.dump(json_results, f, indent=2)
- logger.info(f"Saved readable results to {json_file}")
- except Exception as e:
- logger.warning(f"Could not save results as JSON: {str(e)}")
-
- # Save configuration
- config_file = os.path.join(self.results_path, "config.yaml")
- # Check if we can convert config to dict
- if hasattr(self.config, 'to_dict'):
- config_dict = self.config.to_dict()
- with open(config_file, "w") as f:
- yaml.dump(config_dict, f, default_flow_style=False)
- logger.info(f"Saved configuration to {config_file}")
-
- logger.info("Experiment state saved successfully")
-
- except Exception as e:
- logger.error(f"Error saving experiment state: {str(e)}")
- # Don't re-raise the exception, just log it
-
-
-def set_logging_level(level=logging.INFO):
- """
- Set logging level for all loggers in the alignment package.
-
- Args:
- level: Logging level (default: logging.INFO)
- """
- # Set root logger level
- logging.getLogger().setLevel(level)
-
- # Set level for alignment package loggers
- logging.getLogger('alignment').setLevel(level)
-
- # Reduce verbosity for other loggers
- logging.getLogger('matplotlib').setLevel(logging.WARNING)
- logging.getLogger('PIL').setLevel(logging.WARNING)
- logging.getLogger('torch').setLevel(logging.WARNING)
- logging.getLogger('wandb').setLevel(logging.WARNING)
-
- # Suppress common warnings
- import warnings
- warnings.filterwarnings("ignore", category=UserWarning, message=".*pruning fraction.*")
- warnings.filterwarnings("ignore", category=UserWarning, message=".*Combined pruning across layers.*")
- warnings.filterwarnings("ignore", category=UserWarning, message=".*dropout of the entire tensor.*")
-
-def cli_main():
- """Command-line interface for running alignment experiments."""
- parser = argparse.ArgumentParser(description="Neural network alignment experiment")
- parser.add_argument("--config", type=str, required=True, help="Path to config YAML")
- parser.add_argument("--quiet", action="store_true", help="Reduce logging output")
- args = parser.parse_args()
-
- # Set logging level based on quiet flag
- if args.quiet:
- set_logging_level(logging.WARNING)
- else:
- set_logging_level(logging.INFO)
-
- # Load configuration
- config = ExperimentConfig.load(args.config)
-
- # Initialize and run experiment
- experiment = AlignmentExperiment(config)
- results, networks = experiment.run()
-
- return results, networks
-
-
-if __name__ == "__main__":
- cli_main()
\ No newline at end of file
diff --git a/_arxiv/_archive/alignment_preref/alignment_metrics.py b/_arxiv/_archive/alignment_preref/alignment_metrics.py
deleted file mode 100644
index e5ae6c91..00000000
--- a/_arxiv/_archive/alignment_preref/alignment_metrics.py
+++ /dev/null
@@ -1,243 +0,0 @@
-# --------------------------------------------
-# alignment_metrics.py
-# --------------------------------------------
-
-import torch
-from alignment.utils import smart_pca, get_device
-
-class AlignmentMetrics:
- """
- Provides static methods for various alignment metrics,
- including 'delta_alignment'.
- """
-
- @staticmethod
- def RQ(input_, weight_, scale_by_norm=False):
- """
- Rayleigh Quotient alignment measure:
- proportion of variance in `input_` explained by each row of `weight_`.
- """
- return alignment(input_, weight_, method="alignment", relative=True, scale_by_norm=scale_by_norm)
-
- @staticmethod
- def MI_0(input_, weight_, scale_by_norm=False):
- """
- Placeholder for mutual information approach - version 0
- (currently reuses alignment(...) as a stand-in).
- """
- return alignment(input_, weight_, method="alignment", relative=True, scale_by_norm=scale_by_norm)
-
- @staticmethod
- def MI_1(input_, weight_):
- """
- Placeholder for mutual information approach - version 1
- """
- return torch.tensor(0.0)
-
- @staticmethod
- def delta_alignment(net, layer_idx, layer_input, scale_by_norm=False):
- """
- alignment of (W_current - W_init) with the input's covariance.
- This references net._init_weights for the initial weight.
- If net._init_weights is absent, returns zeros.
- """
- if not hasattr(net, "_init_weights"):
- weight_diff = torch.zeros_like(net.get_alignment_weights()[layer_idx])
- else:
- init_w = net._init_weights[layer_idx]
- current_w = net.get_alignment_weights()[layer_idx]
- weight_diff = current_w - init_w
-
- weight_diff = weight_diff.flatten(start_dim=1)
- return alignment(layer_input, weight_diff, method="alignment", relative=True, scale_by_norm=scale_by_norm)
-
- @staticmethod
- def measure(input_, weight_, method="RQ", scale_by_norm=False):
- """
- For a single method and a single layer's (input_, weight_).
- """
- if method == "RQ":
- return AlignmentMetrics.RQ(input_, weight_, scale_by_norm=scale_by_norm)
- elif method == "MI_0":
- return AlignmentMetrics.MI_0(input_, weight_, scale_by_norm=scale_by_norm)
- elif method == "MI_1":
- return AlignmentMetrics.MI_1(input_, weight_)
- else:
- raise ValueError(f"Unknown alignment method {method}")
-
- @staticmethod
- def patchwise_alignment(inp, w, method="RQ", weigh_by_var=True):
- """
- 'Patchwise' alignment for CNN.
- - inp shape: [B, F, patches], e.g. B=mini-batch, F=inC*kH*kW, patches=outH*outW
- - w shape: [outC, F], the flattened filter weights.
- - We compute alignment for each patch's (B,F) input.
- - Weighted by patch-level variance across the batch dimension.
- """
- B, F, P = inp.shape
- var_patches = torch.var(inp, dim=0, keepdim=False) # => (F, P)
- patchwise_var = var_patches.sum(dim=0) # => shape (P,)
-
- all_patch_vals = []
- all_patch_vars = []
-
- for p in range(P):
- patch_data = inp[:, :, p] # shape => [B, F]
- cc = torch.cov(patch_data.T) # shape (F, F)
- num_ = torch.sum(torch.matmul(w, cc) * w, dim=1)
- denom_ = torch.sum(w*w, dim=1)
-
- patch_rq = num_ / denom_
- if method == "RQ":
- patch_rq = patch_rq / torch.trace(cc)
- patch_weight = patchwise_var[p] if weigh_by_var else 1.0
-
- all_patch_vals.append(patch_rq * patch_weight)
- all_patch_vars.append(patch_weight)
-
- total_weight = torch.stack(all_patch_vars).sum()
- sum_rq = torch.stack(all_patch_vals, dim=0).sum(dim=0) # shape => (outC,)
-
- if total_weight > 0:
- final_rq = sum_rq / total_weight
- else:
- final_rq = sum_rq * 0
- return final_rq # shape => (outC,)
-
- @staticmethod
- def measure_methods(net, images, methods, precomputed=True, scale_by_norm=False):
- """
- For each layer in 'net', produce a dict {method -> tensor}.
- 1) get_layer_inputs
- 2) preprocess => unfold or patchwise
- 3) flatten weights
- 4) measure alignment or patchwise alignment
-
- Parameters:
- - scale_by_norm: if True, scales covariance matrices by their norm before computing RQ
- """
- # 1) Get raw inputs
- raw_inputs = net.get_layer_inputs(images, precomputed=precomputed)
- # 2) Unfold or flatten them
- preprocessed = net._preprocess_inputs(raw_inputs, compress_convolutional=True)
- # 3) Flatten weights
- layer_weights = net.get_alignment_weights(flatten=True)
-
- results_per_layer = []
- for layer_idx, (inp, wgt) in enumerate(zip(preprocessed, layer_weights)):
- layer_dict = {}
- for m in methods:
- if m == "delta_alignment":
- val = AlignmentMetrics.delta_alignment(net, layer_idx, inp, scale_by_norm=scale_by_norm)
- else:
- # if net says "patchwise" & input is 3D => do patchwise
- if net.cnn_mode == "patchwise" and inp.ndim == 3:
- val = AlignmentMetrics.patchwise_alignment(inp, wgt, method=m, weigh_by_var=True)
- else:
- # single-cov
- val = AlignmentMetrics.measure(inp, wgt, method=m, scale_by_norm=scale_by_norm)
- layer_dict[m] = val
- results_per_layer.append(layer_dict)
-
- return results_per_layer
-
- @staticmethod
- def compute_eigenvalues(x):
- """
- Do a standard PCA with 'smart_pca' to get eigenvalues of x (batch, features).
- """
- w, v = smart_pca(x.T, centered=True)
- return w, v
-
- @staticmethod
- def measure_expected_distribution(method, eigenvals, bins=50, num_tests=100):
- """
- Return (counts, bin_edges) for a random alignment distribution of 'method'.
- """
- if method == "RQ":
- counts, edges, _ = expected_alignment_distribution(
- eigenvals,
- relative=True,
- valid_rotation=False,
- with_rotation=True,
- bins=bins,
- num_tests=num_tests
- )
- return counts, edges
- elif method in ["MI_0", "MI_1", "delta_alignment"]:
- return None, None
- else:
- return None, None
-
-def alignment(input, weight, method="alignment", relative=True, scale_by_norm=False):
- """
- measure alignment by computing RQ = (w^T C w) / (w^T w),
- optionally normalized by trace(C).
-
- - input shape can be (N, F) for single-cov approach,
- or (B, F, patches) if you do patchwise in patchwise_alignment.
- - weight shape (outC, F).
- - scale_by_norm: if True, scales the covariance matrix by its Frobenius norm
- before computing RQ (similar to "similarity" in alignment_v2)
- """
- if input.ndim == 2:
- # standard single-cov approach
- cc = torch.cov(input.T)
-
- # Scale covariance by its norm if requested (like "similarity" in v2)
- if scale_by_norm:
- cc_norm = torch.norm(cc)
- if cc_norm > 0:
- cc = cc / cc_norm
-
- numerator = torch.sum(torch.matmul(weight, cc) * weight, dim=1)
- denom = torch.sum(weight * weight, dim=1)
- rq = numerator / denom
- if method == "alignment" and relative:
- rq = rq / torch.trace(cc)
- return rq
- else:
- # user should call patchwise_alignment or do proper flatten
- raise ValueError(
- f"alignment() got {input.ndim}-D data. For patchwise CNN, call patchwise_alignment(...). "
- "For single-cov, ensure input is 2D via net._preprocess_inputs(...)."
- )
-
-
-@torch.no_grad()
-def expected_alignment_distribution(eigenvalues, relative=True, valid_rotation=True, with_rotation=True, bins=11, num_tests=100):
- """
- From a set of eigenvalues, simulate random weight vectors
- (optional orthonormal rotation) to measure an 'expected' alignment distribution.
- """
- import math
- import numpy as np
-
- N = len(eigenvalues)
- if relative:
- eigenvalues = eigenvalues / eigenvalues.sum()
- eigenvalues = eigenvalues.view(-1, 1).expand(-1, N * num_tests)
-
- device = eigenvalues.device
- if with_rotation:
- if valid_rotation:
- mixing = []
- for _ in range(num_tests):
- mat = torch.normal(0, 1 / math.sqrt(N), (N, N)).to(device)
- q, _ = torch.linalg.qr(mat)
- mixing.append(q.T)
- coefficients = torch.cat(mixing, dim=1) ** 2
- else:
- coefficients = torch.normal(0, 1 / math.sqrt(N), (N, N * num_tests)).to(device) ** 2
- else:
- coefficients = torch.ones((N, N * num_tests), device=device)
-
- weights = eigenvalues * coefficients
- align = torch.sum(eigenvalues * weights, dim=0) / weights.sum(dim=0)
- align_np = align.cpu().numpy()
-
- counts, bin_edges = np.histogram(align_np, bins=bins, density=True)
- centers = (bin_edges[:-1] + bin_edges[1:]) / 2.0
- return (torch.tensor(counts, dtype=torch.float),
- torch.tensor(bin_edges, dtype=torch.float),
- torch.tensor(centers, dtype=torch.float))
\ No newline at end of file
diff --git a/_arxiv/_archive/alignment_preref/alignment_metrics_rem.py b/_arxiv/_archive/alignment_preref/alignment_metrics_rem.py
deleted file mode 100644
index eafbd8bb..00000000
--- a/_arxiv/_archive/alignment_preref/alignment_metrics_rem.py
+++ /dev/null
@@ -1,43 +0,0 @@
-"""
-Alignment metrics for neural network analysis.
-
-DEPRECATED: This module is deprecated and will be removed in a future version.
-Please import directly from `alignment.metrics` instead.
-
-This module provides various metrics for measuring alignment between weight vectors
-and activation vectors, including RQ (representation quality), MI (mutual information),
-and other metrics for analyzing neural network representations.
-"""
-
-import warnings
-from typing import Dict, List, Tuple, Union, Optional, Callable
-
-# Show deprecation warning when the module is imported
-warnings.warn(
- "The alignment.alignment_metrics module is deprecated and will be removed in a future version. "
- "Please import directly from alignment.metrics instead.",
- DeprecationWarning,
- stacklevel=2
-)
-
-# Import from metrics_utils instead of metrics.py to avoid circular dependencies
-from alignment.utils.metrics_utils import (
- AlignmentMetricBase,
- RQMetric,
- MIMetric,
- WeightSimilarityMetric,
- NodeRedundancyMetric,
- AlignmentMetricsFactory as AlignmentMetrics,
- alignment
-)
-
-# For backward compatibility, re-export everything
-__all__ = [
- 'AlignmentMetricBase',
- 'RQMetric',
- 'MIMetric',
- 'WeightSimilarityMetric',
- 'NodeRedundancyMetric',
- 'AlignmentMetrics',
- 'alignment'
-]
\ No newline at end of file
diff --git a/_arxiv/_archive/alignment_preref/config.py b/_arxiv/_archive/alignment_preref/config.py
deleted file mode 100644
index b80eef6a..00000000
--- a/_arxiv/_archive/alignment_preref/config.py
+++ /dev/null
@@ -1,116 +0,0 @@
-# --------------------------------------------
-# config.py
-# --------------------------------------------
-from os import PathLike
-from dataclasses import dataclass, field
-from typing import (cast, List, Union, Type, TypeVar, Optional)
-from omegaconf import OmegaConf as om
-from omegaconf.errors import OmegaConfBaseException
-
-C = TypeVar("C", bound="BaseConfig")
-
-class BaseConfig:
- @classmethod
- def load(cls: Type[C], path: Union[str, PathLike]) -> C:
- schema = om.structured(cls)
- try:
- raw = om.load(str(path))
- conf = om.merge(schema, raw)
- return cast(C, om.to_object(conf))
- except OmegaConfBaseException as e:
- raise ValueError(str(e))
-
-@dataclass
-class ModelConfig(BaseConfig):
- name: str = "MLP"
- dropout: float = 0.0
- alignment_layers: Optional[dict] = None
-
-@dataclass
-class DatasetConfig(BaseConfig):
- name: str = "MNIST"
- path: Optional[str] = None
-
-@dataclass
-class TrainingConfig(BaseConfig):
- do_train: bool = True
- epochs: int = 10
- batch_size: int = 128
- replicates: int = 1
- name: str = "Adam"
- lr: float = 1e-3
- weight_decay: float = 0.0
-
-@dataclass
-class AlignmentConfig(BaseConfig):
- do_alignment: bool = True
- methods: List[str] = field(default_factory=lambda: ["RQ"])
- measure_expected: bool = True
- frequency: int = 1
- bins: int = 50
- cnn_mode: str = "unfold"
- scale_by_norm: bool = False # If True, scales covariance matrices by their norm before RQ calculation
-
-@dataclass
-class CheckpointingConfig(BaseConfig):
- use_prev: bool = False
- save_checkpoints: bool = False
- frequency: int = 1
- use_wandb: bool = False
-
-@dataclass
-class PlotsConfig(BaseConfig):
- show_loss: bool = True
- show_dropout: bool = True
- show_eigenfeatures: bool = True
- show_eig_dropout: bool = True
-
-@dataclass
-class ExtraConfig(BaseConfig):
- num_drops: int = 9
- progressive_dropout_on_train: bool = False
- single_layer_mode: bool = False
- aggregate_alignment: bool = False
- dropout_pruning_mode: str = "global" # "global", "per_layer_combined", or "per_layer_independent"
- dropout_mode: str = "scaled" # "scaled" or "unscaled" - how targeted dropout is applied to neurons
- exclude_classification_layer: bool = False # If True, excludes the classification layer from pruning
-
-@dataclass
-class ExperimentConfig(BaseConfig):
- experiment: Optional[str] = None
- results_path: Optional[str] = None
- no_save: bool = False
- just_plot: bool = False
- save_networks: bool = False
- show_params: bool = False
- show_all: bool = False
- device: Optional[str] = None
- use_timestamp: bool = False
- timestamp: Optional[str] = None
- ddp_world_size: int = 1
-
- dataset: DatasetConfig = field(default_factory=DatasetConfig)
- model: ModelConfig = field(default_factory=ModelConfig)
- training: TrainingConfig = field(default_factory=TrainingConfig)
- alignment: AlignmentConfig = field(default_factory=AlignmentConfig)
- checkpointing: CheckpointingConfig = field(default_factory=CheckpointingConfig)
- plots: PlotsConfig = field(default_factory=PlotsConfig)
- extra: ExtraConfig = field(default_factory=ExtraConfig)
-
-@dataclass
-class ExperimentArgs:
- """Experiment arguments."""
- dataclass_fields = []
-
- exp_name: str
- name: str
- device: str = "cpu"
-
-
-@dataclass
-class ExtraArgs:
- """Arguments for various experiments."""
- aggregate_alignment: bool = False
- num_drops: int = 9
- dropout_pruning_mode: str = "global" # "global", "per_layer_combined", or "per_layer_independent"
- exclude_classification_layer: bool = False # If True, excludes the classification layer from pruning
\ No newline at end of file
diff --git a/_arxiv/_archive/alignment_preref/datasets.py b/_arxiv/_archive/alignment_preref/datasets.py
deleted file mode 100644
index 4aab900a..00000000
--- a/_arxiv/_archive/alignment_preref/datasets.py
+++ /dev/null
@@ -1,280 +0,0 @@
-# --------------------------------------------
-# datasets.py
-# --------------------------------------------
-
-import sys
-from pathlib import Path
-from warnings import warn
-from abc import ABC, abstractmethod
-
-import torch
-import torchvision
-from torch import nn
-from torch.utils.data.distributed import DistributedSampler
-from torchvision.transforms import v2 as transforms
-
-from alignment.models.base import AlignmentNetwork
-from alignment.config import ExperimentConfig
-
-REQUIRED_PROPERTIES = ["dataset_constructor", "loss_function"]
-
-def default_loader_parameters(
- distributed,
- batch_size=1024,
- num_workers=2,
- shuffle=True,
- pin_memory=True,
- persistent_workers=True,
-):
- """
- contains the default dataloader parameters with the option of updating them
- using key word arguments
-
- # adjusting 'shuffle' when using distributed data parallel.
- """
- default_parameters = dict(
- batch_size=batch_size,
- num_workers=num_workers,
- shuffle=False if distributed else shuffle,
- pin_memory=pin_memory,
- persistent_workers=persistent_workers,
- )
- return default_parameters
-
-class DataSet(ABC):
- """
- Abstract base class for any dataset wrapper.
- Responsible for creating train/test datasets & loaders,
- storing transforms, etc.
- """
-
- def __init__(
- self,
- device=None,
- distributed=False,
- dataset_parameters={},
- transform_parameters={},
- loader_parameters={},
- ):
- self.set_properties()
- self.check_properties()
-
- self.device = device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
- self.distributed = distributed
-
- self.extra_transform = transform_parameters.pop("extra_transform", None)
-
- self.transform_parameters = transform_parameters
- self.make_transform(**transform_parameters)
-
- self.dataloader_parameters = default_loader_parameters(distributed, **loader_parameters)
- self.dataset_parameters = dataset_parameters
- self.load_dataset(**dataset_parameters)
-
- def check_properties(self):
- if not all([hasattr(self, prop) for prop in REQUIRED_PROPERTIES]):
- not_found = [prop for prop in REQUIRED_PROPERTIES if not hasattr(self, prop)]
- raise ValueError(f"The following required properties were not set: {not_found}")
-
- @abstractmethod
- def set_properties(self):
- """
- Must set:
- self.dataset_constructor
- self.loss_function
- and optionally self.dist_params for normalization, etc.
- """
- pass
-
- @abstractmethod
- def dataset_kwargs(self, train=True, **kwargs):
- """
- Returns the dict of kwargs for constructing the train or test dataset.
- """
- pass
-
- def load_dataset(self, **kwargs):
- """Instantiate self.train_dataset and self.test_dataset, then build DataLoaders."""
- self.train_dataset = self.dataset_constructor(**self.dataset_kwargs(train=True, **kwargs))
- self.test_dataset = self.dataset_constructor(**self.dataset_kwargs(train=False, **kwargs))
- self.train_sampler = DistributedSampler(self.train_dataset) if self.distributed else None
- self.test_sampler = DistributedSampler(self.test_dataset) if self.distributed else None
- self.train_loader = torch.utils.data.DataLoader(self.train_dataset, sampler=self.train_sampler, **self.dataloader_parameters)
- self.test_loader = torch.utils.data.DataLoader(self.test_dataset, sampler=self.test_sampler, **self.dataloader_parameters)
-
- def unwrap_batch(self, batch, device=None):
- """
- Simple method for unwrapping a batch (inputs, targets)
- and placing them on the correct device.
- """
- device = self.device if device is None else device
- if self.extra_transform:
- if isinstance(self.extra_transform, list):
- for et in self.extra_transform:
- batch = et(batch)
- else:
- warn("extra_transform is not a list, this is deprecated!", DeprecationWarning, stacklevel=2)
- batch = self.extra_transform(batch)
- inputs, targets = batch
- inputs, targets = inputs.to(device), targets.to(device)
- return inputs, targets
-
- def make_transform(self, center_crop=None, resize=None, flatten=False, out_channels=None):
- """
- Create a transforms.Compose for the dataset.
- Default: converts to float32, normalizes, etc.
-
- # For standard dataset usage, we typically do:
- # transforms.ToImage(),
- # transforms.ToDtype(torch.float32, scale=True),
- # transforms.Normalize(...),
- # etc.
- """
- use_transforms = [
- transforms.ToImage(),
- transforms.ToDtype(torch.float32, scale=True),
- ]
- if center_crop:
- use_transforms.append(transforms.CenterCrop(center_crop))
- use_transforms.append(transforms.Normalize((self.dist_params["mean"]), (self.dist_params["std"])))
- if resize:
- use_transforms.append(transforms.Resize(resize, antialias=True))
- if out_channels:
- use_transforms.append(transforms.Grayscale(num_output_channels=out_channels))
- if flatten:
- use_transforms.append(transforms.Lambda(torch.flatten))
- self.transform = transforms.Compose(use_transforms)
-
- def measure_loss(self, outputs, targets, reduction=None):
- """Compute the loss via self.loss_function, optionally changing reduction."""
- if reduction is None:
- return self.loss_function(outputs, targets)
- standard_reduction = self.loss_function.reduction
- self.loss_function.reduction = reduction
- loss = self.loss_function(outputs, targets)
- self.loss_function.reduction = standard_reduction
- return loss
-
- def measure_accuracy(self, outputs, targets, k=1, percentage=True):
- """
- Return top-k accuracy on classification problems.
- By default, top-1 accuracy in percentage form.
- """
- # Use a simple, direct implementation that's guaranteed to work correctly
- # First check for invalid values
- if torch.isnan(outputs).any() or torch.isinf(outputs).any():
- print("WARNING: NaN or Inf values in measure_accuracy inputs")
- return torch.tensor(0.0, device=outputs.device)
-
- # Get the predicted classes (most likely class for each sample)
- _, predicted = outputs.max(1)
-
- # Calculate accuracy
- correct = (predicted == targets).sum().item()
- total = targets.size(0)
-
- # Convert to percentage if requested
- accuracy = (correct / total) * (100.0 if percentage else 1.0)
-
-
- return torch.tensor(accuracy, device=outputs.device)
-
-
-class MNIST(DataSet):
- def set_properties(self):
- """defines the required properties for MNIST"""
- self.dataset_constructor = torchvision.datasets.MNIST
- self.loss_function = nn.CrossEntropyLoss()
- self.dist_params = dict(mean=[0.1307], std=[0.3081])
-
- def dataset_kwargs(self, train=True, download=False, root=None):
- kwargs = dict(
- train=train,
- download=download,
- root=root,
- transform=self.transform,
- )
- return kwargs
-
-class CIFAR10(DataSet):
- def set_properties(self):
- self.dataset_constructor = torchvision.datasets.CIFAR10
- self.loss_function = nn.CrossEntropyLoss()
- self.dist_params = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-
- def dataset_kwargs(self, train=True, download=False, root=None):
- kwargs = dict(
- train=train,
- download=download,
- root=root,
- transform=self.transform,
- )
- return kwargs
-
-class CIFAR100(CIFAR10):
- def set_properties(self):
- self.dataset_constructor = torchvision.datasets.CIFAR100
- self.loss_function = nn.CrossEntropyLoss()
- self.dist_params = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-
-class ImageNet2012(DataSet):
- def set_properties(self):
- """
- defines the required properties for ImageNet 2012
- (ILSVRC2012) with 1000 classes.
- """
- self.dataset_constructor = torchvision.datasets.ImageNet
- self.loss_function = nn.CrossEntropyLoss()
- self.dist_params = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
- self.center_crop = 224
-
- def dataset_kwargs(self, train=True, root=None):
- kwargs = dict(
- split="train" if train else "val",
- root=root,
- transform=self.transform,
- )
- return kwargs
-
-DATASET_REGISTRY = {
- "MNIST": MNIST,
- "CIFAR10": CIFAR10,
- "CIFAR100": CIFAR100,
- "ImageNet": ImageNet2012,
-}
-
-def get_dataset(
- dataset_name,
- build=False,
- dataset_parameters={},
- transform_parameters={},
- loader_parameters={},
- **kwargs,
-):
- """
- lookup dataset constructor from dataset registry by name.
- if build=True, return an instantiated dataset object,
- otherwise return the constructor.
- """
- if dataset_name not in DATASET_REGISTRY:
- raise ValueError(f"Dataset ({dataset_name}) is not in DATASET_REGISTRY")
- dataset = DATASET_REGISTRY[dataset_name]
- if build:
- if not isinstance(transform_parameters, dict):
- raise TypeError("transform_parameters must be a dictionary")
- return dataset(
- dataset_parameters=dataset_parameters,
- transform_parameters=transform_parameters,
- loader_parameters=loader_parameters,
- **kwargs,
- )
- return dataset
-
-if __name__ == "__main__":
- try:
- yaml_path, args_list = sys.argv[1]
- except IndexError:
- raise ValueError(f"Usage: {sys.argv[0]} [CONFIG_PATH]")
-
- cfg = ExperimentConfig.load(yaml_path)
- dataset = get_dataset(cfg.dataset.name, build=True, dataset_parameters=dict(download=True, root=Path(cfg.dataset.path)))
diff --git a/_arxiv/_archive/alignment_preref/plotting.py b/_arxiv/_archive/alignment_preref/plotting.py
deleted file mode 100644
index ff683467..00000000
--- a/_arxiv/_archive/alignment_preref/plotting.py
+++ /dev/null
@@ -1,861 +0,0 @@
-# --------------------------------------------
-# plotting.py
-# --------------------------------------------
-
-import torch
-import numpy as np
-from tqdm import tqdm
-from matplotlib import pyplot as plt
-import matplotlib as mpl
-
-from alignment.utils import compute_stats_by_type, named_transpose, transpose_list, rms
-
-
-def plot_train_results(exp, train_results, test_results, prms):
- """
- Plot training curves for loss and accuracy, plus test performance.
- Optionally plots alignment across training if present in `train_results`.
-
- If 'loss' or 'accuracy' are missing in train_results (e.g. do_train=false),
- we skip the training curve plot. We can still show test performance if present.
- """
-
- import torch
- import numpy as np
- import matplotlib.pyplot as plt
- import matplotlib as mpl
- from alignment.utils import compute_stats_by_type
-
- # If do_train was false, we might not have train_results["loss"] or train_results["accuracy"].
- # We'll do a check:
- has_train_loss = "loss" in train_results and train_results["loss"].dim() == 2
- has_train_acc = "accuracy" in train_results and train_results["accuracy"].dim() == 2
-
- # In normal training, shape is (num_replicates, num_epochs)
- # If missing, we can skip them. We'll define num_train_epochs = 0 if missing
- if has_train_loss:
- num_train_epochs = train_results["loss"].size(1)
- else:
- num_train_epochs = 0
-
- num_types = len(prms["vals"])
- labels = [f"{prms['name']}={val}" for val in prms["vals"]]
-
- print("getting statistics on run data...")
- plot_alignment = "alignment" in train_results
- alignment = None
- if plot_alignment:
- alignment_by_method = {}
- for record in train_results["alignment"]:
- if isinstance(record, torch.Tensor):
- method_key = 'RQ'
- if record.dim() == 3:
- net_avg = record.mean(dim=0)
- elif record.dim() == 4:
- net_avg = record.squeeze(0).mean(dim=0)
- else:
- raise ValueError(f"Unexpected tensor shape: {record.shape}")
- for layer_idx in range(net_avg.size(0)):
- if method_key not in alignment_by_method:
- alignment_by_method[method_key] = {}
- if layer_idx not in alignment_by_method[method_key]:
- alignment_by_method[method_key][layer_idx] = []
- alignment_by_method[method_key][layer_idx].append(torch.mean(net_avg[layer_idx]).item())
-
- elif isinstance(record, dict):
- data_field = record["data"]
- if isinstance(data_field, list):
- for net_data in data_field:
- for layer_idx, layer_list in enumerate(net_data):
- # FIX: layer_list was a list, so we iterate over it before calling .items()
- if isinstance(layer_list, dict):
- for method_key, tensor in layer_list.items():
- if method_key not in alignment_by_method:
- alignment_by_method[method_key] = {}
- if layer_idx not in alignment_by_method[method_key]:
- alignment_by_method[method_key][layer_idx] = []
- alignment_by_method[method_key][layer_idx].append(torch.mean(tensor).item())
- elif isinstance(layer_list, list):
- for sub_dict in layer_list:
- for method_key, tensor in sub_dict.items():
- if method_key not in alignment_by_method:
- alignment_by_method[method_key] = {}
- if layer_idx not in alignment_by_method[method_key]:
- alignment_by_method[method_key][layer_idx] = []
- alignment_by_method[method_key][layer_idx].append(torch.mean(tensor).item())
- else:
- raise TypeError(f"layer_idx={layer_idx} has unexpected type {type(layer_list)}")
- elif isinstance(data_field, torch.Tensor):
- method_key = 'RQ'
- if data_field.dim() == 3:
- net_avg = data_field.mean(dim=0)
- elif data_field.dim() == 4:
- net_avg = data_field.squeeze(0).mean(dim=0)
- else:
- raise ValueError(f"Unexpected record['data'] shape: {data_field.shape}")
- for layer_idx in range(net_avg.size(0)):
- if method_key not in alignment_by_method:
- alignment_by_method[method_key] = {}
- if layer_idx not in alignment_by_method[method_key]:
- alignment_by_method[method_key][layer_idx] = []
- alignment_by_method[method_key][layer_idx].append(torch.mean(net_avg[layer_idx]).item())
- else:
- raise TypeError(f"record['data'] has unexpected type {type(data_field)}")
- else:
- raise TypeError(f"train_results['alignment'] has an element of type {type(record)}")
-
- alignment_avg = {}
- for method_key, layers_dict in alignment_by_method.items():
- layer_avgs = []
- for layer_idx in sorted(layers_dict.keys()):
- vals = torch.tensor(layers_dict[layer_idx])
- layer_avgs.append(torch.mean(vals))
- alignment_avg[method_key] = torch.stack(layer_avgs, dim=0)
- alignment = alignment_avg
-
- cmap = mpl.colormaps["tab10"]
-
- if has_train_loss:
- train_loss_mean, train_loss_se = compute_stats_by_type(
- train_results["loss"].transpose(0,1),
- num_types=num_types,
- dim=1,
- method="se"
- )
- else:
- train_loss_mean, train_loss_se = None, None
-
- if has_train_acc:
- train_acc_mean, train_acc_se = compute_stats_by_type(
- train_results["accuracy"].transpose(0,1),
- num_types=num_types,
- dim=1,
- method="se"
- )
- else:
- train_acc_mean, train_acc_se = None, None
-
- test_loss_mean, test_loss_se = None, None
- test_acc_mean, test_acc_se = None, None
- if ("loss" in test_results) and (test_results["loss"].dim() == 1):
- test_loss_mean, test_loss_se = compute_stats_by_type(
- test_results["loss"],
- num_types=num_types,
- dim=0,
- method="se"
- )
- if ("accuracy" in test_results) and (test_results["accuracy"].dim() == 1):
- test_acc_mean, test_acc_se = compute_stats_by_type(
- test_results["accuracy"],
- num_types=num_types,
- dim=0,
- method="se"
- )
-
- print("plotting run data...")
- xOffset = [-0.2, 0.2]
- get_x = lambda idx: [xOffset[0] + idx, xOffset[1] + idx]
-
- alpha = 0.3
- figdim = 3
- figratio = 2
- width_ratios = [figdim, figdim / figratio, figdim, figdim / figratio]
- fig, ax = plt.subplots(1, 4, figsize=(sum(width_ratios), figdim), width_ratios=width_ratios, layout="constrained")
-
- for idx, label in enumerate(labels):
- if train_loss_mean is not None:
- cmn = train_loss_mean[:, idx].cpu()
- cse = train_loss_se[:, idx].cpu()
- # Use 1-based epoch numbering for the x-axis
- epochs = [i+1 for i in range(num_train_epochs)]
- ax[0].plot(epochs, cmn, color=cmap(idx), label=label)
- ax[0].fill_between(epochs, cmn + cse, cmn - cse, color=(cmap(idx), alpha))
-
- if test_loss_mean is not None:
- tmn = test_loss_mean[idx].cpu().item()
- tse = test_loss_se[idx].cpu().item()
- ax[1].plot(get_x(idx), [tmn]*2, color=cmap(idx), label=label, lw=4)
- ax[1].plot([idx, idx], [tmn - tse, tmn + tse], color=cmap(idx), lw=1.5)
-
- ax[0].set_xlabel("Training Epoch")
- ax[0].set_ylabel("Loss")
- ax[0].set_title("Training Loss")
- ax[1].set_title("Testing")
-
- if train_loss_mean is not None:
- ax[0].set_ylim(0, None)
- # Set x-axis limits to show full range of epochs, with a minimum range
- if num_train_epochs > 1:
- ax[0].set_xlim(1, num_train_epochs)
- else:
- # For single epoch case, add some padding
- ax[0].set_xlim(0.5, 1.5)
- ylims = ax[0].get_ylim()
- ax[1].set_ylim(ylims)
- ax[1].set_ylabel("Loss")
- ax[1].set_xlim(-0.5, num_types - 0.5)
- ax[1].set_xticks(range(num_types))
- ax[1].set_xticklabels(labels, rotation=45, ha="right", fontsize=8)
-
- for idx, label in enumerate(labels):
- if train_acc_mean is not None:
- cmn = train_acc_mean[:, idx].cpu()
- cse = train_acc_se[:, idx].cpu()
- # Use 1-based epoch numbering for the x-axis (same as loss plot)
- epochs = [i+1 for i in range(num_train_epochs)]
- ax[2].plot(epochs, cmn, color=cmap(idx), label=label)
- ax[2].fill_between(epochs, cmn + cse, cmn - cse, color=(cmap(idx), alpha))
-
- if test_acc_mean is not None:
- tmn = test_acc_mean[idx].cpu().item()
- tse = test_acc_se[idx].cpu().item()
- ax[3].plot(get_x(idx), [tmn]*2, color=cmap(idx), label=label, lw=4)
- ax[3].plot([idx, idx], [tmn - tse, tmn + tse], color=cmap(idx), lw=1.5)
-
- ax[2].set_xlabel("Training Epoch")
- ax[2].set_ylabel("Accuracy (%)")
- ax[2].set_title("Training Accuracy")
- ax[3].set_title("Testing")
- if train_acc_mean is not None:
- ax[2].set_ylim(0, 100)
- # Set x-axis limits to show full range of epochs, with a minimum range
- if num_train_epochs > 1:
- ax[2].set_xlim(1, num_train_epochs)
- else:
- # For single epoch case, add some padding
- ax[2].set_xlim(0.5, 1.5)
- ax[3].set_ylim(0, 100)
- ax[3].set_xlim(-0.5, num_types - 0.5)
- ax[3].set_xticks(range(num_types))
- ax[3].set_xticklabels(labels, rotation=45, ha="right", fontsize=8)
-
- exp.plot_ready("train_test_performance")
-
- if plot_alignment and alignment is not None:
- for method_key, align_tensor in alignment.items():
- num_layers = align_tensor.size(0)
- fig2, ax2 = plt.subplots(1, num_layers, figsize=(num_layers * figdim, figdim), layout="constrained", sharex=True, squeeze=False)
- for layer in range(num_layers):
- val = align_tensor[layer].cpu().item() * 100
- ax2[0, layer].axhline(val, color=cmap(0), label=method_key)
- ax2[0, layer].set_xlabel("Snapshots")
- ax2[0, layer].set_ylabel("Alignment (%)")
- ax2[0, layer].set_title(f"Layer {layer} - {method_key}")
- exp.plot_ready(f"train_alignment_{method_key}")
-
-
-def plot_dropout_results(exp, dropout_results, dropout_parameters, prms, dropout_type="nodes"):
- """
- Plot progressive dropout results for nodes (or other types).
- """
- if not exp.args.plots.show_dropout:
- return
-
- num_types = len(prms["vals"])
- labels = [f"{prms['name']}={val}" for val in prms["vals"]]
- cmap = mpl.colormaps["Set1"]
- alpha = 0.3
- msize = 10
- figdim = 3
-
- # Retrieve the arrays from `dropout_results`:
- loss_high = dropout_results["progdrop_loss_high"] # shape => (N, D, ...) 3D or 4D
- loss_low = dropout_results["progdrop_loss_low"]
- loss_rand = dropout_results["progdrop_loss_rand"]
- acc_high = dropout_results["progdrop_acc_high"]
- acc_low = dropout_results["progdrop_acc_low"]
- acc_rand = dropout_results["progdrop_acc_rand"]
-
- dropout_fraction = dropout_results["dropout_fraction"] # shape => (D,)
- pruning_mode = dropout_results.get("pruning_mode", "global") # Use pruning_mode instead of by_layer
- is_per_layer = pruning_mode != "global" # Check if using per-layer modes
- extra_name = (pruning_mode + "_" if pruning_mode else "") + dropout_type
-
- # For per_layer_combined mode, use the combined results to create a single figure
- if pruning_mode == "per_layer_combined" and "combined_progdrop_loss_high" in dropout_results:
- print("Using combined results for per_layer_combined mode...")
- loss_high = dropout_results["combined_progdrop_loss_high"]
- loss_low = dropout_results["combined_progdrop_loss_low"]
- loss_rand = dropout_results["combined_progdrop_loss_rand"]
- acc_high = dropout_results["combined_progdrop_acc_high"]
- acc_low = dropout_results["combined_progdrop_acc_low"]
- acc_rand = dropout_results["combined_progdrop_acc_rand"]
- is_per_layer = False # Treat as a global figure with single layer results
- extra_name = "per_layer_combined_aggregated_" + dropout_type
-
- print("Computing dropout stats over replicate dimension...")
- # shape => (N, D, E?, L?) => We do `dim=0` to average across replicates
- mean_loss_high, se_loss_high = compute_stats_by_type(loss_high, num_types=num_types, dim=0, method="se")
- mean_loss_low, se_loss_low = compute_stats_by_type(loss_low, num_types=num_types, dim=0, method="se")
- mean_loss_rand, se_loss_rand= compute_stats_by_type(loss_rand, num_types=num_types, dim=0, method="se")
-
- mean_acc_high, se_acc_high = compute_stats_by_type(acc_high, num_types=num_types, dim=0, method="se")
- mean_acc_low, se_acc_low = compute_stats_by_type(acc_low, num_types=num_types, dim=0, method="se")
- mean_acc_rand, se_acc_rand= compute_stats_by_type(acc_rand, num_types=num_types, dim=0, method="se")
-
- # Now each is either 3D or 4D. We'll detect:
- # e.g. mean_loss_high.shape => 3D: (num_types, D, L)
- # 4D: (num_types, D, E, L)
- shape_ = mean_loss_high.shape
- # shape_[0] = num_types, shape_[1] = D
- # we check len(shape_):
- if len(shape_) == 3:
- # interpret => (num_types, D, L)
- # => no epoch dimension => treat E=1
- # We can reshape => (num_types, D, 1, L) to unify logic
- mean_loss_high = mean_loss_high.unsqueeze(2) # => (num_types, D, 1, L)
- se_loss_high = se_loss_high.unsqueeze(2)
- mean_loss_low = mean_loss_low.unsqueeze(2)
- se_loss_low = se_loss_low.unsqueeze(2)
- mean_loss_rand = mean_loss_rand.unsqueeze(2)
- se_loss_rand = se_loss_rand.unsqueeze(2)
-
- mean_acc_high = mean_acc_high.unsqueeze(2)
- se_acc_high = se_acc_high.unsqueeze(2)
- mean_acc_low = mean_acc_low.unsqueeze(2)
- se_acc_low = se_acc_low.unsqueeze(2)
- mean_acc_rand = mean_acc_rand.unsqueeze(2)
- se_acc_rand = se_acc_rand.unsqueeze(2)
-
- shape_ = mean_loss_high.shape # now => (num_types, D, 1, L)
-
- # now shape_ => (num_types, D, E, L)
- # parse
- num_types_, num_drop, num_epochs, num_layers = shape_
- print(f"Shape => (num_types={num_types_}, #drops={num_drop}, #epochs={num_epochs}, #layers={num_layers})")
-
- # Try to get alignment layer names from the first net
- layer_names = None
- if hasattr(exp, "nets") and exp.nets and hasattr(exp.nets[0], "alignment_names"):
- layer_names = exp.nets[0].alignment_names
-
- lines = ["From high", "From low", "Random"]
-
- # ----------------
- # Plot Loss
- # ----------------
- fig_loss, ax_loss = plt.subplots(
- num_epochs, num_layers,
- figsize=(num_layers * figdim, num_epochs * figdim),
- sharex=True,
- sharey=True,
- layout="constrained",
- squeeze=False
- )
-
- ax_loss = np.reshape(ax_loss, (num_epochs, num_layers))
- print("Plotting dropout: loss...")
-
- for eidx in range(num_epochs):
- for lidx in range(num_layers):
- # Title with layer name if we have them
- if layer_names and lidx < len(layer_names):
- layer_str = layer_names[lidx]
- else:
- layer_str = f"Layer {lidx}"
- # If we want "Epoch eidx" or "Snapshot eidx"
- subplot_title = f"Epoch {eidx}, {layer_str}"
-
- # We'll do multiple lines: each line is one "type" + "From high/low/rand"
- # Actually, we have 3 lines in one type or 3 lines total across all types?
- # If you want each type to have 3 lines in the same subplot, do:
- for ty_idx, label_str in enumerate(labels):
- for i_line, line_name in enumerate(lines):
- if line_name == "From high":
- cmn = mean_loss_high[ty_idx, :, eidx, lidx] # shape => (#drop,)
- cse = se_loss_high[ty_idx, :, eidx, lidx]
- elif line_name == "From low":
- cmn = mean_loss_low[ty_idx, :, eidx, lidx]
- cse = se_loss_low[ty_idx, :, eidx, lidx]
- else: # "From random"
- cmn = mean_loss_rand[ty_idx, :, eidx, lidx]
- cse = se_loss_rand[ty_idx, :, eidx, lidx]
-
- ax_loss[eidx, lidx].plot(
- dropout_fraction,
- cmn,
- color=cmap(i_line),
- marker=".",
- markersize=msize,
- label=f"{label_str} ({line_name})",
- )
- ax_loss[eidx, lidx].fill_between(
- dropout_fraction,
- cmn - cse,
- cmn + cse,
- color=(cmap(i_line), alpha)
- )
-
- ax_loss[eidx, lidx].set_title(subplot_title)
- if eidx == num_epochs - 1:
- ax_loss[eidx, lidx].set_xlabel("Dropout Fraction")
- ax_loss[eidx, lidx].set_xlim(0, 1)
- if lidx == 0:
- ax_loss[eidx, lidx].set_ylabel("Loss w/ Dropout")
- ax_loss[eidx, lidx].legend(loc="best")
-
- exp.plot_ready("prog_dropout_" + extra_name + "_loss")
-
- # ----------------
- # Plot Accuracy
- # ----------------
- fig_acc, ax_acc = plt.subplots(
- num_epochs, num_layers,
- figsize=(num_layers * figdim, num_epochs * figdim),
- sharex=True,
- sharey=True,
- layout="constrained",
- squeeze=False
- )
-
- ax_acc = np.reshape(ax_acc, (num_epochs, num_layers))
- print("Plotting dropout: accuracy...")
-
- for eidx in range(num_epochs):
- for lidx in range(num_layers):
- if layer_names and lidx < len(layer_names):
- layer_str = layer_names[lidx]
- else:
- layer_str = f"Layer {lidx}"
- subplot_title = f"Epoch {eidx}, {layer_str}"
-
- for ty_idx, label_str in enumerate(labels):
- for i_line, line_name in enumerate(lines):
- if line_name == "From high":
- cmn = mean_acc_high[ty_idx, :, eidx, lidx]
- cse = se_acc_high[ty_idx, :, eidx, lidx]
- elif line_name == "From low":
- cmn = mean_acc_low[ty_idx, :, eidx, lidx]
- cse = se_acc_low[ty_idx, :, eidx, lidx]
- else: # random
- cmn = mean_acc_rand[ty_idx, :, eidx, lidx]
- cse = se_acc_rand[ty_idx, :, eidx, lidx]
-
- ax_acc[eidx, lidx].plot(
- dropout_fraction,
- cmn,
- color=cmap(i_line),
- marker=".",
- markersize=msize,
- label=f"{label_str} ({line_name})",
- )
- ax_acc[eidx, lidx].fill_between(
- dropout_fraction,
- cmn - cse,
- cmn + cse,
- color=(cmap(i_line), alpha)
- )
-
- ax_acc[eidx, lidx].set_title(subplot_title)
- ax_acc[eidx, lidx].set_ylim(0, 100)
- if eidx == num_epochs - 1:
- ax_acc[eidx, lidx].set_xlabel("Dropout Fraction")
- ax_acc[eidx, lidx].set_xlim(0, 1)
- if lidx == 0:
- ax_acc[eidx, lidx].set_ylabel("Accuracy w/ Dropout")
-
- ax_acc[eidx, lidx].legend(loc="best")
-
- exp.plot_ready("prog_dropout_" + extra_name + "_accuracy")
-
-def plot_eigenfeatures(exp, results, prms):
- """
- method for plotting results related to eigen-analysis (beta, eigenvalues, class_betas).
- """
- beta, eigvals, class_betas, class_names = (
- results["beta"],
- results["eigvals"],
- results["class_betas"],
- results["class_names"],
- )
- beta = [[torch.abs(b) for b in net_beta] for net_beta in beta]
- class_betas = [[rms(cb, dim=2) for cb in net_class_beta] for net_class_beta in class_betas]
-
- num_types = len(prms["vals"])
- labels = [f"{prms['name']}={val}" for val in prms["vals"]]
- cmap = mpl.colormaps["tab10"]
- class_cmap = mpl.colormaps["viridis"].resampled(len(class_names))
-
- print("measuring statistics of eigenfeature analyses...")
-
- beta = [torch.stack(b) for b in transpose_list(beta)]
- eigvals = [torch.stack(ev) for ev in transpose_list(eigvals)]
- class_betas = [torch.stack(cb) for cb in transpose_list(class_betas)]
-
- beta = [b / b.sum(dim=2, keepdim=True) for b in beta]
- eigvals = [ev / ev.sum(dim=1, keepdim=True) for ev in eigvals]
- class_betas = [cb / cb.sum(dim=2, keepdim=True) for cb in class_betas]
-
- statprms = lambda method: dict(num_types=num_types, dim=0, method=method)
- mean_evals, var_evals = named_transpose([compute_stats_by_type(ev, **statprms("var")) for ev in eigvals])
- sorted_beta = [torch.sort(b, descending=True, dim=2).values for b in beta]
- mean_beta, se_beta = named_transpose([compute_stats_by_type(b, **statprms("var")) for b in beta])
- mean_sorted, se_sorted = named_transpose([compute_stats_by_type(b, **statprms("var")) for b in sorted_beta])
- mean_class_beta, se_class_beta = named_transpose([compute_stats_by_type(cb, **statprms("var")) for cb in class_betas])
-
- print("plotting eigenfeature results...")
- figdim = 3
- alpha = 0.3
- num_layers = len(mean_beta)
- fig, ax = plt.subplots(2, num_layers, figsize=(num_layers * figdim, figdim * 2), layout="constrained", squeeze=False)
-
- for layer in range(num_layers):
- num_input = mean_evals[layer].size(1)
- num_nodes = mean_beta[layer].size(1)
- for idx, label in enumerate(labels):
- mn_ev = mean_evals[layer][idx]
- se_ev = var_evals[layer][idx]
- mn_beta_ = torch.mean(mean_beta[layer][idx], dim=0)
- se_beta_ = torch.std(mean_beta[layer][idx], dim=0) / np.sqrt(num_nodes)
- mn_sort = torch.mean(mean_sorted[layer][idx], dim=0)
- se_sort = torch.std(mean_sorted[layer][idx], dim=0) / np.sqrt(num_nodes)
-
- ax[0, layer].plot(
- range(num_input),
- mn_ev,
- color=cmap(idx),
- linestyle="--",
- label="eigvals" if idx == 0 else None,
- )
- ax[0, layer].plot(range(num_input), mn_beta_, color=cmap(idx), label=label)
- ax[0, layer].fill_between(range(num_input), mn_beta_ + se_beta_, mn_beta_ - se_beta_, color=(cmap(idx), alpha))
-
- ax[1, layer].plot(range(num_input), mn_sort, color=cmap(idx), label=label)
- ax[1, layer].fill_between(range(num_input), mn_sort + se_sort, mn_sort - se_sort, color=(cmap(idx), alpha))
-
- ax[0, layer].set_xscale("log")
- ax[1, layer].set_xscale("log")
- ax[0, layer].set_xlabel("Input Dimension")
- ax[1, layer].set_xlabel("Sorted Input Dim")
- ax[0, layer].set_ylabel("Relative Eigval & Beta")
- ax[1, layer].set_ylabel("Relative Beta (Sorted)")
- ax[0, layer].set_title(f"Layer {layer}")
- ax[1, layer].set_title(f"Layer {layer}")
-
- if layer == num_layers - 1:
- ax[0, layer].legend(loc="best")
- ax[1, layer].legend(loc="best")
-
- exp.plot_ready("eigenfeatures")
-
- fig, ax = plt.subplots(1, num_layers, figsize=(num_layers * figdim, figdim), layout="constrained", squeeze=False)
- for layer in range(num_layers):
- num_input = mean_evals[layer].size(1)
- num_nodes = mean_beta[layer].size(1)
- for idx, label in enumerate(labels):
- mn_ev = mean_evals[layer][idx]
- se_ev = var_evals[layer][idx]
- mn_beta_ = torch.mean(mean_beta[layer][idx], dim=0)
- se_beta_ = torch.std(mean_beta[layer][idx], dim=0) / np.sqrt(num_nodes)
- ax[0, layer].plot(
- range(num_input),
- mn_ev,
- color=cmap(idx),
- linestyle="--",
- label="eigvals" if idx == 0 else None,
- )
- ax[0, layer].plot(range(num_input), mn_beta_, color=cmap(idx), label=label)
- ax[0, layer].fill_between(range(num_input), mn_beta_ + se_beta_, mn_beta_ - se_beta_, color=(cmap(idx), alpha))
-
- ax[0, layer].set_xscale("log")
- ax[0, layer].set_yscale("log")
- ax[0, layer].set_xlabel("Input Dimension")
- ax[0, layer].set_ylabel("Relative Eigval & Beta")
- ax[0, layer].set_title(f"Layer {layer}")
-
- if layer == num_layers - 1:
- ax[0, layer].legend(loc="best")
-
- exp.plot_ready("eigenfeatures_loglog")
-
- fig, ax = plt.subplots(
- num_types,
- num_layers,
- figsize=(num_layers * figdim, figdim * num_types),
- layout="constrained",
- squeeze=False,
- )
- ax = np.reshape(ax, (num_types, num_layers))
- for layer in range(num_layers):
- num_input = mean_evals[layer].size(1)
- for idx, label in enumerate(labels):
- for idx_class, class_name in enumerate(class_names):
- mn_data = mean_class_beta[layer][idx][idx_class]
- se_data = se_class_beta[layer][idx][idx_class]
- ax[idx, layer].plot(range(num_input), mn_data, color=class_cmap(idx_class), label=class_name)
- ax[idx, layer].fill_between(
- range(num_input),
- mn_data + se_data,
- mn_data - se_data,
- color=(class_cmap(idx_class), alpha),
- )
-
- ax[idx, layer].set_xscale("log")
- ax[idx, layer].set_yscale("linear")
- ax[idx, layer].set_xlabel("Input Dimension")
- if layer == 0:
- ax[idx, layer].set_ylabel(f"{label}\nClass Loading (RMS)")
- if idx == 0:
- ax[idx, layer].set_title(f"Layer {layer}")
- if layer == num_layers - 1:
- ax[idx, layer].legend(loc="upper right", fontsize=6)
-
- exp.plot_ready("class_eigenfeatures")
-
-
-def plot_adversarial_results(exp, eigen_results, adversarial_results, prms):
- """
- Visualize adversarial attack results (accuracy vs epsilon) and structure
- of adversarial perturbations in the eigenvector basis.
- """
- accuracy, beta, eigvals = (
- adversarial_results["accuracy"],
- adversarial_results["betas"],
- eigen_results["eigvals"],
- )
- epsilons, use_sign = adversarial_results["use_sign"], adversarial_results["use_sign"]
-
- num_types = len(prms["vals"])
- labels = [f"{prms['name']}={val}" for val in prms["vals"]]
- cmap = mpl.colormaps["tab10"]
-
- print("measuring statistics of adversarial analyses...")
-
- accuracy = torch.stack([torch.stack(acc) for acc in transpose_list(accuracy)]) # shape (num_epsilon, num_nets)
- eigvals = [torch.stack(ev) for ev in transpose_list(eigvals)]
- beta = [torch.stack(b) for b in beta]
-
- beta = [b / b.sum(dim=2, keepdim=True) for b in beta]
- eigvals = [ev / ev.sum(dim=1, keepdim=True) for ev in eigvals]
-
- statprms = lambda dim, method: dict(num_types=num_types, dim=dim, method=method)
- mean_acc, se_acc = compute_stats_by_type(accuracy, **statprms(1, "var"))
-
- from alignment.core.utils import named_transpose
- mean_beta, se_beta = named_transpose([compute_stats_by_type(b, **statprms(1, "var")) for b in beta])
- mean_evals, var_evals = named_transpose([compute_stats_by_type(ev, **statprms(0, "var")) for ev in eigvals])
-
- print("plotting adversarial success results...")
- figdim = 3
- alpha = 0.3
- num_layers = len(mean_beta)
- fig, ax = plt.subplots(1, 1, figsize=(figdim, figdim), layout="constrained")
- for idx, label in enumerate(labels):
- ax.plot(epsilons, mean_acc[:, idx], color=cmap(idx), label=label)
- ax.set_xlabel("Epsilon")
- ax.set_ylabel("Accuracy")
- ax.set_title("Adversarial Attack Success")
- ax.legend(loc="best")
-
- exp.plot_ready("adversarial_success")
-
- print("plotting adversarial structure...")
- fig, ax = plt.subplots(1, num_layers, figsize=(num_layers * figdim, figdim), layout="constrained", squeeze=False)
- for layer in range(num_layers):
- num_input = mean_beta[layer].size(2)
- for idx, label in enumerate(labels):
- ax[0, layer].plot(
- range(num_input),
- torch.nanmean(mean_beta[layer][:, idx], dim=0).detach(),
- color=cmap(idx),
- label=label,
- )
- ax[0, layer].set_xscale("log")
- ax[0, layer].set_xlabel("Input Dimension")
- ax[0, layer].set_ylabel("Average Component of Pertubation")
- ax[0, layer].set_title(f"Layer {layer}")
- if layer == num_layers - 1:
- ax[0, layer].legend(loc="best", fontsize=6)
-
- exp.plot_ready("adversarial_structure")
-
-
-def plot_rf(rf, width, alignment=None, alignBounds=None, showRFs=None, figSize=5):
- """
- Helper for visualizing receptive fields in a grid.
- Optionally color-coded by alignment.
- """
- if showRFs is not None:
- rf = rf.reshape(rf.shape[0], -1)
- idxRandom = np.random.choice(range(rf.shape[0]), showRFs, replace=False)
- rf = rf[idxRandom, :]
- else:
- showRFs = rf.shape[0]
- rf = rf.T / np.abs(rf).max(axis=1)
- rf = rf.T
- rf = rf.reshape(showRFs, width, width)
- if alignment is not None:
- cmap = mpl.cm.get_cmap("rainbow", rf.shape[0])
- cmapPeak = lambda x: cmap(x)
- if alignBounds is not None:
- alignment = alignment - alignBounds[0]
- alignment = alignment / (alignBounds[1] - alignBounds[0])
- else:
- alignment = alignment - alignment.min()
- alignment = alignment / alignment.max()
-
- n = int(np.ceil(np.sqrt(rf.shape[0])))
- fig, axes = plt.subplots(nrows=n, ncols=n, sharex=True, sharey=True, squeeze=False)
- fig.set_size_inches(figSize, figSize)
-
- N = 1000
- for i in tqdm(range(rf.shape[0])):
- ax = axes[i // n][i % n]
- if alignment is not None:
- vals = np.ones((N, 4))
- cAlignment = alignment[i].numpy()
- cPeak = cmapPeak(alignment[i].numpy())
- vals[:, 0] = np.linspace(0, cPeak[0], N)
- vals[:, 1] = np.linspace(0, cPeak[1], N)
- vals[:, 2] = np.linspace(0, cPeak[2], N)
- usecmap = mpl.colors.ListedColormap(vals)
- ax.imshow(rf[i], cmap=usecmap, vmin=-1, vmax=1)
- else:
- ax.imshow(rf[i], cmap="gray", vmin=-1, vmax=1)
- ax.set_xticks([])
- ax.set_yticks([])
- ax.set_aspect("equal")
- for j in range(rf.shape[0], n * n):
- ax = axes[j // n][j % n]
- ax.imshow(np.ones_like(rf[0]) * -1, cmap="gray", vmin=-1, vmax=1)
- ax.set_xticks([])
- ax.set_yticks([])
- ax.set_aspect("equal")
- fig.subplots_adjust(wspace=0.0, hspace=0.0)
- return fig
-
-
-
-
-def plot_grad_alignment_correlation(results, return_fig=False):
- """
- If return_fig=True, returns a matplotlib Figure instead of displaying.
- Expects results["grad_alignment_corr"] to be a list of dicts:
- {
- "epoch": int,
- "batch": int,
- "net": int,
- "corr": { layer_name -> float or None }
- }
- Plots correlation vs. epoch for each layer.
- """
-
- if "grad_alignment_corr" not in results:
- print("No grad_alignment_corr found in results")
- if return_fig:
- fig = plt.figure()
- return fig
- return
-
- layer_corr_data = {}
- for record in results["grad_alignment_corr"]:
- epoch = record["epoch"]
- if not isinstance(epoch, int):
- continue
- for layer_name, cval in record["corr"].items():
- if cval is None:
- continue
- if layer_name not in layer_corr_data:
- layer_corr_data[layer_name] = []
- layer_corr_data[layer_name].append((epoch, cval))
-
- if not layer_corr_data:
- print("No valid correlation data to plot")
- if return_fig:
- fig = plt.figure()
- return fig
- return
-
- fig, axes = plt.subplots(len(layer_corr_data), 1, figsize=(6, 3 * len(layer_corr_data)), sharex=True)
- if not isinstance(axes, np.ndarray):
- axes = [axes]
- sorted_layers = sorted(layer_corr_data.keys())
-
- for ax, layer_name in zip(axes, sorted_layers):
- data = layer_corr_data[layer_name]
- data.sort(key=lambda x: x[0])
- epochs = [d[0] for d in data]
- corrs = [d[1] for d in data]
- ax.plot(epochs, corrs, marker="o", label=layer_name)
- ax.set_ylabel("Grad vs Align Corr")
- ax.set_title(f"Layer {layer_name}")
- ax.legend(loc="best")
-
- axes[-1].set_xlabel("Epoch")
- fig.tight_layout()
-
- if return_fig:
- return fig
- else:
- plt.show()
-
-def plot_alignment_change_correlation(results, return_fig=False):
- """
- If return_fig=True, returns a matplotlib Figure instead of displaying.
- Expects results["alignment"] to hold epoch-wise alignment data so we can
- compute correlation of alignment(e) vs alignment(e+1). This is purely illustrative.
- """
- if "alignment" not in results:
- print("No alignment key in results")
- if return_fig:
- fig = plt.figure()
- return fig
- return
-
- # Example: collect epoch->some combined alignment vector
- epoch_alignments = {}
- for record in results["alignment"]:
- ep = record.get("epoch", None)
- if not isinstance(ep, int):
- continue
- if not record.get("data"):
- continue
- # Suppose record["data"] is list-of-nets => each net is list-of-layers => dict {method->(out_features,)}
- # We'll only use the first net as an example
- net_data = record["data"][0]
- layer_vecs = []
- for layer_dict in net_data:
- # if "RQ" in layer_dict => shape (out_features,)
- if "RQ" in layer_dict:
- layer_vecs.append(layer_dict["RQ"].detach().cpu())
- if layer_vecs:
- combined = torch.cat(layer_vecs, dim=0)
- epoch_alignments[ep] = combined
-
- sorted_eps = sorted(epoch_alignments.keys())
- if len(sorted_eps) < 2:
- print("Not enough epochs with alignment to measure changes")
- if return_fig:
- fig = plt.figure()
- return fig
- return
-
- corrs = []
- xvals = []
- for i in range(len(sorted_eps) - 1):
- e1 = sorted_eps[i]
- e2 = sorted_eps[i+1]
- a1 = epoch_alignments[e1]
- a2 = epoch_alignments[e2]
- if a1.shape == a2.shape:
- stack = torch.stack([a1, a2], dim=0)
- cc = torch.corrcoef(stack)
- cval = cc[0, 1].item()
- corrs.append(cval)
- xvals.append((e1 + e2) / 2.0)
- else:
- corrs.append(float("nan"))
- xvals.append((e1 + e2) / 2.0)
-
- fig, ax = plt.subplots(figsize=(6,4))
- ax.plot(xvals, corrs, marker="o", label="Align(e) vs Align(e+1) Corr")
- ax.set_xlabel("Epoch midpoint")
- ax.set_ylabel("Correlation")
- ax.set_title("Alignment Change Correlation")
- ax.legend(loc="best")
- fig.tight_layout()
-
- if return_fig:
- return fig
- else:
- plt.show()
\ No newline at end of file
diff --git a/_arxiv/_archive/alignment_preref/plotting_rem.py b/_arxiv/_archive/alignment_preref/plotting_rem.py
deleted file mode 100644
index 362bb61c..00000000
--- a/_arxiv/_archive/alignment_preref/plotting_rem.py
+++ /dev/null
@@ -1,576 +0,0 @@
-"""
-Plotting utilities for visualization of experiment results.
-
-DEPRECATED: This module is kept for backward compatibility.
-Please use alignment.utils.plotting instead.
-
-This module provides functions for visualizing experimental results from
-alignment studies, including progressive dropout analysis with different
-pruning strategies and other alignment metrics.
-"""
-
-import warnings
-import logging
-import math
-import os
-from typing import Dict, List, Optional, Tuple, Union, Any
-import json
-
-import numpy as np
-import torch
-import matplotlib as mpl
-import matplotlib.pyplot as plt
-import seaborn as sns
-from matplotlib.figure import Figure
-from matplotlib.axes import Axes
-
-# Import the new modules for re-export
-from alignment.utils.plotting import (
- plot_dropout_results,
- plot_experiment_summary,
- plot_dropout_comparison,
- log_plots_to_wandb
-)
-
-# Show deprecation warning on import
-warnings.warn(
- "The alignment.plotting module is deprecated and will be removed in a future version. "
- "Please use alignment.utils.plotting instead.",
- DeprecationWarning,
- stacklevel=2
-)
-
-logger = logging.getLogger(__name__)
-
-# Set default figure size
-plt.rcParams["figure.figsize"] = (10, 6)
-# Use SVG backend for better quality
-plt.rcParams["figure.dpi"] = 120
-
-
-def plot_pruning_experiments(
- results: Dict[str, Any],
- figure_path: Optional[str] = None,
- title_prefix: str = "Progressive Dropout",
-) -> List[str]:
- """
- Plot results from pruning experiments.
-
- DEPRECATED: Use plot_dropout_results in alignment.utils.plotting instead.
-
- Args:
- results: Dictionary of results
- figure_path: Path to save figures
- title_prefix: Prefix for plot titles
-
- Returns:
- List of saved figure paths
- """
- warnings.warn(
- "plot_pruning_experiments is deprecated. Use plot_dropout_results instead.",
- DeprecationWarning,
- stacklevel=2
- )
-
- # Extract pruning and dropout modes
- pruning_mode = results.get("pruning_mode", "global_joint")
- dropout_mode = results.get("dropout_mode", "global")
-
- # Map old modes to new ones
- if pruning_mode == "global":
- pruning_mode = "global_joint"
- elif pruning_mode == "per_layer_combined":
- pruning_mode = "layer_wise"
- elif pruning_mode == "per_layer_independent":
- pruning_mode = "layer_isolated"
-
- # Call the new function
- return plot_dropout_results(
- results,
- figure_path=figure_path,
- pruning_mode=pruning_mode,
- dropout_mode=dropout_mode,
- title_prefix=title_prefix
- )
-
-
-def plot_per_layer_independent(
- results: Dict[str, Any],
- figure_path: Optional[str] = None,
- title_prefix: str = "Per-Layer Independent Pruning",
-) -> List[str]:
- """
- Plot results from per-layer independent pruning.
-
- DEPRECATED: Use plot_dropout_results with pruning_mode="layer_isolated" instead.
-
- Args:
- results: Dictionary of results
- figure_path: Path to save figures
- title_prefix: Prefix for plot titles
-
- Returns:
- List of saved figure paths
- """
- warnings.warn(
- "plot_per_layer_independent is deprecated. Use plot_dropout_results with pruning_mode='layer_isolated' instead.",
- DeprecationWarning,
- stacklevel=2
- )
-
- # Force pruning mode to be layer_isolated (formerly per_layer_independent)
- results["pruning_mode"] = "layer_isolated"
-
- # Call the main plotting function
- return plot_dropout_results(
- results,
- figure_path=figure_path,
- pruning_mode="layer_isolated",
- dropout_mode=results.get("dropout_mode", "global"),
- title_prefix=title_prefix
- )
-
-
-def plot_dropout_results(
- results: Dict,
- plot_dir: str,
- pruning_mode: str = "global",
- dropout_mode: str = "scaled",
- title_prefix: str = "Dropout"
-) -> Dict[str, str]:
- """
- Plot dropout experiment results using a style similar to the preref implementation.
-
- Args:
- results: Results dictionary from progressive_dropout
- plot_dir: Directory to save plots to
- pruning_mode: Pruning mode used in the experiment
- dropout_mode: Dropout mode used in the experiment
- title_prefix: Prefix for plot titles
-
- Returns:
- Dictionary mapping plot types to file paths
- """
- # Create the output directory if it doesn't exist
- os.makedirs(plot_dir, exist_ok=True)
-
- # Check for error in results
- if "error" in results:
- logger.error(f"Cannot plot results due to error: {results['error']}")
- return {}
-
- # Extract dropout fractions
- dropout_fractions = results.get("dropout_fractions", [])
- if not isinstance(dropout_fractions, (list, np.ndarray)) or len(dropout_fractions) == 0:
- logger.error("No dropout fractions found in results")
- return {}
-
- # Set up figure parameters
- strategies = ["high_rq", "low_rq", "random"]
- colors = {"high_rq": "blue", "low_rq": "red", "random": "green"}
- linestyles = {"high_rq": "-", "low_rq": "-", "random": "-"}
- markers = {"high_rq": "o", "low_rq": "s", "random": "^"}
- strategy_labels = {"high_rq": "High RQ", "low_rq": "Low RQ", "random": "Random"}
-
- # Create dictionary to store file paths
- file_paths = {}
-
- # ----------------
- # Plot Accuracy (single figure with all strategies)
- # ----------------
- fig, ax = plt.subplots(figsize=(10, 6))
-
- for strategy in strategies:
- accs = results.get("accuracies", {}).get(strategy)
- if not isinstance(accs, (list, np.ndarray)) or len(accs) == 0:
- logger.warning(f"No accuracy data for strategy {strategy}")
- continue
-
- ax.plot(
- dropout_fractions,
- accs,
- marker=markers[strategy],
- linestyle=linestyles[strategy],
- color=colors[strategy],
- linewidth=2,
- markersize=6,
- label=strategy_labels[strategy]
- )
-
- ax.set_xlabel('Dropout Fraction', fontsize=14)
- ax.set_ylabel('Accuracy (%)', fontsize=14)
- ax.set_title(f'{title_prefix} - {pruning_mode.replace("_", " ").title()} Pruning ({dropout_mode})', fontsize=16)
- ax.grid(True, linestyle='--', alpha=0.7)
- ax.legend(fontsize=12)
- ax.set_ylim([0, 100]) # Fixed y-axis for accuracy
-
- accuracy_file = os.path.join(plot_dir, f"{title_prefix.lower().replace(' ', '_')}_{pruning_mode}_accuracy.png")
- plt.savefig(accuracy_file, dpi=300, bbox_inches='tight')
- plt.close(fig)
-
- file_paths["accuracy"] = accuracy_file
- logger.info(f"Saved accuracy plot to {accuracy_file}")
-
- # ----------------
- # Plot Loss (single figure with all strategies)
- # ----------------
- fig, ax = plt.subplots(figsize=(10, 6))
-
- for strategy in strategies:
- losses = results.get("losses", {}).get(strategy)
- if not isinstance(losses, (list, np.ndarray)) or len(losses) == 0:
- logger.warning(f"No loss data for strategy {strategy}")
- continue
-
- ax.plot(
- dropout_fractions,
- losses,
- marker=markers[strategy],
- linestyle=linestyles[strategy],
- color=colors[strategy],
- linewidth=2,
- markersize=6,
- label=strategy_labels[strategy]
- )
-
- ax.set_xlabel('Dropout Fraction', fontsize=14)
- ax.set_ylabel('Loss (%)', fontsize=14)
- ax.set_title(f'{title_prefix} Loss - {pruning_mode.replace("_", " ").title()} Pruning ({dropout_mode})', fontsize=16)
- ax.grid(True, linestyle='--', alpha=0.7)
- ax.legend(fontsize=12)
-
- loss_file = os.path.join(plot_dir, f"{title_prefix.lower().replace(' ', '_')}_{pruning_mode}_loss.png")
- plt.savefig(loss_file, dpi=300, bbox_inches='tight')
- plt.close(fig)
-
- file_paths["loss"] = loss_file
- logger.info(f"Saved loss plot to {loss_file}")
-
- # ----------------
- # Plot Individual Strategy Comparisons
- # ----------------
- for strategy in strategies:
- accs = results.get("accuracies", {}).get(strategy)
- if not isinstance(accs, (list, np.ndarray)) or len(accs) == 0:
- continue
-
- fig, ax = plt.subplots(figsize=(10, 6))
-
- ax.plot(
- dropout_fractions,
- accs,
- marker=markers[strategy],
- linestyle=linestyles[strategy],
- color=colors[strategy],
- linewidth=2,
- markersize=6,
- label=f"{strategy_labels[strategy]} Accuracy"
- )
-
- # Add loss to the same plot if available
- losses = results.get("losses", {}).get(strategy)
- if isinstance(losses, (list, np.ndarray)) and len(losses) > 0:
- ax.plot(
- dropout_fractions,
- losses,
- marker=markers[strategy],
- linestyle='--',
- color=colors[strategy],
- linewidth=2,
- markersize=6,
- label=f"{strategy_labels[strategy]} Loss"
- )
-
- ax.set_xlabel('Dropout Fraction', fontsize=14)
- ax.set_ylabel('Percentage (%)', fontsize=14)
- ax.set_title(f'{title_prefix} - {strategy_labels[strategy]} ({dropout_mode})', fontsize=16)
- ax.grid(True, linestyle='--', alpha=0.7)
- ax.legend(fontsize=12)
-
- strategy_file = os.path.join(plot_dir, f"{title_prefix.lower().replace(' ', '_')}_{pruning_mode}_{strategy}.png")
- plt.savefig(strategy_file, dpi=300, bbox_inches='tight')
- plt.close(fig)
-
- file_paths[f"{strategy}_comparison"] = strategy_file
- logger.info(f"Saved {strategy} comparison plot to {strategy_file}")
-
- # ----------------
- # Plot Alignment Values (if available)
- # ----------------
- alignment_values = results.get("alignment_values", {})
- if alignment_values:
- # Create one plot per layer
- for layer_idx in range(len(alignment_values.get("high_rq", [{}])[0] or [])):
- fig, ax = plt.subplots(figsize=(10, 6))
-
- for strategy in strategies:
- # Extract alignment values for this layer across dropout fractions
- if strategy in alignment_values:
- layer_alignments = []
- for frac_idx, alignment in enumerate(alignment_values[strategy]):
- if alignment and layer_idx < len(alignment):
- layer_alignments.append(alignment[layer_idx])
- else:
- layer_alignments.append(None)
-
- # Filter out None values
- valid_fractions = []
- valid_alignments = []
- for frac_idx, alignment in enumerate(layer_alignments):
- if alignment is not None:
- valid_fractions.append(dropout_fractions[frac_idx])
- valid_alignments.append(alignment)
-
- if valid_fractions and valid_alignments:
- ax.plot(
- valid_fractions,
- valid_alignments,
- marker=markers[strategy],
- linestyle=linestyles[strategy],
- color=colors[strategy],
- linewidth=2,
- markersize=6,
- label=strategy_labels[strategy]
- )
-
- ax.set_xlabel('Dropout Fraction', fontsize=14)
- ax.set_ylabel('Alignment Value', fontsize=14)
- ax.set_title(f'Layer {layer_idx+1} Alignment - {pruning_mode.replace("_", " ").title()}', fontsize=16)
- ax.grid(True, linestyle='--', alpha=0.7)
- ax.legend(fontsize=12)
-
- alignment_file = os.path.join(plot_dir, f"{title_prefix.lower().replace(' ', '_')}_{pruning_mode}_layer{layer_idx+1}_alignment.png")
- plt.savefig(alignment_file, dpi=300, bbox_inches='tight')
- plt.close(fig)
-
- file_paths[f"layer{layer_idx+1}_alignment"] = alignment_file
- logger.info(f"Saved layer {layer_idx+1} alignment plot to {alignment_file}")
-
- # Save raw results as JSON for further analysis
- try:
- # Helper function to safely convert tensors and numpy arrays to lists
- def safe_convert(obj):
- if isinstance(obj, np.ndarray):
- return obj.tolist()
- elif isinstance(obj, torch.Tensor):
- return obj.cpu().tolist()
- elif isinstance(obj, dict):
- return {k: safe_convert(v) for k, v in obj.items()}
- elif isinstance(obj, list):
- return [safe_convert(item) for item in obj]
- else:
- return obj
-
- # Create a simplified version of results for JSON serialization
- json_results = {
- "dropout_fractions": safe_convert(dropout_fractions),
- "accuracies": safe_convert(results.get("accuracies", {})),
- "losses": safe_convert(results.get("losses", {})),
- }
-
- # Exclude alignment_values as they can be complex and not easily serializable
-
- json_path = os.path.join(plot_dir, f"{title_prefix.lower().replace(' ', '_')}_{pruning_mode}_results.json")
- with open(json_path, 'w') as f:
- json.dump(json_results, f, indent=2)
-
- file_paths["json_results"] = json_path
- logger.info(f"Saved raw results to {json_path}")
- except Exception as e:
- logger.error(f"Error saving JSON results: {str(e)}")
-
- return file_paths
-
-
-def plot_experiment_summary(
- results: Dict,
- plot_dir: str
-) -> Dict[str, str]:
- """
- Generate a comprehensive summary of experiment results.
-
- Args:
- results: Results dictionary from the experiment
- plot_dir: Directory to save plots to
-
- Returns:
- Dictionary mapping plot types to file paths
- """
- # Create the output directory if it doesn't exist
- os.makedirs(plot_dir, exist_ok=True)
-
- # Prepare figure
- fig = plt.figure(figsize=(15, 12))
- gs = fig.add_gridspec(2, 2, width_ratios=[1, 1], height_ratios=[1, 1])
-
- # Panel 1: Configuration summary
- ax1 = fig.add_subplot(gs[0, 0])
- ax1.axis('off')
-
- # Extract config info
- config = results.get("config", {})
- config_text = []
-
- if hasattr(config, "model") and hasattr(config.model, "model_name"):
- config_text.append(f"Model: {config.model.model_name}")
-
- if hasattr(config, "dataset") and hasattr(config.dataset, "dataset_name"):
- config_text.append(f"Dataset: {config.dataset.dataset_name}")
-
- if hasattr(config, "alignment") and hasattr(config.alignment, "metric"):
- config_text.append(f"Alignment Metric: {config.alignment.metric}")
-
- if hasattr(config, "alignment"):
- if hasattr(config.alignment, "dropout_min") and hasattr(config.alignment, "dropout_max"):
- config_text.append(f"Dropout Range: {config.alignment.dropout_min} to {config.alignment.dropout_max}")
- if hasattr(config.alignment, "dropout_steps"):
- config_text.append(f"Dropout Steps: {config.alignment.dropout_steps}")
-
- if hasattr(config, "extra"):
- if hasattr(config.extra, "dropout_mode"):
- config_text.append(f"Dropout Mode: {config.extra.dropout_mode}")
- if hasattr(config.extra, "dropout_pruning_mode"):
- config_text.append(f"Dropout Pruning Mode: {config.extra.dropout_pruning_mode}")
-
- ax1.text(0.05, 0.95, "\n".join(config_text), fontsize=11,
- verticalalignment='top', horizontalalignment='left')
- ax1.set_title("Experiment Configuration", fontsize=14)
-
- # Panel 2: Progressive Dropout results if available
- ax2 = fig.add_subplot(gs[0, 1])
- prog_results = results.get("progressive_dropout", {})
-
- if "accuracies" in prog_results and "dropout_fractions" in prog_results:
- fractions = prog_results["dropout_fractions"]
- for strategy, color in [("high_rq", "blue"), ("low_rq", "red"), ("random", "green")]:
- if strategy in prog_results["accuracies"]:
- accs = prog_results["accuracies"][strategy]
- ax2.plot(
- fractions,
- accs,
- marker='o',
- linestyle='-',
- color=color,
- label=strategy.replace('_', ' ').title()
- )
-
- ax2.set_xlabel("Dropout Fraction", fontsize=12)
- ax2.set_ylabel("Accuracy (%)", fontsize=12)
- ax2.set_title("Progressive Dropout Results", fontsize=14)
- ax2.grid(True, alpha=0.3)
- ax2.set_ylim([0, 105])
- ax2.legend()
- else:
- ax2.text(0.5, 0.5, "No Progressive Dropout Results",
- fontsize=12, ha='center', va='center')
- ax2.axis('off')
-
- # Panel 3: Eigenvector Dropout results if available
- ax3 = fig.add_subplot(gs[1, 0])
- eig_results = results.get("eigenvector_dropout", {})
-
- if "accuracies" in eig_results and "dropout_fractions" in eig_results:
- fractions = eig_results["dropout_fractions"]
- if "eigenvector" in eig_results["accuracies"]:
- accs = eig_results["accuracies"]["eigenvector"]
- ax3.plot(
- fractions,
- accs,
- marker='o',
- linestyle='-',
- color='purple',
- label="Eigenvector"
- )
-
- # Also add the high_rq from progressive dropout for comparison if available
- if "accuracies" in prog_results and "high_rq" in prog_results["accuracies"]:
- prog_fracs = prog_results["dropout_fractions"]
- prog_accs = prog_results["accuracies"]["high_rq"]
- # Only plot if the fractions match
- if len(prog_fracs) == len(fractions) and all(a == b for a, b in zip(prog_fracs, fractions)):
- ax3.plot(
- fractions,
- prog_accs,
- marker='s',
- linestyle='--',
- color='blue',
- label="High RQ"
- )
-
- ax3.set_xlabel("Dropout Fraction", fontsize=12)
- ax3.set_ylabel("Accuracy (%)", fontsize=12)
- ax3.set_title("Eigenvector Dropout Results", fontsize=14)
- ax3.grid(True, alpha=0.3)
- ax3.set_ylim([0, 105])
- ax3.legend()
- else:
- ax3.text(0.5, 0.5, "No Eigenvector Dropout Results",
- fontsize=12, ha='center', va='center')
- ax3.axis('off')
-
- # Panel 4: Alignment comparison or other metrics
- ax4 = fig.add_subplot(gs[1, 1])
-
- # Check both progressive and eigenvector results for alignment values
- alignment_data = None
- if "alignment_values" in prog_results and "high_rq" in prog_results["alignment_values"]:
- alignment_data = prog_results["alignment_values"]["high_rq"][0]
- elif "alignment_values" in eig_results and "eigenvector" in eig_results["alignment_values"]:
- alignment_data = eig_results["alignment_values"]["eigenvector"][0]
-
- if alignment_data and len(alignment_data) > 0:
- # Extract alignment values as floats
- alignment_values = []
- for val in alignment_data:
- if isinstance(val, torch.Tensor):
- alignment_values.append(val.item())
- else:
- alignment_values.append(float(val))
-
- # Create bar chart of alignment by layer
- x = np.arange(len(alignment_values))
- bars = ax4.bar(x, alignment_values, width=0.6, alpha=0.7)
-
- # Add value labels on top of bars
- for i, bar in enumerate(bars):
- height = bar.get_height()
- ax4.text(
- bar.get_x() + bar.get_width()/2.,
- height + 0.01,
- f'{alignment_values[i]:.3f}',
- ha='center',
- va='bottom',
- rotation=0,
- fontsize=9
- )
-
- ax4.set_xlabel("Layer", fontsize=12)
- ax4.set_ylabel("Alignment Value", fontsize=12)
- ax4.set_title("Alignment by Layer", fontsize=14)
- ax4.set_xticks(x)
- ax4.set_xticklabels([f"Layer {i+1}" for i in range(len(alignment_values))])
- ax4.grid(True, alpha=0.3, axis='y')
- else:
- ax4.text(0.5, 0.5, "No Alignment Data",
- fontsize=12, ha='center', va='center')
- ax4.axis('off')
-
- # Add an overall title
- title = "Experiment Summary"
- if hasattr(config, "model") and hasattr(config.model, "model_name"):
- if hasattr(config, "dataset") and hasattr(config.dataset, "dataset_name"):
- title += f": {config.model.model_name} on {config.dataset.dataset_name}"
- else:
- title += f": {config.model.model_name}"
-
- fig.suptitle(title, fontsize=16, y=0.98)
- plt.tight_layout(rect=[0, 0, 1, 0.96]) # Adjust for suptitle
-
- # Save the figure
- summary_file = os.path.join(plot_dir, "experiment_summary.png")
- plt.savefig(summary_file, dpi=300, bbox_inches='tight')
- plt.close(fig)
-
- logger.info(f"Saved experiment summary to {summary_file}")
-
- return {"summary": summary_file}
\ No newline at end of file
diff --git a/_arxiv/_archive/alignment_preref/processing.py b/_arxiv/_archive/alignment_preref/processing.py
deleted file mode 100644
index a7186a89..00000000
--- a/_arxiv/_archive/alignment_preref/processing.py
+++ /dev/null
@@ -1,869 +0,0 @@
-# processing.py
-
-import os
-import torch
-from tqdm import tqdm
-
-from alignment.utils import load_checkpoints, test_nets
-from Code.alignment.src.alignment_preref.alignment_metrics_rem import AlignmentMetrics
-from alignment.train import train, test
-
-def parse_alignment_to_tensor(alignment_list, aggregate=True, by_layer=False):
- """
- Convert a list of alignment records into a structure suitable for dropout sorting.
- If aggregator=True => you may have multiple records per epoch or per batch.
- If aggregator=False => typically one record per epoch with flattened data.
-
- by_layer=False => (#nets, total_nodes)
- by_layer=True => list of (#nets, node_count) per layer.
- """
-
- if len(alignment_list) == 0:
- raise ValueError("parse_alignment_to_tensor: empty alignment_list")
-
- if not by_layer:
- all_records_tensors = []
- for record in alignment_list:
- if "data" not in record:
- continue
- netwise_tensors = []
- # record["data"] => (#nets, #layers)
- # or if aggregator=false => might be (#nets, #layers) but each layer is partial avg
- for net_i_data in record["data"]:
- node_tensors = []
- for layer_dict in net_i_data:
- if "RQ" not in layer_dict:
- raise ValueError("Expected 'RQ' in layer_dict")
- node_tensors.append(layer_dict["RQ"].flatten())
- cat_tsr = torch.cat(node_tensors, dim=0)
- netwise_tensors.append(cat_tsr)
- if len(netwise_tensors) == 0:
- continue
- netwise_tsr = torch.stack(netwise_tensors, dim=0) # (#nets, total_nodes)
- all_records_tensors.append(netwise_tsr)
- if len(all_records_tensors) == 0:
- raise ValueError("No valid alignment data found")
-
- bigstack = torch.stack(all_records_tensors, dim=0) # (#records, #nets, total_nodes)
- if aggregate:
- return bigstack.mean(dim=0) # => (#nets, total_nodes)
- else:
- # return all records => shape (#records, #nets, total_nodes)
- # The progressive_dropout code typically expects (#nets, total_nodes),
- # so we might do a final average anyway:
- return bigstack.mean(dim=0)
-
- else:
- max_layers_found = 0
- for record in alignment_list:
- if "data" not in record:
- continue
- for net_i_data in record["data"]:
- if len(net_i_data) > max_layers_found:
- max_layers_found = len(net_i_data)
-
- layer_storage = [[] for _ in range(max_layers_found)]
-
- for record in alignment_list:
- if "data" not in record:
- continue
- net_list = record["data"] # (#nets, #layers)
- for layer_i in range(max_layers_found):
- layer_nodevals = []
- for net_i_data in net_list:
- if layer_i < len(net_i_data):
- lay_dict = net_i_data[layer_i]
- if "RQ" not in lay_dict:
- raise ValueError("Expected 'RQ' in lay_dict")
- layer_nodevals.append(lay_dict["RQ"].flatten())
- else:
- layer_nodevals.append(None)
-
- valid_vals = [v for v in layer_nodevals if v is not None]
- if len(valid_vals) == 0:
- continue
- node_count = valid_vals[0].numel()
- netwise_tensor = []
- for val in layer_nodevals:
- if val is not None:
- netwise_tensor.append(val)
- netwise_tensor = torch.stack(netwise_tensor, dim=0)
- layer_storage[layer_i].append(netwise_tensor)
-
- final_layer_list = []
- for layer_i in range(max_layers_found):
- if len(layer_storage[layer_i]) == 0:
- final_layer_list.append(None)
- continue
- stacked = torch.stack(layer_storage[layer_i], dim=0) # (#records, #nets, node_count)
- if aggregate:
- final_layer_list.append(stacked.mean(dim=0)) # => (#nets, node_count)
- else:
- # for progressive_dropout we expect (#nets, node_count)
- # so again do a final average:
- final_layer_list.append(stacked.mean(dim=0))
-
- return final_layer_list
-
-
-def train_networks(exp, nets, optimizers, dataset, **special_parameters):
- do_alignment_train = exp.args.alignment.do_alignment
- methods = exp.args.alignment.methods
- measure_freq = exp.args.alignment.frequency
-
- params = dict(
- train_set=True,
- num_epochs=exp.args.training.epochs,
- alignment=do_alignment_train,
- methods=methods,
- frequency=measure_freq,
- measure_expected=exp.args.alignment.measure_expected,
- results=None,
- verbose=True,
- )
- params.update(**special_parameters)
-
- if exp.args.checkpointing.use_prev and os.path.isfile(exp.get_checkpoint_path()):
- nets, optimizers, cresults = load_checkpoints(nets, optimizers, exp.args.device, exp.get_checkpoint_path())
- for net in nets:
- net.train()
- params["num_complete"] = cresults["epoch"] + 1
- params["results"] = cresults
- print("loaded networks from previous checkpoint")
-
- if exp.args.checkpointing.save_checkpoints:
- params["save_checkpoints"] = (
- True,
- exp.args.checkpointing.frequency,
- exp.get_checkpoint_path(),
- exp.args.device,
- )
-
- print("training networks...")
- train_results = train(nets, optimizers, dataset, **params)
-
- do_alignment_infer = exp.args.alignment.do_alignment
- params["train_set"] = False
- params["alignment"] = do_alignment_infer
- print("testing networks (inference)...")
- test_results = test(nets, dataset, **params)
-
- return train_results, test_results
-
-
-def test_networks(exp, nets, dataset):
- do_align = exp.args.alignment.do_alignment
- methods = exp.args.alignment.methods
- freq = exp.args.alignment.frequency
-
- test_params = dict(
- train_set=False,
- alignment=do_align,
- methods=methods,
- frequency=freq,
- measure_expected=exp.args.alignment.measure_expected,
- bins=exp.args.alignment.bins,
- results=None,
- verbose=True
- )
- print("testing networks (no training)...")
- test_results = test(nets, dataset, **test_params)
- return test_results
-
-
-@torch.no_grad()
-def get_dropout_indices(idx_alignment, fraction):
- """
- Convenience method for getting a fraction of dropout indices from each layer.
- This is the same implementation as in alignment_v2.
- """
- num_nets = idx_alignment[0].size(0)
- num_nodes = [idx.size(1) for idx in idx_alignment]
- num_drop = [int(nodes * fraction) for nodes in num_nodes]
- idx_high = [
- torch.sort(idx[:, -drop:], dim=1).values
- for idx, drop in zip(idx_alignment, num_drop)
- ]
- idx_low = [
- torch.sort(idx[:, :drop], dim=1).values
- for idx, drop in zip(idx_alignment, num_drop)
- ]
- idx_rand = [
- torch.sort(idx[:, torch.randperm(idx.size(1))[:drop]], dim=1).values
- for idx, drop in zip(idx_alignment, num_drop)
- ]
- return idx_high, idx_low, idx_rand
-
-
-@test_nets
-@torch.no_grad()
-def progressive_dropout(nets, dataset, alignment=None, **parameters):
- """
- Main progressive dropout function that calls parse_alignment_to_tensor,
- then does targeted_dropout on each fraction.
-
- pruning_mode options:
- - "global": Prune X% of all nodes sorted together across all layers (original v1 behavior)
- - "per_layer_combined": Prune X% of each layer but apply to all layers at once (v2 behavior)
- - "per_layer_independent": Prune one layer at a time, creating separate results for each layer
-
- dropout_mode options:
- - "scaled": Zeroes neurons and applies scaling to maintain signal magnitude (default)
- - "unscaled": Zeroes neurons but doesn't apply compensation scaling
-
- Parameters:
- - exclude_classification_layer: If True, excludes the final classification layer from pruning
- """
- if not isinstance(nets, list):
- nets = [nets]
-
- if alignment is None:
- alignment = test(nets, dataset, alignment=True, methods=["RQ"], **parameters)["alignment"]
-
- aggregator = parameters.get("aggregate_alignment", False)
- pruning_mode = parameters.get("pruning_mode", "global")
- dropout_mode = parameters.get("dropout_mode", "scaled")
- exclude_classification_layer = parameters.get("exclude_classification_layer", False)
-
- parsed = parse_alignment_to_tensor(alignment, aggregate=aggregator, by_layer=(pruning_mode != "global"))
-
- num_drops = parameters.get("num_drops", 9)
- drop_fraction = torch.linspace(0, 1, num_drops + 2)[1:-1]
- use_train = parameters.get("train_set", False)
- dataloader = dataset.train_loader if use_train else dataset.test_loader
- print(f"Progressive Dropout (mode: {pruning_mode}, dropout_mode: {dropout_mode}):")
-
- num_nets = len(nets)
-
- if pruning_mode == "global":
- # Original v1 behavior - prune X% of all nodes together
- if not isinstance(parsed, torch.Tensor) or parsed.dim() != 2:
- raise ValueError("Expected shape (#nets, total_nodes) for pruning_mode='global'")
-
- # Get all layer indices for targeting
- target_layer_indices = list(range(len(nets[0].alignment_layers)))
-
- # Exclude classification layer if requested
- if exclude_classification_layer and nets[0].is_classification_layer_included():
- # Remove last layer which is the classification layer
- target_layer_indices = target_layer_indices[:-1]
- print(f"Excluding classification layer from pruning")
-
- # Get the alignment values for each layer separately to handle indices correctly
- # This matches alignment_v2's approach of having per-layer indices
- layer_alignments = []
- for i in target_layer_indices:
- # Get alignment values for this layer
- layer_data = []
- for record in alignment:
- if "data" not in record:
- continue
- netwise_layer_data = []
- for net_i_data in record["data"]:
- if i < len(net_i_data):
- if "RQ" in net_i_data[i]:
- netwise_layer_data.append(net_i_data[i]["RQ"].flatten())
- else:
- # Handle case where layer isn't present (shouldn't happen)
- netwise_layer_data.append(None)
-
- # Filter valid data
- valid_data = [d for d in netwise_layer_data if d is not None]
- if len(valid_data) == 0:
- continue
-
- # Stack data for this layer across nets
- layer_tensor = torch.stack(valid_data, dim=0)
- layer_data.append(layer_tensor)
-
- # Aggregate across records if needed
- if len(layer_data) == 0:
- # No valid data for this layer, use dummy
- layer_alignments.append(None)
- else:
- bigstack = torch.stack(layer_data, dim=0)
- if aggregator:
- layer_alignments.append(bigstack.mean(dim=0))
- else:
- layer_alignments.append(bigstack.mean(dim=0))
-
- # Initialize result tensors
- progdrop_loss_high = torch.zeros((num_nets, num_drops, 1), device="cpu")
- progdrop_loss_low = torch.zeros((num_nets, num_drops, 1), device="cpu")
- progdrop_loss_rand = torch.zeros((num_nets, num_drops, 1), device="cpu")
- progdrop_acc_high = torch.zeros((num_nets, num_drops, 1), device="cpu")
- progdrop_acc_low = torch.zeros((num_nets, num_drops, 1), device="cpu")
- progdrop_acc_rand = torch.zeros((num_nets, num_drops, 1), device="cpu")
-
- # Remove any None layers
- valid_alignments = []
- valid_layer_indices = []
- for i, align in enumerate(layer_alignments):
- if align is not None:
- valid_alignments.append(align)
- valid_layer_indices.append(target_layer_indices[i])
-
- # Get sorted indices for each layer
- idx_alignment = [torch.argsort(align, dim=1) for align in valid_alignments]
-
- num_batches = 0
- for batch in tqdm(dataloader):
- images, labels = dataset.unwrap_batch(batch)
- num_batches += 1
-
- for dropidx, fraction in enumerate(drop_fraction):
- # Get dropout indices for each layer (using alignment_v2 approach)
- idx_high, idx_low, idx_rand = get_dropout_indices(idx_alignment, fraction)
-
- # Process each network
- out_high, out_low, out_rand = [], [], []
- for i_net, net in enumerate(nets):
- # Get indices to dropout for this network
- high_indices = [drop_high[i_net] for drop_high in idx_high]
- low_indices = [drop_low[i_net] for drop_low in idx_low]
- rand_indices = [drop_rand[i_net] for drop_rand in idx_rand]
-
- # Apply dropout to all layers
- oh, _ = net.forward_targeted_dropout(images, high_indices, valid_layer_indices, dropout_mode=dropout_mode)
- ol, _ = net.forward_targeted_dropout(images, low_indices, valid_layer_indices, dropout_mode=dropout_mode)
- or_, _ = net.forward_targeted_dropout(images, rand_indices, valid_layer_indices, dropout_mode=dropout_mode)
-
- out_high.append(oh)
- out_low.append(ol)
- out_rand.append(or_)
-
- # Measure performance
- lh, ll, lr = [], [], []
- ah, al, ar = [], [], []
- for i_net in range(num_nets):
- # Losses after pruning
- lv_h = float(dataset.measure_loss(out_high[i_net], labels).cpu())
- lv_l = float(dataset.measure_loss(out_low[i_net], labels).cpu())
- lv_r = float(dataset.measure_loss(out_rand[i_net], labels).cpu())
- lh.append(lv_h)
- ll.append(lv_l)
- lr.append(lv_r)
-
- # Accuracy after pruning
- av_h = float(dataset.measure_accuracy(out_high[i_net], labels).cpu())
- av_l = float(dataset.measure_accuracy(out_low[i_net], labels).cpu())
- av_r = float(dataset.measure_accuracy(out_rand[i_net], labels).cpu())
- ah.append(av_h)
- al.append(av_l)
- ar.append(av_r)
-
- # Accumulate results
- progdrop_loss_high[:, dropidx, 0] += torch.tensor(lh, device="cpu")
- progdrop_loss_low[:, dropidx, 0] += torch.tensor(ll, device="cpu")
- progdrop_loss_rand[:, dropidx, 0] += torch.tensor(lr, device="cpu")
- progdrop_acc_high[:, dropidx, 0] += torch.tensor(ah, device="cpu")
- progdrop_acc_low[:, dropidx, 0] += torch.tensor(al, device="cpu")
- progdrop_acc_rand[:, dropidx, 0] += torch.tensor(ar, device="cpu")
-
- # Normalize by number of batches
- progdrop_loss_high /= num_batches
- progdrop_loss_low /= num_batches
- progdrop_loss_rand /= num_batches
- progdrop_acc_high /= num_batches
- progdrop_acc_low /= num_batches
- progdrop_acc_rand /= num_batches
-
- results = {
- "progdrop_loss_high": progdrop_loss_high,
- "progdrop_loss_low": progdrop_loss_low,
- "progdrop_loss_rand": progdrop_loss_rand,
- "progdrop_acc_high": progdrop_acc_high,
- "progdrop_acc_low": progdrop_acc_low,
- "progdrop_acc_rand": progdrop_acc_rand,
- "dropout_fraction": drop_fraction,
- "pruning_mode": pruning_mode,
- "dropout_mode": dropout_mode,
- "idx_dropout_layers": valid_layer_indices,
- }
- return results
- else:
- # Per-layer parsing for both "per_layer_combined" and "per_layer_independent" modes
- valid_layers = []
- layer_indices = []
- for i, layer_tsr in enumerate(parsed):
- if layer_tsr is not None:
- # Skip classification layer if requested
- if exclude_classification_layer and i == len(parsed) - 1 and nets[0].is_classification_layer_included():
- print(f"Excluding classification layer (index {i}) from pruning")
- continue
-
- valid_layers.append(layer_tsr)
- layer_indices.append(i)
-
- num_layers = len(valid_layers)
- if num_layers == 0:
- raise ValueError("No valid layers found in alignment data when using per-layer pruning")
-
- if pruning_mode == "per_layer_combined":
- # Initialize a class to hold combined results
- class CombinedResults:
- def __init__(self, num_nets, num_drops):
- self.loss_high = torch.zeros((num_nets, num_drops), device="cpu")
- self.loss_low = torch.zeros((num_nets, num_drops), device="cpu")
- self.loss_rand = torch.zeros((num_nets, num_drops), device="cpu")
- self.acc_high = torch.zeros((num_nets, num_drops), device="cpu")
- self.acc_low = torch.zeros((num_nets, num_drops), device="cpu")
- self.acc_rand = torch.zeros((num_nets, num_drops), device="cpu")
- self.count = 0
-
- pruning_combined_results = CombinedResults(num_nets, num_drops)
- else:
- pruning_combined_results = None
-
- progdrop_loss_high = torch.zeros((num_nets, num_drops, num_layers), device="cpu")
- progdrop_loss_low = torch.zeros((num_nets, num_drops, num_layers), device="cpu")
- progdrop_loss_rand = torch.zeros((num_nets, num_drops, num_layers), device="cpu")
- progdrop_acc_high = torch.zeros((num_nets, num_drops, num_layers), device="cpu")
- progdrop_acc_low = torch.zeros((num_nets, num_drops, num_layers), device="cpu")
- progdrop_acc_rand = torch.zeros((num_nets, num_drops, num_layers), device="cpu")
-
- num_batches = 0
- for batch in tqdm(dataloader):
- images, labels = dataset.unwrap_batch(batch)
- num_batches += 1
-
- for dropidx, fraction in enumerate(drop_fraction):
- for lyr_idx, layer_tsr in enumerate(valid_layers):
- idx_sorted = torch.argsort(layer_tsr, dim=1)
- node_count = idx_sorted.size(1)
- dn = int(node_count * fraction)
- if dn > 0:
- # nodes_to_drop_low: Prune lowest RQ nodes (first dn values in sorted order)
- # nodes_to_drop_high: Prune highest RQ nodes (last dn values in sorted order)
- # Variable names reflect which nodes are being DROPPED (zeroed out)
- nodes_to_drop_low = idx_sorted[:, :dn]
- nodes_to_drop_high = idx_sorted[:, node_count-dn:]
- nodes_to_drop_rand = []
- for i_net in range(num_nets):
- perm = torch.randperm(node_count, device=idx_sorted.device)
- nodes_to_drop_rand.append(perm[:dn])
- nodes_to_drop_rand = torch.stack(nodes_to_drop_rand, dim=0)
- else:
- nodes_to_drop_high = idx_sorted[:, :0]
- nodes_to_drop_low = idx_sorted[:, :0]
- nodes_to_drop_rand = idx_sorted[:, :0]
-
- # Determine which layers to apply pruning to
- if pruning_mode == "per_layer_independent":
- # Just prune the current layer
- target_layers = [layer_indices[lyr_idx]]
-
- # Process outputs for current layer only
- out_high, out_low, out_rand = [], [], []
- for i_net, net in enumerate(nets):
- # Forward pass with pruning:
- # out_low = output when DROPPING LOW alignment nodes (keeping HIGH alignment nodes)
- # out_high = output when DROPPING HIGH alignment nodes (keeping LOW alignment nodes)
- ol, _ = net.forward_targeted_dropout(images, [nodes_to_drop_low[i_net]], target_layers, dropout_mode=dropout_mode)
- oh, _ = net.forward_targeted_dropout(images, [nodes_to_drop_high[i_net]], target_layers, dropout_mode=dropout_mode)
- or_, _= net.forward_targeted_dropout(images, [nodes_to_drop_rand[i_net]], target_layers, dropout_mode=dropout_mode)
- out_low.append(ol) # Networks with LOW alignment nodes removed
- out_high.append(oh) # Networks with HIGH alignment nodes removed
- out_rand.append(or_)
-
- # Record metrics for this layer
- lh, ll, lr = [], [], []
- ah, al, ar = [], [], []
- for i_net in range(num_nets):
- # Losses after pruning (high = after pruning high nodes, etc.)
- lv_h = float(dataset.measure_loss(out_high[i_net], labels).cpu())
- lv_l = float(dataset.measure_loss(out_low[i_net], labels).cpu())
- lv_r = float(dataset.measure_loss(out_rand[i_net], labels).cpu())
- lh.append(lv_h)
- ll.append(lv_l)
- lr.append(lv_r)
-
- # Accuracy after pruning
- av_h = float(dataset.measure_accuracy(out_high[i_net], labels).cpu())
- av_l = float(dataset.measure_accuracy(out_low[i_net], labels).cpu())
- av_r = float(dataset.measure_accuracy(out_rand[i_net], labels).cpu())
- ah.append(av_h)
- al.append(av_l)
- ar.append(av_r)
-
- progdrop_loss_high[:, dropidx, lyr_idx] += torch.tensor(lh, device="cpu")
- progdrop_loss_low[:, dropidx, lyr_idx] += torch.tensor(ll, device="cpu")
- progdrop_loss_rand[:, dropidx, lyr_idx] += torch.tensor(lr, device="cpu")
- progdrop_acc_high[:, dropidx, lyr_idx] += torch.tensor(ah, device="cpu")
- progdrop_acc_low[:, dropidx, lyr_idx] += torch.tensor(al, device="cpu")
- progdrop_acc_rand[:, dropidx, lyr_idx] += torch.tensor(ar, device="cpu")
-
- else: # "per_layer_combined"
- # For per_layer_combined, we prune X% from EVERY layer and apply all prunings at once
- # This matches alignment_v2's by_layer=True code branch
-
- # Get the sorted indices for each layer
- idx_sorted_layers = [torch.argsort(layer_tsr, dim=1) for layer_tsr in valid_layers]
-
- # Get dropout indices for each layer using alignment_v2's approach
- idx_high, idx_low, idx_rand = get_dropout_indices(idx_sorted_layers, fraction)
-
- # Apply pruning to all layers simultaneously, exactly like alignment_v2 does
- out_high, out_low, out_rand = [], [], []
- for i_net, net in enumerate(nets):
- # Get indices for this network from each layer
- # This is how alignment_v2 extracts the indices for each network
- high_indices = [drop[i_net, :] for drop in idx_high]
- low_indices = [drop[i_net, :] for drop in idx_low]
- rand_indices = [drop[i_net, :] for drop in idx_rand]
-
-
- # Apply dropout to all layers at once (consistent with alignment_v2)
- oh, _ = net.forward_targeted_dropout(images, high_indices, layer_indices, dropout_mode=dropout_mode)
- ol, _ = net.forward_targeted_dropout(images, low_indices, layer_indices, dropout_mode=dropout_mode)
- or_, _ = net.forward_targeted_dropout(images, rand_indices, layer_indices, dropout_mode=dropout_mode)
-
- # Store network outputs
- out_high.append(oh)
- out_low.append(ol)
- out_rand.append(or_)
-
- # Record metrics for this pruning
- lh, ll, lr = [], [], []
- ah, al, ar = [], [], []
- for i_net in range(num_nets):
- # Losses after pruning
- lv_h = float(dataset.measure_loss(out_high[i_net], labels).cpu())
- lv_l = float(dataset.measure_loss(out_low[i_net], labels).cpu())
- lv_r = float(dataset.measure_loss(out_rand[i_net], labels).cpu())
- lh.append(lv_h)
- ll.append(lv_l)
- lr.append(lv_r)
-
-
- # Get the predicted classes
- _, pred_high = out_high[i_net].max(1)
- _, pred_low = out_low[i_net].max(1)
- _, pred_rand = out_rand[i_net].max(1)
-
- # Calculate accuracy manually
- manual_acc_high = 100 * (pred_high == labels).float().mean().item()
- manual_acc_low = 100 * (pred_low == labels).float().mean().item()
- manual_acc_rand = 100 * (pred_rand == labels).float().mean().item()
-
- # Accuracy after pruning - convert to float to ensure proper handling
- av_h = float(dataset.measure_accuracy(out_high[i_net], labels).cpu())
- av_l = float(dataset.measure_accuracy(out_low[i_net], labels).cpu())
- av_r = float(dataset.measure_accuracy(out_rand[i_net], labels).cpu())
-
- ah.append(av_h)
- al.append(av_l)
- ar.append(av_r)
-
- # Convert lists to tensors
- lh_tensor = torch.tensor(lh, device="cpu")
- ll_tensor = torch.tensor(ll, device="cpu")
- lr_tensor = torch.tensor(lr, device="cpu")
- ah_tensor = torch.tensor(ah, device="cpu")
- al_tensor = torch.tensor(al, device="cpu")
- ar_tensor = torch.tensor(ar, device="cpu")
-
- # Initialize the combined results storage if first batch
- if pruning_combined_results.count == 0:
- num_actual_nets = lh_tensor.size(0)
- pruning_combined_results.loss_high = torch.zeros((num_actual_nets, num_drops), device="cpu")
- pruning_combined_results.loss_low = torch.zeros((num_actual_nets, num_drops), device="cpu")
- pruning_combined_results.loss_rand = torch.zeros((num_actual_nets, num_drops), device="cpu")
- pruning_combined_results.acc_high = torch.zeros((num_actual_nets, num_drops), device="cpu")
- pruning_combined_results.acc_low = torch.zeros((num_actual_nets, num_drops), device="cpu")
- pruning_combined_results.acc_rand = torch.zeros((num_actual_nets, num_drops), device="cpu")
-
- # Store the combined results
- pruning_combined_results.loss_high[:, dropidx] += lh_tensor
- pruning_combined_results.loss_low[:, dropidx] += ll_tensor
- pruning_combined_results.loss_rand[:, dropidx] += lr_tensor
- pruning_combined_results.acc_high[:, dropidx] += ah_tensor
- pruning_combined_results.acc_low[:, dropidx] += al_tensor
- pruning_combined_results.acc_rand[:, dropidx] += ar_tensor
- pruning_combined_results.count += 1
-
- # Also store in the per-layer results for visualization
- # Duplicate results across layers for backward compatibility
- for lyr_idx in range(len(valid_layers)):
- progdrop_loss_high[:, dropidx, lyr_idx] += lh_tensor
- progdrop_loss_low[:, dropidx, lyr_idx] += ll_tensor
- progdrop_loss_rand[:, dropidx, lyr_idx] += lr_tensor
- progdrop_acc_high[:, dropidx, lyr_idx] += ah_tensor
- progdrop_acc_low[:, dropidx, lyr_idx] += al_tensor
- progdrop_acc_rand[:, dropidx, lyr_idx] += ar_tensor
-
-
- progdrop_loss_high /= num_batches
- progdrop_loss_low /= num_batches
- progdrop_loss_rand /= num_batches
- progdrop_acc_high /= num_batches
- progdrop_acc_low /= num_batches
- progdrop_acc_rand /= num_batches
-
- results = {
- "progdrop_loss_high": progdrop_loss_high,
- "progdrop_loss_low": progdrop_loss_low,
- "progdrop_loss_rand": progdrop_loss_rand,
- "progdrop_acc_high": progdrop_acc_high,
- "progdrop_acc_low": progdrop_acc_low,
- "progdrop_acc_rand": progdrop_acc_rand,
- "dropout_fraction": drop_fraction,
- "pruning_mode": pruning_mode,
- "dropout_mode": dropout_mode,
- "idx_dropout_layers": layer_indices,
- }
-
- # For per_layer_combined mode, add the properly calculated combined results
- # instead of averaging duplicate data
- if pruning_mode == "per_layer_combined" and hasattr(pruning_combined_results, 'count'):
- # Normalize by the number of batches
- pruning_combined_results.loss_high /= pruning_combined_results.count
- pruning_combined_results.loss_low /= pruning_combined_results.count
- pruning_combined_results.loss_rand /= pruning_combined_results.count
- pruning_combined_results.acc_high /= pruning_combined_results.count
- pruning_combined_results.acc_low /= pruning_combined_results.count
- pruning_combined_results.acc_rand /= pruning_combined_results.count
-
-
- # Check for flat zero accuracy
- if torch.all(pruning_combined_results.acc_high == 0) and torch.all(pruning_combined_results.acc_low == 0):
-
- # Apply safety measure - add tiny epsilon to first few fractions
- # This ensures plots will show a curve instead of a flat line
- for i in range(min(3, pruning_combined_results.acc_high.shape[1])):
- epsilon = 0.1 * (3 - i) # Small decreasing values: 0.3, 0.2, 0.1
- pruning_combined_results.acc_high[:, i] += epsilon
- pruning_combined_results.acc_low[:, i] += epsilon
- pruning_combined_results.acc_rand[:, i] += epsilon
-
- print(f"After safety measure - acc_high: {pruning_combined_results.acc_high}")
-
- # Add the combined results to the results dictionary
- results["combined_progdrop_loss_high"] = pruning_combined_results.loss_high.unsqueeze(2)
- results["combined_progdrop_loss_low"] = pruning_combined_results.loss_low.unsqueeze(2)
- results["combined_progdrop_loss_rand"] = pruning_combined_results.loss_rand.unsqueeze(2)
- results["combined_progdrop_acc_high"] = pruning_combined_results.acc_high.unsqueeze(2)
- results["combined_progdrop_acc_low"] = pruning_combined_results.acc_low.unsqueeze(2)
- results["combined_progdrop_acc_rand"] = pruning_combined_results.acc_rand.unsqueeze(2)
-
- # Add safety check for all result tensors
- result_keys = list(results.keys())
- for key in result_keys:
- if key.startswith("progdrop_acc"):
- if torch.all(results[key] == 0):
-
- # Add tiny values that decrease with dropout fraction
- result_shape = results[key].shape
- for i in range(result_shape[1]): # For each dropout fraction
- # Add decreasing values for smaller dropout fractions
- epsilon = max(0.1 * (result_shape[1] - i) / result_shape[1], 0.01)
- results[key][:, i, :] += epsilon
-
- return results
-
-
-def progressive_dropout_experiment(exp, nets, dataset, alignment=None, train_set=False):
- """
- Run progressive dropout experiment using the provided experiment config.
- """
- print("\nRunning progressive dropout experiment...")
- # Support both exp.args and exp.cfg patterns for backward compatibility
- config = getattr(exp, 'cfg', getattr(exp, 'args', None))
- if config is None:
- raise ValueError("Experiment object must have either 'cfg' or 'args' attribute with configuration")
-
- dropout_params = {
- "num_drops": config.extra.num_drops,
- "pruning_mode": config.extra.dropout_pruning_mode,
- "dropout_mode": config.extra.dropout_mode,
- "exclude_classification_layer": config.extra.exclude_classification_layer,
- "train_set": train_set,
- }
-
- # Add optional parameters if available
- if hasattr(config, 'alignment') and hasattr(config.alignment, 'scale_by_norm'):
- dropout_params["scale_by_norm"] = config.alignment.scale_by_norm
-
- if hasattr(config.extra, 'aggregate_alignment'):
- dropout_params["aggregate_alignment"] = config.extra.aggregate_alignment
-
- dropout_results = progressive_dropout(nets, dataset, alignment=alignment, **dropout_params)
- return dropout_results, dropout_params
-
-
-def measure_eigenfeatures(exp, nets, dataset, train_set=False):
- from tqdm import tqdm
- beta, eigvals, eigvecs, class_betas = [], [], [], []
- for net in tqdm(nets):
- inputs, labels = net._process_collect_activity(
- dataset,
- train_set=train_set,
- with_updates=False,
- use_training_mode=False,
- )
- efeatures = net.measure_eigenfeatures(inputs, with_updates=False)
- cls_betas = net.measure_class_eigenfeatures(inputs, labels, efeatures[2], rms=False, with_updates=False)
- beta.append(efeatures[0])
- eigvals.append(efeatures[1])
- eigvecs.append(efeatures[2])
- class_betas.append(cls_betas)
- class_names = getattr(dataset.train_loader if train_set else dataset.test_loader, "dataset").classes
- return dict(
- beta=beta,
- eigvals=eigvals,
- eigvecs=eigvecs,
- class_betas=class_betas,
- class_names=class_names,
- )
-
-
-@test_nets
-@torch.no_grad()
-def eigenvector_dropout(nets, dataset, eigenvalues, eigenvectors, **parameters):
- num_nets = len(nets)
- align_layer_indices = list(range(len(nets[0].alignment_layers)))
- pruning_mode = parameters.get("pruning_mode", "global")
- is_per_layer = pruning_mode != "global"
- num_layers = len(align_layer_indices) if is_per_layer else 1
- num_drops = parameters.get("num_drops", 9)
- drop_fraction = torch.linspace(0, 1, num_drops + 2)[1:-1]
- idx_eigenvalue = []
- for net_i in range(num_nets):
- layer_idxs = []
- for evec_j in eigenvectors[net_i]:
- dim = evec_j.size(1)
- layer_idxs.append(torch.arange(dim - 1, -1, -1).unsqueeze(0))
- idx_eigenvalue.append(layer_idxs)
-
- progdrop_loss_high = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_loss_low = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_loss_rand = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_acc_high = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_acc_low = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_acc_rand = torch.zeros((num_nets, num_drops, num_layers))
-
- use_test = not parameters.get("train_set", True)
- dataloader = dataset.test_loader if use_test else dataset.train_loader
- num_batches = 0
-
- from alignment.utils import check_iterable
-
- for batch in tqdm(dataloader):
- images, labels = dataset.unwrap_batch(batch)
- num_batches += 1
-
- for dropidx, fraction in enumerate(drop_fraction):
- # get_dropout_indices logic
- high_list, low_list, rand_list = [], [], []
- for net_i_idx, layer_idx_list in enumerate(idx_eigenvalue):
- # layer_idx_list => list of shape (#layers, 1, dim)
- h_layers, l_layers, r_layers = [], [], []
- for l_i, idx_sorted in enumerate(layer_idx_list):
- device_of_idx = idx_sorted.device
- drop_num = int(idx_sorted.size(1) * fraction)
- if drop_num > 0:
- hi = idx_sorted[:, idx_sorted.size(1)-drop_num:]
- lo = idx_sorted[:, :drop_num]
- rr = []
- for _n in range(1): # we have 1 row => expand if needed
- perm = torch.randperm(idx_sorted.size(1), device=device_of_idx)
- rr.append(perm[:drop_num])
- rr = torch.stack(rr, dim=0)
- else:
- hi = idx_sorted[:, :0]
- lo = idx_sorted[:, :0]
- rr = idx_sorted[:, :0]
- h_layers.append(hi)
- l_layers.append(lo)
- r_layers.append(rr)
- high_list.append(h_layers)
- low_list.append(l_layers)
- rand_list.append(r_layers)
-
- for layer_i in range(num_layers):
- net_out_high, net_out_low, net_out_rand = [], [], []
- for i_net, net in enumerate(nets):
- drop_layer = [align_layer_indices[layer_i]] if is_per_layer else align_layer_indices
- high_idxs = [high_list[i_net][layer_i][0]] # shape => ([indices])
- low_idxs = [low_list[i_net][layer_i][0]]
- rand_idxs = [rand_list[i_net][layer_i][0]]
- oh, _ = net.forward_eigenvector_dropout(images, eigenvalues[i_net], eigenvectors[i_net], high_idxs, drop_layer)
- ol, _ = net.forward_eigenvector_dropout(images, eigenvalues[i_net], eigenvectors[i_net], low_idxs, drop_layer)
- or_, _= net.forward_eigenvector_dropout(images, eigenvalues[i_net], eigenvectors[i_net], rand_idxs, drop_layer)
- net_out_high.append(oh)
- net_out_low.append(ol)
- net_out_rand.append(or_)
-
- lh, ll, lr = [], [], []
- ah, al, ar = [], [], []
- for i_net in range(num_nets):
- lv_h = float(dataset.measure_loss(net_out_high[i_net], labels).detach().cpu())
- lv_l = float(dataset.measure_loss(net_out_low[i_net], labels).detach().cpu())
- lv_r = float(dataset.measure_loss(net_out_rand[i_net], labels).detach().cpu())
- lh.append(lv_h)
- ll.append(lv_l)
- lr.append(lv_r)
-
- av_h = float(dataset.measure_accuracy(net_out_high[i_net], labels).detach().cpu())
- av_l = float(dataset.measure_accuracy(net_out_low[i_net], labels).detach().cpu())
- av_r = float(dataset.measure_accuracy(net_out_rand[i_net], labels).detach().cpu())
- ah.append(av_h)
- al.append(av_l)
- ar.append(av_r)
-
- progdrop_loss_high[:, dropidx, layer_i] += torch.tensor(lh)
- progdrop_loss_low[:, dropidx, layer_i] += torch.tensor(ll)
- progdrop_loss_rand[:, dropidx, layer_i] += torch.tensor(lr)
- progdrop_acc_high[:, dropidx, layer_i] += torch.tensor(ah)
- progdrop_acc_low[:, dropidx, layer_i] += torch.tensor(al)
- progdrop_acc_rand[:, dropidx, layer_i] += torch.tensor(ar)
-
- progdrop_loss_high /= num_batches
- progdrop_loss_low /= num_batches
- progdrop_loss_rand /= num_batches
- progdrop_acc_high /= num_batches
- progdrop_acc_low /= num_batches
- progdrop_acc_rand /= num_batches
-
- return {
- "progdrop_loss_high": progdrop_loss_high,
- "progdrop_loss_low": progdrop_loss_low,
- "progdrop_loss_rand": progdrop_loss_rand,
- "progdrop_acc_high": progdrop_acc_high,
- "progdrop_acc_low": progdrop_acc_low,
- "progdrop_acc_rand": progdrop_acc_rand,
- "dropout_fraction": drop_fraction,
- "pruning_mode": pruning_mode,
- }
-
-
-def eigenvector_dropout_experiment(exp, nets, dataset, eigen_results, train_set=False):
- evec_params = dict(
- num_drops=exp.args.extra.num_drops,
- pruning_mode=exp.args.extra.dropout_pruning_mode,
- train_set=train_set,
- )
- evec_dropout_results = eigenvector_dropout(
- nets,
- dataset,
- eigen_results["eigvals"],
- eigen_results["eigvecs"],
- **evec_params
- )
- return evec_dropout_results, evec_params
-
-def evaluate_pretrained_model(net, dataset):
- net.eval()
- device = next(net.parameters()).device
- total_correct = 0
- total_samples = 0
- with torch.no_grad():
- for batch in dataset.test_loader:
- images, labels = dataset.unwrap_batch(batch, device=device)
- outputs = net(images)
- predictions = outputs.argmax(dim=1)
- total_correct += (predictions == labels).sum().item()
- total_samples += labels.size(0)
- accuracy = 100.0 * total_correct / total_samples if total_samples > 0 else 0.0
- print(f"Pretrained Model Accuracy: {accuracy:.2f}%")
- return accuracy
\ No newline at end of file
diff --git a/_arxiv/_archive/alignment_preref/train.py b/_arxiv/_archive/alignment_preref/train.py
deleted file mode 100644
index b0187735..00000000
--- a/_arxiv/_archive/alignment_preref/train.py
+++ /dev/null
@@ -1,424 +0,0 @@
-# train.py
-
-import torch
-import numpy as np
-from copy import deepcopy
-from tqdm import tqdm
-
-from alignment.utils import train_nets, test_nets, load_checkpoints, save_checkpoint
-from Code.alignment.src.alignment_preref.alignment_metrics_rem import AlignmentMetrics
-
-@train_nets
-def train(nets, optimizers, dataset, **parameters):
- """
- A single function for supervised training with batch/epoch aggregator logic
- and alignment distribution computations.
-
- Args:
- nets (List[nn.Module]): replicate networks
- optimizers (List[torch.optim.Optimizer]): match length #nets
- dataset: dataset wrapper with .train_loader/.test_loader
- alignment (bool): if True, measure alignment
- aggregate_alignment (bool): if True => store alignment each batch as a separate record
- if False => accumulate alignment stats across batches, then produce a single record
- methods (List[str]): alignment methods, e.g. ["RQ"]
- frequency (int): how often (epochs) to measure alignment
- measure_expected (bool): if True => measure expected distribution via PCA
- bins (int): # bins for histogram
- num_epochs (int): total epochs
- results (dict or None): existing results or None
- save_checkpoints (tuple): (bool do_ckpt, ckpt_freq, ckpt_path, device)
- train_set (bool): True => use dataset.train_loader
- verbose (bool): print progress
- use_wandb (bool): if True => wandb logging
- ...
- Returns:
- dict with keys "loss", "accuracy", "alignment", "alignment_distribution", "expected_distribution", ...
- """
-
- do_align = parameters.get("alignment", False)
- methods = parameters.get("methods", ["RQ"])
- freq = parameters.get("frequency", 1)
- measure_expected = parameters.get("measure_expected", True)
- bins = parameters.get("bins", 50)
- aggregate = parameters.get("aggregate_alignment", False)
-
- num_epochs = parameters["num_epochs"]
- results = parameters.get("results", None)
- use_wandb = parameters.get("use_wandb", False)
-
- try:
- import wandb
- wandb_run = wandb.run
- except ImportError:
- wandb_run = None
- wandb_inited = (wandb_run is not None)
-
- if not isinstance(results, dict):
- results = {}
-
- if "alignment" not in results:
- results["alignment"] = []
- if "alignment_distribution" not in results:
- results["alignment_distribution"] = []
- if "expected_distribution" not in results:
- results["expected_distribution"] = []
- if "grad_alignment_corr" not in results:
- results["grad_alignment_corr"] = []
-
- num_replicates = len(nets)
- if "loss" not in results:
- results["loss"] = torch.zeros(num_replicates, num_epochs, dtype=torch.float)
- if "accuracy" not in results:
- results["accuracy"] = torch.zeros(num_replicates, num_epochs, dtype=torch.float)
-
- save_ckpt_info = parameters.get("save_checkpoints", (False, 1, "", ""))
- do_ckpt, ckpt_freq, ckpt_path, dev = save_ckpt_info
- start_epoch = parameters.get("num_complete", 0)
-
- use_train = parameters.get("train_set", True)
- dataloader = dataset.train_loader if use_train else dataset.test_loader
- verbose = parameters.get("verbose", True)
-
- if verbose:
- print(f"Starting training loop with epochs={num_epochs}, alignment={do_align}, "
- f"aggregate_alignment={aggregate}, methods={methods}")
-
- for epoch in range(start_epoch, num_epochs):
- replicate_loss_sums = [0.0] * num_replicates
- replicate_loss_counts = [0] * num_replicates
- replicate_acc_sums = [0.0] * num_replicates
- replicate_acc_counts = [0] * num_replicates
-
- # If aggregator=False, we accumulate alignment in memory to produce a single record
- # If aggregator=True, we store each batch as a separate record
- epoch_align_batches = []
- epoch_dist_batches = []
- epoch_exp_batches = []
-
- epoch_rq_values = []
-
- loop = tqdm(dataloader, desc=f"Train Epoch {epoch+1}", leave=False) if verbose else dataloader
-
- for batch_idx, batch in enumerate(loop):
- images, labels = dataset.unwrap_batch(batch)
-
- for idx_rep, (net, opt) in enumerate(zip(nets, optimizers)):
- opt.zero_grad()
- out = net(images, store_hidden=True)
- loss_val = dataset.measure_loss(out, labels)
- loss_val.backward()
-
- grad_norms_by_layer = {}
- for layer_name, layer in zip(net.alignment_names, net.alignment_layers):
- if layer.weight.grad is not None:
- g = layer.weight.grad.view(layer.weight.shape[0], -1).norm(dim=1)
- grad_norms_by_layer[layer_name] = g.detach().cpu()
-
- opt.step()
-
- replicate_loss_sums[idx_rep] += loss_val.item()
- replicate_loss_counts[idx_rep] += 1
-
- acc_val = dataset.measure_accuracy(out, labels)
- replicate_acc_sums[idx_rep] += float(acc_val)
- replicate_acc_counts[idx_rep] += 1
-
- if do_align and (epoch % freq == 0):
- alignment_data = net.measure_alignment_methods(images, methods=["RQ"], precomputed=False)
- correlation_by_layer = {}
- for layer_idx, layer_n in enumerate(net.alignment_names):
- node_alignment = alignment_data[layer_idx]["RQ"].cpu()
- if layer_n in grad_norms_by_layer:
- node_gradnorm = grad_norms_by_layer[layer_n]
- if node_alignment.shape == node_gradnorm.shape:
- stack = torch.stack([node_alignment, node_gradnorm], dim=0)
- corr_mat = torch.corrcoef(stack)
- correlation_by_layer[layer_n] = corr_mat[0, 1].item()
- else:
- correlation_by_layer[layer_n] = None
- else:
- correlation_by_layer[layer_n] = None
- results["grad_alignment_corr"].append({
- "epoch": epoch,
- "batch": batch_idx,
- "net": idx_rep,
- "corr": correlation_by_layer
- })
-
- if do_align and (epoch % freq == 0):
- batch_align_data = []
- for net in nets:
- layer_metrics = AlignmentMetrics.measure_methods(net, images, methods=methods, precomputed=False)
- batch_align_data.append(layer_metrics)
-
- dist_data = []
- for net_layer_list in batch_align_data:
- layer_dists = []
- for layer_dict in net_layer_list:
- method_dists = {}
- for m, val_tensor in layer_dict.items():
- val_cpu = val_tensor.detach().cpu()
- c, e = torch.histogram(val_cpu, bins=bins, density=True)
- method_dists[m] = (c, e)
- layer_dists.append(method_dists)
- dist_data.append(layer_dists)
-
- exp_data = []
- if measure_expected:
- for net in nets:
- layer_inps = net.get_layer_inputs(images, precomputed=False)
- layer_exp_list = []
- for inp in layer_inps:
- if inp.ndim == 4:
- inp = inp.flatten(start_dim=1)
- w_vals, _ = AlignmentMetrics.compute_eigenvalues(inp)
- method_exp = {}
- for m in methods:
- ccounts, cedges = AlignmentMetrics.measure_expected_distribution(m, w_vals, bins=bins)
- method_exp[m] = (ccounts, cedges)
- layer_exp_list.append(method_exp)
- exp_data.append(layer_exp_list)
-
- # aggregator logic
- if aggregate:
- # store each batch separately
- results["alignment"].append({
- "epoch": epoch,
- "batch": batch_idx,
- "data": batch_align_data
- })
- results["alignment_distribution"].append({
- "epoch": epoch,
- "batch": batch_idx,
- "data": dist_data
- })
- if measure_expected:
- results["expected_distribution"].append({
- "epoch": epoch,
- "batch": batch_idx,
- "data": exp_data
- })
- else:
- # store them in memory for now
- epoch_align_batches.append(batch_align_data)
- epoch_dist_batches.append(dist_data)
- epoch_exp_batches.append(exp_data)
-
- if batch_align_data and batch_align_data[0]:
- first_net_first_layer = batch_align_data[0][0]
- if "RQ" in first_net_first_layer:
- rq_val = first_net_first_layer["RQ"].mean().item()
- epoch_rq_values.append(rq_val)
-
- for idx_rep in range(num_replicates):
- if replicate_loss_counts[idx_rep] > 0:
- avg_loss = replicate_loss_sums[idx_rep] / replicate_loss_counts[idx_rep]
- else:
- avg_loss = 0.0
- results["loss"][idx_rep, epoch] = avg_loss
- if replicate_acc_counts[idx_rep] > 0:
- avg_acc = replicate_acc_sums[idx_rep] / replicate_acc_counts[idx_rep]
- else:
- avg_acc = 0.0
- results["accuracy"][idx_rep, epoch] = avg_acc
-
- if do_align and (epoch % freq == 0) and not aggregate:
- # we produce a single record for the entire epoch
- if epoch_align_batches:
- # Flatten or unify them
- # We'll combine all batch_align_data into one large net-layers structure
- # We do so by concatenating node-level alignment across batches
- # Then average node-level alignment
- # This can replicate aggregator=True final average
-
- # shape => list of (#batches) elements, each => [net_i_data], net_i_data => list of layer_dict
- # we'll unify them as if we had large data
- combined_net_data = []
- for _net_i in range(num_replicates):
- combined_net_data.append([]) # each net => layers
-
- # for each batch => batch_align_data is shape (#nets, #layers)
- for batch_align_data in epoch_align_batches:
- for net_idx, layer_list in enumerate(batch_align_data):
- # layer_list => list of dicts, each => method->Tensor
- combined_net_data[net_idx].append(layer_list)
-
- # now combined_net_data[net_idx] => list of (#batches) layer_list
- # we unify them layer by layer, node by node
- final_epoch_align_data = []
- for net_idx in range(num_replicates):
- # gather all batch-layers => flatten into per-layer accum
- all_layers = list(zip(*combined_net_data[net_idx]))
- # each element of all_layers => list of dicts from each batch
- net_layer_list = []
- for layer_items in all_layers:
- # layer_items => each batch a dict like {"RQ": Tensor(...)}
- # we unify them across batches => cat their Tensors => average
- # for safety, do so for each method
- methods_dict = {}
- for m in layer_items[0].keys():
- cat_list = [li[m].flatten() for li in layer_items]
- big_cat = torch.cat(cat_list, dim=0)
- mean_cat = big_cat.view(-1).mean(dim=0, keepdim=False)
- # store shape => (some_nodes,) => we store as 1D
- # or we can keep as single scalar => depends on usage
- # typical aggregator => node-level? we might want to keep node-level means => but here
- # we produce a single average => lose node granularity
- # If we want node granularity => cat them without mean => but aggregator=False => unify them
- # For this example => let's store node-level average if you want
- # We'll keep as big_cat => or we do node-level? let's do full node-level cat
- # That might produce big memory. We'll do a single average for each node across all batches:
- # big_cat => shape (#batches * node_count,) => we can keep it or average =>
- # aggregator=False => we want final single "RQ" per node? => we'd do stack or cat?
- # We'll do node-level average => big_cat is large => We'll keep the same # of nodes as 1 batch
- # We can't guess node_count if each batch might have a different # of samples =>
- # but alignment is node-based, does not depend on samples, only on the node dimension =>
- # Actually alignment is node-based. We'll produce a single average:
- # average across all batches => same shape as layer_items[0][m]
- # We do a stack approach =>
- shapes = [li[m].shape for li in layer_items]
- # assume all same shape => we stack
- stacked = torch.stack([li[m] for li in layer_items], dim=0) # (#batches, node_count)
- mean_per_node = stacked.mean(dim=0) # shape => (node_count,)
- methods_dict[m] = mean_per_node
- net_layer_list.append(methods_dict)
- final_epoch_align_data.append(net_layer_list)
-
- # final_epoch_align_data => shape (#nets, #layers)
- # store as a single record
- results["alignment"].append({
- "epoch": epoch,
- "batch": "aggregated",
- "data": final_epoch_align_data
- })
-
- # we do the same for distribution & exp_data if needed
- # skip for brevity => or replicate same logic
- # or produce a single record "alignment_distribution" at epoch
- # ignoring for shortness
-
- # optional distribution code
- if epoch_dist_batches:
- # we unify them in a simpler approach => just pick last batch or do an average
- # for shortness, skip or do your logic
- pass
-
- if measure_expected and epoch_exp_batches:
- pass
-
- mean_loss_ep = float(torch.mean(results["loss"][:, epoch]))
- mean_acc_ep = float(torch.mean(results["accuracy"][:, epoch]))
- mean_rq = float(np.mean(epoch_rq_values)) if len(epoch_rq_values) > 0 else 0.0
-
- if use_wandb and wandb_inited:
- import wandb
- wandb.log({
- "epoch": epoch,
- "train_loss": mean_loss_ep,
- "train_acc": mean_acc_ep,
- "train_alignment_RQ_epoch": mean_rq
- })
-
- if do_ckpt and (epoch % ckpt_freq == 0):
- cpy_res = deepcopy(results)
- cpy_res["epoch"] = epoch
- cpy_res["device"] = dev
- cpy_res["prms"] = parameters
- load_checkpoints(nets, optimizers, dev, None)
- save_checkpoint(nets, optimizers, cpy_res, ckpt_path)
-
- if verbose:
- print(f"Epoch {epoch+1}/{num_epochs} => loss={mean_loss_ep:.4f}, acc={mean_acc_ep:.2f}")
-
- results["loss"] = results["loss"].detach()
- results["accuracy"] = results["accuracy"].detach()
- return results
-
-
-@test_nets
-def test(nets, dataset, **parameters):
- """
- A single function for testing/evaluation, possibly measuring alignment too.
- Returns a dict with "loss", "accuracy", possibly "alignment" etc.
- """
-
- do_align = parameters.get("alignment", False)
- methods = parameters.get("methods", ["RQ"])
- measure_expected = parameters.get("measure_expected", True)
- bins = parameters.get("bins", 50)
- train_set = parameters.get("train_set", False)
- results = parameters.get("results", {})
- if not isinstance(results, dict):
- results = {}
-
- num_reps = len(nets)
- device = dataset.device
- loader = dataset.train_loader if train_set else dataset.test_loader
-
- loss_vec = torch.zeros(num_reps)
- acc_vec = torch.zeros(num_reps)
-
- total_correct = torch.zeros(num_reps, device=device)
- total_samples = 0
- total_loss = torch.zeros(num_reps, device=device)
-
- for batch in loader:
- images, labels = dataset.unwrap_batch(batch, device=device)
- for idx_rep, net in enumerate(nets):
- out = net(images)
- loss_val = dataset.measure_loss(out, labels, reduction="sum")
- total_loss[idx_rep] += loss_val.detach()
- pred = out.argmax(dim=1)
- total_correct[idx_rep] += (pred == labels).sum()
- total_samples += labels.size(0)
-
- for idx_rep in range(num_reps):
- if total_samples > 0:
- loss_vec[idx_rep] = total_loss[idx_rep].cpu().item() / float(total_samples)
- acc_vec[idx_rep] = (total_correct[idx_rep].cpu().item() / total_samples) * 100.0
-
- results["loss"] = loss_vec
- results["accuracy"] = acc_vec
-
- if do_align:
- images, labels = next(iter(loader))
- images, labels = dataset.unwrap_batch((images, labels), device=device)
- align_data = []
- dist_data = []
- exp_data = []
- for net in nets:
- net.forward(images, store_hidden=True)
- metrics = AlignmentMetrics.measure_methods(net, images, methods=methods, precomputed=False)
- align_data.append(metrics)
-
- layer_dists = []
- for layer_dict in metrics:
- m_d = {}
- for m, val_tensor in layer_dict.items():
- val_cpu = val_tensor.detach().cpu()
- c, e = torch.histogram(val_cpu, bins=bins, density=True)
- m_d[m] = (c, e)
- layer_dists.append(m_d)
- dist_data.append(layer_dists)
-
- if measure_expected:
- net_inps = net.get_layer_inputs(images, precomputed=False)
- layer_exp_list = []
- for inp in net_inps:
- if inp.ndim == 4:
- inp = inp.flatten(start_dim=1)
- wvals, _ = AlignmentMetrics.compute_eigenvalues(inp)
- method_exp = {}
- for m in methods:
- ccounts, cedges = AlignmentMetrics.measure_expected_distribution(m, wvals, bins=bins)
- method_exp[m] = (ccounts, cedges)
- layer_exp_list.append(method_exp)
- exp_data.append(layer_exp_list)
-
- results["alignment"] = [{"epoch":"test", "batch":"all", "data": align_data}]
- results["alignment_distribution"] = [{"epoch":"test", "batch":"all", "data": dist_data}]
- if measure_expected:
- results["expected_distribution"] = [{"epoch":"test", "batch":"all", "data": exp_data}]
-
- return results
\ No newline at end of file
diff --git a/_arxiv/_archive/alignment_preref/utils.py b/_arxiv/_archive/alignment_preref/utils.py
deleted file mode 100644
index 8d6cbd1f..00000000
--- a/_arxiv/_archive/alignment_preref/utils.py
+++ /dev/null
@@ -1,631 +0,0 @@
-# --------------------------------------------
-# utils.py
-# --------------------------------------------
-
-import os
-import math
-import zipfile
-from typing import List
-from warnings import warn
-from contextlib import contextmanager
-from functools import wraps
-from natsort import natsorted
-from gitignore_parser import parse_gitignore
-
-import torch
-import numpy as np
-from scipy.linalg import null_space
-from sklearn.decomposition import IncrementalPCA
-
-@contextmanager
-def no_grad(no_grad=True):
- if no_grad:
- with torch.no_grad():
- yield
- else:
- yield
-
-def test_nets(func):
- @wraps(func)
- def wrapper(nets, *args, **kwargs):
- # get original training mode and set to eval
- in_training_mode = [set_net_mode(net, training=False) for net in nets]
- func_outputs = func(nets, *args, **kwargs)
- # return networks to whatever mode they used to be in
- for train_mode, net in zip(in_training_mode, nets):
- set_net_mode(net, training=train_mode)
- return func_outputs
- return wrapper
-
-def train_nets(func):
- @wraps(func)
- def wrapper(nets, *args, **kwargs):
- # get original training mode and set to train
- in_training_mode = [set_net_mode(net, training=True) for net in nets]
- func_outputs = func(nets, *args, **kwargs)
- # return networks to whatever mode they used to be in
- for train_mode, net in zip(in_training_mode, nets):
- set_net_mode(net, training=train_mode)
- return func_outputs
- return wrapper
-
-def set_net_mode(net, training=True):
- """
- Helper for toggling train/eval mode of a network.
- """
- in_training_mode = net.training
- # set to training mode or evaluation mode
- if training:
- net.train()
- else:
- net.eval()
- # return original mode of network
- return in_training_mode
-
-def get_device(obj):
- """
- Returns 'cuda' if the module or tensor is on GPU, else 'cpu'.
- """
- if isinstance(obj, torch.nn.Module):
- return next(obj.parameters()).device.type
- elif isinstance(obj, torch.Tensor):
- return "cuda" if obj.is_cuda else "cpu"
- else:
- raise ValueError("get_device: object must be nn.Module or torch.Tensor")
-
-def check_iterable(val):
- """duck-type check if val is iterable"""
- try:
- _ = iter(val)
- except:
- return False
- else:
- return True
-
-def remove_by_idx(input, idx, dim):
- """
- remove part of input along dimension 'dim' for the indices in 'idx'
- """
- idx_keep = [i for i in range(input.size(dim)) if i not in idx]
- return torch.index_select(input, dim, torch.tensor(idx_keep).to(input.device))
-
-def get_eval_transform_by_cutoff(cutoff):
- """
- Return a function that zeros out eigenvalues below a fraction 'cutoff'.
- """
- def eval_transform(evals):
- assert torch.all(evals >= 0), "found negative eigenvalues, doesn't work for 'cutoff' eval_transform"
- evals = evals / torch.sum(evals)
- return 1.0 * (evals > cutoff)
- return eval_transform
-
-def fractional_histogram(*args, **kwargs):
- """wrapper of np.histogram() with relative counts instead of total or density"""
- counts, bins = np.histogram(*args, **kwargs)
- counts = counts / np.sum(counts)
- return counts, bins
-
-def edge2center(edges):
- """from a list of edges of bins (e.g. for torch.histogram()), return the centers between the edges"""
- assert edges.ndim == 1, "edges must be a 1-d array"
- return edges[:-1] + np.diff(edges) / 2
-
-def smartcorr(input):
- """
- Wraps torch.corrcoef but zeros out rows/cols that have zero variance.
- """
- idx_zeros = torch.var(input, dim=1) == 0
- cc = torch.corrcoef(input)
- cc[idx_zeros, :] = 0
- cc[:, idx_zeros] = 0
- return cc
-
-def batch_cov(input, centered=True, correction=True):
- """
- Batched covariance.
- If input.ndim==3, shape is (batch, dim, samples).
- If input.ndim==2, shape is (dim, samples).
- """
- assert (input.ndim == 2) or (input.ndim == 3), "input must be a 2D or 3D tensor"
- assert isinstance(correction, bool), "correction must be bool"
-
- no_batch = input.ndim == 2
- if no_batch:
- input = input.unsqueeze(0)
-
- S = input.size(2)
- if centered:
- input = input - input.mean(dim=2, keepdim=True)
- bcov = torch.bmm(input, input.transpose(1, 2))
- bcov /= (S - 1.0 * correction)
-
- if no_batch:
- bcov = bcov.squeeze(0)
- return bcov
-
-def smart_pcaOLD(input, centered=True, use_rank=True, correction=True):
- """
- smart algorithm for pca optimized for speed
-
- input should either have shape (batch, dim, samples) or (dim, samples)
- if dim > samples, will use svd and if samples < dim will use covariance/eigh method
-
- will center data when centered=True
-
- if it fails, will fall back on performing sklearns IncrementalPCA whenever forcetry=True
- """
- assert (input.ndim == 2) or (input.ndim == 3), "input should be a matrix or batched matrices"
- assert isinstance(correction, bool), "correction should be a boolean"
-
- if input.ndim == 2:
- no_batch = True
- input = input.unsqueeze(0) # create batch dimension for uniform code
- else:
- no_batch = False
-
- _, D, S = input.size()
- if D > S:
- # subtract mean if doing centered covariance
- if centered:
- input = input - input.mean(dim=2, keepdim=True)
- # if more dimensions than samples, it's more efficient to run svd
- v, w, _ = named_transpose([torch.linalg.svd(inp) for inp in input])
- # convert singular values to eigenvalues
- w = [ww**2 / (S - 1.0 * correction) for ww in w]
- # append zeros because svd returns w in R**k where k = min(D, S)
- w = [torch.concatenate((ww, torch.zeros(D - S))) for ww in w]
-
- else:
- # if more samples than dimensions, it's more efficient to run eigh
- bcov = batch_cov(input, centered=centered, correction=correction)
- w, v = named_transpose([eigendecomposition(C, use_rank=use_rank) for C in bcov])
-
- # return to stacked tensor across batch dimension
- w = torch.stack(w)
- v = torch.stack(v)
-
- # if no batch originally provided, squeeze out batch dimension
- if no_batch:
- w = w.squeeze(0)
- v = v.squeeze(0)
-
- # return eigenvalues and eigenvectors
- return w, v
-
-def smart_pca(input, centered=True, use_rank=True, correction=True):
- """
- Efficient PCA using either SVD or eigen-decomposition depending on shape.
- Falls back on a known method (like sklearn) if it fails to converge.
-
- Args:
- input (torch.Tensor): shape could be (D, S) or (B, D, S).
- - If 2D, treat it as a single (D, S) case (no batch).
- - If 3D, treat it as (batch, D, S).
- centered (bool): subtract mean along samples dimension if True
- use_rank (bool): zero out small eigenvalues beyond the matrix rank
- correction (bool): if True, use (S-1) denominator
- Returns:
- w (torch.Tensor): eigen/singular values, shape is either (D,) or (B, D)
- v (torch.Tensor): eigen/singular vectors, shape is either (D, D) or (B, D, D)
- """
- # 1) We only handle 2D or 3D
- assert input.ndim in (2, 3), "smart_pca: input must be 2D or 3D (D,S or B,D,S)."
- assert isinstance(correction, bool), "correction must be bool (True/False)."
-
- # 2) Convert to batch form if 2D
- no_batch = (input.ndim == 2)
- if no_batch:
- input = input.unsqueeze(0) # shape => (1, D, S)
-
- B, D, S = input.shape
-
- # 3) If #dims > #samples, use SVD on (D, S) for each batch
- # else use covariance+eigendecomposition
- if D > S:
- # Possibly center the data
- if centered:
- mean_ = input.mean(dim=2, keepdim=True) # shape (B, D, 1)
- input = input - mean_
-
- # We'll collect eigenvalues and eigenvectors for each batch
- w_list = []
- v_list = []
- for b in range(B):
- inp_ = input[b] # shape (D, S)
- # Using SVD => inp_ = U * S_ * V^T
- # torch.linalg.svd => (U, S, Vh)
- U, Svals, Vh = torch.linalg.svd(inp_, full_matrices=False)
- # Convert singular values to eigenvalues
- w_ = Svals**2 / (S - (1.0 if correction else 0.0))
- v_ = Vh.transpose(0, 1) # from shape (S, D) => (D, S)
-
- w_list.append(w_)
- v_list.append(v_)
- # Stack to shape => (B, D) or (B, S) … but in standard PCA we want dimension = D
- w = torch.stack(w_list, dim=0) # shape => (B, k)
- v = torch.stack(v_list, dim=0) # shape => (B, D, S)
- # If we truly want (B, D, D) we might need to pad or handle the mismatch if D != S
-
- else:
- # Use covariance => (D x D), then eigh
- bcov = batch_cov(input, centered=centered, correction=correction)
- # bcov shape => (B, D, D)
- w_list = []
- v_list = []
- for b in range(B):
- C = bcov[b] # shape (D, D)
- w_, v_ = eigendecomposition(C, use_rank=use_rank)
- w_list.append(w_)
- v_list.append(v_)
- w = torch.stack(w_list, dim=0) # => shape (B, D)
- v = torch.stack(v_list, dim=0) # => shape (B, D, D)
-
- # 4) If we started with no batch, remove batch dim
- if no_batch:
- w = w.squeeze(0) # => (D,)
- v = v.squeeze(0) # => (D, D) or possibly (D, S) in the SVD branch
-
- return w, v
-
-def eigendecomposition(C, use_rank=True):
- """
- helper for getting eigenvalues and eigenvectors of covariance matrix
-
- will measure eigenvalues and eigenvectors with torch.linalg.eigh()
- the output will be sorted from highest to lowest eigenvalue (& eigenvector)
-
- if use_rank=True, will measure the rank of the covariance matrix and zero
- out any eigenvalues beyond the rank (that are usually nonzero numerical errors)
- """
- try:
- w, v = torch.linalg.eigh(C)
- except torch._C._LinAlgError as error:
- # this happens if the algorithm failed to converge
- # try with sklearn's incrementalPCA algorithm
- return sklearn_pca(C, use_rank=use_rank)
- except Exception as error:
- raise error
- w_idx = torch.argsort(-w)
- w = w[w_idx]
- v = v[:, w_idx]
- # iff use_rank=True, will set eigenvalues to 0 for probable numerical errors
- if use_rank:
- crank = torch.linalg.matrix_rank(C)
- w[crank:] = 0
- return w, v
-
-def sklearn_pca(input, use_rank=True, rank=None):
- """
- sklearn incrementalPCA algorithm serving as a replacement for eigh when it fails
-
- input should be a tensor with shape (num_samples, num_features) or it can be a
- covariance matrix with (num_features, num_features)
-
- if use_rank=True, will set num_components to the rank of input and then fill out the
- rest of the components with random orthogonal components in the null space of the true
- components and set the eigenvalues to 0
-
- if use_rank=False, will attempt to fit all the components
- if rank is not None, will attempt to fit #=rank components without measuring the rank directly
- (will ignore "rank" if use_rank=False)
-
- returns w, v where w is eigenvalues and v is eigenvectors sorted from highest to lowest
- """
- # Move tensor to CPU if it's on GPU
- if input.is_cuda:
- input = input.cpu()
-
- num_samples, num_features = input.shape
- rank = None if not use_rank else (rank if rank is not None else fast_rank(input))
- ipca = IncrementalPCA(n_components=rank).fit(input)
- v = ipca.components_
- w = ipca.singular_values_**2 / num_samples
- # if v is a subspace of input (e.g. not a full basis, fill it out)
- if v.shape[0] < num_features:
- msg = "this condition should always be true, and if not we have to find out why"
- assert w.shape[0] == v.shape[0], msg
- v_kernel = null_space(v).T
- v = np.vstack((v, v_kernel))
- w = np.concatenate((w, np.zeros(v_kernel.shape[0])))
- return torch.tensor(w, dtype=torch.float), torch.tensor(v, dtype=torch.float).T
-
-def fast_rank(input):
- """uses transpose to speed up rank computation, otherwise normal"""
- # Move tensor to CPU if it's on GPU
- if input.is_cuda:
- input = input.cpu()
-
- if input.size(-2) < input.size(-1):
- input = torch.transpose(input, -2, -1)
- return int(torch.linalg.matrix_rank(input))
-
-def get_maximum_strides(h_input, w_input, layer):
- """
- Helper for computing the number of strides h_max, w_max after
- convolution with given kernel/stride/padding/dilation.
- """
- h_max = int(np.floor((h_input + 2 * layer.padding[0] - layer.dilation[0] * (layer.kernel_size[0] - 1) - 1) / layer.stride[0] + 1))
- w_max = int(np.floor((w_input + 2 * layer.padding[1] - layer.dilation[1] * (layer.kernel_size[1] - 1) - 1) / layer.stride[1] + 1))
- return h_max, w_max
-
-def get_unfold_params(layer):
- return dict(stride=layer.stride, padding=layer.padding, dilation=layer.dilation)
-
-@torch.no_grad()
-def cvPCA(X1, X2):
- """X1, X2 are both (dimensions x samples)"""
- D, B = X1.shape
- assert X2.shape == (D, B), "shape mismatch"
- _, u = smart_pca(X1)
- cproj0 = X1.T @ u
- cproj1 = X2.T @ u
- ss = (cproj0 * cproj1).mean(axis=0)
- return ss
-
-def get_num_components(nc, shape):
- return nc if nc is not None else min(shape)
-
-@torch.no_grad()
-def shuff_cvPCA(X1, X2, nshuff=5, cvmethod=cvPCA):
- D, B = X1.shape
- assert X2.shape == (D, B), "shape mismatch"
- nc = get_num_components(None, (D, B))
- ss = torch.zeros((nshuff, nc))
- X = torch.stack((X1, X2))
- for k in range(nshuff):
- iflip = 1 * (torch.rand(B) > 0.5)
- X1c = torch.gather(X, 0, iflip.view(1, 1, -1).expand(1, D, -1)).squeeze(0)
- X2c = torch.gather(X, 0, -(iflip - 1).view(1, 1, -1).expand(1, D, -1)).squeeze(0)
- ss[k] = cvmethod(X1c, X2c)
- return ss
-
-def avg_value_by_layer(full):
- """
- Return average value per layer across a list of epochs or minibatches.
-
- **full** is a list of lists where the outer list is each snapshot through training or
- minibatch etc and each inner list is the value for each node in the network across layers
- of a particular measurement
-
- For example:
- num_epochs = 1000
- nodes_per_layer = [50, 40, 30, 20]
- len(full) == 1000
- len(full[i]) == 4 ... for all i
- [f.shape for f in full[i]] = [50, 40, 30, 20] ... for all i
-
- this method will return a tensor of size (num_layers, num_epochs) of the average value (for
- whatever value is in **full**) for each list/list of values in **full**
- """
- num_epochs = len(full)
- num_layers = len(full[0])
- avg_full = torch.zeros((num_layers, num_epochs))
- for layer in range(num_layers):
- avg_full[layer, :] = torch.tensor([torch.mean(f[layer]) for f in full])
- return avg_full.cpu()
-
-def value_by_layer(full: List[List[torch.Tensor]], layer: int) -> torch.Tensor:
- """
- return all value measurements for a particular layer from **full**
-
- **full** is a list of lists where the outer list is each snapshot through training or
- minibatch etc and each inner list is the value for each node in the network across layers
-
- this method will return just the part of **full** corresponding to the layer indexed
- by **layer** as a tensor of shape (num_epochs, num_nodes)
-
- see ``avg_value_by_layer`` for a little more explanation
- """
- return torch.cat([f[layer].view(1, -1) for f in full], dim=0).cpu()
-
-def condense_values(full: List[List[List[torch.Tensor]]]) -> List[torch.Tensor]:
- """
- condense List[List[List[Tensor]]] -> list of #=num_layers Tensors
- shape: (num_networks, num_batches, num_nodes_per_layer)
-
- returns list of #=num_layers tensors, where each tensor has shape (num_networks, num_batches, num_nodes_per_layer)
-
- full should be a list of list of lists
- the first list should have length = number of networks
- the second list should have length = number of batches
- the third list should have length = number of layers in the network (this has to be the same for each network!)
- the tensor should have shape = number of nodes in this layer (also must be the same for each network) (or can be anything as long as consistent across layers)
- """
- num_layers = len(full[0][0])
- return [torch.stack([value_by_layer(value, layer) for value in full]) for layer in range(num_layers)]
-
-def transpose_list(list_of_lists):
- """helper function for transposing the order of a list of lists"""
- return list(map(list, zip(*list_of_lists)))
-
-def named_transpose(list_of_lists, reduction=None):
- """
- helper function for transposing lists without forcing the output to be a list like transpose_list
-
- for example, if list_of_lists contains 10 copies of lists that each have 3 iterable elements you
- want to name "A", "B", and "C", then write:
- A, B, C = named_transpose(list_of_lists)
-
- if reduction is used, it will be applied to each output, otherwise will make them lists
- """
- if reduction is not None:
- return map(reduction, zip(*list_of_lists))
- return map(list, zip(*list_of_lists))
-
-def ptp(tensor, dim=None, keepdim=False):
- """
- simple method for measuring range of tensor on requested dimension or on all data
- """
- if dim is None:
- return tensor.max() - tensor.min()
- return tensor.max(dim, keepdim).values - tensor.min(dim, keepdim).values
-
-def rms(tensor, dim=None, keepdim=False):
- """simple method for measuring root-mean-square on requested dimension or on all data in tensor"""
- if dim is None:
- return torch.sqrt(torch.mean(tensor**2))
- return torch.sqrt(torch.mean(tensor**2, dim=dim, keepdim=keepdim))
-
-def compute_stats_by_type(tensor, num_types, dim, method="var"):
- """
- helper method for returning the mean and variance across a certain dimension
- where multiple types are concatenated on that dimension
-
- for example, suppose we trained 2 networks each with 3 sets of parameters
- and concatenated the loss in a tensor like [set1-loss-net1, set1-loss-net2, set2-loss-net1, ...]
- then this would contract across the nets from each set and return the mean and variance
- """
- num_on_dim = tensor.size(dim)
- num_per_type = int(num_on_dim / num_types)
- tensor_by_type = tensor.unsqueeze(dim)
- expand_shape = list(tensor_by_type.shape)
- expand_shape[dim + 1] = num_per_type
- expand_shape[dim] = num_types
- tensor_by_type = tensor_by_type.view(expand_shape)
- type_means = torch.mean(tensor_by_type, dim=dim + 1)
- if method == "var":
- type_dev = torch.var(tensor_by_type, dim=dim + 1)
- elif method == "std":
- type_dev = torch.std(tensor_by_type, dim=dim + 1)
- elif method == "se":
- type_dev = torch.std(tensor_by_type, dim=dim + 1) / np.sqrt(num_per_type)
- elif method == "range":
- type_dev = ptp(tensor_by_type, dim=dim + 1)
- else:
- raise ValueError(f"Method ({method}) not recognized.")
- return type_means, type_dev
-
-
-def weighted_average(data, weights, dim, keepdim=False, ignore_nan=False):
- """
- take the weighted average of **data** on a certain dimension with **weights**
-
- weights should be a nonnegative vector that broadcasts into data
- avg = sum_i(data_i * weight_i, dim) / sum_i(weight_i, dim)
-
- if ignore_nan=True, (default=False), will ignore nans in weighted average
- """
- assert data.ndim == weights.ndim, "data and weights must have same number of dimensions"
- assert torch.all(weights[~torch.isnan(weights)] >= 0), "weights must be nonnegative"
-
- for d in dim if check_iterable(dim) else [dim]:
- assert data.size(d) == weights.size(
- d
- ), f"data and weights must have same size in averaging dimensions (data.size({d})={data.size(d)}, (weight.size({d})={weights.size(d)}))"
-
- # use normal sum if not ignore nan, otherwise use nansum
- sum = torch.nansum if ignore_nan else torch.sum
-
- # make sure nans are in the same place in weights and data for accurate division by total weight
- if ignore_nan:
- weights = weights.expand(data.size())
- weights = torch.masked_fill(weights, torch.isnan(data), torch.nan)
-
- # numerator & denominator of weighted average
- numerator = sum(data * weights, dim=dim, keepdim=keepdim)
- denominator = sum(weights, dim=dim, keepdim=keepdim)
-
- # return weighted average
- return numerator / denominator
-
-
-def fgsm_attack(image, epsilon, data_grad, transform, sign):
- """update an image with fast-gradient sign method"""
- warn("fgsm_attack is only going to be in utils temporarily!", DeprecationWarning, stacklevel=2)
- # Collect the element-wise sign of the data gradient
- if sign:
- data_grad = data_grad.sign()
- else:
- data_grad = data_grad.clone()
- # Create the perturbed image by adjusting each pixel of the input image
- perturbed_image = image + epsilon * data_grad
- # Adding clipping to maintain [0,1] range
- perturbed_image = transform(perturbed_image)
- # Return the perturbed image
- return perturbed_image
-
-def str2bool(str):
- if isinstance(str, bool):
- return str
- if str.lower() in ("true", "1"):
- return True
- elif str.lower() in ("false", "0"):
- return False
- else:
- raise TypeError("Boolean type expected")
-
-def save_checkpoint(nets, optimizers, results, path):
- """
- Method for saving checkpoints for networks throughout training.
- """
- multi_model_ckpt = {f"model_state_dict_{i}": net.state_dict() for i, net in enumerate(nets)}
- multi_optimizer_ckpt = {f"optimizer_state_dict_{i}": opt.state_dict() for i, opt in enumerate(optimizers)}
- checkpoint = results | multi_model_ckpt | multi_optimizer_ckpt
- torch.save(checkpoint, path)
-
-def load_checkpoints(nets, optimizers, device, path):
- """
- Method for loading presaved checkpoint during training.
- """
- if device == "cpu":
- checkpoint = torch.load(path, map_location=device)
- elif device == "cuda":
- checkpoint = torch.load(path)
-
- net_ids = natsorted([key for key in checkpoint if key.startswith("model_state_dict")])
- opt_ids = natsorted([key for key in checkpoint if key.startswith("optimizer_state_dict")])
- assert all(
- [oi.split("_")[-1] == ni.split("_")[-1] for oi, ni in zip(opt_ids, net_ids)]
- ), "nets and optimizers cannot be matched up from checkpoint"
-
- [net.load_state_dict(checkpoint.pop(net_id)) for net, net_id in zip(nets, net_ids)]
- [opt.load_state_dict(checkpoint.pop(opt_id)) for opt, opt_id in zip(optimizers, opt_ids)]
-
- if device == "cuda":
- [net.to(device) for net in nets]
- return nets, optimizers, checkpoint
-
-def match_git(path):
- """simple method for determining if a path is a git-related file or directory"""
- if ".git" in path:
- return True
- return False
-
-def compress_directory(output_path, directory_path=None):
- """
- Utility to compress entire directory into a zip while respecting .gitignore and ignoring .git.
- """
- if directory_path is None:
- directory_path = os.path.dirname(os.path.abspath(__file__)) + "/../.."
- gitignore_path = os.path.join(directory_path, ".gitignore")
- matches = parse_gitignore(gitignore_path)
-
- files_to_copy = []
- archive_names = []
- for dirpath, dirnames, files in os.walk(directory_path):
- if matches(dirpath) or match_git(dirpath):
- dirnames[:] = []
- else:
- # Filter files based on .gitignore rules (and don't save any .git files)
- keep_files = [f for f in files if not matches(f) and not match_git(f)]
- full_files = [os.path.join(dirpath, f) for f in keep_files]
- for file in full_files:
- files_to_copy.append(file)
- archive_names.append(os.path.relpath(file, directory_path))
-
- with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zipf:
- for file, name in zip(files_to_copy, archive_names):
- zipf.write(file, arcname=name)
-
-
-def condense_values(al_list):
- data_list = [item["data"] for item in al_list]
- aggregated = []
- for net_i in range(len(data_list[0])):
- net_snapshots = [d[net_i] for d in data_list]
- net_layers = []
- for layer_i in range(len(net_snapshots[0])):
- layer_values = [snap[layer_i] for snap in net_snapshots]
- net_layers.append(torch.stack(layer_values, dim=0).mean(dim=0))
- aggregated.append(net_layers)
- return aggregated
\ No newline at end of file
diff --git a/_arxiv/_archive/alignment_preref/utils_rem.py b/_arxiv/_archive/alignment_preref/utils_rem.py
deleted file mode 100644
index e933d1e2..00000000
--- a/_arxiv/_archive/alignment_preref/utils_rem.py
+++ /dev/null
@@ -1,57 +0,0 @@
-"""
-Utilities for alignment research.
-
-DEPRECATED: This module is kept for backward compatibility.
-Please use alignment.utils.core and alignment.utils.math instead.
-
-This module contains various utility functions for the alignment package.
-"""
-
-import warnings
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
-
-# Import from the new modules for re-export
-from alignment.utils.core import (
- setup_logging,
- timer,
- debug,
- to_numpy,
- to_tensor,
- check_iterable,
- ensure_device,
- timed
-)
-
-from alignment.utils.math import (
- orthogonalize,
- compute_correlation_matrix,
- matrix_angles,
- project_to_subspace
-)
-
-# Show deprecation warning on import
-warnings.warn(
- "The alignment.utils module is deprecated and will be removed in a future version. "
- "Please use alignment.utils.core and alignment.utils.math instead.",
- DeprecationWarning,
- stacklevel=2
-)
-
-# Re-export everything from the imported modules
-__all__ = [
- # Core utilities
- 'setup_logging',
- 'timer',
- 'debug',
- 'to_numpy',
- 'to_tensor',
- 'check_iterable',
- 'ensure_device',
- 'timed',
-
- # Math utilities
- 'orthogonalize',
- 'compute_correlation_matrix',
- 'matrix_angles',
- 'project_to_subspace'
-]
\ No newline at end of file
diff --git a/_arxiv/_archive/alignment_stats.py b/_arxiv/_archive/alignment_stats.py
deleted file mode 100644
index 1f01906c..00000000
--- a/_arxiv/_archive/alignment_stats.py
+++ /dev/null
@@ -1,97 +0,0 @@
-import torch
-
-
-from alignment_v2.datasets import get_dataset
-from alignment_v2.models.registry import get_model, get_transform_parameters
-from alignment_v2 import processing
-from alignment_v2.config import ExperimentConfig
-from alignment_v2.experiments.experiment import Experiment
-from alignment_v2 import plotting
-
-class AlignmentStatistics(Experiment):
- def get_basename(self):
- return "alignment_stats"
-
- def prepare_path(self):
- return [self.args.model.name, self.args.dataset.name, self.args.optimizer.name]
-
- def create_networks(self):
- """
- method for creating networks
- """
- if self.args.optimizer.name == "Adam":
- optim = torch.optim.Adam
- elif self.args.optimizer.name == "SGD":
- optim = torch.optim.SGD
- else:
- raise ValueError(f"optimizer ({self.args.optimizer.name}) not recognized")
-
- nets = [
- get_model(
- self.args.model.name,
- alignment_layer_names=self.args.model.alignment_layers,
- build=True,
- dataset=self.args.dataset.name,
- dropout=self.args.model.dropout,
- )
- for _ in range(self.args.training.replicates)
- ]
- nets = [net.to(self.device) for net in nets]
-
- optimizers = [optim(net.parameters(), lr=self.args.optimizer.lr, weight_decay=self.args.optimizer.weight_decay) for net in nets]
-
- prms = {
- "vals": [self.args.model.name],
- "name": "network",
- "dataset": self.args.dataset.name,
- "dropout": self.args.model.dropout,
- "lr": self.args.optimizer.lr,
- "weight_decay": self.args.optimizer.weight_decay,
- }
- return nets, optimizers, prms
-
- def main(self):
- nets, optimizers, prms = self.create_networks()
-
- dataset = self.prepare_dataset(get_transform_parameters(self.args.model.name, self.args.dataset.name))
-
- train_results, test_results = processing.train_networks(self, nets, optimizers, dataset)
-
- dropout_results, dropout_parameters = processing.progressive_dropout_experiment(
- self, nets, dataset, alignment=test_results.get("alignment", None), train_set=False
- )
-
- eigen_results = processing.measure_eigenfeatures(self, nets, dataset, train_set=False)
-
- evec_dropout_results, evec_dropout_parameters = processing.eigenvector_dropout(self, nets, dataset, eigen_results, train_set=False)
-
- results = dict(
- prms=prms,
- train_results=train_results,
- test_results=test_results,
- dropout_results=dropout_results,
- dropout_parameters=dropout_parameters,
- eigen_results=eigen_results,
- evec_dropout_results=evec_dropout_results,
- evec_dropout_parameters=evec_dropout_parameters,
- )
-
- return results, nets
-
- def plot(self, results):
- plotting.plot_train_results(self, results["train_results"], results["test_results"], results["prms"])
- plotting.plot_dropout_results(
- self,
- results["dropout_results"],
- results["dropout_parameters"],
- results["prms"],
- dropout_type="nodes",
- )
- plotting.plot_eigenfeatures(self, results["eigen_results"], results["prms"])
- plotting.plot_dropout_results(
- self,
- results["evec_dropout_results"],
- results["evec_dropout_parameters"],
- results["prms"],
- dropout_type="eigenvectors",
- )
\ No newline at end of file
diff --git a/_arxiv/_archive/alignment_v2/__init__.py b/_arxiv/_archive/alignment_v2/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/_arxiv/_archive/alignment_v2/datasets.py b/_arxiv/_archive/alignment_v2/datasets.py
deleted file mode 100644
index 5f57be31..00000000
--- a/_arxiv/_archive/alignment_v2/datasets.py
+++ /dev/null
@@ -1,322 +0,0 @@
-from warnings import warn
-from abc import ABC, abstractmethod
-
-import torch
-import torchvision
-from torch import nn
-from torch.utils.data.distributed import DistributedSampler
-from torchvision.transforms import v2 as transforms
-
-from alignment_v2 import files
-from alignment_v2.models.base import AlignmentNetwork
-
-REQUIRED_PROPERTIES = ["dataset_path", "dataset_constructor", "loss_function"]
-
-
-def default_loader_parameters(
- distributed,
- batch_size=1024*2,
- num_workers=2,
- shuffle=True,
- pin_memory=True,
- persistent_workers=True,
-):
- """
- contains the default dataloader parameters with the option of updating them
- using key word argument
- """
- default_parameters = dict(
- batch_size=batch_size,
- num_workers=num_workers, # usually 2 workers is appropriate for swapping loading during batch processing
- shuffle=False if distributed else shuffle, # can't use shuffle=True if using DDP
- pin_memory=pin_memory,
- persistent_workers=persistent_workers,
- )
- return default_parameters
-
-
-class DataSet(ABC):
- def __init__(
- self,
- device=None,
- distributed=False,
- dataset_parameters={},
- transform_parameters={},
- loader_parameters={},
- ):
- # set properties of dataset and check that all required properties are defined
- self.set_properties()
- self.check_properties()
-
- # define device for dataloading
- self.device = device if device is not None else ("cuda" if torch.cuda.is_available() else "cpu")
- self.distributed = distributed
-
- # define extra transform (should be a callable method or None) for any transformations that
- # can't go in the torchvision.transforms.Compose(...), hopefully this won't be needed later
- # when that issue is resolved (grayscale to RGB transform isn't working in Compose right now)
- self.extra_transform = transform_parameters.pop("extra_transform", None)
-
- # create transform for dataloader
- self.transform_parameters = transform_parameters
- self.make_transform(**transform_parameters)
-
- # define the dataloader parameters
- self.dataloader_parameters = default_loader_parameters(distributed, **loader_parameters) # get dataloader parameters
-
- # load the dataset and create the dataloaders
- self.dataset_parameters = dataset_parameters
- self.load_dataset(**dataset_parameters)
-
- def check_properties(self):
- """
- DataSet objects have a few properties that are required but need to be set
- by the children. For simplicity, I want the properties to be attributes,
- instead of @property methods, but I want to make sure that all the required
- properties are loaded. Hence this method.
- """
- if not all([hasattr(self, prop) for prop in REQUIRED_PROPERTIES]):
- not_found = [prop for prop in REQUIRED_PROPERTIES if not hasattr(self, prop)]
- raise ValueError(f"The following required properties were not set: {not_found}")
-
- @abstractmethod
- def set_properties(self):
- """
- DataSets have a few required properties (listed in the **REQUIRED_PROPERTIES**
- global variable). This method is used to set them all. There is a check implemented
- in the __init__ method to make sure all required properties are set.
-
- required
- --------
- dataset_path: string, defines the local file location containing the relevant dataset
- files. in particular, whatever filepath is returned by this method will
- be passed into the "root" input for torch datasets that is required to
- load one of the standard datasets (like MNIST, CIFAR, ImageNet...)
-
- dataset_constructor: callable method, defines the constructor object for the dataset
- loss_function: callable method, defines how to evaluate the loss of the output and target
- """
- pass
-
- @abstractmethod
- def dataset_kwargs(self, train=True, **kwargs):
- """
- keyword arguments passed into the torch dataset constructor
-
- different datasets have different kwarg requirements, including
- the way train vs test is defined by the kwargs. This class calls
- the train dataset and the test dataset using dataset_kwargs(train=True)
- or train=False, so the children need to define whatever kwargs go
- into the dataset_kwargs appropriately to be determined by the kwarg
- train
- """
- pass
-
- def load_dataset(self, **kwargs):
- """load dataset using the established path and parameters"""
- self.train_dataset = self.dataset_constructor(**self.dataset_kwargs(train=True, **kwargs))
- self.test_dataset = self.dataset_constructor(**self.dataset_kwargs(train=False, **kwargs))
- self.train_sampler = DistributedSampler(self.train_dataset) if self.distributed else None
- self.test_sampler = DistributedSampler(self.test_dataset) if self.distributed else None
- self.train_loader = torch.utils.data.DataLoader(self.train_dataset, sampler=self.train_sampler, **self.dataloader_parameters)
- self.test_loader = torch.utils.data.DataLoader(self.test_dataset, sampler=self.test_sampler, **self.dataloader_parameters)
-
- def unwrap_batch(self, batch, device=None):
- """simple method for unwrapping batch for simple training loops"""
- device = self.device if device is None else device
- if self.extra_transform:
- if type(self.extra_transform) == list:
- for et in self.extra_transform:
- batch = et(batch)
- else:
- warn("extra_transform is not a list, this is deprecated!", DeprecationWarning, stacklevel=2)
- batch = self.extra_transform(batch)
- inputs, targets = batch
- inputs, targets = inputs.to(device), targets.to(device)
- return inputs, targets
-
- def make_transform(self, center_crop=None, resize=None, flatten=False, out_channels=None):
- """
- create transform for dataloader
- resize is the new (H, W) shape of the image for the transforms.Resize transform (or None)
- flatten is a boolean indicating whether to flatten the image, (i.e. for a linear input layer)
- """
- # default transforms
- use_transforms = [
- # Convert PIL Image to PyTorch Tensor
- transforms.ToImage(),
- transforms.ToDtype(torch.float32, scale=True),
- ]
-
- if center_crop:
- use_transforms.append(transforms.CenterCrop(center_crop))
-
- # Normalize inputs to canonical distribution
- use_transforms.append(transforms.Normalize((self.dist_params["mean"]), (self.dist_params["std"])))
-
- # extra transforms depending on network
- if resize:
- use_transforms.append(transforms.Resize(resize, antialias=True))
- if out_channels:
- use_transforms.append(transforms.Grayscale(num_output_channels=out_channels))
- if flatten:
- use_transforms.append(transforms.Lambda(torch.flatten))
-
- # store composed transformation
- self.transform = transforms.Compose(use_transforms)
-
- def measure_loss(self, outputs, targets, reduction=None):
- """simple method for measuring loss with stored loss function"""
- if reduction is None:
- return self.loss_function(outputs, targets)
-
- standard_reduction = self.loss_function.reduction
- self.loss_function.reduction = reduction
- loss = self.loss_function(outputs, targets)
- self.loss_function.reduction = standard_reduction
- return loss
-
- def measure_accuracy(self, outputs, targets, k=1, percentage=True):
- """
- simple method for measuring accuracy on a classification problem
-
- default output is top1 percentage, but k can set the top-k accuracy
- and if percentage=False then returns the number correct (by topk)
- """
- topk = outputs.topk(k, dim=1, sorted=True, largest=True)[1] # get topk indices
- num_correct = torch.sum(torch.any(topk == targets.view(-1, 1), dim=1)) # num correct
- if percentage:
- return 100 * num_correct / outputs.size(0) # percentage
- else:
- return num_correct
-
-
-class MNIST(DataSet):
- def set_properties(self):
- """defines the required properties for MNIST"""
- self.dataset_path = files.dataset_path("MNIST")
- self.dataset_constructor = torchvision.datasets.MNIST
- self.loss_function = nn.CrossEntropyLoss()
- self.dist_params = dict(mean=[0.1307], std=[0.3081])
-
- def dataset_kwargs(self, train=True, download=False):
- """set data constructor kwargs for MNIST"""
- kwargs = dict(
- train=train,
- root=self.dataset_path,
- download=download,
- transform=self.transform,
- )
- return kwargs
-
-
-class CIFAR10(DataSet):
- def set_properties(self):
- """defines the required properties for CIFAR10"""
- self.dataset_path = files.dataset_path("CIFAR10")
- self.dataset_constructor = torchvision.datasets.CIFAR10
- self.loss_function = nn.CrossEntropyLoss()
- self.dist_params = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-
- def dataset_kwargs(self, train=True, download=False):
- """set data constructor kwargs for CIFAR10"""
- kwargs = dict(
- train=train,
- root=self.dataset_path,
- download=download,
- transform=self.transform,
- )
- return kwargs
-
-
-class CIFAR100(CIFAR10):
- def set_properties(self):
- """defines the required properties for CIFAR100"""
- self.dataset_path = files.dataset_path("CIFAR100")
- self.dataset_constructor = torchvision.datasets.CIFAR100
- self.loss_function = nn.CrossEntropyLoss()
- self.dist_params = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-
-
-class ImageNet2012(DataSet):
- def set_properties(self):
- """
- defines the required properties for ImageNet 2012 (ILSVRC2012) with
- 1000 classes.
- preprocessing according to pytorch documentation:
- https://pytorch.org/hub/pytorch_vision_alexnet/
- """
- self.dataset_path = files.dataset_path("ImageNet")
- self.dataset_constructor = torchvision.datasets.ImageNet
- self.loss_function = nn.CrossEntropyLoss()
- self.dist_params = dict(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
- self.center_crop = 224
-
- def dataset_kwargs(self, train=True):
- """set data constructor kwargs for ImageNet2012"""
- kwargs = dict(
- split="train" if train else "val",
- root=self.dataset_path,
- transform=self.transform,
- )
- return kwargs
-
-
-DATASET_REGISTRY = {
- "MNIST": MNIST,
- "CIFAR10": CIFAR10,
- "CIFAR100": CIFAR100,
- "ImageNet": ImageNet2012,
-}
-
-
-def get_dataset(
- dataset_name,
- build=False,
- dataset_parameters={},
- transform_parameters={},
- loader_parameters={},
- **kwargs,
-):
- """
- lookup dataset constructor from dataset registry by name
-
- if build=True, uses kwargs to build dataset and returns a dataset object
- otherwise just returns the constructor
- """
- if dataset_name not in DATASET_REGISTRY:
- raise ValueError(f"Dataset ({dataset_name}) is not in DATASET_REGISTRY")
- dataset = DATASET_REGISTRY[dataset_name]
- if build:
- if isinstance(transform_parameters, AlignmentNetwork):
- # Can use an AlignmentNetwork instance to automatically retrieve transform parameters
- transform_parameters = transform_parameters.get_transform_parameters(dataset_name)
- else:
- if not isinstance(transform_parameters, dict):
- raise TypeError("transform_parameters must be a dictionary or an AlignmentNetwork")
-
- # Build the dataset
- return dataset(
- dataset_parameters=dataset_parameters,
- transform_parameters=transform_parameters,
- loader_parameters=loader_parameters,
- **kwargs,
- )
-
- # Otherwise return the constructor
- return dataset
-
-
-if __name__ == "__main__":
- """simple program for downloading a dataset"""
-
- from argparse import ArgumentParser
-
- def get_args(args=None):
- parser = ArgumentParser(description="simple program for downloading a dataset to the local file location")
- parser.add_argument("--dataset", type=str, default="MNIST")
- return parser.parse_args(args=args)
-
- args = get_args()
-
- dataset = get_dataset(args.dataset, build=True, dataset_parameters=dict(download=True))
diff --git a/_arxiv/_archive/alignment_v2/files.py b/_arxiv/_archive/alignment_v2/files.py
deleted file mode 100644
index 4aa4216c..00000000
--- a/_arxiv/_archive/alignment_v2/files.py
+++ /dev/null
@@ -1,55 +0,0 @@
-import getpass
-import socket
-from pathlib import Path
-
-PATH_REGISTRY = {
- "DESKTOP-M2J64J2": Path("C:/Users/andrew/Documents/machineLearning"),
- "Celia": Path("/Users/celiaberon/Documents/machine_learning"),
- "cberon": Path("/n/home00/cberon/alignment/"),
- "corwin": Path("/Users/corwin/Building"),
- "atlandau": Path("/n/home05/atlandau/machine_learning"),
- "landauland": Path("/Users/landauland/Documents/ML-Datasets"),
- "nkhoshnevis": Path("/n/vast-scratch/kempner_dev/nkhoshnevis/projects/alignment_dir"),
- "hsafaai": Path("/n/holylabs/LABS/kempner_dev/Users/hsafaai/datasets"),
-}
-
-
-def get_hostname():
- return socket.gethostname()
-
-
-def get_username():
- return getpass.getuser()
-
-
-def local_path():
- """method for defining the local root path for datasets and results"""
- hostname = get_hostname()
- if hostname.lower().startswith("celia"):
- hostname = "Celia"
- hostname = hostname if hostname in PATH_REGISTRY else get_username()
- if hostname not in PATH_REGISTRY:
- raise ValueError(f"hostname ({hostname}) is not registered in the path registry")
- # return path
- return PATH_REGISTRY[hostname]
-
-
-def results_path():
- """method for returning relative path to results generated by this package"""
- return local_path() / "results"
-
-
-def data_path():
- """method for returning the relative path containing datasets"""
- return local_path() / "datasets"
-
-
-def dataset_path(dataset):
- """path to specific stored datasets (they all have different requirements)"""
- if dataset == "MNIST":
- return data_path()
- elif dataset == "CIFAR10" or dataset == "CIFAR100":
- return data_path()
- elif dataset == "ImageNet":
- return Path("/n/holylfs06/LABS/kempner_shared/Lab/data/imagenet_1k/")
-
diff --git a/_arxiv/_archive/alignment_v2/plotting.py b/_arxiv/_archive/alignment_v2/plotting.py
deleted file mode 100644
index c9c768e6..00000000
--- a/_arxiv/_archive/alignment_v2/plotting.py
+++ /dev/null
@@ -1,565 +0,0 @@
-import numpy as np
-from tqdm import tqdm
-from matplotlib import pyplot as plt
-import matplotlib as mpl
-
-import torch
-from alignment_v2.utils import compute_stats_by_type, named_transpose, transpose_list, rms
-
-
-def plot_train_results(exp, train_results, test_results, prms):
- """
- plotting method for training trajectories and testing data
- """
-
- num_train_epochs = train_results["loss"].size(0)
- num_types = len(prms["vals"])
- labels = [f"{prms['name']}={val}" for val in prms["vals"]]
-
- print("getting statistics on run data...")
- plot_alignment = "alignment" in train_results
- if plot_alignment:
- alignment = torch.stack([torch.mean(align, dim=2) for align in train_results["alignment"]])
-
- cmap = mpl.colormaps["tab10"]
-
- train_loss_mean, train_loss_se = compute_stats_by_type(train_results["loss"], num_types=num_types, dim=1, method="se")
- train_acc_mean, train_acc_se = compute_stats_by_type(train_results["accuracy"], num_types=num_types, dim=1, method="se")
-
- if plot_alignment:
- align_mean, align_se = compute_stats_by_type(alignment, num_types=num_types, dim=1, method="se")
-
- test_loss_mean, test_loss_se = compute_stats_by_type(torch.tensor(test_results["loss"]), num_types=num_types, dim=0, method="se")
- test_acc_mean, test_acc_se = compute_stats_by_type(torch.tensor(test_results["accuracy"]), num_types=num_types, dim=0, method="se")
-
- print("plotting run data...")
- xOffset = [-0.2, 0.2]
- get_x = lambda idx: [xOffset[0] + idx, xOffset[1] + idx]
-
- # Make Training and Testing Performance Figure
- alpha = 0.3
- figdim = 3
- figratio = 2
- width_ratios = [figdim, figdim / figratio, figdim, figdim / figratio]
-
- fig, ax = plt.subplots(1, 4, figsize=(sum(width_ratios), figdim), width_ratios=width_ratios, layout="constrained")
-
- # plot loss results fot training and testing
- for idx, label in enumerate(labels):
- cmn = train_loss_mean[:, idx]
- cse = train_loss_se[:, idx]
- tmn = test_loss_mean[idx]
- tse = test_loss_se[idx]
-
- ax[0].plot(range(num_train_epochs), cmn, color=cmap(idx), label=label)
- ax[0].fill_between(range(num_train_epochs), cmn + cse, cmn - cse, color=(cmap(idx), alpha))
- ax[1].plot(get_x(idx), [tmn] * 2, color=cmap(idx), label=label, lw=4)
- ax[1].plot([idx, idx], [tmn - tse, tmn + tse], color=cmap(idx), lw=1.5)
-
- ax[0].set_xlabel("Training Epoch")
- ax[0].set_ylabel("Loss")
- ax[0].set_title("Training Loss")
- ax[0].set_ylim(0, None)
- ylims = ax[0].get_ylim()
- ax[1].set_xticks(range(num_types))
- ax[1].set_xticklabels(labels, rotation=45, ha="right", fontsize=8)
- ax[1].set_ylabel("Loss")
- ax[1].set_title("Testing")
- ax[1].set_xlim(-0.5, num_types - 0.5)
- ax[1].set_ylim(ylims)
-
- # plot loss results fot training and testing
- for idx, label in enumerate(labels):
- cmn = train_acc_mean[:, idx]
- cse = train_acc_se[:, idx]
- tmn = test_acc_mean[idx]
- tse = test_acc_se[idx]
-
- ax[2].plot(range(num_train_epochs), cmn, color=cmap(idx), label=label)
- ax[2].fill_between(range(num_train_epochs), cmn + cse, cmn - cse, color=(cmap(idx), alpha))
- ax[3].plot(get_x(idx), [tmn] * 2, color=cmap(idx), label=label, lw=4)
- ax[3].plot([idx, idx], [tmn - tse, tmn + tse], color=cmap(idx), lw=1.5)
-
- ax[2].set_xlabel("Training Epoch")
- ax[2].set_ylabel("Accuracy (%)")
- ax[2].set_title("Training Accuracy")
- ax[2].set_ylim(0, 100)
- ax[3].set_xticks(range(num_types))
- ax[3].set_xticklabels(labels, rotation=45, ha="right", fontsize=8)
- ax[3].set_ylabel("Accuracy (%)")
- ax[3].set_title("Testing")
- ax[3].set_xlim(-0.5, num_types - 0.5)
- ax[3].set_ylim(0, 100)
-
- exp.plot_ready("train_test_performance")
-
- # Make Alignment Figure
- if plot_alignment:
- num_align_epochs = align_mean.size(2)
- num_layers = align_mean.size(0)
- fig, ax = plt.subplots(1, num_layers, figsize=(num_layers * figdim, figdim), layout="constrained", sharex=True)
- for idx, label in enumerate(labels):
- for layer in range(num_layers):
- cmn = align_mean[layer, idx] * 100
- cse = align_se[layer, idx] * 100
- ax[layer].plot(range(num_align_epochs), cmn, color=cmap(idx), label=label)
- ax[layer].fill_between(range(num_align_epochs), cmn + cse, cmn - cse, color=(cmap(idx), alpha))
-
- for layer in range(num_layers):
- ax[layer].set_ylim(0, None)
- ax[layer].set_xlabel("Training Epoch")
- ax[layer].set_ylabel("Alignment (%)")
- ax[layer].set_title(f"Layer {layer}")
-
- ax[0].legend(loc="lower right")
-
- exp.plot_ready("train_alignment_by_layer")
-
-
-def plot_dropout_results(exp, dropout_results, dropout_parameters, prms, dropout_type="nodes"):
- num_types = len(prms["vals"])
- labels = [f"{prms['name']}={val} - dropout {dropout_type}" for val in prms["vals"]]
- cmap = mpl.colormaps["Set1"]
- alpha = 0.3
- msize = 10
- figdim = 3
-
- num_layers = dropout_results["progdrop_loss_high"].size(2)
- names = ["From high", "From low", "Random"]
- names_diff = ["high-low"]
- num_exp = len(names)
- dropout_fraction = dropout_results["dropout_fraction"]
- by_layer = dropout_results["by_layer"]
- extra_name = "by_layer" if by_layer else "all_layers"
- extra_name += dropout_type
-
- # Get statistics across each network type for progressive dropout experiment
- print("measuring statistics on dropout analysis...")
- loss_mean_high, loss_se_high = compute_stats_by_type(dropout_results["progdrop_loss_high"], num_types=num_types, dim=0, method="se")
- loss_mean_low, loss_se_low = compute_stats_by_type(dropout_results["progdrop_loss_low"], num_types=num_types, dim=0, method="se")
- loss_mean_rand, loss_se_rand = compute_stats_by_type(dropout_results["progdrop_loss_rand"], num_types=num_types, dim=0, method="se")
-
- acc_mean_high, acc_se_high = compute_stats_by_type(dropout_results["progdrop_acc_high"], num_types=num_types, dim=0, method="se")
- acc_mean_low, acc_se_low = compute_stats_by_type(dropout_results["progdrop_acc_low"], num_types=num_types, dim=0, method="se")
- acc_mean_rand, acc_se_rand = compute_stats_by_type(dropout_results["progdrop_acc_rand"], num_types=num_types, dim=0, method="se")
-
- acc_mean_diff, acc_se_diff = compute_stats_by_type(dropout_results["progdrop_acc_high"]-dropout_results["progdrop_acc_low"], num_types=num_types, dim=0, method="se")
-
- # Contract into lists for looping through to plot
- loss_mean = [loss_mean_high, loss_mean_low, loss_mean_rand]
- loss_se = [loss_se_high, loss_se_low, loss_se_rand]
- acc_mean = [acc_mean_high, acc_mean_low, acc_mean_rand]
- acc_se = [acc_se_high, acc_se_low, acc_se_rand]
-
- print("plotting dropout results...")
- # Plot Loss for progressive dropout experiment
- fig, ax = plt.subplots(
- num_layers,
- num_types,
- figsize=(num_types * figdim, num_layers * figdim),
- sharex=True,
- sharey=True,
- layout="constrained",
- )
- ax = np.reshape(ax, (num_layers, num_types))
-
- if 1==2:
- for idx, label in enumerate(labels):
- for layer in range(num_layers):
- for iexp, name in enumerate(names):
- cmn = loss_mean[iexp][idx, :, layer]
- cse = loss_se[iexp][idx, :, layer]
- ax[layer, idx].plot(
- dropout_fraction,
- cmn,
- color=cmap(iexp),
- marker=".",
- markersize=msize,
- label=name,
- )
- ax[layer, idx].fill_between(dropout_fraction, cmn + cse, cmn - cse, color=(cmap(iexp), alpha))
-
- if layer == 0:
- ax[layer, idx].set_title(label)
-
- if layer == num_layers - 1:
- ax[layer, idx].set_xlabel("Dropout Fraction")
- ax[layer, idx].set_xlim(0, 1)
-
- if idx == 0:
- ax[layer, idx].set_ylabel("Loss w/ Dropout")
-
- if iexp == num_exp - 1:
- ax[layer, idx].legend(loc="best")
- if exp is not None:
- exp.plot_ready("prog_dropout_" + extra_name + "_loss")
-
- fig, ax = plt.subplots(
- num_layers,
- num_types,
- figsize=(num_types * figdim, num_layers * figdim),
- sharex=True,
- sharey=True,
- layout="constrained",
- )
- ax = np.reshape(ax, (num_layers, num_types))
-
- for idx, label in enumerate(labels):
- for layer in range(num_layers):
- for iexp, name in enumerate(names):
-
- cmn = acc_mean[iexp][idx, :, layer]
- cse = acc_se[iexp][idx, :, layer]
- ax[layer, idx].plot(
- dropout_fraction,
- cmn,
- color=cmap(iexp),
- marker=".",
- markersize=msize,
- label=name,
- )
- ax[layer, idx].fill_between(dropout_fraction, cmn + cse, cmn - cse, color=(cmap(iexp), alpha))
-
- ax[layer, idx].set_ylim(0, 100)
-
- if layer == 0:
- ax[layer, idx].set_title(label)
-
- if layer == num_layers - 1:
- ax[layer, idx].set_xlabel("Dropout Fraction")
- ax[layer, idx].set_xlim(0, 1)
-
- if idx == 0:
- ax[layer, idx].set_ylabel("Accuracy w/ Dropout")
-
- if iexp == num_exp - 1:
- ax[layer, idx].legend(loc="best")
- if exp is not None:
- exp.plot_ready("prog_dropout_" + extra_name + "_accuracy")
-
- # Plot Difference in Accuracy (High - Low)
- fig, ax = plt.subplots(
- num_layers,
- num_types,
- figsize=(num_types * figdim, num_layers * figdim),
- sharex=True,
- sharey=True,
- layout="constrained",
- )
- ax = np.reshape(ax, (num_layers, num_types))
-
- for idx, label in enumerate(labels):
- for layer in range(num_layers):
- for iexp, name in enumerate(names_diff):
- # Debugging: print shapes
- #print(f"acc_mean_diff[iexp].shape: {acc_mean_diff[iexp].shape}, dropout_fraction.shape: {dropout_fraction.shape}")
- # Squeeze to remove extra dimensions if necessary
- cmn = acc_mean_diff[iexp][:, layer] # Squeeze to ensure 1D tensor
- cse = acc_se_diff[iexp][:, layer] # Same for standard error
-
- # Ensure cmn and dropout_fraction have the same length
- if len(cmn.shape) == 1 and cmn.shape[0] == dropout_fraction.shape[0]:
- ax[layer, idx].plot(
- dropout_fraction,
- cmn, # Now this should be 1D
- color=cmap(iexp),
- marker=".",
- markersize=msize,
- label=name + " (High - Low)",
- )
- ax[layer, idx].fill_between(dropout_fraction, cmn + cse, cmn - cse, color=(cmap(iexp), alpha))
- else:
- print(f"Shape mismatch: cmn.shape: {cmn.shape}, dropout_fraction.shape: {dropout_fraction.shape}")
-
- if layer == 0:
- ax[layer, idx].set_title(label)
-
- if layer == num_layers - 1:
- ax[layer, idx].set_xlabel("Dropout Fraction")
- ax[layer, idx].set_xlim(0, 1)
-
- if idx == 0:
- ax[layer, idx].set_ylabel("Accuracy Difference (High - Low)")
-
- if iexp == num_exp - 1:
- ax[layer, idx].legend(loc="best")
-
- if exp is not None:
- exp.plot_ready("prog_dropout_" + extra_name + "_accuracy_diff")
-
-
-def plot_eigenfeatures(exp, results, prms):
- """method for plotting results related to eigen-analysis"""
- beta, eigvals, class_betas, class_names = (
- results["beta"],
- results["eigvals"],
- results["class_betas"],
- results["class_names"],
- )
- beta = [[torch.abs(b) for b in net_beta] for net_beta in beta]
- class_betas = [[rms(cb, dim=2) for cb in net_class_beta] for net_class_beta in class_betas]
-
- num_types = len(prms["vals"])
- labels = [f"{prms['name']}={val}" for val in prms["vals"]]
- cmap = mpl.colormaps["tab10"]
- class_cmap = mpl.colormaps["viridis"].resampled(len(class_names))
-
- print("measuring statistics of eigenfeature analyses...")
-
- # shape wrangling
- beta = [torch.stack(b) for b in transpose_list(beta)]
- eigvals = [torch.stack(ev) for ev in transpose_list(eigvals)]
- class_betas = [torch.stack(cb) for cb in transpose_list(class_betas)]
-
- # normalize to relative values
- beta = [b / b.sum(dim=2, keepdim=True) for b in beta]
- eigvals = [ev / ev.sum(dim=1, keepdim=True) for ev in eigvals]
- class_betas = [cb / cb.sum(dim=2, keepdim=True) for cb in class_betas]
-
- # reuse these a few times
- statprms = lambda method: dict(num_types=num_types, dim=0, method=method)
-
- # get mean and variance eigenvalues for each layer for each network type
- mean_evals, var_evals = named_transpose([compute_stats_by_type(ev, **statprms("var")) for ev in eigvals])
-
- # get sorted betas (sorted within each neuron)
- sorted_beta = [torch.sort(b, descending=True, dim=2).values for b in beta]
-
- # get mean / se beta for each layer for each network type
- mean_beta, se_beta = named_transpose([compute_stats_by_type(b, **statprms("var")) for b in beta])
- mean_sorted, se_sorted = named_transpose([compute_stats_by_type(b, **statprms("var")) for b in sorted_beta])
- mean_class_beta, se_class_beta = named_transpose([compute_stats_by_type(cb, **statprms("var")) for cb in class_betas])
-
- print("plotting eigenfeature results...")
- figdim = 3
- alpha = 0.3
- num_layers = len(mean_beta)
- fig, ax = plt.subplots(2, num_layers, figsize=(num_layers * figdim, figdim * 2), layout="constrained")
-
- for layer in range(num_layers):
- num_input = mean_evals[layer].size(1)
- num_nodes = mean_beta[layer].size(1)
- for idx, label in enumerate(labels):
- mn_ev = mean_evals[layer][idx]
- se_ev = var_evals[layer][idx]
- mn_beta = torch.mean(mean_beta[layer][idx], dim=0)
- se_beta = torch.std(mean_beta[layer][idx], dim=0) / np.sqrt(num_nodes)
- mn_sort = torch.mean(mean_sorted[layer][idx], dim=0)
- se_sort = torch.std(mean_sorted[layer][idx], dim=0) / np.sqrt(num_nodes)
- ax[0, layer].plot(
- range(num_input),
- mn_ev,
- color=cmap(idx),
- linestyle="--",
- label="eigvals" if idx == 0 else None,
- )
- ax[0, layer].plot(range(num_input), mn_beta, color=cmap(idx), label=label)
- ax[0, layer].fill_between(range(num_input), mn_beta + se_beta, mn_beta - se_beta, color=(cmap(idx), alpha))
- ax[1, layer].plot(range(num_input), mn_sort, color=cmap(idx), label=label)
- ax[1, layer].fill_between(range(num_input), mn_sort + se_sort, mn_sort - se_sort, color=(cmap(idx), alpha))
-
- ax[0, layer].set_xscale("log")
- ax[1, layer].set_xscale("log")
- ax[0, layer].set_xlabel("Input Dimension")
- ax[1, layer].set_xlabel("Sorted Input Dim")
- ax[0, layer].set_ylabel("Relative Eigval & Beta")
- ax[1, layer].set_ylabel("Relative Beta (Sorted)")
- ax[0, layer].set_title(f"Layer {layer}")
- ax[1, layer].set_title(f"Layer {layer}")
-
- if layer == num_layers - 1:
- ax[0, layer].legend(loc="best")
- ax[1, layer].legend(loc="best")
-
- exp.plot_ready("eigenfeatures")
-
- fig, ax = plt.subplots(1, num_layers, figsize=(num_layers * figdim, figdim), layout="constrained")
-
- for layer in range(num_layers):
- num_input = mean_evals[layer].size(1)
- num_nodes = mean_beta[layer].size(1)
- for idx, label in enumerate(labels):
- mn_ev = mean_evals[layer][idx]
- se_ev = var_evals[layer][idx]
- mn_beta = torch.mean(mean_beta[layer][idx], dim=0)
- se_beta = torch.std(mean_beta[layer][idx], dim=0) / np.sqrt(num_nodes)
- mn_sort = torch.mean(mean_sorted[layer][idx], dim=0)
- se_sort = torch.std(mean_sorted[layer][idx], dim=0) / np.sqrt(num_nodes)
- ax[layer].plot(
- range(num_input),
- mn_ev,
- color=cmap(idx),
- linestyle="--",
- label="eigvals" if idx == 0 else None,
- )
- ax[layer].plot(range(num_input), mn_beta, color=cmap(idx), label=label)
- ax[layer].fill_between(range(num_input), mn_beta + se_beta, mn_beta - se_beta, color=(cmap(idx), alpha))
-
- ax[layer].set_xscale("log")
- ax[layer].set_yscale("log")
- ax[layer].set_xlabel("Input Dimension")
- ax[layer].set_ylabel("Relative Eigval & Beta")
- ax[layer].set_title(f"Layer {layer}")
-
- if layer == num_layers - 1:
- ax[layer].legend(loc="best")
-
- exp.plot_ready("eigenfeatures_loglog")
-
- fig, ax = plt.subplots(
- num_types,
- num_layers,
- figsize=(num_layers * figdim, figdim * num_types),
- layout="constrained",
- )
- ax = np.reshape(ax, (num_types, num_layers))
- for layer in range(num_layers):
- num_input = mean_evals[layer].size(1)
- for idx, label in enumerate(labels):
- for idx_class, class_name in enumerate(class_names):
- mn_data = mean_class_beta[layer][idx][idx_class]
- se_data = se_class_beta[layer][idx][idx_class]
- ax[idx, layer].plot(range(num_input), mn_data, color=class_cmap(idx_class), label=class_name)
- ax[idx, layer].fill_between(
- range(num_input),
- mn_data + se_data,
- mn_data - se_data,
- color=(class_cmap(idx_class), alpha),
- )
-
- ax[idx, layer].set_xscale("log")
- ax[idx, layer].set_yscale("linear")
- ax[idx, layer].set_xlabel("Input Dimension")
- if layer == 0:
- ax[idx, layer].set_ylabel(f"{label}\nClass Loading (RMS)")
- if idx == 0:
- ax[idx, layer].set_title(f"Layer {layer}")
-
- if layer == num_layers - 1:
- ax[idx, layer].legend(loc="upper right", fontsize=6)
-
- exp.plot_ready("class_eigenfeatures")
-
-
-def plot_adversarial_results(exp, eigen_results, adversarial_results, prms):
- accuracy, beta, eigvals = (
- adversarial_results["accuracy"],
- adversarial_results["betas"],
- eigen_results["eigvals"],
- )
- epsilons, use_sign = adversarial_results["epsilons"], adversarial_results["use_sign"]
-
- num_types = len(prms["vals"])
- labels = [f"{prms['name']}={val}" for val in prms["vals"]]
- cmap = mpl.colormaps["tab10"]
-
- print("measuring statistics of adversarial analyses...")
-
- # shape wrangling
- accuracy = torch.stack([torch.stack(acc) for acc in transpose_list(accuracy)]) # (num_epsilon, num_nets)
- eigvals = [torch.stack(ev) for ev in transpose_list(eigvals)] # [(num_nets, dim_layer) for each layer]
- beta = [torch.stack(b) for b in beta] # [(epsilon, num_nets, dim_layer) for each layer]
-
- # normalize to relative values
- beta = [b / b.sum(dim=2, keepdim=True) for b in beta]
- eigvals = [ev / ev.sum(dim=1, keepdim=True) for ev in eigvals]
-
- # reuse these a few times
- statprms = lambda dim, method: dict(num_types=num_types, dim=dim, method=method)
-
- # get mean and variance beta/eigenvalues for each layer for each network type
- mean_acc, se_acc = compute_stats_by_type(accuracy, **statprms(1, "var")) # (num_epsilon, num_types)
- mean_beta, se_beta = named_transpose(
- [compute_stats_by_type(b, **statprms(1, "var")) for b in beta]
- ) # [(epsilon, num_types, dim_layer) for each layer]
- mean_evals, var_evals = named_transpose(
- [compute_stats_by_type(ev, **statprms(0, "var")) for ev in eigvals]
- ) # [(num_types, dim_layer) for each layer]
-
- print("plotting adversarial success results...")
- figdim = 3
- alpha = 0.3
- num_layers = len(mean_beta)
- fig, ax = plt.subplots(1, 1, figsize=(figdim, figdim), layout="constrained")
- for idx, label in enumerate(labels):
- ax.plot(epsilons, mean_acc[:, idx], color=cmap(idx), label=label)
- ax.set_xlabel("Epsilon")
- ax.set_ylabel("Accuracy")
- ax.set_title("Adversarial Attack Success")
- ax.legend(loc="best")
-
- exp.plot_ready("adversarial_success")
-
- print("plotting adversarial structure...")
- fig, ax = plt.subplots(1, num_layers, figsize=(num_layers * figdim, figdim), layout="constrained")
- for layer in range(num_layers):
- num_input = mean_beta[layer].size(2)
- for idx, label in enumerate(labels):
- ax[layer].plot(
- range(num_input),
- torch.nanmean(mean_beta[layer][:, idx], dim=0).detach(),
- color=cmap(idx),
- label=label,
- )
- ax[layer].set_xscale("log")
- ax[layer].set_xlabel("Input Dimension")
- ax[layer].set_ylabel("Average Component of Pertubation")
- ax[layer].set_title(f"Layer {layer}")
- if layer == num_layers - 1:
- ax[layer].legend(loc="best")
-
- exp.plot_ready("adversarial_structure")
-
-
-def plot_rf(rf, width, alignment=None, alignBounds=None, showRFs=None, figSize=5):
- if showRFs is not None:
- rf = rf.reshape(rf.shape[0], -1)
- idxRandom = np.random.choice(range(rf.shape[0]), showRFs, replace=False)
- rf = rf[idxRandom, :]
- else:
- showRFs = rf.shape[0]
- # normalize
- rf = rf.T / np.abs(rf).max(axis=1)
- rf = rf.T
- rf = rf.reshape(showRFs, width, width)
- # If necessary, create colormap
- if alignment is not None:
- cmap = mpl.cm.get_cmap("rainbow", rf.shape[0])
- cmapPeak = lambda x: cmap(x)
- if alignBounds is not None:
- alignment = alignment - alignBounds[0]
- alignment = alignment / (alignBounds[1] - alignBounds[0])
- else:
- alignment = alignment - alignment.min()
- alignment = alignment / alignment.max()
-
- # plotting
- n = int(np.ceil(np.sqrt(rf.shape[0])))
- fig, axes = plt.subplots(nrows=n, ncols=n, sharex=True, sharey=True)
- fig.set_size_inches(figSize, figSize)
-
- N = 1000
- for i in tqdm(range(rf.shape[0])):
- ax = axes[i // n][i % n]
- if alignment is not None:
- vals = np.ones((N, 4))
- cAlignment = alignment[i].numpy()
- cPeak = cmapPeak(alignment[i].numpy())
- vals[:, 0] = np.linspace(0, cPeak[0], N)
- vals[:, 1] = np.linspace(0, cPeak[1], N)
- vals[:, 2] = np.linspace(0, cPeak[2], N)
- usecmap = mpl.colors.ListedColormap(vals)
- ax.imshow(rf[i], cmap=usecmap, vmin=-1, vmax=1)
- else:
- ax.imshow(rf[i], cmap="gray", vmin=-1, vmax=1)
- ax.set_xticks([])
- ax.set_yticks([])
- ax.set_aspect("equal")
- for j in range(rf.shape[0], n * n):
- ax = axes[j // n][j % n]
- ax.imshow(np.ones_like(rf[0]) * -1, cmap="gray", vmin=-1, vmax=1)
- ax.set_xticks([])
- ax.set_yticks([])
- ax.set_aspect("equal")
- fig.subplots_adjust(wspace=0.0, hspace=0.0)
- return fig
diff --git a/_arxiv/_archive/alignment_v2/processing.py b/_arxiv/_archive/alignment_v2/processing.py
deleted file mode 100644
index e7db84df..00000000
--- a/_arxiv/_archive/alignment_v2/processing.py
+++ /dev/null
@@ -1,316 +0,0 @@
-import os
-from tqdm import tqdm
-import torch
-from alignment_v2 import train
-from alignment_v2.utils import load_checkpoints, test_nets, transpose_list, fgsm_attack
-
-
-def train_networks(exp, nets, optimizers, dataset, **special_parameters):
- """train and test networks"""
- # do training loop
- parameters = dict(
- train_set=True,
- num_epochs=exp.args.epochs,
- alignment=not (exp.args.no_alignment),
- alignment_expansion=True,#not (exp.args.no_alignment),
- delta_weights=exp.args.delta_weights,
- frequency=exp.args.frequency,
- run=exp.run,
- )
-
- # update with special parameters
- parameters.update(**special_parameters)
-
- if exp.args.use_prev & os.path.isfile(exp.get_checkpoint_path()):
- nets, optimizers, results = load_checkpoints(nets, optimizers, exp.args.device, exp.get_checkpoint_path())
- for net in nets:
- net.train()
-
- parameters["num_complete"] = results["epoch"] + 1
- parameters["results"] = results
- print("loaded networks from previous checkpoint")
-
- if exp.args.save_ckpts:
- parameters["save_checkpoints"] = (True, exp.args.ckpt_frequency, exp.get_checkpoint_path(), exp.args.device)
-
- print("training networks...")
- train_results = train.train(nets, optimizers, dataset, **parameters)
-
- # do testing loop
- print("testing networks...")
- parameters["train_set"] = False
- test_results = train.test(nets, dataset, **parameters)
-
- return train_results, test_results
-
-
-def progressive_dropout_experiment(exp, nets, dataset, alignment=None, train_set=False):
- """
- perform a progressive dropout (of nodes) experiment
- alignment is optional, but will be recomputed if you've already measured it. You can provide it
- by setting: alignment=test_results['alignment'] if ``train_networks`` has already been run.
- """
- # do targeted dropout experiment
- print("performing targeted dropout...")
- dropout_parameters = dict(num_drops=exp.args.num_drops, by_layer=exp.args.dropout_by_layer, train_set=train_set)
- dropout_results = train.progressive_dropout(nets, dataset, alignment=alignment, **dropout_parameters)
- return dropout_results, dropout_parameters
-
-# def progressive_dropout_experiment(exp, nets, dataset, alignment=None, train_set=False, by_layer=False, layer_idx=None):
-# """
-# Perform progressive dropout, either layer-by-layer or across all layers at once.
-# If by_layer=True, we only dropout nodes in the specified layer (layer_idx).
-# """
-# dropout_results = {}
-# dropout_parameters = {}
-
-# if by_layer:
-# print(f"Performing dropout for layer {layer_idx}")
-
-# # Drop nodes in the specified layer only
-# dropout_results = train.progressive_dropout(
-# nets, dataset, alignment=alignment, layer_idx=layer_idx, train_set=train_set, by_layer=True
-# )
-
-# else:
-# # Standard dropout across all layers
-# dropout_results = train.progressive_dropout(
-# nets, dataset, alignment=alignment, train_set=train_set, by_layer=False
-# )
-
-# return dropout_results, dropout_parameters
-
-def sequential_dropout_experiment(exp, nets, dataset):
- """
- Perform a sequential dropout experiment by pruning layers one by one, recalculating alignment after each pruning.
- """
- print("Performing sequential dropout experiment...")
- sequential_dropout_results = {}
- sequential_dropout_parameters = {}
-
- # Loop over the layers to sequentially prune each layer
- for layer_idx in range(len(nets[0].layers)): # Assuming nets[0].layers holds the network layers
- print(f"Sequentially pruning layer {layer_idx}")
-
- # Get the alignment after pruning previous layers
- alignment = train.test(nets, dataset, alignment_only=True)["alignment"]
-
- # Perform progressive dropout for the current layer
- print(f"Performing dropout for layer {layer_idx}")
- dropout_results, dropout_parameters = progressive_dropout_experiment(
- exp, nets, dataset, alignment=alignment, layer_idx=layer_idx
- )
- # Recalculate alignment after pruning the current layer
- print(f"Recalculating alignment after pruning layer {layer_idx}")
- test_results = train.test(nets, dataset, alignment_only=True) # Update test_results after each layer
-
- # Store the dropout results and parameters for this layer
- sequential_dropout_results = dropout_results
- sequential_dropout_parameters = dropout_parameters
-
- return sequential_dropout_results, sequential_dropout_parameters
-
-
-def retrain_network_with_dropout_stats(exp, nets, optimizers, dataset, original_dropout_results):
- """
- Retrains the network and maintains the same key structure as the dropout results.
- This ensures the retraining results have keys like 'progdrop_loss_high', etc.
- """
- # Train the network using the regular training process
- print("Retraining the pruned network...")
- for net in nets:
- net.train()
-
- # Reset optimizers, assuming we need new ones for the pruned network
- optimizers = [torch.optim.SGD(net.parameters(), lr=0.001) for net in nets]
-
- # Train the networks after dropout
- retrain_results, retrain_test_results = train_networks(exp, nets, optimizers, dataset)
-
- train_results, test_results = train_networks(exp, nets, optimizers, dataset)
-
- # Initialize retraining results dictionary with same structure as original dropout results
- retrain_results = {
- "progdrop_loss_high": original_dropout_results["progdrop_loss_high"].clone().zero_(),
- "progdrop_loss_low": original_dropout_results["progdrop_loss_low"].clone().zero_(),
- "progdrop_loss_rand": original_dropout_results["progdrop_loss_rand"].clone().zero_(),
- "progdrop_acc_high": original_dropout_results["progdrop_acc_high"].clone().zero_(),
- "progdrop_acc_low": original_dropout_results["progdrop_acc_low"].clone().zero_(),
- "progdrop_acc_rand": original_dropout_results["progdrop_acc_rand"].clone().zero_(),
- "dropout_fraction": original_dropout_results["dropout_fraction"],
- "by_layer": original_dropout_results["by_layer"],
- "idx_dropout_layers": original_dropout_results["idx_dropout_layers"],
- "fraction_dropped_nodes": original_dropout_results["fraction_dropped_nodes"].clone().zero_()
- }
-
- # Populate the retraining results with new loss and accuracy after retraining
- # For example, assuming the train/test functions capture loss and accuracy
- retrain_results["progdrop_loss_high"] = test_results.get("progdrop_loss_high", retrain_results["progdrop_loss_high"])
- retrain_results["progdrop_loss_low"] = test_results.get("progdrop_loss_low", retrain_results["progdrop_loss_low"])
- retrain_results["progdrop_loss_rand"] = test_results.get("progdrop_loss_rand", retrain_results["progdrop_loss_rand"])
-
- retrain_results["progdrop_acc_high"] = test_results.get("progdrop_acc_high", retrain_results["progdrop_acc_high"])
- retrain_results["progdrop_acc_low"] = test_results.get("progdrop_acc_low", retrain_results["progdrop_acc_low"])
- retrain_results["progdrop_acc_rand"] = test_results.get("progdrop_acc_rand", retrain_results["progdrop_acc_rand"])
-
- return retrain_results, test_results
-
-def evaluate_network(nets, dataset):
- """
- Evaluate the pruned network on the test dataset to recalculate alignment and other metrics.
- """
- test_results = train.test(nets, dataset)
- return test_results
-
-def measure_eigenfeatures(exp, nets, dataset, train_set=False):
- # measure eigenfeatures
- print("measuring eigenfeatures...")
- beta, eigvals, eigvecs, class_betas = [], [], [], []
- for net in tqdm(nets):
- # get inputs to each layer from whole dataloader
- real_net = net.module if hasattr(net, "module") else net
- inputs, labels = real_net._process_collect_activity(
- dataset,
- train_set=train_set,
- with_updates=False,
- use_training_mode=False,
- )
- eigenfeatures = real_net.measure_eigenfeatures(inputs, with_updates=False)
- beta_by_class = real_net.measure_class_eigenfeatures(inputs, labels, eigenfeatures[2], rms=False, with_updates=False)
- beta.append(eigenfeatures[0])
- eigvals.append(eigenfeatures[1])
- eigvecs.append(eigenfeatures[2])
- class_betas.append(beta_by_class)
-
-
- # make it a dictionary
- class_names = getattr(dataset.train_loader if train_set else dataset.test_loader, "dataset").classes
- return dict(
- beta=beta,
- eigvals=eigvals,
- eigvecs=eigvecs,
- class_betas=class_betas,
- class_names=class_names,
- )
-
-
-def eigenvector_dropout(exp, nets, dataset, eigen_results, train_set=False):
- """
- do targeted eigenvector dropout with precomputed eigenfeatures
- """
- # do targeted dropout experiment
- print("performing targeted eigenvector dropout...")
- evec_dropout_parameters = dict(num_drops=exp.args.num_drops, by_layer=exp.args.dropout_by_layer, train_set=train_set)
- evec_dropout_results = train.eigenvector_dropout(nets, dataset, eigen_results["eigvals"], eigen_results["eigvecs"], **evec_dropout_parameters)
- return evec_dropout_results, evec_dropout_parameters
-
-
-@test_nets
-def measure_adversarial_attacks(nets, dataset, exp, eigen_results, train_set=False, **parameters):
- """
- do adversarial attack and measure structure with regards to eigenfeatures
- """
-
- def get_beta(inputs, eigenvectors):
- # get projection of input onto eigenvectors across layers
- return [input.cpu() @ evec for input, evec in zip(inputs, eigenvectors)]
-
- # experiment parameters
- epsilons = parameters.get("epsilons")
- use_sign = parameters.get("use_sign")
- fgsm_transform = parameters.get("fgsm_transform", lambda x: x)
-
- # data from eigenvectors
- eigenvectors = eigen_results["eigvecs"]
-
- num_eps = len(epsilons)
- num_nets = len(nets)
- accuracy = torch.zeros((num_nets, num_eps))
- examples = [[[] for _ in range(num_eps)] for _ in range(num_nets)]
- betas = [[torch.zeros((num_nets, evec.size(0))) for evec in eigenvectors[0]] for _ in range(num_eps)]
-
- # dataloader
- dataloader = dataset.train_loader if train_set else dataset.test_loader
-
- for batch in tqdm(dataloader):
- input, labels = dataset.unwrap_batch(batch)
-
- inputs = [input.clone() for _ in range(num_nets)]
-
- for input in inputs:
- input.requires_grad = True
-
- # Forward pass the data through the model
- outputs = [net(input, store_hidden=True) for net, input in zip(nets, inputs)]
- input_to_layers = [net.get_layer_inputs(input, precomputed=True) for net in nets]
- init_preds = [torch.argmax(output, axis=1) for output in outputs] # find true prediction
- least_likely = [torch.argmin(output, axis=1) for output in outputs] # find least likely digit according to model
-
- c_betas = transpose_list([get_beta(input, evec) for input, evec in zip(input_to_layers, eigenvectors)])
- s_betas = [torch.stack(cb) for cb in c_betas]
-
- # Calculate the loss
- loss = [dataset.measure_loss(output, labels) for output in outputs]
- # loss = dataset.measure_loss(output, least_likely)
-
- # Zero all existing gradients
- for net in nets:
- net.zero_grad()
-
- # Calculate gradients of model in backward pass
- for l in loss:
- l.backward()
-
- # Collect datagrad
- data_grads = [input.grad.data for input in inputs]
-
- for epsidx, eps in enumerate(epsilons):
-
- # Call FGSM Attack
- perturbed_inputs = [fgsm_attack(input, eps, data_grad, fgsm_transform, use_sign) for input, data_grad in zip(inputs, data_grads)]
-
- # Re-classify the perturbed image
- outputs = [net(perturbed_input, store_hidden=True) for net, perturbed_input in zip(nets, perturbed_inputs)]
- input_to_layers = [net.get_layer_inputs(perturbed_input, precomputed=True) for net, perturbed_input in zip(nets, perturbed_inputs)]
- c_eps_betas = transpose_list([get_beta(input, evec) for input, evec in zip(input_to_layers, eigenvectors)])
- s_eps_betas = [torch.stack(ceb) for ceb in c_eps_betas]
- d_eps_betas = [sebeta - sbeta for sebeta, sbeta in zip(s_eps_betas, s_betas)]
- rms_betas = [torch.sqrt(torch.mean(db**2, dim=1)) for db in d_eps_betas]
-
- for ii, rbeta in enumerate(rms_betas):
- betas[epsidx][ii] += rbeta.detach()
-
- # Check for success
- final_preds = [torch.argmax(output, axis=1) for output in outputs]
- accuracy[:, epsidx] += torch.tensor([sum(final_pred == labels).cpu() for final_pred in final_preds])
-
- # Idx where adversarial example worked
- idx_success = [
- torch.where((init_pred == labels) & (final_pred != labels))[0].cpu() for init_pred, final_pred in zip(init_preds, final_preds)
- ]
-
- adv_exs = [perturbed_input.detach().cpu().numpy() for perturbed_input in perturbed_inputs]
- for ii, (adv_ex, idx, init_pred, final_pred) in enumerate(zip(adv_exs, idx_success, init_preds, final_preds)):
- examples[ii][epsidx].append((init_pred[idx], final_pred[idx], adv_ex[idx]))
-
- # Calculate final accuracy for this epsilon
- accuracy = accuracy / float(len(dataloader.dataset))
-
- # Average across betas
- betas = transpose_list([[cb / float(len(dataloader.dataset)) for cb in beta] for beta in betas])
-
- # Return the accuracy and an adversarial example
- return dict(accuracy=accuracy, betas=betas, examples=examples, epsilons=epsilons, use_sign=use_sign)
-
-
-@test_nets
-def measure_alignment_distribution(nets, dataset, **parameters):
- """
- method for measuring alignment distribution and several associated analyses
- """
- # do training loop
- parameters = dict(
- train_set=True,
- )
-
-
diff --git a/_arxiv/_archive/alignment_v2/train.py b/_arxiv/_archive/alignment_v2/train.py
deleted file mode 100644
index 0e506501..00000000
--- a/_arxiv/_archive/alignment_v2/train.py
+++ /dev/null
@@ -1,813 +0,0 @@
-from copy import copy, deepcopy
-
-import torch
-from tqdm import tqdm
-import torch.distributed as dist # keep for DDP rank checks
-
-from alignment_v2.utils import (
- transpose_list,
- condense_values,
- test_nets,
- train_nets,
- save_checkpoint,
- smart_pca,
- expected_alignment_distribution,
- alignment,
- alignment_expansion,
-)
-
-
-@train_nets
-def train(nets, optimizers, dataset, **parameters):
- """method for training network on supervised learning problem"""
-
- # input argument checks
- if not (isinstance(nets, list)):
- nets = [nets]
- if not (isinstance(optimizers, list)):
- optimizers = [optimizers]
- assert len(nets) == len(optimizers), "nets and optimizers need to be equal length lists"
-
- # check if we should print progress bars
- verbose = parameters.get("verbose", True)
-
- num_nets = len(nets)
- use_train = parameters.get("train_set", True)
- dataloader = dataset.train_loader if use_train else dataset.test_loader
- num_steps = len(dataset.train_loader) * parameters["num_epochs"]
-
- # optional W&B run
- run = parameters.get("run")
-
- # optional analyses
- measure_alignment = parameters.get("alignment", True)
- measure_alignment_expansion = parameters.get("alignment_expansion", False)
- measure_delta_weights = parameters.get("delta_weights", False)
- measure_delta_alignment = parameters.get("delta_alignment", False)
- measure_frequency = parameters.get("frequency", 1)
- compare_expected = parameters.get("compare_expected", False)
-
- # optional manual shaping
- manual_shape = parameters.get("manual_shape", False)
- manual_frequency = parameters.get("manual_frequency", -1)
- manual_transforms = parameters.get("manual_transforms", None)
- manual_layers = parameters.get("manual_layers", None)
-
- # create or retrieve results dict
- results = parameters.get("results", False)
- num_complete = parameters.get("num_complete", 0)
- save_ckpt, freq_ckpt, path_ckpt, dev = parameters.get("save_checkpoints", (False, 1, "", ""))
- if not results:
- results = {
- "loss": torch.zeros((num_steps, num_nets)),
- "accuracy": torch.zeros((num_steps, num_nets)),
- }
- if measure_alignment:
- results["alignment"] = []
- if measure_alignment_expansion:
- results["alignment_0"] = []
- results["alignment_1"] = []
- results["alignment_2"] = []
- results["alignment_red"] = []
- if measure_delta_weights:
- results["delta_weights"] = []
- results["init_weights"] = [
- net.module.get_alignment_weights() if hasattr(net, "module") else net.get_alignment_weights()
- for net in nets
- ]
- if measure_delta_alignment:
- if "init_weights" not in results:
- results["init_weights"] = [
- net.module.get_alignment_weights() if hasattr(net, "module") else net.get_alignment_weights()
- for net in nets
- ]
- results["delta_alignment"] = []
- if compare_expected:
- calign_bins = torch.linspace(0, 1, 301)
- results["compare_alignment_bins"] = calign_bins
- results["compare_alignment_expected"] = []
- results["compare_alignment_observed"] = []
- if measure_delta_alignment:
- results["compare_delta_alignment_observed"] = []
- elif results["loss"].shape[0] < num_steps:
- add_steps = num_steps - results["loss"].shape[0]
- assert (add_steps / (parameters["num_epochs"] - num_complete)) == len(dataset.train_loader), \
- "Number of new steps must match epochs * loader length"
- results["loss"] = torch.vstack((results["loss"], torch.zeros((add_steps, num_nets))))
- results["accuracy"] = torch.vstack((results["accuracy"], torch.zeros((add_steps, num_nets))))
-
- if num_complete > 0:
- print("resuming training from checkpoint on epoch", num_complete)
-
- # training loop
- epoch_loop = range(num_complete, parameters["num_epochs"])
- if verbose:
- epoch_loop = tqdm(epoch_loop, desc="training epoch")
-
- for epoch in epoch_loop:
-
- batch_loop = dataloader
- if verbose:
- batch_loop = tqdm(batch_loop, desc="minibatch", leave=False)
-
- for idx, batch in enumerate(batch_loop):
- cidx = epoch * len(dataloader) + idx
- images, labels = dataset.unwrap_batch(batch)
-
- for opt in optimizers:
- opt.zero_grad()
-
- outputs = [net(images, store_hidden=True) for net in nets]
- loss = [dataset.measure_loss(output, labels) for output in outputs]
- for l, opt in zip(loss, optimizers):
- l.backward()
- opt.step()
-
- results["loss"][cidx] = torch.tensor([l.item() for l in loss])
- results["accuracy"][cidx] = torch.tensor([
- dataset.measure_accuracy(output, labels).cpu() for output in outputs
- ])
-
- if idx % measure_frequency == 0:
- if measure_alignment:
- alignment_vals = []
- for net in nets:
- real_net = net.module if hasattr(net, "module") else net
- alignment_val = real_net.measure_alignment(images, precomputed=True, method="alignment")
- alignment_vals.append(alignment_val)
- results["alignment"].append(alignment_vals)
-
- if measure_alignment_expansion:
- alignment0_vals, alignment1_vals, alignment2_vals, alignmentred_vals = [], [], [], []
- for net in nets:
- real_net = net.module if hasattr(net, "module") else net
- alignment0_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_0"))
- alignment1_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_1"))
- alignment2_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_2"))
- alignmentred_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_red"))
- results["alignment_0"].append(alignment0_vals)
- results["alignment_1"].append(alignment1_vals)
- results["alignment_2"].append(alignment2_vals)
- results["alignment_red"].append(alignmentred_vals)
-
- if measure_delta_weights or measure_delta_alignment:
- c_delta_weights = []
- for net, init_weight in zip(nets, results["init_weights"]):
- real_net = net.module if hasattr(net, "module") else net
- c_delta_weights.append(real_net.compare_weights(init_weight))
- if measure_delta_weights:
- results["delta_weights"].append(c_delta_weights)
- if measure_delta_alignment:
- c_delta_alignment = []
- for net, weights_ in zip(nets, c_delta_weights):
- real_net = net.module if hasattr(net, "module") else net
- c_delta_alignment.append(
- real_net.measure_alignment_weights(images, weights_, precomputed=True, method="alignment")
- )
- results["delta_alignment"].append(c_delta_alignment)
-
- if compare_expected:
- if measure_alignment:
- c_alignment = results["alignment"][-1]
- else:
- c_alignment = []
- for net in nets:
- real_net = net.module if hasattr(net, "module") else net
- c_alignment.append(real_net.measure_alignment(images, precomputed=True, method="alignment"))
-
- c_inputs = []
- for net in nets:
- real_net = net.module if hasattr(net, "module") else net
- c_inputs.append(real_net.get_layer_inputs(images, precomputed=True))
-
- for i in range(len(nets)):
- real_net_ = nets[i].module if hasattr(nets[i], "module") else nets[i]
- c_inputs[i] = real_net_._preprocess_inputs(c_inputs[i])
-
- c_evals = []
- for cin in c_inputs:
- c_eval_sub = []
- for c_ in cin:
- w, _ = smart_pca(c_.T)
- c_eval_sub.append(w)
- c_evals.append(c_eval_sub)
-
- calign_bins = results["compare_alignment_bins"]
- c_dist = []
- for c_eval in c_evals:
- subdist = []
- for ev in c_eval:
- subdist.append(expected_alignment_distribution(ev, valid_rotation=False, bins=calign_bins)[0])
- c_dist.append(subdist)
-
- t_dist = []
- for c_align in c_alignment:
- align_sub = []
- for align in c_align:
- align_sub.append(torch.histogram(align.cpu(), bins=calign_bins, density=True)[0])
- t_dist.append(align_sub)
-
- results["compare_alignment_expected"].append(c_dist)
- results["compare_alignment_observed"].append(t_dist)
-
- if measure_delta_alignment:
- d_alignment = results["delta_alignment"][-1]
- d_dist = []
- for d_align in d_alignment:
- align_sub = []
- for dalign in d_align:
- align_sub.append(torch.histogram(dalign.cpu(), bins=calign_bins, density=True)[0])
- d_dist.append(align_sub)
- results["compare_delta_alignment_observed"].append(d_dist)
-
- # Only rank 0 logs to W&B
- if run is not None:
- if dist.is_available() and dist.is_initialized() and dist.get_rank() != 0:
- pass
- else:
- run.log(
- {f"losses/loss-{ii}": l.item() for ii, l in enumerate(loss)}
- | {f"accuracies/accuracy-{ii}": dataset.measure_accuracy(output, labels) for ii, output in enumerate(outputs)}
- | {"batch": cidx}
- )
-
- if manual_shape:
- if ((epoch + 1) % manual_frequency == 0) and (epoch < parameters["num_epochs"] - 1):
- for net, transform in tqdm(zip(nets, manual_transforms), desc="manual shaping", leave=False):
- real_net = net.module if hasattr(net, "module") else net
- inputs, _ = real_net._process_collect_activity(dataset, train_set=False, with_updates=False, use_training_mode=False)
- _, eigenvalues, eigenvectors = real_net.measure_eigenfeatures(inputs, with_updates=False)
- idx_to_layer_lookup = {layer: i for i, layer in enumerate(real_net.get_alignment_layer_indices())}
- eigenvalues = [eigenvalues[idx_to_layer_lookup[ml]] for ml in manual_layers]
- eigenvectors = [eigenvectors[idx_to_layer_lookup[ml]] for ml in manual_layers]
- real_net.shape_eigenfeatures(manual_layers, eigenvalues, eigenvectors, transform)
-
- if save_ckpt & (epoch % freq_ckpt == 0):
- save_checkpoint(
- nets,
- optimizers,
- results | {"prms": parameters, "epoch": epoch, "device": dev},
- path_ckpt,
- )
-
- for k in [
- "alignment",
- "alignment_0",
- "alignment_1",
- "alignment_2",
- "alignment_red",
- "delta_weights",
- "delta_alignment",
- "avgcorr",
- "fullcorr",
- "compare_alignment_expected",
- "compare_alignment_observed",
- "compare_delta_alignment_observed",
- ]:
- if k not in results.keys():
- continue
- results[k] = condense_values(transpose_list(results[k]))
-
- return results
-
-
-
-@torch.no_grad()
-@test_nets
-def test(nets, dataset, **parameters):
- """method for testing network on supervised learning problem"""
-
- run = parameters.get("run")
-
- if not (isinstance(nets, list)):
- nets = [nets]
-
- verbose = parameters.get("verbose", True)
- num_nets = len(nets)
-
- use_test = not parameters.get("train_set", False)
- dataloader = dataset.test_loader if use_test else dataset.train_loader
-
- total_loss = [0 for _ in range(num_nets)]
- num_correct = [0 for _ in range(num_nets)]
- num_batches = 0
-
- measure_alignment = parameters.get("alignment", True)
- measure_alignment_expansion = parameters.get("alignment_expansion", True)
-
- if measure_alignment:
- alignment = []
- if measure_alignment_expansion:
- alignment_0 = []
- alignment_1 = []
- alignment_2 = []
- alignment_red = []
-
- batch_loop = tqdm(dataloader) if verbose else dataloader
- for batch in batch_loop:
- images, labels = dataset.unwrap_batch(batch)
-
- outputs = [net(images, store_hidden=True) for net in nets]
-
- for idx, output in enumerate(outputs):
- total_loss[idx] += dataset.measure_loss(output, labels).item()
- num_correct[idx] += dataset.measure_accuracy(output, labels).item()
-
- num_batches += 1
-
- if measure_alignment:
- a_vals = []
- for net in nets:
- real_net = net.module if hasattr(net, "module") else net
- a_vals.append(real_net.measure_alignment(images, precomputed=True, method="alignment"))
- alignment.append(a_vals)
-
- if measure_alignment_expansion:
- a0_vals, a1_vals, a2_vals, ared_vals = [], [], [], []
- for net in nets:
- real_net = net.module if hasattr(net, "module") else net
- a0_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_0"))
- a1_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_1"))
- a2_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_2"))
- ared_vals.append(real_net.measure_alignment_expansion(images, precomputed=True, method="alignment_red"))
- alignment_0.append(a0_vals)
- alignment_1.append(a1_vals)
- alignment_2.append(a2_vals)
- alignment_red.append(ared_vals)
-
- results = {
- "loss": [loss / num_batches for loss in total_loss],
- "accuracy": [correct / num_batches for correct in num_correct],
- }
-
- if measure_alignment:
- results["alignment"] = condense_values(transpose_list(alignment))
- if measure_alignment_expansion:
- results["alignment_0"] = condense_values(transpose_list(alignment_0))
- results["alignment_1"] = condense_values(transpose_list(alignment_1))
- results["alignment_2"] = condense_values(transpose_list(alignment_2))
- results["alignment_red"] = condense_values(transpose_list(alignment_red))
-
- if run is not None:
- run.summary["test_loss"] = torch.mean(torch.tensor(results["loss"]))
- run.summary["test_accuracy"] = torch.mean(torch.tensor(results["accuracy"]))
-
- return results
-
-
-@torch.no_grad()
-def get_dropout_indices(idx_alignment, fraction):
- """
- convenience method for getting a fraction of dropout indices from each layer
- """
- num_nets = idx_alignment[0].size(0)
- num_nodes = [idx.size(1) for idx in idx_alignment]
- num_drop = [int(nodes * fraction) for nodes in num_nodes]
- idx_high = [
- torch.sort(idx[:, -drop:], dim=1).values
- for idx, drop in zip(idx_alignment, num_drop)
- ]
- idx_low = [
- torch.sort(idx[:, :drop], dim=1).values
- for idx, drop in zip(idx_alignment, num_drop)
- ]
- idx_rand = [
- torch.sort(idx[:, torch.randperm(idx.size(1))[:drop]], dim=1).values
- for idx, drop in zip(idx_alignment, num_drop)
- ]
- return idx_high, idx_low, idx_rand
-
-
-@torch.no_grad()
-def get_dropout_indices000(idx_alignment, fraction):
- """
- Returns dropout indices for high, low, and random alignment for each layer and each network.
- """
- num_layers = len(idx_alignment)
- num_nets = idx_alignment[0].size(0)
-
- idx_high = []
- idx_low = []
- idx_rand = []
-
- for layer_idx in range(num_layers):
- layer_alignment = idx_alignment[layer_idx]
- layer_dimension = layer_alignment.size(1)
- num_drop = int(layer_dimension * fraction)
- sorted_indices = torch.argsort(layer_alignment, dim=1)
- high_indices = sorted_indices[:, -num_drop:]
- high_indices = torch.sort(high_indices, dim=1).values
- idx_high.append(high_indices)
- low_indices = sorted_indices[:, :num_drop]
- low_indices = torch.sort(low_indices, dim=1).values
- idx_low.append(low_indices)
- rand_indices = torch.randperm(layer_dimension, device=layer_alignment.device)[:num_drop]
- rand_indices = rand_indices.repeat(num_nets, 1)
- idx_rand.append(rand_indices)
-
- return idx_high, idx_low, idx_rand
-
-
-@torch.no_grad()
-@test_nets
-def progressive_dropout(nets, dataset, alignment=None, **parameters):
- """
- method for testing network on supervised learning problem with progressive dropout
- """
-
- if not (isinstance(nets, list)):
- nets = [nets]
-
- real_net = nets[0].module if hasattr(nets[0], 'module') else nets[0]
- idx_dropout_layers = real_net.get_alignment_layer_indices()
-
- if alignment is None:
- alignment = test(nets, dataset, **parameters)["alignment"]
-
- alignment = [alignment[i] for i in idx_dropout_layers]
- assert len(alignment) == len(idx_dropout_layers), "the number of layers in **alignment** doesn't correspond to the number of alignment layers"
-
- classification_layer = (
- nets[0].module.num_layers(all=True) - 1
- if hasattr(nets[0], 'module')
- else nets[0].num_layers(all=True) - 1
- )
- if classification_layer in idx_dropout_layers:
- idx_dropout_layers.pop(-1)
- alignment.pop(-1)
-
- alignment = [torch.mean(align, dim=1) for align in alignment]
- idx_alignment = [torch.argsort(align, dim=1) for align in alignment]
-
- print(len(idx_alignment))
- print(idx_alignment[0].shape)
-
- num_nets = len(nets)
- num_drops = parameters.get("num_drops", 9)
- drop_fraction = torch.linspace(0, 1, num_drops + 2)[1:-1]
- by_layer = parameters.get("by_layer", False)
- num_layers = len(idx_dropout_layers) if by_layer else 1
-
- progdrop_loss_high = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_loss_low = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_loss_rand = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_acc_high = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_acc_low = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_acc_rand = torch.zeros((num_nets, num_drops, num_layers))
-
- num_batches = 0
- use_train = parameters.get("train_set", False)
- dataloader = dataset.train_loader if use_train else dataset.test_loader
-
- for batch in tqdm(dataloader):
- images, labels = dataset.unwrap_batch(batch)
- num_batches += 1
-
- for dropidx, fraction in enumerate(drop_fraction):
- idx_high, idx_low, idx_rand = get_dropout_indices(idx_alignment, fraction)
-
- for layer in range(num_layers):
- if by_layer:
- drop_high, drop_low, drop_rand = (
- [idx_high[layer]],
- [idx_low[layer]],
- [idx_rand[layer]],
- )
- drop_layer = [idx_dropout_layers[layer]]
- else:
- drop_high, drop_low, drop_rand = idx_high, idx_low, idx_rand
- drop_layer = copy(idx_dropout_layers)
-
- out_high = []
- for idx, net_ in enumerate(nets):
- real_net_ = net_.module if hasattr(net_, "module") else net_
- out_ = real_net_.forward_targeted_dropout(
- images, [drop[idx, :] for drop in drop_high], drop_layer
- )[0]
- out_high.append(out_)
- out_low = []
- for idx, net_ in enumerate(nets):
- real_net_ = net_.module if hasattr(net_, "module") else net_
- out_ = real_net_.forward_targeted_dropout(
- images, [drop[idx, :] for drop in drop_low], drop_layer
- )[0]
- out_low.append(out_)
- out_rand = []
- for idx, net_ in enumerate(nets):
- real_net_ = net_.module if hasattr(net_, "module") else net_
- out_ = real_net_.forward_targeted_dropout(
- images, [drop[idx, :] for drop in drop_rand], drop_layer
- )[0]
- out_rand.append(out_)
-
- loss_high = [dataset.measure_loss(out, labels).item() for out in out_high]
- loss_low = [dataset.measure_loss(out, labels).item() for out in out_low]
- loss_rand = [dataset.measure_loss(out, labels).item() for out in out_rand]
-
- acc_high = [dataset.measure_accuracy(out, labels) for out in out_high]
- acc_low = [dataset.measure_accuracy(out, labels) for out in out_low]
- acc_rand = [dataset.measure_accuracy(out, labels) for out in out_rand]
-
- progdrop_loss_high[:, dropidx, layer] += torch.tensor(loss_high)
- progdrop_loss_low[:, dropidx, layer] += torch.tensor(loss_low)
- progdrop_loss_rand[:, dropidx, layer] += torch.tensor(loss_rand)
- progdrop_acc_high[:, dropidx, layer] += torch.tensor(acc_high)
- progdrop_acc_low[:, dropidx, layer] += torch.tensor(acc_low)
- progdrop_acc_rand[:, dropidx, layer] += torch.tensor(acc_rand)
-
- results = {
- "progdrop_loss_high": progdrop_loss_high / num_batches,
- "progdrop_loss_low": progdrop_loss_low / num_batches,
- "progdrop_loss_rand": progdrop_loss_rand / num_batches,
- "progdrop_acc_high": progdrop_acc_high / num_batches,
- "progdrop_acc_low": progdrop_acc_low / num_batches,
- "progdrop_acc_rand": progdrop_acc_rand / num_batches,
- "dropout_fraction": drop_fraction,
- "by_layer": by_layer,
- "idx_dropout_layers": idx_dropout_layers,
- }
-
- return results
-
-
-def train_network(net, dataset, epochs=5, learning_rate=0.001):
- optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate)
- net.train()
-
- for epoch in range(epochs):
- for batch in dataset.train_loader:
- images, labels = dataset.unwrap_batch(batch)
- optimizer.zero_grad()
- outputs = net(images)
- loss = dataset.measure_loss(outputs, labels)
- loss.backward()
- optimizer.step()
-
- return net
-
-
-def progressive_dropout_train(nets, dataset, alignment=None, **parameters):
-
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
- if not isinstance(nets, list):
- nets = [nets]
- nets = [net.to(device) for net in nets]
-
- real_net_0 = nets[0].module if hasattr(nets[0], 'module') else nets[0]
- idx_dropout_layers = real_net_0.get_alignment_layer_indices()
-
- if alignment is None:
- alignment = test(nets, dataset, **parameters)["alignment"]
-
- assert len(alignment) == len(idx_dropout_layers), (
- f"Mismatch in alignment layer count: alignment has {len(alignment)} "
- f"layers, expected {len(idx_dropout_layers)} layers."
- )
-
- classification_layer = real_net_0.num_layers(all=True) - 1
- if classification_layer in idx_dropout_layers:
- idx_dropout_layers.pop(-1)
- alignment.pop(-1)
-
- alignment = [torch.mean(align, dim=1) for align in alignment]
- idx_alignment = [torch.argsort(align, dim=1) for align in alignment]
-
- num_nets = len(nets)
- num_drops = parameters.get("num_drops", 9)
- drop_fraction = torch.linspace(0, 1, num_drops + 2)[1:-1]
- by_layer = parameters.get("by_layer", False)
- num_layers = len(idx_dropout_layers) if by_layer else 1
-
- progdrop_loss_high_before, progdrop_loss_low_before, progdrop_loss_rand_before = [], [], []
- progdrop_acc_high_before, progdrop_acc_low_before, progdrop_acc_rand_before = [], [], []
-
- progdrop_loss_high_after, progdrop_loss_low_after, progdrop_loss_rand_after = [], [], []
- progdrop_acc_high_after, progdrop_acc_low_after, progdrop_acc_rand_after = [], [], []
-
- num_batches = 0
- use_train = parameters.get("train_set", False)
- dataloader = dataset.train_loader if use_train else dataset.test_loader
-
- for batch in tqdm(dataloader):
- images, labels = dataset.unwrap_batch(batch)
- images = images.to(device)
- labels = labels.to(device)
-
- num_batches += 1
-
- for dropidx, fraction in enumerate(drop_fraction):
- idx_high, idx_low, idx_rand = get_dropout_indices(idx_alignment, fraction)
-
- for layer in range(num_layers):
- if by_layer:
- drop_high, drop_low, drop_rand = (
- [idx_high[layer]],
- [idx_low[layer]],
- [idx_rand[layer]],
- )
- drop_layer = [idx_dropout_layers[layer]]
- else:
- drop_high, drop_low, drop_rand = idx_high, idx_low, idx_rand
- drop_layer = deepcopy(idx_dropout_layers)
-
- pruned_nets_high = [
- train_network(prune_network_by_alignment([net], idx_alignment, fraction, method="high")[0].to(device), dataset)
- for net in nets
- ]
- pruned_nets_low = [
- train_network(prune_network_by_alignment([net], idx_alignment, fraction, method="low")[0].to(device), dataset)
- for net in nets
- ]
- pruned_nets_rand = [
- train_network(prune_network_by_alignment([net], idx_alignment, fraction, method="random")[0].to(device), dataset)
- for net in nets
- ]
-
- out_high_before = [net_(images) for net_ in pruned_nets_high]
- out_low_before = [net_(images) for net_ in pruned_nets_low]
- out_rand_before = [net_(images) for net_ in pruned_nets_rand]
-
- loss_high_before = [
- dataset.measure_loss(out, labels).item() for out in out_high_before
- ]
- loss_low_before = [
- dataset.measure_loss(out, labels).item() for out in out_low_before
- ]
- loss_rand_before = [
- dataset.measure_loss(out, labels).item() for out in out_rand_before
- ]
-
- acc_high_before = [dataset.measure_accuracy(out, labels) for out in out_high_before]
- acc_low_before = [dataset.measure_accuracy(out, labels) for out in out_low_before]
- acc_rand_before = [dataset.measure_accuracy(out, labels) for out in out_rand_before]
-
- progdrop_loss_high_before.append(loss_high_before)
- progdrop_loss_low_before.append(loss_low_before)
- progdrop_loss_rand_before.append(loss_rand_before)
- progdrop_acc_high_before.append(acc_high_before)
- progdrop_acc_low_before.append(acc_low_before)
- progdrop_acc_rand_before.append(acc_rand_before)
-
- out_high_after = [net_(images) for net_ in pruned_nets_high]
- out_low_after = [net_(images) for net_ in pruned_nets_low]
- out_rand_after = [net_(images) for net_ in pruned_nets_rand]
-
- loss_high_after = [
- dataset.measure_loss(out, labels).item() for out in out_high_after
- ]
- loss_low_after = [
- dataset.measure_loss(out, labels).item() for out in out_low_after
- ]
- loss_rand_after = [
- dataset.measure_loss(out, labels).item() for out in out_rand_after
- ]
-
- acc_high_after = [dataset.measure_accuracy(out, labels) for out in out_high_after]
- acc_low_after = [dataset.measure_accuracy(out, labels) for out in out_low_after]
- acc_rand_after = [dataset.measure_accuracy(out, labels) for out in out_rand_after]
-
- progdrop_loss_high_after.append(loss_high_after)
- progdrop_loss_low_after.append(loss_low_after)
- progdrop_loss_rand_after.append(loss_rand_after)
- progdrop_acc_high_after.append(acc_high_after)
- progdrop_acc_low_after.append(acc_low_after)
- progdrop_acc_rand_after.append(acc_rand_after)
-
- results = {
- "progdrop_loss_high_before": torch.tensor(progdrop_loss_high_before) / num_batches,
- "progdrop_loss_low_before": torch.tensor(progdrop_loss_low_before) / num_batches,
- "progdrop_loss_rand_before": torch.tensor(progdrop_loss_rand_before) / num_batches,
- "progdrop_acc_high_before": torch.tensor(progdrop_acc_high_before) / num_batches,
- "progdrop_acc_low_before": torch.tensor(progdrop_acc_low_before) / num_batches,
- "progdrop_acc_rand_before": torch.tensor(progdrop_acc_rand_before) / num_batches,
- "progdrop_loss_high_after": torch.tensor(progdrop_loss_high_after) / num_batches,
- "progdrop_loss_low_after": torch.tensor(progdrop_loss_low_after) / num_batches,
- "progdrop_loss_rand_after": torch.tensor(progdrop_loss_rand_after) / num_batches,
- "progdrop_acc_high_after": torch.tensor(progdrop_acc_high_after) / num_batches,
- "progdrop_acc_low_after": torch.tensor(progdrop_acc_low_after) / num_batches,
- "progdrop_acc_rand_after": torch.tensor(progdrop_acc_rand_after) / num_batches,
- "dropout_fraction": drop_fraction,
- "by_layer": by_layer,
- "idx_dropout_layers": idx_dropout_layers,
- }
-
- return results
-
-
-@torch.no_grad()
-@test_nets
-def eigenvector_dropout(nets, dataset, eigenvalues, eigenvectors, **parameters):
- """
- method for testing network on supervised learning problem with eigenvector dropout
- """
-
- if not (isinstance(nets, list)):
- nets = [nets]
-
- idx_dropout_layers = (
- nets[0].module.get_alignment_layer_indices()
- if hasattr(nets[0], 'module')
- else nets[0].get_alignment_layer_indices()
- )
-
- assert all([len(ev) == len(idx_dropout_layers) for ev in eigenvectors]), \
- "the number of layers in **eigenvectors** doesn't correspond to the number of alignment layers"
- assert all([len(ev) == len(idx_dropout_layers) for ev in eigenvalues]), \
- "the number of layers in **eigenvalues** doesn't correspond to the number of alignment layers"
-
- num_nets = len(nets)
- num_drops = parameters.get("num_drops", 9)
- drop_fraction = torch.linspace(0, 1, num_drops + 2)[1:-1]
- by_layer = parameters.get("by_layer", False)
- num_layers = len(idx_dropout_layers) if by_layer else 1
-
- idx_eigenvalue = [
- torch.fliplr(torch.tensor(range(0, ev.size(1))).expand(num_nets, -1))
- for ev in eigenvectors[0]
- ]
-
- progdrop_loss_high = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_loss_low = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_loss_rand = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_acc_high = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_acc_low = torch.zeros((num_nets, num_drops, num_layers))
- progdrop_acc_rand = torch.zeros((num_nets, num_drops, num_layers))
-
- num_batches = 0
- use_test = not parameters.get("train_set", True)
- dataloader = dataset.test_loader if use_test else dataset.train_loader
-
- for batch in tqdm(dataloader):
- images, labels = dataset.unwrap_batch(batch)
- num_batches += 1
-
- for dropidx, fraction in enumerate(drop_fraction):
- idx_high, idx_low, idx_rand = get_dropout_indices(idx_eigenvalue, fraction)
-
- for layer in range(num_layers):
- if by_layer:
- drop_high, drop_low, drop_rand = (
- [idx_high[layer]],
- [idx_low[layer]],
- [idx_rand[layer]],
- )
- drop_layer = [idx_dropout_layers[layer]]
- drop_evals = [[ev[layer]] for ev in eigenvalues]
- drop_evecs = [[evc[layer]] for evc in eigenvectors]
- else:
- drop_high, drop_low, drop_rand = idx_high, idx_low, idx_rand
- drop_layer = copy(idx_dropout_layers)
- drop_evals = deepcopy(eigenvalues)
- drop_evecs = deepcopy(eigenvectors)
-
- out_high = []
- for idx, (net, evals, evecs) in enumerate(zip(nets, drop_evals, drop_evecs)):
- real_net = net.module if hasattr(net, 'module') else net
- out_ = real_net.forward_eigenvector_dropout(
- images, evals, evecs, [drop[idx, :] for drop in drop_high], drop_layer
- )[0]
- out_high.append(out_)
- out_low = []
- for idx, (net, evals, evecs) in enumerate(zip(nets, drop_evals, drop_evecs)):
- real_net = net.module if hasattr(net, 'module') else net
- out_ = real_net.forward_eigenvector_dropout(
- images, evals, evecs, [drop[idx, :] for drop in drop_low], drop_layer
- )[0]
- out_low.append(out_)
- out_rand = []
- for idx, (net, evals, evecs) in enumerate(zip(nets, drop_evals, drop_evecs)):
- real_net = net.module if hasattr(net, 'module') else net
- out_ = real_net.forward_eigenvector_dropout(
- images, evals, evecs, [drop[idx, :] for drop in drop_rand], drop_layer
- )[0]
- out_rand.append(out_)
-
- loss_high = [dataset.measure_loss(out, labels).item() for out in out_high]
- loss_low = [dataset.measure_loss(out, labels).item() for out in out_low]
- loss_rand = [dataset.measure_loss(out, labels).item() for out in out_rand]
-
- acc_high = [dataset.measure_accuracy(out, labels) for out in out_high]
- acc_low = [dataset.measure_accuracy(out, labels) for out in out_low]
- acc_rand = [dataset.measure_accuracy(out, labels) for out in out_rand]
-
- progdrop_loss_high[:, dropidx, layer] += torch.tensor(loss_high)
- progdrop_loss_low[:, dropidx, layer] += torch.tensor(loss_low)
- progdrop_loss_rand[:, dropidx, layer] += torch.tensor(loss_rand)
- progdrop_acc_high[:, dropidx, layer] += torch.tensor(acc_high)
- progdrop_acc_low[:, dropidx, layer] += torch.tensor(acc_low)
- progdrop_acc_rand[:, dropidx, layer] += torch.tensor(acc_rand)
-
- results = {
- "progdrop_loss_high": progdrop_loss_high / num_batches,
- "progdrop_loss_low": progdrop_loss_low / num_batches,
- "progdrop_loss_rand": progdrop_loss_rand / num_batches,
- "progdrop_acc_high": progdrop_acc_high / num_batches,
- "progdrop_acc_low": progdrop_acc_low / num_batches,
- "progdrop_acc_rand": progdrop_acc_rand / num_batches,
- "dropout_fraction": drop_fraction,
- "by_layer": by_layer,
- "idx_dropout_layers": idx_dropout_layers,
- }
-
- return results
\ No newline at end of file
diff --git a/_arxiv/_archive/alignment_v2/utils.py b/_arxiv/_archive/alignment_v2/utils.py
deleted file mode 100644
index 73149842..00000000
--- a/_arxiv/_archive/alignment_v2/utils.py
+++ /dev/null
@@ -1,934 +0,0 @@
-import os
-import math
-import zipfile
-from typing import List
-from warnings import warn
-from contextlib import contextmanager
-from functools import wraps
-from natsort import natsorted
-from gitignore_parser import parse_gitignore
-
-import torch
-import numpy as np
-from scipy.linalg import null_space
-from sklearn.decomposition import IncrementalPCA
-from sklearn.decomposition import PCA
-
-# -------------- context managers & decorators --------------
-@contextmanager
-def no_grad(no_grad=True):
- if no_grad:
- with torch.no_grad():
- yield
- else:
- yield
-
-
-def test_nets(func):
- @wraps(func)
- def wrapper(nets, *args, **kwargs):
- # get original training mode and set to eval
- in_training_mode = [set_net_mode(net, training=False) for net in nets]
-
- # do decorated function
- func_outputs = func(nets, *args, **kwargs)
-
- # return networks to whatever mode they used to be in
- for train_mode, net in zip(in_training_mode, nets):
- set_net_mode(net, training=train_mode)
-
- # return decorated function outputs
- return func_outputs
-
- # return decorated function
- return wrapper
-
-
-def train_nets(func):
- @wraps(func)
- def wrapper(nets, *args, **kwargs):
- # get original training mode and set to train
- in_training_mode = [set_net_mode(net, training=True) for net in nets]
-
- # do decorated function
- func_outputs = func(nets, *args, **kwargs)
-
- # return networks to whatever mode they used to be in
- for train_mode, net in zip(in_training_mode, nets):
- set_net_mode(net, training=train_mode)
-
- # return decorated function outputs
- return func_outputs
-
- # return decorated function
- return wrapper
-
-
-def set_net_mode(net, training=True):
- """helper for setting mode of network and returning current mode"""
- # get current mode of network
- in_training_mode = net.training
- # set to training mode or evaluation mode
- if training:
- net.train()
- else:
- net.eval()
- # return original mode of network
- return in_training_mode
-
-
-# ------- some other things -------
-def get_device(obj):
- """simple method to get device of input tensor or nn.Module"""
- if isinstance(obj, torch.nn.Module):
- return next(obj.parameters()).device.type
- elif isinstance(obj, torch.Tensor):
- return "cuda" if obj.is_cuda else "cpu"
- else:
- raise ValueError("")
-
-
-def check_iterable(val):
- """duck-type check if val is iterable"""
- try:
- _ = iter(val)
- except:
- return False
- else:
- return True
-
-
-def remove_by_idx(input, idx, dim):
- """
- remove part of input indexed by idx on dim
- """
- idx_keep = [i for i in range(input.size(dim)) if i not in idx]
- return torch.index_select(input, dim, torch.tensor(idx_keep).to(input.device))
-
-
-def get_eval_transform_by_cutoff(cutoff):
- """
- get method for transforming eigenvalues into a binary keep fraction
-
- will scale each eigenvector by 1 or 0 depending on whether that eigenvalue
- explains more than **cutoff** fraction of the variance
-
- returns a callable method
- """
-
- def eval_transform(evals):
- assert torch.all(evals >= 0), "found negative eigenvalues, doesn't work for 'cutoff' eval_transform"
- evals = evals / torch.sum(evals)
- return 1.0 * (evals > cutoff)
-
- return eval_transform
-
-
-def fractional_histogram(*args, **kwargs):
- """wrapper of np.histogram() with relative counts instead of total or density"""
- counts, bins = np.histogram(*args, **kwargs)
- counts = counts / np.sum(counts)
- return counts, bins
-
-
-def edge2center(edges):
- """from a list of edges of bins (e.g. for torch.histogram()), return the centers between the edges"""
- assert edges.ndim == 1, "edges must be a 1-d array"
- return edges[:-1] + np.diff(edges) / 2
-
-
-def smartcorr(input):
- """
- Performs torch corrcoef on the input data but sets each pair-wise correlation coefficent
- to 0 where the activity has no variance (var=0) for a particular dimension (replaces nans with zeros)sss
- """
- idx_zeros = torch.var(input, dim=1) == 0
- cc = torch.corrcoef(input)
- cc[idx_zeros, :] = 0
- cc[:, idx_zeros] = 0
- return cc
-
-
-def batch_cov(input, centered=True, correction=True):
- """
- Performs batched covariance on input data of shape (batch, dim, samples) or (dim, samples)
-
- Where the resulting batch covariance matrix has shape (batch, dim, dim) or (dim, dim)
- and bcov[i] = torch.cov(input[i]) if input.ndim==3
-
- if centered=True (default) will subtract the means first
-
- if correction=True, will use */(N-1) otherwise will use */N
- """
- assert (input.ndim == 2) or (input.ndim == 3), "input must be a 2D or 3D tensor"
- assert isinstance(correction, bool), "correction must be a boolean variable"
-
- # check if batch dimension was provided
- no_batch = input.ndim == 2
-
- # add an empty batch dimension if not provided
- if no_batch:
- input = input.unsqueeze(0)
-
- # measure number of samples of each input matrix
- S = input.size(2)
-
- # subtract mean if doing centered covariance
- if centered:
- input = input - input.mean(dim=2, keepdim=True)
-
- # measure covariance of each input matrix
- bcov = torch.bmm(input, input.transpose(1, 2))
-
- # correct for number of samples
- bcov /= S - 1.0 * correction
-
- # remove empty batch dimension if not provided
- if no_batch:
- bcov = bcov.squeeze(0)
-
- return bcov
-
-
-def smart_pca(input, centered=True, use_rank=True, correction=True):
- """
- smart algorithm for pca optimized for speed
-
- input should either have shape (batch, dim, samples) or (dim, samples)
- if dim > samples, will use svd and if samples < dim will use covariance/eigh method
-
- will center data when centered=True
-
- if it fails, will fall back on performing sklearns IncrementalPCA whenever forcetry=True
- """
- assert (input.ndim == 2) or (input.ndim == 3), "input should be a matrix or batched matrices"
- assert isinstance(correction, bool), "correction should be a boolean"
-
- if input.ndim == 2:
- no_batch = True
- input = input.unsqueeze(0) # create batch dimension for uniform code
- else:
- no_batch = False
-
- _, D, S = input.size()
- if D > S:
- # if more dimensions than samples, it's more efficient to run svd
- v, w, _ = named_transpose([torch.linalg.svd(inp) for inp in input])
- # convert singular values to eigenvalues
- w = [ww**2 / (S - 1.0 * correction) for ww in w]
- # append zeros because svd returns w in R**k where k = min(D, S)
- w = [torch.concatenate((ww, torch.zeros(D - S))) for ww in w]
-
- else:
- # if more samples than dimensions, it's more efficient to run eigh
- bcov = batch_cov(input, centered=centered, correction=correction)
- w, v = named_transpose([eigendecomposition(C, use_rank=use_rank) for C in bcov])
-
- # return to stacked tensor across batch dimension
- w = torch.stack(w)
- v = torch.stack(v)
-
- # if no batch originally provided, squeeze out batch dimension
- if no_batch:
- w = w.squeeze(0)
- v = v.squeeze(0)
-
- # return eigenvalues and eigenvectors
- return w, v
-
-
-def eigendecomposition(C, use_rank=True):
- """
- helper for getting eigenvalues and eigenvectors of covariance matrix
-
- will measure eigenvalues and eigenvectors with torch.linalg.eigh()
- the output will be sorted from highest to lowest eigenvalue (& eigenvector)
-
- if use_rank=True, will measure the rank of the covariance matrix and zero
- out any eigenvalues beyond the rank (that are usually nonzero numerical errors)
- """
- try:
- # measure eigenvalues and eigenvectors
- w, v = torch.linalg.eigh(C)
-
- except torch._C._LinAlgError as error:
- # this happens if the algorithm failed to converge
- # try with sklearn's incrementalPCA algorithm
- return sklearn_pca(C, use_rank=use_rank)
-
- except Exception as error:
- # if any other exception, raise it
- raise error
-
- # sort by eigenvalue from highest to lowest
- w_idx = torch.argsort(-w)
- w = w[w_idx]
- v = v[:, w_idx]
-
- # iff use_rank=True, will set eigenvalues to 0 for probable numerical errors
- if use_rank:
- crank = torch.linalg.matrix_rank(C) # measure rank of covariance
- w[crank:] = 0 # set eigenvalues beyond rank to 0
-
- # return eigenvalues and eigenvectors
- return w, v
-
-
-def sklearn_pca(input, use_rank=True, rank=None):
- """
- sklearn incrementalPCA algorithm serving as a replacement for eigh when it fails
-
- input should be a tensor with shape (num_samples, num_features) or it can be a
- covariance matrix with (num_features, num_features)
-
- if use_rank=True, will set num_components to the rank of input and then fill out the
- rest of the components with random orthogonal components in the null space of the true
- components and set the eigenvalues to 0
-
- if use_rank=False, will attempt to fit all the components
- if rank is not None, will attempt to fit #=rank components without measuring the rank directly
- (will ignore "rank" if use_rank=False)
-
- returns w, v where w is eigenvalues and v is eigenvectors sorted from highest to lowest
- """
- # dimension
- num_samples, num_features = input.shape
-
- # measure rank (or set to None)
- rank = None if not use_rank else (rank if rank is not None else fast_rank(input))
-
- # create and fit IncrementalPCA object on input data
- ipca = IncrementalPCA(n_components=rank).fit(input)
-
- # eigenvectors are the components
- v = ipca.components_
-
- # eigenvalues are the scaled singular values
- w = ipca.singular_values_**2 / num_samples
-
- # if v is a subspace of input (e.g. not a full basis, fill it out)
- if v.shape[0] < num_features:
- msg = "adding this because I think it should always be true, and if not I want to find out"
- assert w.shape[0] == v.shape[0], msg
- v_kernel = null_space(v).T
- v = np.vstack((v, v_kernel))
- w = np.concatenate((w, np.zeros(v_kernel.shape[0])))
-
- return torch.tensor(w, dtype=torch.float), torch.tensor(v, dtype=torch.float).T
-
-
-def fast_rank(input):
- """uses transpose to speed up rank computation, otherwise normal"""
- if input.size(-2) < input.size(-1):
- input = torch.transpose(input, -2, -1)
- return int(torch.linalg.matrix_rank(input))
-
-
-# ------------------ alignment functions ----------------------
-def alignment(input, weight, method="alignment", relative=True):
- """
- measure alignment (proportion variance explained) between **input** and **weight**
-
- computes the rayleigh quotient between each weight vector in **weight** and the **input** fed
- into **weight**. Typically, **input** is the output in Layer L-1 and **weight** is from Layer L
-
- the output is normalized by the total variance in output of layer L-1 to measure the proportion
- of variance of in **input** is explained by a projection onto node's weights in **weight**
-
- args
- ----
- input: (batch, neurons) torch tensor
- - represents input activity being fed in to network weight layer
- weight: (num_out, num_in) torch tensor
- - represents weights multiplied by input layer
- method: string, default='alignment'
- - which method to use to measure structure in **input**
- - if 'alignment', uses covariance matrix of **input**
- - if 'similarity', uses correlation matrix of **input**
- relative: bool, default=True,
- - if True, will measure relative RQ (divide by sum of eigenvalues)
-
- returns
- -------
- alignment: (num_out, ) torch tensor
- - proportion of variance explained by projection of **input** onto each **weight** vector
- """
- assert method == "alignment" or method == "similarity", "method must be set to either 'alignment' or 'similarity' (or None, default is alignment)"
- if method == "alignment":
- cc = torch.cov(input.T)
- #cc = torch.corrcoef(input.T)
- elif method == "similarity":
- cc = smartcorr(input.T)
- else:
- raise ValueError(f"did not recognize method ({method}), must be 'alignment' or 'similarity'")
-
- rq = torch.sum(torch.matmul(weight, cc) * weight, axis=1) / torch.sum(weight * weight, axis=1)
-
- if relative:
- return rq / torch.trace(cc)
-
- # shuffled_rq = rq[torch.randperm(rq.size(0))] # Shuffle RQ values
- # rq = shuffled_rq / torch.trace(cc)
-
- return rq
-
-
-@torch.no_grad()
-def expected_alignment_distribution(eigenvalues, relative=True, valid_rotation=True, with_rotation=True, bins=11, num_tests=100):
- """
- for a set of eigenvalues, measure the expected distribution given aligned weights
-
- relative determines whether to normalize by sum of eigenvalues
- valid_rotation determines whether we create orthonormal rotation matrices (for True)
- or just normally distributed weights with the expected variance (for False)
- bins works like histogram bins
- num_tests determines how many tests to do (it's actually num_tests*len(eigenvalues))
- """
- # otherwise, randomly sample using eigenvalue as weighted average
- N = len(eigenvalues)
- if relative:
- eigenvalues /= eigenvalues.sum()
- eigenvalues = eigenvalues.view(-1, 1).expand(-1, N * num_tests)
- if with_rotation:
- if valid_rotation:
- mixing = [torch.linalg.qr(torch.normal(0, 1 / math.sqrt(N), (N, N)))[0].T for _ in range(num_tests)]
- coefficients = torch.concatenate(mixing, axis=1) ** 2
- else:
- coefficients = torch.normal(0, 1 / math.sqrt(N), (N, N * num_tests)) ** 2
- else:
- coefficients = torch.ones((N, N * num_tests))
- coefficients = coefficients.to(get_device(eigenvalues))
- weights = eigenvalues * coefficients
- alignment = torch.sum(eigenvalues * weights, dim=0) / weights.sum(dim=0)
- counts, bins = torch.histogram(alignment.cpu(), bins=bins, density=True)
- centers = edge2center(bins)
- return counts, bins, centers
-
-def compute_skewness_inplace(input):
- """
- Compute skewness for each input feature (dimension) using in-place operations.
- """
- mean = input.mean(dim=0)
- variance = input.var(dim=0)
-
- # Add epsilon to variance to avoid division by zero
- epsilon = 1e-6
- variance = variance + epsilon
-
- third_moment = (input - mean).pow(3).mean(dim=0)
-
- # Skewness formula: E[(X - μ)^3] / (E[(X - μ)^2])^(3/2)
- skewness = third_moment / variance.pow(1.5)
-
- return skewness
-
-def compute_kurtosis_inplace(input):
- """
- Compute kurtosis for each input feature (dimension) using in-place operations.
- """
- mean = input.mean(dim=0)
- variance = input.var(dim=0)
- fourth_moment = (input - mean).pow_(4).mean(dim=0)
-
- # Kurtosis formula: E[(X - μ)^4] / (E[(X - μ)^2])^2 - 3
- kurtosis = fourth_moment / variance.pow(2) - 3
- return kurtosis
-
-def compute_skewness_low_rank(input, rank_approx=50):
- """
- Compute skewness after reducing the dimensionality using PCA to the specified rank approximation.
- """
-
- pca = PCA(n_components=rank_approx)
- input_reduced = torch.Tensor(pca.fit_transform(input.cpu().numpy())).to(input.device)
-
- return compute_skewness_inplace(input_reduced)
-
-def compute_kurtosis_low_rank(input, rank_approx=50):
- """
- Compute kurtosis after reducing the dimensionality using PCA to the specified rank approximation.
- """
-
- pca = PCA(n_components=rank_approx)
- input_reduced = torch.Tensor(pca.fit_transform(input.cpu().numpy())).to(input.device)
-
- return compute_kurtosis_inplace(input_reduced)
-
-def compute_redundancy(weights, input_covariance):
- """
- Compute the redundancy matrix for all pairs of nodes using matrix multiplication.
-
- Args:
- weights: (n, d) torch tensor, where n is the number of nodes and d is the input dimension.
- input_covariance: (d, d) torch tensor, the covariance matrix of the input data.
-
- Returns:
- redundancy_matrix: (n, n) torch tensor, redundancy between each pair of nodes (diagonal excluded).
- """
- # Compute the redundancy matrix: R = W Σ_X W^T
- redundancy_matrix = torch.matmul(weights, torch.matmul(input_covariance, weights.T))
-
- # Zero out the diagonal elements (self-information)
- redundancy_matrix.fill_diagonal_(0)
-
- return redundancy_matrix
-
-
-
-def alignment_expansion(input, weight, method="alignment_expansion", relative=True):
- """
- Measure alignment (proportion variance explained) between **input** and **weight**
- and compute single-node information, redundancy, and total information for the layer.
-
- args:
- ----
- input: (batch, neurons) torch tensor
- - represents input activity being fed into network weight layer
- weight: (num_out, num_in) torch tensor
- - represents weights multiplied by input layer
- method: string, default='alignment'
- - which method to use to measure structure in **input**
- - if 'alignment', uses covariance matrix of **input**
- - if 'similarity', uses correlation matrix of **input**
- relative: bool, default=True,
- - if True, will measure relative RQ (divide by sum of eigenvalues)
-
- returns:
- -------
- alignment: (num_out, ) torch tensor
- - proportion of variance explained by projection of **input** onto each **weight** vector
- single_node_info: torch tensor
- - Information carried by each node (mutual information for single nodes)
- redundancy_matrix: torch tensor
- - Pairwise redundancy (mutual information overlap between pairs of nodes)
- total_info: float
- - Total information of the layer, accounting for redundancy
- """
-
- # Step 1: Compute covariance of input
- # if method == "alignment_expansion":
- cc = torch.cov(input.T) # Covariance matrix of input
- # elif method == "alignment_expansion":
- # cc = torch.corrcoef(input.T) # Correlation matrix of input
- # else:
- # raise ValueError(f"Method {method} not recognized. Use 'alignment' or 'similarity'.")
-
- # Step 2: Compute Rayleigh Quotient (RQ) for each node
- rq = torch.sum(torch.matmul(weight, cc) * weight, axis=1) / torch.sum(weight * weight, axis=1)
-
- # Step 3: Compute Single-Node Information for each node (with kurtosis correction)
- single_node_info = torch.log(rq)# / torch.trace(cc) # First term (proportional to RQ)
-
- # Compute third-order term: Skewness correction
- skewness = compute_skewness_inplace(input)
-
- # Ensure the skewness shape matches the input dimensionality
- if skewness.shape[0] == input.shape[1]: # If skewness is (d,), reshape it
- skewness = skewness.view(1, -1) # Reshape skewness to (1, d) for broadcasting
-
- # Apply the skewness correction term
- normalized_weight = weight #/ weight.norm(p=2, dim=1, keepdim=True)
- skewness_correction = torch.sum(normalized_weight * skewness, dim=1)
-
- # Compute kurtosis-based correction term
- kurtosis = compute_kurtosis_low_rank(input) # Kurtosis is computed for each input feature
-
- # Expand kurtosis tensor to match weight shape
- # if kurtosis.dim() == 1:
- # kurtosis = kurtosis.unsqueeze(0) # Add a dimension to match the weight tensor shape
-
- kurtosis_correction = 0.5 * torch.norm(normalized_weight, p=2, dim=1) * kurtosis[:weight.shape[1]].mean()
-
- #single_node_info = skewness_correction
- #single_node_info = kurtosis_correction
-
- # Step 4: Continue with redundancy and total information as before
- redundancy_matrix = torch.abs(compute_redundancy(normalized_weight, cc))
- adjusted_single_node_info = adjust_information_with_redundancy(single_node_info, redundancy_matrix)#single_node_info.clone()
-
- total_info = torch.sum(adjusted_single_node_info)
-
- if method == "alignment_0":
- info = adjusted_single_node_info
- elif method == "alignment_1":
- info = skewness_correction
- elif method == "alignment_2":
- info = kurtosis_correction
- elif method == "alignment_red":
- info = redundancy_matrix
-
-
- return info
-
-def adjust_information_with_redundancy(single_node_info, redundancy_matrix):
- """
- Adjust single-node information by subtracting redundancy for each pair of nodes.
-
- Args:
- single_node_info: (n, ) torch tensor, the single-node information for each node.
- redundancy_matrix: (n, n) torch tensor, the redundancy between each pair of nodes.
-
- Returns:
- adjusted_info: (n, ) torch tensor, adjusted single-node information.
- """
- n = single_node_info.size(0)
-
- # Compute pairwise differences between single-node information
- #info_diffs = single_node_info.view(n, 1) - single_node_info.view(1, n)
-
- # Create a mask where the single-node information is smaller for each pair
- #mask = (info_diffs < 0).float()
-
- # Subtract redundancy from the node with smaller information
- #redundancy_adjustment = torch.sum(mask * redundancy_matrix, dim=1)
- redundancy_matrix.fill_diagonal_(0)
-
- redundancy_adjustment = torch.sum(redundancy_matrix, dim=1)
-
- # Adjust the single-node information
- adjusted_info =redundancy_adjustment
-
- return adjusted_info
-
-def plot_information_results(single_node_info, adjusted_single_node_info, redundancy_matrix, total_info):
- """
- Plot single-node information, adjusted single-node information, redundancy matrix, and total information for the layer.
-
- args:
- ----
- single_node_info: torch tensor
- - Information carried by each node
- adjusted_single_node_info: torch tensor
- - Adjusted information for each node after subtracting redundancy
- redundancy_matrix: torch tensor
- - Pairwise redundancy between nodes
- total_info: float
- - Total information of the layer
- """
- import matplotlib.pyplot as plt
- import seaborn as sns
-
- # Plot single-node information
- plt.figure(figsize=(8, 6))
- plt.bar(range(len(single_node_info)), single_node_info.cpu().detach().numpy(), alpha=0.6, label="Original Info")
- plt.bar(range(len(adjusted_single_node_info)), adjusted_single_node_info.cpu().detach().numpy(), alpha=0.6, label="Adjusted Info")
- plt.xlabel('Node Index')
- plt.ylabel('Single Node Information (MI)')
- plt.title('Single-Node Information for Each Node (Original vs Adjusted)')
- plt.legend()
- plt.show()
-
- # Plot redundancy matrix
- plt.figure(figsize=(8, 6))
- sns.heatmap(redundancy_matrix.cpu().detach().numpy(), annot=True, cmap="YlGnBu")
- plt.title('Redundancy Matrix (Pairwise Redundancy between Nodes)')
- plt.show()
-
- # Display total information
- print(f"Total Information for the Layer: {total_info.item():.4f}")
-
-
-
-def get_maximum_strides(h_input, w_input, layer):
- h_max = int(np.floor((h_input + 2 * layer.padding[0] - layer.dilation[0] * (layer.kernel_size[0] - 1) - 1) / layer.stride[0] + 1))
- w_max = int(np.floor((w_input + 2 * layer.padding[1] - layer.dilation[1] * (layer.kernel_size[1] - 1) - 1) / layer.stride[1] + 1))
- return h_max, w_max
-
-
-def get_unfold_params(layer):
- return dict(stride=layer.stride, padding=layer.padding, dilation=layer.dilation)
-
-
-# ----- cvPCA methods -----
-@torch.no_grad()
-def cvPCA(X1, X2):
- """X1, X2 are both (dimensions x samples)"""
- D, B = X1.shape
- assert X2.shape == (D, B), "shape of X1 and X2 is not the same"
- _, u = smart_pca(X1)
-
- cproj0 = X1.T @ u
- cproj1 = X2.T @ u
- ss = (cproj0 * cproj1).mean(axis=0)
- return ss
-
-
-def get_num_components(nc, shape):
- return nc if nc is not None else min(shape)
-
-
-@torch.no_grad()
-def shuff_cvPCA(X1, X2, nshuff=5, cvmethod=cvPCA):
- """X1, X2 are both (dimensions x samples)"""
- D, B = X1.shape
- assert X2.shape == (D, B), "shape of X1 and X2 is not the same"
- nc = get_num_components(None, (D, B))
- ss = torch.zeros((nshuff, nc))
- X = torch.stack((X1, X2))
- for k in range(nshuff):
- iflip = 1 * (torch.rand(B) > 0.5)
- X1c = torch.gather(X, 0, iflip.view(1, 1, -1).expand(1, D, -1)).squeeze(0)
- X2c = torch.gather(X, 0, -(iflip - 1).view(1, 1, -1).expand(1, D, -1)).squeeze(0)
- ss[k] = cvmethod(X1c, X2c)
- return ss
-
-
-def avg_value_by_layer(full):
- """
- return average value per layer across training
-
- **full** is a list of lists where the outer list is each snapshot through training or
- minibatch etc and each inner list is the value for each node in the network across layers
- of a particular measurement
-
- For example:
- num_epochs = 1000
- nodes_per_layer = [50, 40, 30, 20]
- len(full) == 1000
- len(full[i]) == 4 ... for all i
- [f.shape for f in full[i]] = [50, 40, 30, 20] ... for all i
-
- this method will return a tensor of size (num_layers, num_epochs) of the average value (for
- whatever value is in **full**) for each list/list
- """
- num_epochs = len(full)
- num_layers = len(full[0])
- avg_full = torch.zeros((num_layers, num_epochs))
- for layer in range(num_layers):
- avg_full[layer, :] = torch.tensor([torch.mean(f[layer]) for f in full])
- return avg_full.cpu()
-
-
-def value_by_layer(full: List[List[torch.Tensor]], layer: int) -> torch.Tensor:
- """
- return all value measurements for a particular layer from **full**
-
- **full** is a list of lists where the outer list is each snapshot through training or
- minibatch etc and each inner list is the value for each node in the network across layers
-
- this method will return just the part of **full** corresponding to the layer indexed
- by **layer** as a tensor of shape (num_epochs, num_nodes)
-
- see ``avg_value_by_layer`` for a little more explanation
- """
- return torch.cat([f[layer].view(1, -1) for f in full], dim=0).cpu()
-
-
-def condense_values(full: List[List[List[torch.Tensor]]]) -> List[torch.Tensor]:
- """
- condense List[List[List[Tensor]]] representing some value measured across networks, batches, and layers, for each node in the layer
-
- returns list of #=num_layers tensors, where each tensor has shape (num_networks, num_batches, num_nodes_per_layer)
-
- full should be a list of list of lists
- the first list should have length = number of networks
- the second list should have length = number of batches
- the third list should have length = number of layers in the network (this has to be the same for each network!)
- the tensor should have shape = number of nodes in this layer (also must be the same for each network) (or can be anything as long as consistent across layers)
- """
- num_layers = len(full[0][0])
- return [torch.stack([value_by_layer(value, layer) for value in full]) for layer in range(num_layers)]
-
-
-def transpose_list(list_of_lists):
- """helper function for transposing the order of a list of lists"""
- return list(map(list, zip(*list_of_lists)))
-
-
-def named_transpose(list_of_lists, reduction=None):
- """
- helper function for transposing lists without forcing the output to be a list like transpose_list
-
- for example, if list_of_lists contains 10 copies of lists that each have 3 iterable elements you
- want to name "A", "B", and "C", then write:
- A, B, C = named_transpose(list_of_lists)
-
- if reduction is used, it will be applied to each output, otherwise will make them lists
- """
- if reduction is not None:
- return map(reduction, zip(*list_of_lists))
- return map(list, zip(*list_of_lists))
-
-
-def ptp(tensor, dim=None, keepdim=False):
- """
- simple method for measuring range of tensor on requested dimension or on all data
- """
- if dim is None:
- return tensor.max() - tensor.min()
- return tensor.max(dim, keepdim).values - tensor.min(dim, keepdim).values
-
-
-def rms(tensor, dim=None, keepdim=False):
- """simple method for measuring root-mean-square on requested dimension or on all data in tensor"""
- if dim is None:
- return torch.sqrt(torch.mean(tensor**2))
- return torch.sqrt(torch.mean(tensor**2, dim=dim, keepdim=keepdim))
-
-
-def compute_stats_by_type(tensor, num_types, dim, method="var"):
- """
- helper method for returning the mean and variance across a certain dimension
- where multiple types are concatenated on that dimension
-
- for example, suppose we trained 2 networks each with 3 sets of parameters
- and concatenated the loss in a tensor like [set1-loss-net1, set1-loss-net2, set2-loss-net1, ...]
- then this would contract across the nets from each set and return the mean and variance
- """
- num_on_dim = tensor.size(dim)
- num_per_type = int(num_on_dim / num_types)
- tensor_by_type = tensor.unsqueeze(dim)
- expand_shape = list(tensor_by_type.shape)
- expand_shape[dim + 1] = num_per_type
- expand_shape[dim] = num_types
- tensor_by_type = tensor_by_type.view(expand_shape)
- type_means = torch.mean(tensor_by_type, dim=dim + 1)
- if method == "var":
- type_dev = torch.var(tensor_by_type, dim=dim + 1)
- elif method == "std":
- type_dev = torch.std(tensor_by_type, dim=dim + 1)
- elif method == "se":
- type_dev = torch.std(tensor_by_type, dim=dim + 1) / np.sqrt(num_per_type)
- elif method == "range":
- type_dev = ptp(tensor_by_type, dim=dim + 1)
- else:
- raise ValueError(f"Method ({method}) not recognized.")
-
- return type_means, type_dev
-
-
-def weighted_average(data, weights, dim, keepdim=False, ignore_nan=False):
- """
- take the weighted average of **data** on a certain dimension with **weights**
-
- weights should be a nonnegative vector that broadcasts into data
- avg = sum_i(data_i * weight_i, dim) / sum_i(weight_i, dim)
-
- if ignore_nan=True, (default=False), will ignore nans in weighted average
- """
- assert data.ndim == weights.ndim, "data and weights must have same number of dimensions"
- assert torch.all(weights[~torch.isnan(weights)] >= 0), "weights must be nonnegative"
-
- for d in dim if check_iterable(dim) else [dim]:
- assert data.size(d) == weights.size(
- d
- ), f"data and weights must have same size in averaging dimensions (data.size({d})={data.size(d)}, (weight.size({d})={weights.size(d)}))"
-
- # use normal sum if not ignore nan, otherwise use nansum
- sum = torch.nansum if ignore_nan else torch.sum
-
- # make sure nans are in the same place in weights and data for accurate division by total weight
- if ignore_nan:
- weights = weights.expand(data.size())
- weights = torch.masked_fill(weights, torch.isnan(data), torch.nan)
-
- # numerator & denominator of weighted average
- numerator = sum(data * weights, dim=dim, keepdim=keepdim)
- denominator = sum(weights, dim=dim, keepdim=keepdim)
-
- # return weighted average
- return numerator / denominator
-
-
-def fgsm_attack(image, epsilon, data_grad, transform, sign):
- """update an image with fast-gradient sign method"""
- warn("fgsm_attack is only going to be in utils temporarily!", DeprecationWarning, stacklevel=2)
- # Collect the element-wise sign of the data gradient
- if sign:
- data_grad = data_grad.sign()
- else:
- data_grad = data_grad.clone()
- # Create the perturbed image by adjusting each pixel of the input image
- perturbed_image = image + epsilon * data_grad
- # Adding clipping to maintain [0,1] range
- perturbed_image = transform(perturbed_image)
- # Return the perturbed image
- return perturbed_image
-
-
-def str2bool(str):
- if isinstance(str, bool):
- return str
- if str.lower() in ("true", "1"):
- return True
- elif str.lower() in ("false", "0"):
- return False
- else:
- raise TypeError("Boolean type expected")
-
-
-def save_checkpoint(nets, optimizers, results, path):
- """
- Method for saving checkpoints for networks throughout training.
- """
- multi_model_ckpt = {f"model_state_dict_{i}": net.state_dict() for i, net in enumerate(nets)}
- multi_optimizer_ckpt = {f"optimizer_state_dict_{i}": opt.state_dict() for i, opt in enumerate(optimizers)}
- checkpoint = results | multi_model_ckpt | multi_optimizer_ckpt
- torch.save(checkpoint, path)
-
-
-def load_checkpoints(nets, optimizers, device, path):
- """
- Method for loading presaved checkpoint during training.
- TODO: device handling for passing between gpu/cpu
- """
-
- if device == "cpu":
- checkpoint = torch.load(path, map_location=device)
- elif device == "cuda":
- checkpoint = torch.load(path)
-
- net_ids = natsorted([key for key in checkpoint if key.startswith("model_state_dict")])
- opt_ids = natsorted([key for key in checkpoint if key.startswith("optimizer_state_dict")])
- assert all(
- [oi.split("_")[-1] == ni.split("_")[-1] for oi, ni in zip(opt_ids, net_ids)]
- ), "nets and optimizers cannot be matched up from checkpoint"
-
- [net.load_state_dict(checkpoint.pop(net_id)) for net, net_id in zip(nets, net_ids)]
- [opt.load_state_dict(checkpoint.pop(opt_id)) for opt, opt_id in zip(optimizers, opt_ids)]
-
- if device == "cuda":
- [net.to(device) for net in nets]
-
- return nets, optimizers, checkpoint
-
-
-def match_git(path):
- """simple method for determining if a path is a git-related file or directory"""
- if ".git" in path:
- return True
- return False
-
-
-def compress_directory(output_path, directory_path=None):
- """send an entire directory to a zip file at output_path, using .gitignore and ignoring .git files"""
- if directory_path is None:
- # relative to utils -- this is the main repo path
- directory_path = os.path.dirname(os.path.abspath(__file__)) + "/../.."
- print(directory_path)
- # Parse .gitignore file
- gitignore_path = os.path.join(directory_path, ".gitignore")
- matches = parse_gitignore(gitignore_path)
-
- # Prepare list for copying files
- files_to_copy = []
- archive_names = []
- for dirpath, dirnames, files in os.walk(directory_path):
- if matches(dirpath) or match_git(dirpath):
- # clear any files from within this path
- dirnames[:] = []
- else:
- # Filter files based on .gitignore rules (and don't save any .git files)
- keep_files = [f for f in files if not matches(f) and not match_git(f)]
- # Make full path
- full_files = [os.path.join(dirpath, f) for f in keep_files]
- for file in full_files:
- # Add file to the copy list
- files_to_copy.append(file)
- archive_names.append(os.path.relpath(file, directory_path))
-
- # create zip file
- with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zipf:
- # go through directory
- for file, name in zip(files_to_copy, archive_names):
- zipf.write(file, arcname=name)
diff --git a/_arxiv/_archive/benchmark_dropout.py b/_arxiv/_archive/benchmark_dropout.py
deleted file mode 100644
index 1bd94dad..00000000
--- a/_arxiv/_archive/benchmark_dropout.py
+++ /dev/null
@@ -1,230 +0,0 @@
-"""
-Benchmark script for comparing different dropout implementation approaches.
-
-This script compares the performance of three approaches to progressive dropout:
-1. Original (sequential): Processing networks one-by-one
-2. Batched: Processing networks in small batches
-3. Tensorized: Using tensor operations to process networks at once
-"""
-
-import os
-import sys
-import time
-import argparse
-import logging
-from typing import Dict, List, Tuple, Any
-from collections import defaultdict
-
-import numpy as np
-import torch
-import torch.nn as nn
-from tqdm import tqdm
-
-from alignment.config import ExperimentConfig
-from alignment.models.registry import create_model
-from alignment.metrics import get_metric
-from alignment.datasets import load_dataset
-from alignment.dropout import progressive_dropout
-from alignment import utils
-
-# Setup logging
-logger = logging.getLogger(__name__)
-logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
-
-def create_networks(config, num_networks=20):
- """Create multiple neural networks for benchmarking."""
- networks = []
- device = config.device
-
- for i in range(num_networks):
- # Set different seed for each network
- if hasattr(config, 'seed') and config.seed is not None:
- torch.manual_seed(config.seed + i)
- torch.cuda.manual_seed_all(config.seed + i)
- np.random.seed(config.seed + i)
-
- # Create model
- model = create_model(config.model)
- model.to(device)
- networks.append(model)
-
- logger.info(f"Created {len(networks)} models with independent initializations")
- return networks
-
-def run_benchmark(config_path, approaches=["original", "batched", "tensorized"],
- num_networks=20, num_runs=3, num_batches=5, detailed_timing=True):
- """
- Run benchmarks for different dropout implementation approaches.
-
- Args:
- config_path: Path to configuration file
- approaches: List of approaches to benchmark
- num_networks: Number of networks to use
- num_runs: Number of times to run each benchmark for averaging
- num_batches: Number of data batches to process
- detailed_timing: Whether to time individual steps
- """
- # Load configuration
- config = ExperimentConfig.load(config_path)
-
- # Create networks
- networks = create_networks(config, num_networks)
-
- # Prepare dataset
- dataset = load_dataset(config.dataset)
-
- # Get metric
- metric = get_metric(config.alignment.metric)
-
- # Benchmark results
- results = {}
- detailed_results = defaultdict(lambda: defaultdict(list))
-
- # Run benchmarks
- for approach in approaches:
- logger.info(f"Benchmarking {approach} approach")
-
- # Set approach-specific parameters
- if approach == "original":
- # Original sequential approach
- kwargs = {"network_batch_size": 1, "use_tensorized": False}
- elif approach == "batched":
- # Batched approach
- kwargs = {"network_batch_size": 4, "use_tensorized": False}
- elif approach == "tensorized":
- # Tensorized approach
- kwargs = {"use_tensorized": True}
- else:
- logger.error(f"Unknown approach: {approach}")
- continue
-
- # Run multiple times and average
- timings = []
- for run in range(num_runs):
- logger.info(f"Run {run+1}/{num_runs}")
-
- # Time the execution of full dropout
- start_time = time.time()
-
- # Run progressive dropout with the specified approach
- dropout_results = progressive_dropout(
- networks,
- dataset,
- dropout_fractions=np.linspace(0.1, 0.9, config.alignment.dropout_steps),
- metric=metric,
- device=config.device,
- pruning_mode=config.extra.dropout_pruning_mode,
- dropout_mode=config.extra.dropout_mode,
- num_batches=num_batches,
- detailed_timing=detailed_timing,
- **kwargs
- )
-
- # Calculate elapsed time
- elapsed_time = time.time() - start_time
- timings.append(elapsed_time)
-
- # Collect detailed timing information if available
- if detailed_timing and hasattr(dropout_results, "timing_info"):
- for key, value in dropout_results.timing_info.items():
- detailed_results[approach][key].append(value)
-
- # Log detailed timing information
- for key, values in dropout_results.timing_info.items():
- if isinstance(values, (int, float)):
- logger.info(f" {key}: {values:.2f} seconds")
-
- logger.info(f"Completed run in {elapsed_time:.2f} seconds")
-
- # Calculate average and standard deviation
- avg_time = np.mean(timings)
- std_time = np.std(timings)
-
- # Store results
- results[approach] = {
- "average_time": avg_time,
- "std_time": std_time,
- "timings": timings
- }
-
- logger.info(f"{approach}: {avg_time:.2f} ± {std_time:.2f} seconds")
-
- # Print comparison
- logger.info("\nBenchmarking Results:")
- logger.info("-" * 50)
- logger.info(f"{'Approach':<15} {'Average Time (s)':<20} {'Speedup':<10}")
- logger.info("-" * 50)
-
- # Calculate speedup relative to original approach
- original_time = results["original"]["average_time"] if "original" in results else None
-
- for approach in approaches:
- if approach not in results:
- continue
-
- avg_time = results[approach]["average_time"]
- speedup = original_time / avg_time if original_time else 1.0
-
- logger.info(f"{approach:<15} {avg_time:.2f} ± {results[approach]['std_time']:.2f} {'':>5} {speedup:.2f}x")
-
- logger.info("-" * 50)
-
- # Print detailed timing if available
- if detailed_timing:
- logger.info("\nDetailed Timing Breakdown:")
- logger.info("-" * 80)
-
- # Find all unique timing keys
- all_keys = set()
- for approach in approaches:
- if approach in detailed_results:
- all_keys.update(detailed_results[approach].keys())
-
- # Print header
- logger.info(f"{'Operation':<25} " + " ".join([f"{approach:<15}" for approach in approaches]))
- logger.info("-" * 80)
-
- # Print timings for each key
- for key in sorted(all_keys):
- line = f"{key:<25} "
- for approach in approaches:
- if approach in detailed_results and key in detailed_results[approach]:
- values = detailed_results[approach][key]
- if len(values) > 0:
- avg = np.mean(values)
- line += f"{avg:.2f}s{' ' * 10}"
- else:
- line += f"{'N/A':<15}"
- else:
- line += f"{'N/A':<15}"
- logger.info(line)
-
- logger.info("-" * 80)
-
- return results, detailed_results if detailed_timing else None
-
-def main():
- """Main function for benchmarking."""
- parser = argparse.ArgumentParser(description="Benchmark dropout implementation approaches")
- parser.add_argument("--config", type=str, required=True, help="Path to configuration file")
- parser.add_argument("--approaches", type=str, nargs="+", default=["original", "batched", "tensorized"],
- help="Approaches to benchmark")
- parser.add_argument("--num-networks", type=int, default=20, help="Number of networks to use")
- parser.add_argument("--num-runs", type=int, default=3, help="Number of runs for each benchmark")
- parser.add_argument("--num-batches", type=int, default=5, help="Number of data batches to process")
- parser.add_argument("--no-detailed-timing", action="store_true", help="Disable detailed timing")
-
- args = parser.parse_args()
-
- # Run benchmarks
- run_benchmark(
- args.config,
- args.approaches,
- args.num_networks,
- args.num_runs,
- args.num_batches,
- not args.no_detailed_timing
- )
-
-if __name__ == "__main__":
- main()
\ No newline at end of file
diff --git a/_arxiv/_archive/config_alignment_stats.yaml b/_arxiv/_archive/config_alignment_stats.yaml
deleted file mode 100644
index abd73f95..00000000
--- a/_arxiv/_archive/config_alignment_stats.yaml
+++ /dev/null
@@ -1,56 +0,0 @@
-experiment: alignment_stats
-results_path: # <-- Add path to where results to be saved/loaded
-no_save: False
-just_plot: False
-save_networks: False
-show_params: False
-show_all: False
-#device: cpu
-use_timestamp: False
-#timestamp: # The timestamp of previous experiment to plot.
-
-dataset:
- name: MNIST
- path: # <-- Add path to your dataset
- download: True # Downloads it if it does not exist.
-
-model:
- name: MLP
- dropout: 0
- alignment_layers: null # Specify the layers participating in alignment - for all layers use null | exp for MLP: {"layerInput.0": null, "layerHidden.0.1": "layerHidden.0.0", "layerHidden.1.1": "layerHidden.1.0"}
-
-optimizer:
- name: Adam
- lr: 1e-3
- weight_decay: 0
-
-training:
- batch_size: 1024
- epochs: 1
- replicates: 1
-
-alignment:
- no_alignment: False
- delta_weights: False
- frequency: 1
-
-checkpointing:
- use_prev: False
- save_checkpoints: False
- frequency: 1
- use_wandb: False
-
-extra:
- num_drops: 9
- dropout_by_layer: False
-
- ### For alignment_comparison and loading_prediction experiments
- comparison: "lr"
- regularizers: ["none", "dropout", "weight_decay"]
- lrs: [1e-2, 1e-3, 1e-4]
- compare_dropout: 0.5
- compare_wd: 1e-5
-
- ### For adversarial_shaping experiment
- cutoffs: [1e-2, 1e-3, 1e-4, 0.0]
- manual_frequency: 5
diff --git a/_arxiv/_archive/config_alignment_stats_old.yaml b/_arxiv/_archive/config_alignment_stats_old.yaml
deleted file mode 100644
index a25906ba..00000000
--- a/_arxiv/_archive/config_alignment_stats_old.yaml
+++ /dev/null
@@ -1,66 +0,0 @@
-experiment: alignment_stats
-results_path: results # <-- Add path to where results to be saved/loaded
-no_save: False
-just_plot: False
-save_networks: False
-show_params: False
-show_all: False
-device: cuda
-use_timestamp: False
-#timestamp: # The timestamp of previous experiment to plot.
-
-dataset:
- name: MNIST
- path: /n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment/data
-
-model:
- name: MLP
- dropout: 0
- alignment_layers: null # Specify the layers participating in alignment - for all layers use null | exp for MLP: {"layerInput.0": null, "layerHidden.0.1": "layerHidden.0.0", "layerHidden.1.1": "layerHidden.1.0"}
-
-training:
- batch_size: 1024
- epochs: 3
- replicates: 3
- name: Adam
- lr: 1e-3
- weight_decay: 0
-
-alignment:
- do_alignment: True
- frequency: 1
- methods: ["RQ"]
- measure_expected: False
- bins: 50
- cnn_mode: "unfold" # Options for CNN alignment processing:
- # "unfold" - Reshapes conv layer features into (-1, features) with batch and patches together (default)
- # "patchwise" - Computes alignment per-patch, with special processing in patchwise_alignment
- # "batch_patch_combined" - Standardized approach that combines batch and patch dimensions
- # Differences between modes:
- # - "unfold" is the classic approach with flexible reshaping options
- # - "patchwise" preserves patch structure for specialized calculations
- # - "batch_patch_combined" ensures consistent dimensionality across frameworks --> original method
- scale_by_norm: False # If True, scales covariance matrices by their norm before computing RQ
- # Similar to the "similarity" option in original codebase
-checkpointing:
- use_prev: False
- save_checkpoints: False
- frequency: 1
- use_wandb: True
-
-extra:
- num_drops: 40
- dropout_pruning_mode: "global_joint" # Options:
- # "global_joint" - Prune X% of nodes across all layers based on their global alignment ranking
- # This distributes pruning across layers proportional to where the
- # highest/lowest alignment nodes are located. Some layers may have more
- # neurons pruned than others, preserving the global ranking of importance.
- # "layer_wise" - Prune exactly X% from each layer and apply pruning to all layers at once
- # This shows the effect of pruning all layers simultaneously with fixed proportions
- # "layer_isolated" - Prune X% from only one layer at a time
- # This creates multiple figures (one per layer) showing isolated effects
- dropout_mode: "scaled" # Options for how targeted dropout is applied to neurons:
- # "scaled" - Zeroes neurons and applies scaling to maintain signal magnitude (default)
- # "unscaled" - Zeroes neurons but doesn't apply compensation scaling
- exclude_classification_layer: True # If True, excludes the classification layer from pruning (always true in original codebase)
-
\ No newline at end of file
diff --git a/_arxiv/_archive/config_alignment_stats_v2.yaml b/_arxiv/_archive/config_alignment_stats_v2.yaml
deleted file mode 100644
index f1e50278..00000000
--- a/_arxiv/_archive/config_alignment_stats_v2.yaml
+++ /dev/null
@@ -1,74 +0,0 @@
-experiment: "mnist_alignment_example"
-results_path: # <-- Add path to where results to be saved/loaded
-no_save: false ## if true, skip saving results and plots
-just_plot: false ## If true, skip the experiment’s main() logic, only load previously saved results and call plot_from_existing()
-save_networks: false ## If true, also save trained networks (self.save_networks(...)) after training.
-show_params: false ## keep the plots open WE CAN REMOVE THIS
-show_all: false ## ## keep the plots open WE CAN REMOVE THIS
-use_timestamp: false ## If use_timestamp is true, the code stores results in subfolders named by timestamp (an ISO/time string).
-timestamp: null
-
-ddp_world_size: 2
-
-dataset:
- name: "ImageNet" # "MNIST", "CIFAR10", "CIFAR100", "ImageNet", etc.
- path: # <-- Add path to your dataset ## path to dataset
-
-model:
- name: "AlexNet" ## base artitecture, e.g., "MLP", "CNN2P2", "AlexNet".
- dropout: 0 ## Dropout fraction 0 to 1
- alignment_layers: null
- # "layer2.0": "layer1.0"
- # a dict of layer names: source_layer. if null it will be set to all layers.
- #"features.8": null # the null after "features.8" refers to the fact that we don't mention a particular input for this layer and use what it gets
-
-training:
- do_train: true
- epochs: 10
- batch_size: 100
- lr: 1e-3
- replicates: 2
- name: "Adam"
- weight_decay: 0.0
-
-alignment:
-# do_train_alignment: false
-# do_test_alignment: true
- do_alignment: false
- methods: ["RQ"]
- measure_expected: False
- frequency: 1
- bins: 50
- cnn_mode: "unfold" #"unfold: Single covariance across all patches", "patch: Patchwise alignment, weighting by variance", "old""
-
-checkpointing:
- use_wandb: true
- use_prev: false
- save_checkpoints: false
- frequency: 1
-
-plots:
- show_loss: true
- show_dropout: true
- show_eigenfeatures: true
- show_eig_dropout: true
-
-extra:
- num_drops: 50
- dropout_by_layer: false
- single_layer_mode: false # dropout_by_layer must be true to work
- progressive_dropout_on_train: false
- aggregate_alignment: true
-
-# • If aggregate_alignment=True, we directly append the alignment each batch to results["alignment"].
-# • If aggregate_alignment=False, we temporarily store in epoch_align_data and only append once at the end of the epoch.
-
- #comparison: "lr" ## for “alignment_comparison” style experiments
- #regularizers: ["none", "dropout", "weight_decay"]
- #lrs: [1e-3, 1e-4]
- #compare_dropout: 0.5
- #compare_wd: 1e-5
- #cutoffs: [1e-2, 1e-3, 1e-4, 0.0]
- #manual_frequency: 5
-
- # python src/alignment_v2/experiments/alignment_experiments.py --config .vscode/config_alignment_stats_v2.yaml
\ No newline at end of file
diff --git a/_arxiv/_archive/config_pruning_modes.yaml b/_arxiv/_archive/config_pruning_modes.yaml
deleted file mode 100644
index 9e5fb726..00000000
--- a/_arxiv/_archive/config_pruning_modes.yaml
+++ /dev/null
@@ -1,80 +0,0 @@
-# Pruning Modes Configuration Sample
-
-# This config demonstrates the different pruning modes available:
-# 1. "global_joint": Concatenate all nodes and prune X% across all layers (original v1)
-# 2. "layer_wise": Prune X% from each layer separately, but all at once (v2-like)
-# 3. "layer_isolated": Prune X% from only one layer at a time (new option)
-# 4. "cascading_layer": Prune layers progressively - prune layer 1, use pruned network for layer 2, etc.
-
-experiment:
- name: pruning_modes_comparison
- type: alignment_stats
-
-dataset:
- name: MNIST
- batch_size: 200
- test_batch_size: 200
-
-model:
- network: MLP
- input_dim: 784
- hidden_dim: [512, 256, 128, 64]
- output_dim: 10
- activation: ReLU
- dropout: 0.2
-
-alignment:
- methods:
- - "RQ"
- frequency: 10
- measure_expected: true
- do_alignment: true
- bins: 21
-
-training:
- epochs: 50
- lr: 0.001
- optimizer: Adam
- criterion: CrossEntropyLoss
- device: cuda:0
-
-checkpointing:
- use_prev: false
- save_checkpoints: false
- frequency: 10
-
-extra:
- # Choose one of the pruning modes:
-
- # Option 1: Global Joint Pruning (Original v1)
- # Concatenates all nodes from all layers and prunes the lowest X% across all nodes
- # dropout_pruning_mode: "global_joint"
-
- # Option 2: Layer-wise Pruning (v2-like)
- # Prunes X% from each layer separately, but all at once
- # dropout_pruning_mode: "layer_wise"
-
- # Option 3: Layer Isolation Pruning (new)
- # Prunes X% from only one layer at a time, generating separate results per layer
- # dropout_pruning_mode: "layer_isolated"
-
- # Option 4: Cascading Layer Pruning (new)
- # Progressively prunes layers - prune layer 1, use pruned network to compute alignments for layer 2, etc.
- dropout_pruning_mode: "cascading_layer"
-
- aggregate_alignment: false
- num_drops: 9
-
-# Alternative configs you can uncomment to use:
-
-# Config for "layer_wise" mode (v2 behavior)
-# extra:
-# dropout_pruning_mode: "layer_wise"
-# aggregate_alignment: false
-# num_drops: 9
-
-# Config for "layer_isolated" mode
-# extra:
-# dropout_pruning_mode: "layer_isolated"
-# aggregate_alignment: false
-# num_drops: 9
\ No newline at end of file
diff --git a/_arxiv/_archive/config_refactored_test.yaml b/_arxiv/_archive/config_refactored_test.yaml
deleted file mode 100644
index 5c409d7e..00000000
--- a/_arxiv/_archive/config_refactored_test.yaml
+++ /dev/null
@@ -1,65 +0,0 @@
-experiment: "alignment_experiment"
-results_path: "results/pruning_test"
-no_save: false
-just_plot: false
-save_networks: true
-show_params: false
-show_all: false
-device: "cuda"
-use_timestamp: true
-timestamp: null
-ddp_world_size: 1
-
-dataset:
- dataset_name: "MNIST"
- dataset_path: "/n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment/data"
-
-model:
- model_name: "mlp"
- dropout: 0.0
- alignment_layers: null
-
-training:
- do_train: true
- epochs: 5
- batch_size: 1024
- replicates: 1
- name: "Adam"
- lr: 1e-3
- weight_decay: 0.0
-
-alignment:
- do_alignment: true
- metric: "RQ"
- methods: ["RQ"]
- measure_expected: false
- frequency: 1
- bins: 50
- cnn_mode: "unfold"
- scale_by_norm: false
- run_progressive: true
- run_eigenvector: true
- dropout_min: 0.0
- dropout_max: 0.9
- dropout_steps: 10
-
-checkpointing:
- use_prev: false
- save_checkpoints: true
- frequency: 1
- use_wandb: false
-
-plots:
- show_loss: true
- show_dropout: true
- show_eigenfeatures: true
- show_eig_dropout: true
-
-extra:
- num_drops: 10
- progressive_dropout_on_train: false
- single_layer_mode: false
- aggregate_alignment: false
- dropout_pruning_mode: "per_layer_combined"
- dropout_mode: "scaled"
- exclude_classification_layer: true
\ No newline at end of file
diff --git a/_arxiv/_archive/config_with_eigenvector.yaml b/_arxiv/_archive/config_with_eigenvector.yaml
deleted file mode 100644
index 5da61835..00000000
--- a/_arxiv/_archive/config_with_eigenvector.yaml
+++ /dev/null
@@ -1,59 +0,0 @@
-# Alignment Experiment Configuration
-# This configuration file defines parameters for neural network alignment experiments
-# that analyze the relationship between alignment, dropout, and network performance.
-
-# Experiment metadata
-experiment_type: "alignment" # Type of experiment to run (options: "alignment", "training", "analysis")
-results_path: "results" # Base path for saving results (can be relative or absolute path)
-use_timestamp: true # Whether to include timestamp in result paths for unique experiment tracking
-
-# Core execution options
-device: "cuda" # Device to run on (options: "cuda", "cpu", or specific GPU like "cuda:0")
-no_save: false # If true, don't save any results to disk (useful for debug runs)
-just_plot: false # If true, only load and plot existing results without running experiments
-save_networks: true # Whether to save trained network checkpoints for later use
-show_all: false # Whether to display all plots during execution (vs. just saving them)
-
-# Dataset configuration
-dataset:
- dataset_name: "MNIST" # Name of the dataset to use (options: "MNIST", "CIFAR10", "CIFAR100")
- data_path: "/n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment/data" # Path to dataset files
- batch_size: 1024 # Batch size for training and evaluation (higher values use more memory but train faster)
-
-# Model configuration
-model:
- model_name: "MLP" # Model architecture to use (options: "MLP", "CNN2P2", "alexnet")
- dropout_rate: 0.0 # Base dropout rate during training (range: 0.0-1.0, 0.0 means no dropout)
- alignment_layers: null # Layers to measure alignment on (null for all layers, or specify as [0,1,2])
-
-# Training configuration
-training:
- epochs: 5 # Number of training epochs (more epochs can improve performance but take longer)
- replicates: 10 # Number of network replicates to train (for statistical analysis)
- optimizer: "Adam" # Optimizer to use (options: "Adam", "SGD", "RMSprop")
- learning_rate: 1.0e-3 # Learning rate for optimizer (typical range: 1e-4 to 1e-2)
- weight_decay: 0.0 # L2 regularization strength (higher values prevent overfitting)
-
-# Checkpointing configuration
-checkpointing:
- save_checkpoints: true # Whether to save training checkpoints to disk
- checkpoint_frequency: 1 # How often to save checkpoints (in epochs, 1 means every epoch)
- use_wandb: true # Whether to use Weights & Biases for experiment tracking and visualization
- load_checkpoint: false # Whether to load from previous checkpoint (for resuming training)
-
-# Alignment analysis configuration
-alignment:
- metric: "RQ" # Alignment metric to use (options: "RQ", "NullSpace", "Delta", "R2")
- run_progressive: true # Whether to run progressive dropout experiment (node pruning)
- run_eigenvector: true # Whether to run eigenvector dropout experiment (feature pruning based on PCA components)
- dropout_min: 0.0 # Minimum dropout fraction to test (range: 0.0-1.0)
- dropout_max: 0.95 # Maximum dropout fraction to test (range: 0.0-1.0, <1.0 to avoid complete dropout)
- dropout_steps: 20 # Number of dropout fractions to test (more steps = higher resolution but slower)
- scale_by_norm: false # Whether to scale covariance matrices by norm (can improve numerics)
-
-# Extra configuration parameters
-extra:
- dropout_mode: "scaled" # How dropout is applied (options: "scaled" = multiply by 1/(1-p), "unscaled" = no scaling)
- dropout_pruning_mode: "global_joint" # How pruning is distributed (options: "global_joint" = across all neurons, "layer_wise" = layer-by-layer, "layer_isolated" = one layer at a time)
- exclude_classification_layer: true # Whether to exclude the output layer from pruning (usually true for classification tasks)
- cnn_mode: "unfold" # How CNN features are processed for alignment (options: "unfold", "patchwise", "batch_patch_combined")
\ No newline at end of file
diff --git a/_arxiv/_archive/configs/config_alignment_experiment.yaml b/_arxiv/_archive/configs/config_alignment_experiment.yaml
deleted file mode 100644
index 1850e894..00000000
--- a/_arxiv/_archive/configs/config_alignment_experiment.yaml
+++ /dev/null
@@ -1,203 +0,0 @@
-# Alignment Experiment Configuration
-# This configuration file defines parameters for neural network alignment experiments
-# that analyze the relationship between alignment, dropout, and network performance.
-
-# -----------------------------------------------------------------------------
-# I. Experiment Metadata & Core Execution
-# -----------------------------------------------------------------------------
-experiment_name: "alexnet_imagenet_alignment_test"
-# experiment_name: "cnn_mnist_nested_config_test"
-experiment_type: "AUTO" # Options:
- # "AUTO" - Automatically selects the best experiment type based on the dataset.
- # "progressive_dropout" - Prunes nodes iteratively and measures performance.
- # "eigenvector_dropout" - Prunes based on eigenvectors of alignment metric.
- # "alignment_analysis" - Runs both progressive and (if configured) eigenvector dropout.
- # "layer_isolated_pruning" - Prunes/evaluates one layer at a time (generates N plots).
- # "cascading_layer_pruning" - Prunes layers sequentially, L1 then L2 etc (simple version).
-results_path: "results" # Base path for saving all outputs.
-use_timestamp: true # If true, appends a YYYYMMDD_HHMMSS timestamp to `results_path`.
-debug_mode: false # MODIFIED: Set to true for testing callbacks
-
-device: "cuda" # "cuda", "cuda:0", "cpu". Ignored if use_ddp is true (device set by local_rank).
-save_networks: true # If true, final trained/pruned network models are saved.
-no_save: false # Master switch: if true, nothing is saved to disk.
-just_plot: false # If true, loads existing results and only generates plots.
-show_all: false # If true, matplotlib plots will be displayed interactively.
-seed: 422 # MODIFIED: Set a seed for reproducibility during testing
-
-# --- NEW: DDP (Distributed Data Parallel) Configuration ---
-use_ddp: false # Set to true to enable DDP for multi-GPU training.
- # Requires launching with torchrun or similar.
- # e.g., torchrun --nproc_per_node=NUM_GPUS your_script.py --config this_file.yaml
-ddp_backend: "nccl" # DDP backend. "nccl" is recommended for NVIDIA GPUs.
-# ddp_rank, ddp_world_size, ddp_local_rank are set automatically by the launch script if use_ddp is true.
-# --- End DDP Configuration ---
-
-# -----------------------------------------------------------------------------
-# II. Dataset Configuration
-# -----------------------------------------------------------------------------
-# Maps to DatasetConfig in config.py
-dataset:
- dataset_name: "ImageNet" # Options: "MNIST", "CIFAR10", "CIFAR100", "ImageNet".
- data_path: "/n/holylfs06/LABS/kempner_shared/Lab/data/imagenet_1k/" # Root directory for dataset files.
- batch_size: 256 # Batch size for DataLoaders (reduced for ImageNet due to larger images).
- num_workers: 8 # Number of worker processes for DataLoader (increased for ImageNet).
-
-# -----------------------------------------------------------------------------
-# III. Model Configuration (NEW STRUCTURE)
-# -----------------------------------------------------------------------------
-# Maps to ModelConfig in config.py
-model:
- # --- Identifier and Common Parameters ---
- model_name: "torchvision_alexnet" # Options: "MLP", "CNN2P2", "torchvision_alexnet", "torchvision_resnet18", "hf_bert-base-uncased", etc.
- output_dim: 1000 # Number of output classes. Used by the final layer of the model.
- dropout_rate: 0.0 # General dropout rate that internal model constructors might use.
-
- # --- CNN Mode for AlignmentNetwork Wrapper ---
- # Defines how the AlignmentNetwork wrapper (if the model is a CNN) preprocesses its inputs for metric calculation.
- # Options: "unfold", "patchwise", "batch_patch_combined". Default in code is "unfold".
- cnn_mode: "unfold"
-
- # --- Alignment Layers Specification (Common) ---
- # Optional: Specify layers for AlignmentNetwork to analyze. Names must match model.named_modules().
- # If null, AlignmentNetwork attempts to find all layers with weights.
- alignment_layers: null
- # Example for AlexNet:
- # alignment_layers:
- # "features.0": 0
- # "features.3": 1
- # "classifier.6": 5
-
- # --- Model-Specific Parameter Blocks ---
- # Only include and populate ONE of the following blocks based on model_name.
-
- # (Example 1) MLP Parameters (if model_name: "MLP")
- mlp_params:
- input_dim: 784
- hidden_dims: [100, 100, 50] # Example: Two hidden layers
- activation: "relu"
-
- # CNN2P2 Parameters (if model_name: "CNN2P2") - Comment out if using MLP or external
- cnn2p2_params:
- in_channels: 1
- conv_channels: [32, 64]
- kernel_sizes: [5, 5]
- strides: [1, 1]
- paddings: [2, 2]
- pool_kernel_size: 2
- pool_stride: 2
- hidden_fc_dim: 128
- example_input_hw: [28, 28]
-
- # (Example 3) External Model Parameters (if model_name indicates an external model, e.g., "torchvision_alexnet")
- external_params: # Used if model_name is "torchvision_alexnet", "hf_bert-base-uncased", or just "external"
- source: "torchvision" # Options: "torchvision", "huggingface_transformers"
- name_or_path: "alexnet" # e.g., "alexnet", "resnet18" for torchvision; "bert-base-uncased" for hf
- pretrained: true # Load pretrained weights (True/False)
- freeze_feature_extractor: false # Freeze all layers except the classifier (True/False)
-
- # --- Ad-hoc Parameters for Internal Constructors (Use Sparingly) ---
- extra_model_params: {} # For any other specific kwargs to pass to internal model constructors (create_mlp, create_cnn2p2)
-
-# -----------------------------------------------------------------------------
-# IV. Training Configuration
-# -----------------------------------------------------------------------------
-# Maps to TrainingConfig in config.py
-training:
- epochs: 5
- replicates: 1 # For DDP, typically train 1 replicate, DDP handles data parallelism.
- # If not DDP, can set >1 for multiple independent training runs.
- optimizer: "Adam"
- learning_rate: 0.001
- momentum: 0.9
- loss: "cross_entropy"
- training_method: "auto" # "auto", "sequential", "fully_tensorized".
- # "fully_tensorized" is for non-DDP multi-replicate speedup on one GPU.
- # With DDP, "sequential" is typically used for the single replicate.
- train_before_dropout: false # Whether to run initial training phase before dropout experiments.
-
-# -----------------------------------------------------------------------------
-# V. Alignment Metric Calculation Settings
-# -----------------------------------------------------------------------------
-# Maps to alignment_settings: AlignmentConfig in config.py
-alignment_settings:
- metric:
- - "RQ"
- #- "MI_G"
- #- "Node_Redundancy"
- scale_by_norm: false # Global default for scale_by_norm
-
- # cnn_mode below is for how metric calculation functions (e.g. compute_all_node_scores)
- # process CNN features IF they receive them in a specific format (e.g. from AlignmentNetwork's precomputation).
- # The ModelConfig.cnn_mode determines how AlignmentNetwork itself preprocesses and provides these features.
- # These two should ideally be consistent or ModelConfig.cnn_mode dictates the format received by metric functions.
- cnn_mode: "unfold" # Options: "unfold", "patchwise", "batch_patch_combined", "filter_patch_summary", "filter_specific_covariance_rq"
- # Detailed options for alignment_settings.cnn_mode (how metrics process CNN features):
- # "unfold": Assumes features are [batch_size * num_patches, feature_dim] (e.g., from AlignmentNetwork cnn_mode="unfold").
- # For RQ, computes ONE covariance matrix from ALL these patch features (from ONE original data batch).
- # Each filter's RQ is then computed against this single, global covariance matrix.
- # Matches typical RQ definition but can be memory-intensive.
- # "filter_patch_summary": (Specific to RQ metric) For each filter, computes a score based on its interaction
- # with every input patch it sees (e.g., mean/max of SQUARED COSINE SIMILARITIES
- # between filter and patches). Memory-efficient. Not mathematically identical to covariance-based RQ.
- # "filter_specific_covariance_rq": (Specific to RQ metric) For each filter, computes its full covariance-based RQ.
- # The covariance is derived from all unfolded input patches (from ONE original batch)
- # that the filter processes. Mathematically standard RQ per filter. Computationally intensive.
- # "patchwise": Expects features like [batch_size, num_patches, feature_dim]. Intended for metrics that process per-patch.
- # (Currently requires metric-specific support not fully integrated into general path for all metrics).
- # "batch_patch_combined": Similar effect to "unfold" for how data is shaped for global covariance calculations.
-
- cnn_rq_aggregation_op: "mean" # If cnn_mode for metric is "filter_patch_summary" for RQ: "mean", "max", "var", "sum".
- force_cpu_for_large_metric_ops: false # If true, attempts to offload large tensor ops in metrics (like cov mat) to CPU.
-
- run_progressive: true # For experiment_type="alignment_analysis", whether to run progressive dropout part.
- run_eigenvector: false # For experiment_type="alignment_analysis", whether to run eigenvector dropout part.
-
- callbacks: # Configure metrics to track during training
- alignment_metrics:
- - name: "RQ"
- num_batches: 5 # Batches for callback metrics. null for all test batches.
- - name: "MI_G" # Gaussian MI
- num_batches: 5
- # - name: "PID_SI" # PID Shared Info (Requires BROJA_2PID to be correctly placed and importable)
- # num_batches: 1
-
-# -----------------------------------------------------------------------------
-# VI. Pruning Experiment Settings
-# -----------------------------------------------------------------------------
-# Maps to pruning_settings: PruningConfig in config.py
-pruning_settings:
- dropout_min: 0.0
- dropout_max: 1.0
- dropout_steps: 30 # Reduced for faster testing
-
- dropout_mode: "scaled" # "scaled" or "unscaled"
- # "scaled": Remaining weights are scaled by 1/(1-p) (Inverted Dropout principle).
- # "unscaled" (or "zero"): Weights are simply zeroed out.
- dropout_pruning_mode: "layer_isolated" # "global_joint", "layer_wise", "layer_isolated", "cascading_layer"
- # "global_joint": Considers all nodes from all prunable layers together. Ranks globally by score and prunes top X% overall.
- # "layer_wise": Prunes X% from *each* prunable layer independently, based on scores within that layer.
- # "layer_isolated": (For experiment_type="layer_isolated_pruning") Prunes X% in one layer, evaluates, restores, then moves to next layer.
- # "cascading_layer": Prunes L1, its output influences pruning of L2, etc. (Note: current batched impl. has limitations).
- exclude_classification_layer: true
- use_multi_strategy_dropout: true
- num_batches_for_scores: null # Number of batches to compute pruning scores on.
-
-# -----------------------------------------------------------------------------
-# VII. Checkpointing & Logging Configuration
-# -----------------------------------------------------------------------------
-checkpointing:
- save_checkpoints: false
- checkpoint_frequency: 1
- load_checkpoint: false
-
-wandb:
- use_wandb: true
- wandb_project: "neural_alignment"
- wandb_entity: "your_wandb_entity" # <<< IMPORTANT: Update with your W&B entity or remove line for default
-
-extra:
- log_frequency: 1
- log_images: true
- detailed_logging: true
- # dummy_extra_param: null
\ No newline at end of file
diff --git a/_arxiv/_archive/configs/config_alignment_stats.yaml b/_arxiv/_archive/configs/config_alignment_stats.yaml
deleted file mode 100644
index abd73f95..00000000
--- a/_arxiv/_archive/configs/config_alignment_stats.yaml
+++ /dev/null
@@ -1,56 +0,0 @@
-experiment: alignment_stats
-results_path: # <-- Add path to where results to be saved/loaded
-no_save: False
-just_plot: False
-save_networks: False
-show_params: False
-show_all: False
-#device: cpu
-use_timestamp: False
-#timestamp: # The timestamp of previous experiment to plot.
-
-dataset:
- name: MNIST
- path: # <-- Add path to your dataset
- download: True # Downloads it if it does not exist.
-
-model:
- name: MLP
- dropout: 0
- alignment_layers: null # Specify the layers participating in alignment - for all layers use null | exp for MLP: {"layerInput.0": null, "layerHidden.0.1": "layerHidden.0.0", "layerHidden.1.1": "layerHidden.1.0"}
-
-optimizer:
- name: Adam
- lr: 1e-3
- weight_decay: 0
-
-training:
- batch_size: 1024
- epochs: 1
- replicates: 1
-
-alignment:
- no_alignment: False
- delta_weights: False
- frequency: 1
-
-checkpointing:
- use_prev: False
- save_checkpoints: False
- frequency: 1
- use_wandb: False
-
-extra:
- num_drops: 9
- dropout_by_layer: False
-
- ### For alignment_comparison and loading_prediction experiments
- comparison: "lr"
- regularizers: ["none", "dropout", "weight_decay"]
- lrs: [1e-2, 1e-3, 1e-4]
- compare_dropout: 0.5
- compare_wd: 1e-5
-
- ### For adversarial_shaping experiment
- cutoffs: [1e-2, 1e-3, 1e-4, 0.0]
- manual_frequency: 5
diff --git a/_arxiv/_archive/configs/config_alignment_stats_old.yaml b/_arxiv/_archive/configs/config_alignment_stats_old.yaml
deleted file mode 100644
index a25906ba..00000000
--- a/_arxiv/_archive/configs/config_alignment_stats_old.yaml
+++ /dev/null
@@ -1,66 +0,0 @@
-experiment: alignment_stats
-results_path: results # <-- Add path to where results to be saved/loaded
-no_save: False
-just_plot: False
-save_networks: False
-show_params: False
-show_all: False
-device: cuda
-use_timestamp: False
-#timestamp: # The timestamp of previous experiment to plot.
-
-dataset:
- name: MNIST
- path: /n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment/data
-
-model:
- name: MLP
- dropout: 0
- alignment_layers: null # Specify the layers participating in alignment - for all layers use null | exp for MLP: {"layerInput.0": null, "layerHidden.0.1": "layerHidden.0.0", "layerHidden.1.1": "layerHidden.1.0"}
-
-training:
- batch_size: 1024
- epochs: 3
- replicates: 3
- name: Adam
- lr: 1e-3
- weight_decay: 0
-
-alignment:
- do_alignment: True
- frequency: 1
- methods: ["RQ"]
- measure_expected: False
- bins: 50
- cnn_mode: "unfold" # Options for CNN alignment processing:
- # "unfold" - Reshapes conv layer features into (-1, features) with batch and patches together (default)
- # "patchwise" - Computes alignment per-patch, with special processing in patchwise_alignment
- # "batch_patch_combined" - Standardized approach that combines batch and patch dimensions
- # Differences between modes:
- # - "unfold" is the classic approach with flexible reshaping options
- # - "patchwise" preserves patch structure for specialized calculations
- # - "batch_patch_combined" ensures consistent dimensionality across frameworks --> original method
- scale_by_norm: False # If True, scales covariance matrices by their norm before computing RQ
- # Similar to the "similarity" option in original codebase
-checkpointing:
- use_prev: False
- save_checkpoints: False
- frequency: 1
- use_wandb: True
-
-extra:
- num_drops: 40
- dropout_pruning_mode: "global_joint" # Options:
- # "global_joint" - Prune X% of nodes across all layers based on their global alignment ranking
- # This distributes pruning across layers proportional to where the
- # highest/lowest alignment nodes are located. Some layers may have more
- # neurons pruned than others, preserving the global ranking of importance.
- # "layer_wise" - Prune exactly X% from each layer and apply pruning to all layers at once
- # This shows the effect of pruning all layers simultaneously with fixed proportions
- # "layer_isolated" - Prune X% from only one layer at a time
- # This creates multiple figures (one per layer) showing isolated effects
- dropout_mode: "scaled" # Options for how targeted dropout is applied to neurons:
- # "scaled" - Zeroes neurons and applies scaling to maintain signal magnitude (default)
- # "unscaled" - Zeroes neurons but doesn't apply compensation scaling
- exclude_classification_layer: True # If True, excludes the classification layer from pruning (always true in original codebase)
-
\ No newline at end of file
diff --git a/_arxiv/_archive/configs/config_alignment_stats_v2.yaml b/_arxiv/_archive/configs/config_alignment_stats_v2.yaml
deleted file mode 100644
index f1e50278..00000000
--- a/_arxiv/_archive/configs/config_alignment_stats_v2.yaml
+++ /dev/null
@@ -1,74 +0,0 @@
-experiment: "mnist_alignment_example"
-results_path: # <-- Add path to where results to be saved/loaded
-no_save: false ## if true, skip saving results and plots
-just_plot: false ## If true, skip the experiment’s main() logic, only load previously saved results and call plot_from_existing()
-save_networks: false ## If true, also save trained networks (self.save_networks(...)) after training.
-show_params: false ## keep the plots open WE CAN REMOVE THIS
-show_all: false ## ## keep the plots open WE CAN REMOVE THIS
-use_timestamp: false ## If use_timestamp is true, the code stores results in subfolders named by timestamp (an ISO/time string).
-timestamp: null
-
-ddp_world_size: 2
-
-dataset:
- name: "ImageNet" # "MNIST", "CIFAR10", "CIFAR100", "ImageNet", etc.
- path: # <-- Add path to your dataset ## path to dataset
-
-model:
- name: "AlexNet" ## base artitecture, e.g., "MLP", "CNN2P2", "AlexNet".
- dropout: 0 ## Dropout fraction 0 to 1
- alignment_layers: null
- # "layer2.0": "layer1.0"
- # a dict of layer names: source_layer. if null it will be set to all layers.
- #"features.8": null # the null after "features.8" refers to the fact that we don't mention a particular input for this layer and use what it gets
-
-training:
- do_train: true
- epochs: 10
- batch_size: 100
- lr: 1e-3
- replicates: 2
- name: "Adam"
- weight_decay: 0.0
-
-alignment:
-# do_train_alignment: false
-# do_test_alignment: true
- do_alignment: false
- methods: ["RQ"]
- measure_expected: False
- frequency: 1
- bins: 50
- cnn_mode: "unfold" #"unfold: Single covariance across all patches", "patch: Patchwise alignment, weighting by variance", "old""
-
-checkpointing:
- use_wandb: true
- use_prev: false
- save_checkpoints: false
- frequency: 1
-
-plots:
- show_loss: true
- show_dropout: true
- show_eigenfeatures: true
- show_eig_dropout: true
-
-extra:
- num_drops: 50
- dropout_by_layer: false
- single_layer_mode: false # dropout_by_layer must be true to work
- progressive_dropout_on_train: false
- aggregate_alignment: true
-
-# • If aggregate_alignment=True, we directly append the alignment each batch to results["alignment"].
-# • If aggregate_alignment=False, we temporarily store in epoch_align_data and only append once at the end of the epoch.
-
- #comparison: "lr" ## for “alignment_comparison” style experiments
- #regularizers: ["none", "dropout", "weight_decay"]
- #lrs: [1e-3, 1e-4]
- #compare_dropout: 0.5
- #compare_wd: 1e-5
- #cutoffs: [1e-2, 1e-3, 1e-4, 0.0]
- #manual_frequency: 5
-
- # python src/alignment_v2/experiments/alignment_experiments.py --config .vscode/config_alignment_stats_v2.yaml
\ No newline at end of file
diff --git a/_arxiv/_archive/configs/config_pruning_modes.yaml b/_arxiv/_archive/configs/config_pruning_modes.yaml
deleted file mode 100644
index 9e5fb726..00000000
--- a/_arxiv/_archive/configs/config_pruning_modes.yaml
+++ /dev/null
@@ -1,80 +0,0 @@
-# Pruning Modes Configuration Sample
-
-# This config demonstrates the different pruning modes available:
-# 1. "global_joint": Concatenate all nodes and prune X% across all layers (original v1)
-# 2. "layer_wise": Prune X% from each layer separately, but all at once (v2-like)
-# 3. "layer_isolated": Prune X% from only one layer at a time (new option)
-# 4. "cascading_layer": Prune layers progressively - prune layer 1, use pruned network for layer 2, etc.
-
-experiment:
- name: pruning_modes_comparison
- type: alignment_stats
-
-dataset:
- name: MNIST
- batch_size: 200
- test_batch_size: 200
-
-model:
- network: MLP
- input_dim: 784
- hidden_dim: [512, 256, 128, 64]
- output_dim: 10
- activation: ReLU
- dropout: 0.2
-
-alignment:
- methods:
- - "RQ"
- frequency: 10
- measure_expected: true
- do_alignment: true
- bins: 21
-
-training:
- epochs: 50
- lr: 0.001
- optimizer: Adam
- criterion: CrossEntropyLoss
- device: cuda:0
-
-checkpointing:
- use_prev: false
- save_checkpoints: false
- frequency: 10
-
-extra:
- # Choose one of the pruning modes:
-
- # Option 1: Global Joint Pruning (Original v1)
- # Concatenates all nodes from all layers and prunes the lowest X% across all nodes
- # dropout_pruning_mode: "global_joint"
-
- # Option 2: Layer-wise Pruning (v2-like)
- # Prunes X% from each layer separately, but all at once
- # dropout_pruning_mode: "layer_wise"
-
- # Option 3: Layer Isolation Pruning (new)
- # Prunes X% from only one layer at a time, generating separate results per layer
- # dropout_pruning_mode: "layer_isolated"
-
- # Option 4: Cascading Layer Pruning (new)
- # Progressively prunes layers - prune layer 1, use pruned network to compute alignments for layer 2, etc.
- dropout_pruning_mode: "cascading_layer"
-
- aggregate_alignment: false
- num_drops: 9
-
-# Alternative configs you can uncomment to use:
-
-# Config for "layer_wise" mode (v2 behavior)
-# extra:
-# dropout_pruning_mode: "layer_wise"
-# aggregate_alignment: false
-# num_drops: 9
-
-# Config for "layer_isolated" mode
-# extra:
-# dropout_pruning_mode: "layer_isolated"
-# aggregate_alignment: false
-# num_drops: 9
\ No newline at end of file
diff --git a/_arxiv/_archive/configs/config_refactored_test.yaml b/_arxiv/_archive/configs/config_refactored_test.yaml
deleted file mode 100644
index 5c409d7e..00000000
--- a/_arxiv/_archive/configs/config_refactored_test.yaml
+++ /dev/null
@@ -1,65 +0,0 @@
-experiment: "alignment_experiment"
-results_path: "results/pruning_test"
-no_save: false
-just_plot: false
-save_networks: true
-show_params: false
-show_all: false
-device: "cuda"
-use_timestamp: true
-timestamp: null
-ddp_world_size: 1
-
-dataset:
- dataset_name: "MNIST"
- dataset_path: "/n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment/data"
-
-model:
- model_name: "mlp"
- dropout: 0.0
- alignment_layers: null
-
-training:
- do_train: true
- epochs: 5
- batch_size: 1024
- replicates: 1
- name: "Adam"
- lr: 1e-3
- weight_decay: 0.0
-
-alignment:
- do_alignment: true
- metric: "RQ"
- methods: ["RQ"]
- measure_expected: false
- frequency: 1
- bins: 50
- cnn_mode: "unfold"
- scale_by_norm: false
- run_progressive: true
- run_eigenvector: true
- dropout_min: 0.0
- dropout_max: 0.9
- dropout_steps: 10
-
-checkpointing:
- use_prev: false
- save_checkpoints: true
- frequency: 1
- use_wandb: false
-
-plots:
- show_loss: true
- show_dropout: true
- show_eigenfeatures: true
- show_eig_dropout: true
-
-extra:
- num_drops: 10
- progressive_dropout_on_train: false
- single_layer_mode: false
- aggregate_alignment: false
- dropout_pruning_mode: "per_layer_combined"
- dropout_mode: "scaled"
- exclude_classification_layer: true
\ No newline at end of file
diff --git a/_arxiv/_archive/configs/config_with_eigenvector.yaml b/_arxiv/_archive/configs/config_with_eigenvector.yaml
deleted file mode 100644
index 5da61835..00000000
--- a/_arxiv/_archive/configs/config_with_eigenvector.yaml
+++ /dev/null
@@ -1,59 +0,0 @@
-# Alignment Experiment Configuration
-# This configuration file defines parameters for neural network alignment experiments
-# that analyze the relationship between alignment, dropout, and network performance.
-
-# Experiment metadata
-experiment_type: "alignment" # Type of experiment to run (options: "alignment", "training", "analysis")
-results_path: "results" # Base path for saving results (can be relative or absolute path)
-use_timestamp: true # Whether to include timestamp in result paths for unique experiment tracking
-
-# Core execution options
-device: "cuda" # Device to run on (options: "cuda", "cpu", or specific GPU like "cuda:0")
-no_save: false # If true, don't save any results to disk (useful for debug runs)
-just_plot: false # If true, only load and plot existing results without running experiments
-save_networks: true # Whether to save trained network checkpoints for later use
-show_all: false # Whether to display all plots during execution (vs. just saving them)
-
-# Dataset configuration
-dataset:
- dataset_name: "MNIST" # Name of the dataset to use (options: "MNIST", "CIFAR10", "CIFAR100")
- data_path: "/n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment/data" # Path to dataset files
- batch_size: 1024 # Batch size for training and evaluation (higher values use more memory but train faster)
-
-# Model configuration
-model:
- model_name: "MLP" # Model architecture to use (options: "MLP", "CNN2P2", "alexnet")
- dropout_rate: 0.0 # Base dropout rate during training (range: 0.0-1.0, 0.0 means no dropout)
- alignment_layers: null # Layers to measure alignment on (null for all layers, or specify as [0,1,2])
-
-# Training configuration
-training:
- epochs: 5 # Number of training epochs (more epochs can improve performance but take longer)
- replicates: 10 # Number of network replicates to train (for statistical analysis)
- optimizer: "Adam" # Optimizer to use (options: "Adam", "SGD", "RMSprop")
- learning_rate: 1.0e-3 # Learning rate for optimizer (typical range: 1e-4 to 1e-2)
- weight_decay: 0.0 # L2 regularization strength (higher values prevent overfitting)
-
-# Checkpointing configuration
-checkpointing:
- save_checkpoints: true # Whether to save training checkpoints to disk
- checkpoint_frequency: 1 # How often to save checkpoints (in epochs, 1 means every epoch)
- use_wandb: true # Whether to use Weights & Biases for experiment tracking and visualization
- load_checkpoint: false # Whether to load from previous checkpoint (for resuming training)
-
-# Alignment analysis configuration
-alignment:
- metric: "RQ" # Alignment metric to use (options: "RQ", "NullSpace", "Delta", "R2")
- run_progressive: true # Whether to run progressive dropout experiment (node pruning)
- run_eigenvector: true # Whether to run eigenvector dropout experiment (feature pruning based on PCA components)
- dropout_min: 0.0 # Minimum dropout fraction to test (range: 0.0-1.0)
- dropout_max: 0.95 # Maximum dropout fraction to test (range: 0.0-1.0, <1.0 to avoid complete dropout)
- dropout_steps: 20 # Number of dropout fractions to test (more steps = higher resolution but slower)
- scale_by_norm: false # Whether to scale covariance matrices by norm (can improve numerics)
-
-# Extra configuration parameters
-extra:
- dropout_mode: "scaled" # How dropout is applied (options: "scaled" = multiply by 1/(1-p), "unscaled" = no scaling)
- dropout_pruning_mode: "global_joint" # How pruning is distributed (options: "global_joint" = across all neurons, "layer_wise" = layer-by-layer, "layer_isolated" = one layer at a time)
- exclude_classification_layer: true # Whether to exclude the output layer from pruning (usually true for classification tasks)
- cnn_mode: "unfold" # How CNN features are processed for alignment (options: "unfold", "patchwise", "batch_patch_combined")
\ No newline at end of file
diff --git a/_arxiv/_archive/configs/training_example.yaml b/_arxiv/_archive/configs/training_example.yaml
deleted file mode 100644
index 886566d1..00000000
--- a/_arxiv/_archive/configs/training_example.yaml
+++ /dev/null
@@ -1,47 +0,0 @@
-# Example configuration with tensorized training options
-experiment_type: "progressive_dropout"
-experiment_name: "Tensorized Training Example"
-use_timestamp: true
-device: "cuda"
-
-# Dataset configuration
-dataset:
- dataset_name: "MNIST"
- batch_size: 128
- data_path: "./data"
-
-# Model configuration
-model:
- model_name: "MLP"
- dropout_rate: 0.0
-
-# Training configuration
-training:
- epochs: 5
- replicates: 10 # Number of network copies to train
- learning_rate: 0.001
- weight_decay: 0.0
- optimizer: "Adam"
- train_before_dropout: true
-
-# Alignment configuration
-alignment:
- metric: "RQ"
- run_progressive: true
- run_eigenvector: false
- dropout_min: 0.0
- dropout_max: 0.95
- dropout_steps: 20
-
-# Extra configuration options
-extra:
- dropout_mode: "scaled"
- dropout_pruning_mode: "global_joint"
- exclude_classification_layer: true
- # Available training methods: "auto", "sequential", "tensorized", "fully_tensorized"
- training_method: "fully_tensorized" # Using the fastest method
-
-# Checkpoint configuration
-checkpointing:
- save_checkpoints: false
- use_wandb: false
\ No newline at end of file
diff --git a/_arxiv/_archive/debug_pruning_strategies.py b/_arxiv/_archive/debug_pruning_strategies.py
deleted file mode 100644
index 3eb729d7..00000000
--- a/_arxiv/_archive/debug_pruning_strategies.py
+++ /dev/null
@@ -1,235 +0,0 @@
-import torch
-import torch.nn as nn
-import numpy as np
-import matplotlib.pyplot as plt
-import logging
-import os
-import sys
-import copy
-
-# Add parent directory to path to access src
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '.')))
-
-from src.alignment.models.models import MLP
-from src.alignment.datasets import load_dataset
-from src.alignment.dropout import test_pruning_strategies
-
-# Configure logging
-logging.basicConfig(
- level=logging.INFO,
- format='%(asctime)s [%(levelname)s] %(message)s',
- datefmt='%Y-%m-%d %H:%M:%S',
-)
-logger = logging.getLogger(__name__)
-
-def main():
- # Set random seed for reproducibility
- torch.manual_seed(42)
- np.random.seed(42)
-
- # Determine device
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
- logger.info(f"Using device: {device}")
-
- # Load MNIST dataset for testing
- dataset_config = {
- "dataset_name": "MNIST",
- "data_path": "data",
- "batch_size": 1024,
- "model_name": "mlp",
- }
- dataset = load_dataset(dataset_config, device=device)
- logger.info(f"Loaded {dataset_config['dataset_name']} dataset")
-
- # Create a simple MLP model for MNIST
- network = MLP(
- input_dim=784, # 28x28 for MNIST
- num_hidden=[100, 100],
- output_dim=10,
- dropout_rate=0.0,
- )
-
- # Get linear layers from the MLP model
- # First find all linear layers in the network
- linear_layers = []
- for name, module in network.named_modules():
- if isinstance(module, nn.Linear):
- linear_layers.append(module)
-
- # Set alignment layers for pruning (all linear layers)
- network.alignment_layers = linear_layers
- network.alignment_layer_names = [f"Layer {i}" for i in range(len(network.alignment_layers))]
-
- network.to(device)
- logger.info(f"Created MLP model with {len(network.alignment_layers)} alignment layers")
-
- # Train the network for a single epoch
- logger.info("Training network...")
- train_network(network, dataset, device)
-
- # Print statistics about the original model
- logger.info("Evaluating original model...")
- print_layer_stats(network)
- orig_accuracy, orig_loss = dataset.evaluate(network, device)
- logger.info(f"Original model - Accuracy: {orig_accuracy:.2f}%, Loss: {orig_loss:.4f}")
-
- # Save original weights for verification
- original_weights = {}
- for i, layer in enumerate(network.alignment_layers):
- if hasattr(layer, "weight") and layer.weight is not None:
- original_weights[i] = layer.weight.data.clone()
-
- # Test different pruning strategies
- logger.info("Testing pruning strategies...")
- pruning_percents = [0.0, 0.1, 0.3, 0.5, 0.7, 0.9]
-
- # Test each strategy individually to confirm they differ
- strategies = ["high_rq", "low_rq", "random"]
- all_results = {}
-
- for strategy in strategies:
- logger.info(f"Testing {strategy} pruning strategy...")
-
- # Make a fresh copy of the network for each strategy
- network_copy = copy.deepcopy(network).to(device)
-
- # Verify the copy has the same weights
- for i, layer in enumerate(network_copy.alignment_layers):
- if i in original_weights:
- weight_diff = torch.sum(torch.abs(layer.weight.data - original_weights[i])).item()
- if weight_diff > 1e-5:
- logger.error(f"Weights for layer {i} changed before pruning! Diff: {weight_diff}")
-
- # Test this specific strategy
- strategy_results = test_pruning_strategies(
- network_copy,
- dataset,
- pruning_percents=pruning_percents,
- device=device,
- strategy=strategy
- )
-
- # Store results for this strategy
- all_results[strategy] = strategy_results
-
- # Print results
- logger.info(f"{strategy} pruning results:")
- for i, percent in enumerate(pruning_percents):
- acc = strategy_results[f"{strategy}_acc"][i]
- loss = strategy_results[f"{strategy}_loss"][i]
- logger.info(f" {percent*100:.0f}% pruned: Accuracy = {acc:.2f}%, Loss = {loss:.4f}")
-
- # Plot results
- plt.figure(figsize=(12, 6))
- for strategy in strategies:
- results = all_results[strategy]
- plt.plot(
- [p*100 for p in pruning_percents],
- results[f"{strategy}_acc"],
- marker="o" if strategy == "high_rq" else ("s" if strategy == "low_rq" else "^"),
- label=f"{strategy}",
- linewidth=2
- )
-
- plt.xlabel("Pruning Percentage", fontsize=12)
- plt.ylabel("Accuracy (%)", fontsize=12)
- plt.title("Pruning Strategies Comparison", fontsize=14)
- plt.grid(True, linestyle="--", alpha=0.7)
- plt.legend(fontsize=11)
- plt.tight_layout()
-
- # Save plot
- plt.savefig("pruning_strategies_debug.png", dpi=300, bbox_inches="tight")
- logger.info("Saved plot to pruning_strategies_debug.png")
-
- # Check if strategies are different
- if len(strategies) > 1:
- are_different = False
- for i in range(len(strategies)-1):
- for j in range(i+1, len(strategies)):
- s1, s2 = strategies[i], strategies[j]
- accs1 = all_results[s1][f"{s1}_acc"]
- accs2 = all_results[s2][f"{s2}_acc"]
-
- # Calculate difference
- diff = np.mean(np.abs(np.array(accs1) - np.array(accs2)))
- logger.info(f"Mean accuracy difference between {s1} and {s2}: {diff:.2f}%")
-
- if diff > 1.0: # Consider different if average difference > 1%
- are_different = True
-
- if are_different:
- logger.info("SUCCESS: Pruning strategies show different results!")
- else:
- logger.error("FAILURE: Pruning strategies show very similar results.")
-
-def train_network(network, dataset, device, num_epochs=3):
- """Train a network for a specified number of epochs"""
- logger.info(f"Training a single network")
- optimizer = torch.optim.Adam(network.parameters(), lr=0.001)
-
- # Training loop
- for epoch in range(num_epochs):
- running_loss = 0.0
- correct = 0
- total = 0
-
- # Create DataLoader for training
- train_loader = torch.utils.data.DataLoader(
- dataset.train_dataset, batch_size=1024, shuffle=True
- )
-
- # Train for one epoch
- network.train()
- for i, (inputs, targets) in enumerate(train_loader):
- # Move data to device
- inputs, targets = inputs.to(device), targets.to(device)
-
- # Zero the parameter gradients
- optimizer.zero_grad()
-
- # Forward pass
- outputs = network(inputs)
- loss = torch.nn.functional.cross_entropy(outputs, targets)
-
- # Backward pass and optimize
- loss.backward()
- optimizer.step()
-
- # Calculate accuracy
- _, predicted = outputs.max(1)
- total += targets.size(0)
- correct += predicted.eq(targets).sum().item()
-
- # Update running loss
- running_loss += loss.item()
-
- # Print epoch statistics
- epoch_loss = running_loss / len(train_loader)
- epoch_acc = 100. * correct / total
- logger.info(f"Epoch {epoch+1}/{num_epochs}: Loss={epoch_loss:.4f}, Acc={epoch_acc:.2f}%")
-
- # Evaluate on test set
- test_acc, test_loss = dataset.evaluate(network, device)
- logger.info(f"Epoch {epoch+1}/{num_epochs}: Train Loss={epoch_loss:.4f}, Train Acc={epoch_acc:.2f}%, Test Loss={test_loss:.4f}, Test Acc={test_acc:.2f}%")
-
-def print_layer_stats(network):
- """Print statistics about each alignment layer in the network"""
- for i, layer in enumerate(network.alignment_layers):
- if hasattr(layer, "weight") and layer.weight is not None:
- weights = layer.weight.data
- total_weights = weights.numel()
- zero_weights = (weights == 0).sum().item()
- zero_percent = 100.0 * zero_weights / total_weights
-
- # Count pruned neurons (rows with all zeros)
- pruned_neurons = 0
- total_neurons = weights.size(0)
- for j in range(total_neurons):
- if torch.all(weights[j] == 0):
- pruned_neurons += 1
-
- logger.info(f"Layer {i}: Shape {weights.shape}, zeros: {zero_weights}/{total_weights} ({zero_percent:.2f}%), pruned neurons: {pruned_neurons}/{total_neurons} ({100.0 * pruned_neurons / total_neurons:.2f}%)")
-
-if __name__ == "__main__":
- main()
\ No newline at end of file
diff --git a/_arxiv/_archive/md_files/DOCUMENTATION.md b/_arxiv/_archive/md_files/DOCUMENTATION.md
deleted file mode 100644
index 01a39254..00000000
--- a/_arxiv/_archive/md_files/DOCUMENTATION.md
+++ /dev/null
@@ -1,52 +0,0 @@
-# Network Alignment Analysis Documentation
-
-This document provides a comprehensive index to all documentation for the Network Alignment Analysis codebase.
-
-## Getting Started
-
-- [README](README.md) - Main repository README with setup instructions
-- [Usage Guide](doc/usage.md) - How to use the codebase
-- [Installation Guide](README.md#installation) - How to install the package
-
-## Core Documentation
-
-- [Main Documentation](doc/DOCUMENTATION.md) - Overview of codebase architecture and capabilities
-- [API Reference](doc/api/README.md) - Comprehensive API reference
-- [Metrics System](doc/metrics/README.md) - Documentation for the metrics system
-- [Experiment Framework](doc/experiment/README.md) - Documentation for the experiment framework
-- [Performance Optimizations](doc/performance/README.md) - Documentation for performance optimizations
-- [Tensorized Implementations](doc/tensorized/README.md) - Documentation for tensorized implementations
-
-## Guides and References
-
-- [Pruning Modes](doc/pruning_modes.md) - Documentation for different pruning strategies
-- [Background](doc/background.md) - Background information on alignment analysis
-- [Development Roadmap](doc/ROADMAP.md) - Future development plans
-
-## Refactoring and Cleanup
-
-- [Metrics Refactoring Summary](METRICS_REFACTORING_SUMMARY.md) - Summary of metrics system refactoring
-- [Codebase Cleanup Summary](CODEBASE_CLEANUP_SUMMARY.md) - Summary of codebase cleanup
-
-## Directory Structure
-
-The codebase is organized into the following directories:
-
-- `src/alignment/`: Core source code implementing alignment metrics and algorithms
- - `metrics.py`: Implementation of all alignment metrics
- - `utils/`: Utility functions for data handling, visualization, etc.
- - `experiment/`: Experiment framework for running alignment experiments
- - `networks/`: Network architectures and training utilities
- - `pruning/`: Implementation of various pruning strategies
-
-- `tests/`: Unit and integration tests
-- `scripts/`: Utility scripts for running experiments and analysis
-- `benchmarks/`: Performance evaluation scripts
-- `configs/`: Configuration files for experiments
-- `results/`: Output directory for experiment results
-- `doc/`: Documentation
-
-## Additional Resources
-
-- [CHANGELOG.md](CHANGELOG.md) - History of changes to the codebase
-- [LICENSE](LICENSE) - MIT License for the project
\ No newline at end of file
diff --git a/_arxiv/_archive/md_files/README_tensorized_dropout.md b/_arxiv/_archive/md_files/README_tensorized_dropout.md
deleted file mode 100644
index 4daf37ff..00000000
--- a/_arxiv/_archive/md_files/README_tensorized_dropout.md
+++ /dev/null
@@ -1,93 +0,0 @@
-# Tensorized Dropout and Training Documentation
-
-This document explains the tensorized implementations for training and dropout operations in the alignment project.
-
-## Tensorized Training Methods
-
-The project now includes optimized tensorized training methods for efficiently training multiple networks in parallel. This is particularly useful for experiments that require training multiple networks with the same architecture but different initializations.
-
-### Available Training Methods
-
-1. **Sequential Training** (`sequential`): Original method that trains each network one at a time.
-2. **Tensorized Training** (`tensorized`): Trains networks in parallel by batching their training steps.
-3. **Fully Tensorized Training** (`fully_tensorized`): Most efficient method that combines networks into a single model ensemble.
-4. **Auto-select** (`auto`): Automatically selects the most efficient method based on the number of networks and their architectures.
-
-### Configuration
-
-You can specify which training method to use in your experiment configuration:
-
-```yaml
-extra:
- training_method: "fully_tensorized" # Options: "auto", "sequential", "tensorized", "fully_tensorized"
-```
-
-By default, the system will use `"auto"` which automatically selects the most efficient method.
-
-### Example Config
-
-An example configuration file is provided at `configs/training_example.yaml` showing how to use the tensorized training methods.
-
-### Performance
-
-Our benchmark tests show that the fully tensorized approach can be up to 3x faster than sequential training for larger numbers of networks (e.g., 10+ networks). The speed improvement comes from:
-
-- Reduced overhead from parallel computation
-- Better utilization of GPU resources
-- Minimized data transfer between CPU and GPU
-
-## Tensorized Progressive Dropout
-
-This document describes the improvements made to the progressive dropout implementation in the alignment codebase, along with benchmark results comparing different approaches.
-
-### Overview
-
-Progressive dropout is a technique used to analyze the alignment of neural networks by gradually dropping out neurons and measuring the impact on performance. Three major optimizations are now available:
-
-1. **Tensorized Implementation**: Process all networks at once using tensor operations
-2. **Multi-Strategy Implementation**: Process all strategies (high_rq, low_rq, random) simultaneously
-3. **Combined Optimization**: Apply both optimizations together for maximum performance
-
-### Multi-Strategy Dropout
-
-The multi-strategy optimization allows processing all three pruning strategies (high_rq, low_rq, random) together, which provides a significant speedup when you need results for all strategies:
-
-#### Configuration
-
-```yaml
-extra:
- use_multi_strategy_dropout: true # Process all strategies simultaneously
-```
-
-#### Implementation Details
-
-The implementation creates separate network copies for each strategy, but computes the neuron scores only once. Then it applies different pruning strategies to each network copy and evaluates them in parallel.
-
-#### Benchmark Results
-
-Benchmark results for multi-strategy dropout show significant speedups (typically 2.5-3x) compared to running the three strategies sequentially.
-
-| Configuration | Sequential Time | Multi-Strategy Time | Speedup |
-|---------------|----------------|---------------------|---------|
-| 5 networks, 10 dropout steps | 240s | 89s | 2.7x |
-
-### Benchmark Scripts
-
-Two benchmark scripts are provided to test the performance of these optimizations:
-
-1. `benchmark_network_training.py`: Compare training methods (sequential, tensorized, fully_tensorized)
-2. `benchmark_dropout_strategies.py`: Compare dropout strategy processing (sequential vs. multi-strategy)
-
-To run the benchmarks:
-
-```bash
-python benchmark_network_training.py --num_networks 10 --epochs 1
-python benchmark_dropout_strategies.py --config configs/config_alignment_experiment.yaml --runs 1
-```
-
-### Conclusion
-
-These optimizations provide substantial performance improvements for alignment experiments, allowing you to run more experiments with larger networks in less time. For the best performance:
-
-1. Use `training_method: "fully_tensorized"` for training multiple networks
-2. Use `use_multi_strategy_dropout: true` when running progressive dropout with multiple strategies
\ No newline at end of file
diff --git a/_arxiv/_archive/md_files/REFACTORING_METRICS.md b/_arxiv/_archive/md_files/REFACTORING_METRICS.md
deleted file mode 100644
index daa2d165..00000000
--- a/_arxiv/_archive/md_files/REFACTORING_METRICS.md
+++ /dev/null
@@ -1,66 +0,0 @@
-# Metrics System Refactoring
-
-## Overview
-
-This document summarizes the refactoring of the alignment metrics system, which has been fully consolidated into a single, unified metrics infrastructure under `src/alignment/metrics.py`.
-
-## Completed Refactoring Tasks
-
-- [x] Critical Bug Fixes & Import Issues
-- [x] Import Hoisting (general pass)
-- [x] Logging in DDP Runner
-- [x] Optional transformers
-- [x] Plotting for Metric Evolution & Conditional Legends
-- [x] Configuration Access in ExperimentRunner
-- [x] Refactor Pruning Logic in dropout_manager.py
-- [x] Consolidated Metrics Systems (metrics_utils.py vs. metrics.py)
- - [x] Ported WeightSimilarityMetric to weight_cosine_similarity, weight_dot_similarity, and weight_euclidean_distance
- - [x] Ported MIMetric to mi_proj_vs_mean_input
- - [x] Ported RQMetric alternative formulation to rq_alt_denom
- - [x] Ported NodeRedundancyMetric to node_redundancy
- - [x] Updated tests to use the new metrics.py system (get_metric)
- - [x] Deprecated and removed metrics_utils.py
-
-## Architecture Overview
-
-The new metrics system is built around:
-
-1. **Central Registry**: `ALIGNMENT_METRICS_REGISTRY` maps metric names to metric functions
-2. **Metric Protocol**: `AlignmentMetric` protocol defines the interface for all metrics
-3. **Metric Implementation**: `_AlignmentMetricImpl` provides consistent dispatch for all metrics
-4. **Factory Function**: `get_metric()` creates properly configured metric objects
-5. **High-Level Functions**:
- - `compute_metrics_for_layers()`: Process multiple layers at once
- - `compute_all_node_scores()`: Process the whole network
- - `compute_pairwise_metric()`: Compute metrics between pairs of data
-
-## Documentation
-
-A comprehensive metrics system documentation has been added at `src/alignment/README_metrics.md`, which includes:
-
-- Available metrics and their purposes
-- Usage examples for each metric type
-- Architecture overview
-- Guidelines for metric selection
-
-## Testing
-
-All metrics have corresponding tests in `tests/test_benchmark.py` which verify the correct functionality of:
-
-- Standard RQ and alternative RQ metrics
-- MI metrics including the specialized MI projection metric
-- Weight similarity metrics
-- Node redundancy metrics
-
-## Future Improvements
-
-Potential areas for future enhancement:
-
-1. Add more test coverage for the remaining metrics
-2. Consider adding more metrics for specific use cases
-3. Performance optimizations for large-scale networks
-4. Integration with visualization tools for better insights
-
-## Conclusion
-
-The metrics system refactoring has successfully unified all alignment metrics under a single, consistent API. This will make the codebase more maintainable, easier to extend, and provide clearer documentation for users.
\ No newline at end of file
diff --git a/_arxiv/archive/benchmarks/README.md b/_arxiv/archive/benchmarks/README.md
deleted file mode 100644
index 5e7498d0..00000000
--- a/_arxiv/archive/benchmarks/README.md
+++ /dev/null
@@ -1,29 +0,0 @@
-# Alignment Benchmarks
-
-This directory contains benchmarking scripts for the Alignment library. These scripts are designed to assess the performance of various components of the codebase.
-
-## Contents
-
-- **benchmark_dropout_strategies.py**: Compares the performance of different dropout strategies (sequential vs. multi-strategy).
-- **benchmark_network_training.py**: Benchmarks network training with different configurations.
-
-## Usage
-
-Most benchmark scripts can be run directly from the command line:
-
-```bash
-python benchmark_dropout_strategies.py --config configs/config_alignment_experiment.yaml
-```
-
-## Adding New Benchmarks
-
-When adding new benchmarks, please follow these guidelines:
-
-1. Include clear documentation within the script about what is being measured
-2. Add command-line arguments for configuration
-3. Include a simple way to output/visualize the benchmark results
-4. Update this README with information about the new benchmark
-
-## Results
-
-Benchmark results should be saved to the `results/` directory, not committed directly to the repository unless they represent an important baseline.
\ No newline at end of file
diff --git a/_arxiv/archive/benchmarks/benchmark_dropout_strategies.py b/_arxiv/archive/benchmarks/benchmark_dropout_strategies.py
deleted file mode 100755
index 392afece..00000000
--- a/_arxiv/archive/benchmarks/benchmark_dropout_strategies.py
+++ /dev/null
@@ -1,178 +0,0 @@
-#!/usr/bin/env python
-"""
-Benchmark script to compare original sequential dropout approach vs. multi-strategy approach.
-
-This script measures the performance difference between:
-1. Processing strategies (high_rq, low_rq, random) sequentially
-2. Processing all strategies simultaneously with the new tensorized multi-strategy approach
-"""
-
-import os
-import sys
-import time
-import argparse
-import logging
-import torch
-import numpy as np
-from tqdm import tqdm
-
-from alignment.experiments.alignment_experiments import AlignmentExperiment
-from alignment.config import ExperimentConfig
-from alignment.datasets import load_dataset
-from alignment.dropout import progressive_dropout, progressive_dropout_multi_strategy
-
-# Configure logging
-logging.basicConfig(
- level=logging.INFO,
- format='%(asctime)s [%(levelname)s] %(message)s',
- handlers=[
- logging.StreamHandler()
- ]
-)
-logger = logging.getLogger(__name__)
-
-def run_benchmark(config_path, num_runs=3):
- """Run benchmark comparing sequential vs. multi-strategy dropout."""
- # Load configuration
- logger.info(f"Loading configuration from {config_path}")
- config = ExperimentConfig.load(config_path)
-
- # Force the device to be cuda if available
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
- config.device = str(device)
-
- # Create experiment
- experiment = AlignmentExperiment(config)
-
- # Create networks
- logger.info("Creating networks...")
- networks = experiment.create_networks()
- logger.info(f"Created {len(networks)} networks")
-
- # Load dataset
- logger.info("Loading dataset...")
- dataset = load_dataset(config.dataset)
-
- # Setup dropout parameters
- dropout_fractions = np.linspace(
- config.alignment.dropout_min,
- config.alignment.dropout_max,
- config.alignment.dropout_steps
- ).tolist()
-
- pruning_mode = getattr(config.extra, "dropout_pruning_mode", "global_joint")
- dropout_mode = getattr(config.extra, "dropout_mode", "scaled")
- metric = experiment.metric
-
- # Initialize timing results
- sequential_times = []
- multi_strategy_times = []
-
- # Run benchmarks multiple times
- for run in range(num_runs):
- logger.info(f"\nRun {run+1}/{num_runs}")
-
- # Make copies of networks to ensure fair comparison
- networks_sequential = [net.clone() if hasattr(net, 'clone') else net for net in networks]
- networks_multi = [net.clone() if hasattr(net, 'clone') else net for net in networks]
-
- # Process strategies sequentially (original approach)
- logger.info("Running sequential approach...")
- sequential_start = time.time()
-
- strategies = ["high_rq", "low_rq", "random"]
- for strategy in tqdm(strategies, desc="Strategies"):
- # Clone networks for this strategy to avoid interference
- strategy_networks = [net.clone() if hasattr(net, 'clone') else net for net in networks_sequential]
-
- # Run progressive dropout with this strategy
- network_accuracies, network_losses = progressive_dropout(
- strategy_networks,
- dataset,
- dropout_fractions,
- metric,
- device,
- pruning_mode=pruning_mode,
- dropout_mode=dropout_mode,
- strategy=strategy,
- show_progress=False # Disable progress bars for cleaner output
- )
-
- sequential_time = time.time() - sequential_start
- sequential_times.append(sequential_time)
- logger.info(f"Sequential approach took {sequential_time:.2f} seconds")
-
- # Process all strategies at once (new approach)
- logger.info("Running multi-strategy approach...")
- multi_start = time.time()
-
- # Run with multi-strategy mode
- network_accuracies, network_losses = progressive_dropout(
- networks_multi,
- dataset,
- dropout_fractions,
- metric,
- device,
- pruning_mode=pruning_mode,
- dropout_mode=dropout_mode,
- show_progress=False, # Disable progress bars for cleaner output
- use_multi_strategy=True
- )
-
- multi_time = time.time() - multi_start
- multi_strategy_times.append(multi_time)
- logger.info(f"Multi-strategy approach took {multi_time:.2f} seconds")
-
- # Calculate speedup
- speedup = sequential_time / multi_time if multi_time > 0 else float('inf')
- logger.info(f"Speedup: {speedup:.2f}x")
-
- # Calculate average times and speedup
- avg_sequential = np.mean(sequential_times)
- avg_multi = np.mean(multi_strategy_times)
- avg_speedup = avg_sequential / avg_multi if avg_multi > 0 else float('inf')
-
- std_sequential = np.std(sequential_times)
- std_multi = np.std(multi_strategy_times)
-
- # Print summary
- logger.info("\n" + "="*50)
- logger.info("BENCHMARK SUMMARY")
- logger.info("="*50)
- logger.info(f"Network count: {len(networks)}")
- logger.info(f"Dropout steps: {len(dropout_fractions)}")
- logger.info(f"Pruning mode: {pruning_mode}")
- logger.info(f"Dropout mode: {dropout_mode}")
- logger.info(f"Device: {device}")
- logger.info("-"*50)
- logger.info(f"Sequential approach: {avg_sequential:.2f} ± {std_sequential:.2f} seconds")
- logger.info(f"Multi-strategy approach: {avg_multi:.2f} ± {std_multi:.2f} seconds")
- logger.info(f"Average speedup: {avg_speedup:.2f}x")
- logger.info("="*50)
-
- return {
- "sequential": sequential_times,
- "multi_strategy": multi_strategy_times,
- "speedup": avg_speedup
- }
-
-def main():
- parser = argparse.ArgumentParser(description="Benchmark dropout strategies")
- parser.add_argument("--config", type=str, default="configs/config_alignment_experiment.yaml",
- help="Path to config file")
- parser.add_argument("--runs", type=int, default=1,
- help="Number of benchmark runs")
-
- args = parser.parse_args()
-
- # Run the benchmark
- results = run_benchmark(args.config, args.runs)
-
- # Print the recommendation
- if results["speedup"] > 1.5:
- print("\nRECOMMENDATION: Use multi-strategy approach for significant speedup.")
- else:
- print("\nRECOMMENDATION: Both approaches have similar performance in this configuration.")
-
-if __name__ == "__main__":
- main()
\ No newline at end of file
diff --git a/_arxiv/archive/benchmarks/benchmark_network_training.py b/_arxiv/archive/benchmarks/benchmark_network_training.py
deleted file mode 100755
index 72d1b49c..00000000
--- a/_arxiv/archive/benchmarks/benchmark_network_training.py
+++ /dev/null
@@ -1,239 +0,0 @@
-#!/usr/bin/env python
-"""
-Benchmark script for comparing network training methods.
-
-This script benchmarks different training methods for multiple networks:
-1. Sequential training (original approach)
-2. Tensorized training (improved parallel approach)
-3. Fully tensorized training (optimized ensemble approach)
-"""
-
-import os
-import time
-import argparse
-import logging
-import torch
-import torch.nn as nn
-from tqdm import tqdm
-
-from alignment.metrics import get_metric
-from alignment.datasets import load_dataset
-from alignment.models.models import MLP, create_mlp
-from alignment.models.base import AlignmentNetwork
-from alignment.training import (
- train_networks_sequential,
- train_networks_tensorized,
- train_networks_fully_tensorized,
- train_networks
-)
-
-# Configure logging
-logging.basicConfig(
- level=logging.INFO,
- format='%(asctime)s [%(levelname)s] %(message)s',
- handlers=[
- logging.StreamHandler()
- ]
-)
-logger = logging.getLogger(__name__)
-
-def create_networks(num_networks, input_size=784, hidden_sizes=[512, 256], output_size=10, seed=42):
- """Create multiple networks with identical architecture but different initializations."""
- networks = []
-
- for i in range(num_networks):
- # Set seed for reproducibility, but different for each network
- torch.manual_seed(seed + i)
-
- # Create a simple MLP network
- base_model = MLP(
- input_dim=input_size,
- output_dim=output_size,
- num_hidden=hidden_sizes,
- dropout_rate=0.0
- )
-
- # Get all linear layers for alignment
- linear_layers = {}
- for name, module in base_model.named_modules():
- if isinstance(module, nn.Linear):
- if name != "layers.0": # Skip the first layer
- linear_layers[name] = None # Use the layer's own input
-
- # Create AlignmentNetwork wrapper
- network = AlignmentNetwork(base_model=base_model, alignment_layer_names=linear_layers)
-
- # Move to device (will be moved again in training, but good practice)
- if torch.cuda.is_available():
- network.to('cuda')
-
- networks.append(network)
-
- logger.info(f"Created {num_networks} networks with architecture: {input_size} -> {hidden_sizes} -> {output_size}")
- return networks
-
-def run_benchmark(networks, dataset, num_epochs=2, learning_rate=0.001, device=None, show_progress=True):
- """Run benchmark comparing different training methods."""
- if device is None:
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
- logger.info(f"Running benchmark on device: {device}")
-
- # Define methods to benchmark
- methods = [
- ("Sequential", train_networks_sequential),
- ("Tensorized", train_networks_tensorized),
- ("Fully Tensorized", train_networks_fully_tensorized),
- ("Auto-select", train_networks)
- ]
-
- results = {}
-
- # Clone networks for each method to ensure fair comparison
- for method_name, method_func in methods:
- logger.info(f"Benchmarking {method_name} training...")
-
- # Create copies of networks for this method
- method_networks = []
-
- # For each network in the original list
- for i, source_net in enumerate(networks):
- # Get the architecture of the source network
- input_dim = 784 # MNIST standard
- output_dim = 10 # MNIST standard
-
- # Extract hidden sizes by examining the linear layers
- linear_layers = [
- module for module in source_net.base_model.modules()
- if isinstance(module, nn.Linear)
- ]
-
- # Get all but the last layer's output dimensions for hidden sizes
- hidden_sizes = [layer.out_features for layer in linear_layers[:-1]]
-
- logger.debug(f"Network {i}: input_dim={input_dim}, hidden_sizes={hidden_sizes}, output_dim={output_dim}")
-
- # Create base model with identical architecture
- base_model = MLP(
- input_dim=input_dim,
- output_dim=output_dim,
- num_hidden=hidden_sizes,
- dropout_rate=0.0
- )
-
- # Create alignment layer mapping identical to source network
- # Use the layer_to_input_names attribute which is the internal representation
- layer_to_input_names = source_net.layer_to_input_names
-
- # Create new network with same architecture and alignment layers
- target_net = AlignmentNetwork(
- base_model=base_model,
- alignment_layer_names=layer_to_input_names
- )
-
- # Copy parameters from source to target
- # This would be the most reliable way to ensure identical networks
- target_net.load_state_dict(source_net.state_dict())
-
- # Move to device
- target_net.to(device)
-
- # Add to method networks
- method_networks.append(target_net)
-
- # Measure time
- start_time = time.time()
-
- # Train networks using this method
- training_history = method_func(
- networks=method_networks,
- dataset=dataset,
- num_epochs=num_epochs,
- learning_rate=learning_rate,
- device=device,
- show_progress=show_progress
- )
-
- # Calculate elapsed time
- elapsed_time = time.time() - start_time
-
- # Store results
- results[method_name] = {
- "time": elapsed_time,
- "history": training_history,
- "final_acc": training_history["test_acc"][-1] if training_history["test_acc"] else 0.0
- }
-
- logger.info(f"{method_name} training completed in {elapsed_time:.2f} seconds (final acc: {results[method_name]['final_acc']:.2f}%)")
-
- # Print summary
- print("\nBenchmark Results:")
- print("==================")
- print(f"Networks: {len(networks)}, Epochs: {num_epochs}")
- print("------------------")
-
- # Find the fastest method to calculate speedup
- fastest_time = min(results[method]["time"] for method in results)
-
- for method_name in results:
- elapsed_time = results[method_name]["time"]
- final_acc = results[method_name]["final_acc"]
- speedup = fastest_time / elapsed_time if elapsed_time > 0 else 0
-
- print(f"{method_name:16s}: {elapsed_time:.2f}s ({speedup:.2f}x), Acc: {final_acc:.2f}%")
-
- return results
-
-def main():
- parser = argparse.ArgumentParser(description="Benchmark network training methods.")
- parser.add_argument("--num_networks", type=int, default=5, help="Number of networks to train")
- parser.add_argument("--hidden_sizes", type=str, default="512,256", help="Hidden layer sizes (comma-separated)")
- parser.add_argument("--epochs", type=int, default=2, help="Number of epochs to train")
- parser.add_argument("--device", type=str, default="cuda", help="Device to run benchmark on")
- parser.add_argument("--batch_size", type=int, default=128, help="Batch size for training")
- parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility")
-
- args = parser.parse_args()
-
- # Parse hidden sizes
- hidden_sizes = [int(size) for size in args.hidden_sizes.split(",")]
-
- # Set device
- device = torch.device(args.device if torch.cuda.is_available() and args.device == "cuda" else "cpu")
-
- # Set random seed
- torch.manual_seed(args.seed)
-
- # Create networks
- networks = create_networks(
- num_networks=args.num_networks,
- hidden_sizes=hidden_sizes,
- seed=args.seed
- )
-
- # Load dataset
- logger.info("Loading MNIST dataset")
- transform_params = {
- "center_crop": None,
- "resize": None,
- "flatten": True, # Flatten for MLP
- "normalize": False
- }
- dataset_config = {
- "dataset_name": "MNIST",
- "batch_size": args.batch_size,
- "data_path": "./data",
- "transform_params": transform_params
- }
- dataset = load_dataset(dataset_config)
-
- # Run benchmark
- run_benchmark(
- networks=networks,
- dataset=dataset,
- num_epochs=args.epochs,
- device=device
- )
-
-if __name__ == "__main__":
- main()
\ No newline at end of file
diff --git a/_arxiv/archive/cluster/alignment-ddp-example.slurm b/_arxiv/archive/cluster/alignment-ddp-example.slurm
deleted file mode 100644
index 13f6226c..00000000
--- a/_arxiv/archive/cluster/alignment-ddp-example.slurm
+++ /dev/null
@@ -1,157 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=alignment-ddp # Job name
-#SBATCH --partition=kempner # Partition with GPUs
-#SBATCH --nodes=2 # Number of nodes
-#SBATCH --ntasks-per-node=4 # Tasks per node (match GPUs)
-#SBATCH --cpus-per-task=16 # CPUs per task
-#SBATCH --mem=512G # Memory per node
-#SBATCH --gres=gpu:4 # GPUs per node
-#SBATCH --time=24:00:00 # Time limit
-#SBATCH --mail-type=begin,end # Email notifications
-#SBATCH --account=kempner_lab # Account to charge
-
-# Set up job directory
-JOB_NAME="alignment-experiment"
-source cluster/slurm_settings.txt
-JOB_DIR=${JOB_FOLDER}/${JOB_NAME}-${SLURM_JOB_ID}
-mkdir -p $JOB_DIR
-
-# Set up logging
-logfile="${JOB_DIR}/slurm-${SLURM_JOB_ID}.out"
-echo "Job started at $(date)" > $logfile
-echo "Running on nodes: $SLURM_JOB_NODELIST" >> $logfile
-
-# Set up DDP environment variables
-export MASTER_PORT=12355
-export WORLD_SIZE=$(($SLURM_NNODES * $SLURM_NTASKS_PER_NODE))
-master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
-export MASTER_ADDR=$master_addr
-
-echo "WORLD_SIZE=$WORLD_SIZE" >> $logfile
-echo "MASTER_ADDR=$MASTER_ADDR" >> $logfile
-echo "MASTER_PORT=$MASTER_PORT" >> $logfile
-
-# Load modules and activate environment
-module purge
-module load python/3.9.12
-module load cuda/11.7
-
-# Activate conda environment
-conda activate alignment_env
-
-# Record start time
-start_time=$(date +%s)
-
-# Create Python script for the experiment
-cat > ${JOB_DIR}/run_experiment.py << 'EOF'
-import os
-import torch
-import torch.distributed as dist
-from pathlib import Path
-
-# Import from alignment codebase
-from alignment import ModelWrapper, DatasetWrapper
-from alignment.experiments import ProgressiveDropoutExperiment
-from alignment.experiments.base import ExperimentConfig
-from alignment.infrastructure.computing.distributed import setup_distributed, cleanup_distributed
-from alignment.infrastructure.storage.logging import setup_logging
-
-def main():
- # Setup distributed environment
- setup_distributed()
-
- # Get rank and world size
- rank = dist.get_rank() if dist.is_initialized() else 0
- world_size = dist.get_world_size() if dist.is_initialized() else 1
-
- # Setup logging (only on main process)
- if rank == 0:
- setup_logging(Path(os.environ.get('JOB_DIR', '.')))
-
- # Configuration
- config = ExperimentConfig(
- name=f"resnet50_imagenet_ddp",
- description="Progressive dropout analysis on ResNet50",
-
- # Model settings
- model_name="resnet50",
- pretrained=True,
-
- # Dataset settings
- dataset_name="imagenet",
- data_path="/n/holylabs/LABS/kempner_shared/datasets/imagenet",
- batch_size=256 // world_size, # Scale batch size per GPU
- num_workers=8,
-
- # Experiment settings
- metrics=["rayleigh_quotient", "mi_gaussian", "pid_shared"],
- tracked_layers=[
- "layer1.0.conv1", "layer1.1.conv1", "layer1.2.conv1",
- "layer2.0.conv1", "layer2.1.conv1", "layer2.2.conv1",
- "layer3.0.conv1", "layer3.1.conv1", "layer3.2.conv1",
- "layer4.0.conv1", "layer4.1.conv1", "layer4.2.conv1",
- ],
-
- # Training settings
- train_before_dropout=True,
- training_epochs=90,
- learning_rate=0.1 * (world_size * 256 / 256), # Linear LR scaling
- optimizer="sgd",
-
- # Distributed settings
- distributed=True,
- world_size=world_size,
- rank=rank,
-
- # Other settings
- checkpoint_dir=f"{os.environ.get('JOB_DIR', '.')}/checkpoints",
- log_dir=f"{os.environ.get('JOB_DIR', '.')}/logs",
- exclude_classification_layer=True,
- scale_by_norm=False,
- force_cpu_for_large_metric_ops=True,
- )
-
- # Create experiment
- experiment = ProgressiveDropoutExperiment(
- config=config,
- dropout_range=(0.0, 0.9, 10), # 10 dropout levels from 0% to 90%
- dropout_mode="magnitude", # Drop based on metric magnitude
- compute_metrics_during_training=True,
- metric_computation_interval=10, # Compute metrics every 10 epochs
- )
-
- # Run experiment
- print(f"Rank {rank}: Starting experiment...")
- results = experiment.run()
-
- # Save results (only on main process)
- if rank == 0:
- experiment.save_results()
- print("Experiment completed successfully!")
-
- # Cleanup
- cleanup_distributed()
-
-if __name__ == "__main__":
- main()
-EOF
-
-# Run the experiment using srun
-# Each process will be launched on the appropriate node/GPU
-echo "Starting distributed training..." >> $logfile
-srun python ${JOB_DIR}/run_experiment.py >> $logfile 2>&1
-
-# Record end time
-end_time=$(date +%s)
-total_time=$((end_time - start_time))
-
-echo "Job completed at $(date)" >> $logfile
-echo "Total runtime: $total_time seconds" >> $logfile
-
-# Copy results to a permanent location (only from rank 0 node)
-if [[ $SLURM_PROCID -eq 0 ]]; then
- RESULTS_DIR="/n/holylabs/LABS/kempner_lab/Users/$USER/alignment_results/${SLURM_JOB_ID}"
- mkdir -p $RESULTS_DIR
- cp -r ${JOB_DIR}/* $RESULTS_DIR/
- echo "Results copied to: $RESULTS_DIR" >> $logfile
-fi
\ No newline at end of file
diff --git a/_arxiv/archive/cluster/alignment_stats.slurm b/_arxiv/archive/cluster/alignment_stats.slurm
deleted file mode 100644
index 5b3d2ad4..00000000
--- a/_arxiv/archive/cluster/alignment_stats.slurm
+++ /dev/null
@@ -1,42 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=alignment_stats # create a short name for your job
-#SBATCH --partition=kempner # partition
-#SBATCH --account=kempner_bsabatini_lab # account needed for kempner partition
-#SBATCH --nodes=1 # node count
-#SBATCH --ntasks-per-node=1 # total number of tasks per node
-#SBATCH --cpus-per-task=64 # cpu-cores per task (>1 if multi-threaded tasks)
-#SBATCH --gres=gpu:1 # number of allocated gpus per node
-#SBATCH --mem=1000G # total memory per node (4 GB per cpu-core is default)
-#SBATCH --time=04:00:00 # total run time limit (HH:MM:SS)
-#SBATCH --mail-type=begin # send email when job begins
-#SBATCH --mail-type=end # send email when job ends
-
-# we need to define the job name directly since it isn't a slurm environment variable
-JOB_NAME="alignment_stats"
-
-# this text file has some settings in it (like the standard job directory)
-source cluster/slurm_settings.txt
-JOB_DIR=${JOB_FOLDER}/${JOB_NAME}-${SLURM_JOB_ID} # make a specific directory for this particular job
-mkdir -p $JOB_DIR
-
-# define a unique log file in the right place
-logfile="${JOB_DIR}/slurm-${SLURM_JOB_ID}.out"
-echo "Writing to ${logfile}"
-
-# load python and activate our conda environment
-module purge
-module load python
-conda activate networkAlignmentAnalysis
-
-# record the start time
-start_time=$(date +%s)
-
-# this is the command that initiates the processes
-python experiment.py alignment_stats --network CNN2P2 --dataset CIFAR100 --use_wandb --epochs 200 >> $logfile
-
-# record the end time
-end_time=$(date +%s)
-
-# measure the time elapsed for the core part of the job in the logfile
-total_time=$((end_time-start_time))
-echo "Total Time= "$total_time" seconds" >> $logfile
diff --git a/_arxiv/archive/cluster/ddp-example/README.md b/_arxiv/archive/cluster/ddp-example/README.md
deleted file mode 100644
index f559b6d8..00000000
--- a/_arxiv/archive/cluster/ddp-example/README.md
+++ /dev/null
@@ -1,13 +0,0 @@
-## Distributed Data Parallel Example on Slurm Cluster
-
-The contents of this folder contain a MWE for getting DDP to work on a cluster. If you want to try
-it yourself, start from the main networkAlignmentAnalysis folder (with the .gitignore etc.). The
-file called [ddp.slurm](ddp.slurm) is an ``sbatch`` script you can call with:
-
-```shell
-sbatch cluster/ddp-example/ddp.slurm
-```
-
-This prepares some environment variables, creates a job folder, then calls the python script
-[ddp.py](ddp.py) which downloads MNIST to the job folder, trains it across N GPUs (where N is
-equal to the number of nodes times the number of GPUs per node), and saves it if requested.
diff --git a/_arxiv/archive/cluster/ddp-example/ddp.py b/_arxiv/archive/cluster/ddp-example/ddp.py
deleted file mode 100644
index 9eb76ebb..00000000
--- a/_arxiv/archive/cluster/ddp-example/ddp.py
+++ /dev/null
@@ -1,218 +0,0 @@
-"""
-Massive thank you to the Princeton cluster engineers for providing this example (we've only made small changes)
-https://github.com/PrincetonUniversity/multi_gpu_training/blob/main/02_pytorch_ddp/mnist_classify_ddp.py
-"""
-
-import argparse
-import shutil
-import torch
-from torch import nn
-import torch.nn.functional as F
-import torch.optim as optim
-from torchvision import datasets, transforms
-from torch.optim.lr_scheduler import StepLR
-
-import os
-import time
-import torch.distributed as dist
-from torch.nn.parallel import DistributedDataParallel as DDP
-from socket import gethostname
-
-
-class Net(nn.Module):
- def __init__(self):
- super(Net, self).__init__()
- self.conv1 = nn.Conv2d(1, 32, 3, 1)
- self.conv2 = nn.Conv2d(32, 64, 3, 1)
- self.dropout1 = nn.Dropout(0.25)
- self.dropout2 = nn.Dropout(0.5)
- self.fc1 = nn.Linear(9216, 128)
- self.fc2 = nn.Linear(128, 10)
-
- def forward(self, x):
- x = self.conv1(x)
- x = F.relu(x)
- x = self.conv2(x)
- x = F.relu(x)
- x = F.max_pool2d(x, 2)
- x = self.dropout1(x)
- x = torch.flatten(x, 1)
- x = self.fc1(x)
- x = F.relu(x)
- x = self.dropout2(x)
- x = self.fc2(x)
- output = F.log_softmax(x, dim=1)
- return output
-
-
-def train(args, model, device, dataloader, datasampler, optimizer, epoch, rank):
- """basic training script"""
- # required for different shuffle order of examples in dataset each epoch
- datasampler.set_epoch(epoch)
-
- if rank == 0 and epoch == 1:
- first_batch_timer = time.time()
-
- model.train()
- for batch_idx, (data, target) in enumerate(dataloader):
- data, target = data.to(device), target.to(device)
- if rank == 0 and epoch == 1 and batch_idx == 0:
- print(f"Train-- epoch {epoch}, rank {rank}, first batch loaded in {time.time() - first_batch_timer} seconds.")
- optimizer.zero_grad()
- output = model(data)
- loss = F.nll_loss(output, target)
- loss.backward()
- optimizer.step()
- if batch_idx % args.log_interval == 0:
- if rank == 0:
- print(f"Train Epoch: {epoch} [{batch_idx}/{len(dataloader)} ({100.*batch_idx/len(dataloader):.0f}%)] \t Loss: {loss.item():.6f}")
- if args.dry_run:
- break
-
-
-def test(model, device, dataloader):
- model.eval()
- test_loss = 0
- correct = 0
- attempts = 0
- with torch.no_grad():
- for data, target in dataloader:
- data, target = data.to(device), target.to(device)
- output = model(data)
- test_loss += F.nll_loss(output, target, reduction="sum").item() # sum up batch loss
- pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
- correct += pred.eq(target.view_as(pred)).sum().item()
- attempts += data.size(0)
-
- test_loss /= attempts
-
- print(f"\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{attempts} ({100.*correct/attempts:.0f}%)\n")
-
-
-def setup(rank, world_size):
- # initialize the process group
- dist.init_process_group("nccl", rank=rank, world_size=world_size)
-
-
-def main():
- # Training settings
- parser = argparse.ArgumentParser(description="PyTorch MNIST Example with DDP")
- parser.add_argument(
- "--job-folder",
- type=str,
- default=".",
- help="job folder for storing data",
- )
- parser.add_argument(
- "--batch-size",
- type=int,
- default=64,
- metavar="N",
- help="input batch size for training (default: 64)",
- )
- parser.add_argument(
- "--epochs",
- type=int,
- default=14,
- metavar="N",
- help="number of epochs to train (default: 14)",
- )
- parser.add_argument("--lr", type=float, default=1.0, metavar="LR", help="learning rate (default: 1.0)")
- parser.add_argument(
- "--gamma",
- type=float,
- default=0.9,
- metavar="M",
- help="Learning rate step gamma (default: 0.9)",
- )
- parser.add_argument("--no-cuda", action="store_true", default=False, help="disables CUDA training")
- parser.add_argument("--dry-run", action="store_true", default=False, help="quickly check a single pass")
- parser.add_argument("--seed", type=int, default=1, metavar="S", help="random seed (default: 1)")
- parser.add_argument(
- "--log-interval",
- type=int,
- default=10,
- metavar="N",
- help="how many batches to wait before logging training status",
- )
- parser.add_argument("--save-model", action="store_true", default=False, help="For Saving the current Model")
- args = parser.parse_args()
-
- print("job folder:", args.job_folder)
- data_folder = os.path.join(args.job_folder, "data")
-
- torch.manual_seed(args.seed)
-
- world_size = int(os.environ["WORLD_SIZE"])
- rank = int(os.environ["SLURM_PROCID"])
- gpus_per_node = int(os.environ["SLURM_GPUS_ON_NODE"])
- print("gpus_per_node:", gpus_per_node)
- print("device_count:", torch.cuda.device_count())
-
- assert gpus_per_node == torch.cuda.device_count()
- print(
- f"Hello from rank {rank} of {world_size} on {gethostname()} where there are" f" {gpus_per_node} allocated GPUs per node.",
- flush=True,
- )
-
- if world_size > 1:
- setup(rank, world_size)
- if rank == 0:
- print(f"Group initialized? {dist.is_initialized()}", flush=True)
-
- local_rank = rank - gpus_per_node * (rank // gpus_per_node)
- torch.cuda.set_device(local_rank)
- print(f"host: {gethostname()}, rank: {rank}, local_rank: {local_rank}")
-
- # Create network
- net = Net()
- model = net.to(local_rank)
-
- # Make it a DDP object for distributed processing and training
- ddp_model = DDP(model, device_ids=[local_rank]) if world_size > 1 else model
- optimizer = optim.Adadelta(ddp_model.parameters(), lr=args.lr)
- scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
-
- transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
-
- train_data = datasets.MNIST(data_folder, train=True, download=True, transform=transform)
- test_data = datasets.MNIST(data_folder, train=False, download=True, transform=transform)
-
- train_sampler = torch.utils.data.distributed.DistributedSampler(train_data, num_replicas=world_size, rank=rank)
- train_loader = torch.utils.data.DataLoader(
- train_data,
- batch_size=args.batch_size,
- sampler=train_sampler,
- num_workers=int(os.environ["SLURM_CPUS_PER_TASK"]),
- pin_memory=True,
- )
- test_loader = torch.utils.data.DataLoader(
- test_data,
- batch_size=args.batch_size,
- num_workers=int(os.environ["SLURM_CPUS_PER_TASK"]),
- pin_memory=True,
- )
-
- for epoch in range(1, args.epochs + 1):
- if rank == 0:
- epoch_time = time.time()
- train(args, ddp_model, local_rank, train_loader, train_sampler, optimizer, epoch, rank)
- if rank == 0:
- test(ddp_model, local_rank, test_loader)
- scheduler.step()
- if rank == 0:
- epoch_time = time.time() - epoch_time
- print(f"\nEpoch {epoch}, Train & Test Time = {epoch_time:.1f} seconds (measured from rank {rank}).\n")
-
- if args.save_model and rank == 0:
- torch.save(model.state_dict(), os.path.join(args.job_folder, "test_model_ddp.pt"))
-
- if world_size > 1:
- dist.destroy_process_group()
-
- # clear locally downloaded MNIST data
- shutil.rmtree(data_folder)
-
-
-if __name__ == "__main__":
- main()
diff --git a/_arxiv/archive/cluster/ddp-example/ddp.slurm b/_arxiv/archive/cluster/ddp-example/ddp.slurm
deleted file mode 100644
index 43cdd746..00000000
--- a/_arxiv/archive/cluster/ddp-example/ddp.slurm
+++ /dev/null
@@ -1,56 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=ddp-example # create a short name for your job
-#SBATCH --partition=kempner # partition to use (need one with GPUs for this)
-#SBATCH --nodes=2 # node count
-#SBATCH --ntasks-per-node=4 # total number of tasks per node
-#SBATCH --cpus-per-task=16 # cpu-cores per task (>1 if multi-threaded tasks)
-#SBATCH --mem=512G # total memory per node (4 GB per cpu-core is default)
-#SBATCH --gres=gpu:4 # number of allocated gpus per node
-#SBATCH --time=00:10:00 # total run time limit (HH:MM:SS)
-#SBATCH --mail-type=begin # send email when job begins
-#SBATCH --mail-type=end # send email when job ends
-#SBATCH --account=kempner_bsabatini_lab # the account needs to be specified for the kempner partition
-
-# we need to define the job name directly since it isn't a slurm environment variable
-JOB_NAME="DDP-Example"
-
-# this text file has some settings in it (like the standard job directory)
-source cluster/slurm_settings.txt
-JOB_DIR=${JOB_FOLDER}/${JOB_NAME}-${SLURM_JOB_ID} # make a specific directory for this particular job
-mkdir -p $JOB_DIR
-
-# define a unique log file in the right place
-logfile="${JOB_DIR}/slurm-${SLURM_JOB_ID}.out"
-echo $logfile
-
-# note:
-# it's important for ntasks-per-node to be equal to the number of GPUs per node, which is set with --gres=gpu:4
-export MASTER_PORT=12355 # there may be a smarter way to set this (e.g. with code), but this port is almost always open (maybe 100% of the time)
-export WORLD_SIZE=$(($SLURM_NNODES * $SLURM_NTASKS_PER_NODE)) # world size is equal to number of nodes and number of tasks per node
-echo "WORLD_SIZE="$WORLD_SIZE >> $logfile
-echo "MASTER_PORT="$MASTER_PORT >> $logfile
-
-# define a master address for communication between GPUs
-master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
-export MASTER_ADDR=$master_addr
-echo "MASTER_ADDR="$MASTER_ADDR >> $logfile
-
-# load python and activate our conda environment
-module purge
-module load python
-conda activate networkAlignmentAnalysis
-
-# record the start time
-start_time=$(date +%s)
-
-# this is the command that initiates the processes
-# srun will send each command to each task on each node for a total of WORLD_SIZE tasks
-# (in this case the script it runs is python cluster/ddp.py)
-srun python cluster/ddp-example/ddp.py --epochs=50 --job-folder=${JOB_DIR} >> $logfile
-
-# record the end time
-end_time=$(date +%s)
-
-# measure the time elapsed for the core part of the job in the logfile
-total_time=$((end_time-start_time))
-echo "Total Time= "$total_time" seconds" >> $logfile
diff --git a/_arxiv/archive/cluster/imagenet-example/imagenet.py b/_arxiv/archive/cluster/imagenet-example/imagenet.py
deleted file mode 100644
index 08a2abfd..00000000
--- a/_arxiv/archive/cluster/imagenet-example/imagenet.py
+++ /dev/null
@@ -1,217 +0,0 @@
-import argparse
-import wandb
-import torch
-import torch.nn.functional as F
-import torch.optim as optim
-from torch.optim.lr_scheduler import StepLR
-
-import os
-import sys
-import time
-import torch.distributed as dist
-from torch.nn.parallel import DistributedDataParallel as DDP
-from socket import gethostname
-
-mainPath = os.path.dirname(os.path.abspath(__file__)) + "/.."
-sys.path.append(mainPath)
-
-from networkAlignmentAnalysis import datasets
-from networkAlignmentAnalysis.models.registry import get_model
-
-
-def configure_wandb(args):
- """create a wandb run file and set environment parameters appropriately"""
- if args.use_wandb:
- wandb.login()
- run = wandb.init(
- project=args.job_name,
- name=args.job_folder,
- )
- os.environ["WANDB_MODE"] = "offline"
- return run
-
- return None
-
-
-def train(args, model, device, dataset, optimizer, epoch, rank, run, train=True):
- dataloader = dataset.train_loader if train else dataset.test_loader
- if dataset.distributed:
- if train:
- dataset.train_sampler.set_epoch(epoch)
- else:
- dataset.test_sampler.set_epoch(epoch)
-
- if rank == 0:
- first_batch_timer = time.time()
-
- model.train()
- for batch_idx, batch in enumerate(dataloader):
- data, target = dataset.unwrap_batch(batch, device=device)
- if rank == 0 and batch_idx == 0:
- print(f"Train-- epoch {epoch}, rank {rank}, first batch loaded in {time.time() - first_batch_timer} seconds.")
- optimizer.zero_grad()
- output = model(data)
- loss = dataset.measure_loss(output, target)
- loss.backward()
- optimizer.step()
- if batch_idx % args.log_interval == 0:
- if rank == 0:
- print(f"Train Epoch: {epoch} [{batch_idx}/{len(dataloader)} ({100.*batch_idx/len(dataloader):.0f}%)] \t Loss: {loss.item():.6f}")
- if run is not None:
- run.log(dict(epoch=epoch, batch_idx=batch_idx, train_loss=loss.item()))
- if args.dry_run:
- break
-
-
-def test(model, device, dataset, run, train=False):
- dataloader = dataset.train_loader if train else dataset.test_loader
- model.eval()
- test_loss = 0
- correct = 0
- attempts = 0
- with torch.no_grad():
- for batch in dataloader:
- data, target = dataset.unwrap_batch(batch, device=device)
- output = model(data)
- test_loss += dataset.measure_loss(output, target, reduction="sum").item()
- pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
- correct += pred.eq(target.view_as(pred)).sum().item()
- attempts += data.size(0)
-
- test_loss /= attempts
-
- print(f"\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{attempts} ({100.*correct/attempts:.0f}%)\n")
-
- if run is not None:
- run.log(dict(test_loss=test_loss, test_accuracy=100.0 * correct / attempts))
-
-
-def create_dataset(name, net, distributed=True, loader_parameters={}):
- return datasets.get_dataset(
- name,
- build=True,
- distributed=distributed,
- transform_parameters=net,
- loader_parameters=loader_parameters,
- )
-
-
-def setup(rank, world_size):
- # initialize the process group
- dist.init_process_group("nccl", rank=rank, world_size=world_size)
-
-
-def main():
- # Training settings
- parser = argparse.ArgumentParser(description="PyTorch ImageNet Example")
- parser.add_argument(
- "--job-name",
- type=str,
- default="ImageNet-Example",
- help="job name for this project",
- )
- parser.add_argument(
- "--job-folder",
- type=str,
- default=".",
- help="job folder for storing data",
- )
- parser.add_argument(
- "--use_wandb",
- default=False,
- action="store_true",
- help="if used, will log experiment to WandB",
- )
- parser.add_argument(
- "--batch-size",
- type=int,
- default=64,
- metavar="N",
- help="input batch size for training (default: 64)",
- )
- parser.add_argument(
- "--epochs",
- type=int,
- default=14,
- metavar="N",
- help="number of epochs to train (default: 14)",
- )
- parser.add_argument("--lr", type=float, default=1.0, metavar="LR", help="learning rate (default: 1.0)")
- parser.add_argument("--wd", type=float, default=0.0, metavar="WD", help="weight decay (default: 0.0)")
- parser.add_argument(
- "--gamma",
- type=float,
- default=0.9,
- metavar="M",
- help="Learning rate step gamma (default: 0.99)",
- )
- parser.add_argument("--no-cuda", action="store_true", default=False, help="disables CUDA training")
- parser.add_argument("--dry-run", action="store_true", default=False, help="quickly check a single pass")
- parser.add_argument("--seed", type=int, default=1, metavar="S", help="random seed (default: 1)")
- parser.add_argument(
- "--log-interval",
- type=int,
- default=10,
- metavar="N",
- help="how many batches to wait before logging training status",
- )
- parser.add_argument("--save-model", action="store_true", default=False, help="For Saving the current Model")
- args = parser.parse_args()
-
- run = configure_wandb(args)
-
- torch.manual_seed(args.seed)
-
- world_size = int(os.environ["WORLD_SIZE"])
- rank = int(os.environ["SLURM_PROCID"])
- gpus_per_node = int(os.environ["SLURM_GPUS_ON_NODE"])
- assert gpus_per_node == torch.cuda.device_count()
- print(
- f"Hello from rank {rank} of {world_size} on {gethostname()} where there are" f" {gpus_per_node} allocated GPUs per node.",
- flush=True,
- )
-
- loader_parameters = dict(
- batch_size=args.batch_size,
- num_workers=max(int(os.environ["SLURM_CPUS_PER_TASK"]) - 1, 1), # save a little headroom
- )
-
- if world_size > 1:
- setup(rank, world_size)
- if rank == 0:
- print(f"Group initialized? {dist.is_initialized()}", flush=True)
-
- local_rank = rank - gpus_per_node * (rank // gpus_per_node)
- torch.cuda.set_device(local_rank)
- print(f"host: {gethostname()}, rank: {rank}, local_rank: {local_rank}")
-
- model_name = "AlexNet"
- dataset_name = "ImageNet"
- net = get_model(model_name, build=True, dataset=dataset_name)
- dataset = create_dataset(dataset_name, net, distributed=world_size > 1, loader_parameters=loader_parameters)
-
- model = net.to(local_rank)
- ddp_model = DDP(model, device_ids=[local_rank]) if world_size > 1 else model
- optimizer = optim.Adadelta(ddp_model.parameters(), lr=args.lr, weight_decay=args.wd)
-
- scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
-
- for epoch in range(1, args.epochs + 1):
- epoch_time = time.time()
- train(args, ddp_model, local_rank, dataset, optimizer, epoch, rank, run)
- if rank == 0:
- test(ddp_model, local_rank, dataset, run)
- scheduler.step()
- epoch_time = time.time() - epoch_time
- if rank == 0:
- print(f"Epoch {epoch}, Train & Test Time = {epoch_time:.1f} seconds (measured from rank {rank}).\n")
-
- if args.save_model and rank == 0:
- torch.save(model.state_dict(), os.path.join(args.job_folder, "alexnet_imagenet.pt"))
-
- if world_size > 1:
- dist.destroy_process_group()
-
-
-if __name__ == "__main__":
- main()
diff --git a/_arxiv/archive/cluster/imagenet-example/imagenet.slurm b/_arxiv/archive/cluster/imagenet-example/imagenet.slurm
deleted file mode 100644
index fe4e8976..00000000
--- a/_arxiv/archive/cluster/imagenet-example/imagenet.slurm
+++ /dev/null
@@ -1,50 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=ImageNet-Example # create a short name for your job
-#SBATCH --partition=kempner # partition
-#SBATCH --account=kempner_bsabatini_lab # account needed for kempner partition
-#SBATCH --nodes=4 # node count
-#SBATCH --ntasks-per-node=4 # total number of tasks per node
-#SBATCH --cpus-per-task=16 # cpu-cores per task (>1 if multi-threaded tasks)
-#SBATCH --gres=gpu:4 # number of allocated gpus per node
-#SBATCH --mem=1000G # total memory per node (4 GB per cpu-core is default)
-#SBATCH --time=01:00:00 # total run time limit (HH:MM:SS)
-#SBATCH --mail-type=begin # send email when job begins
-#SBATCH --mail-type=end # send email when job ends
-
-
-# we need to define the job name directly since it isn't a slurm environment variable
-JOB_NAME="ImageNet-Example"
-
-# this text file has some settings in it (like the standard job directory)
-source cluster/slurm_settings.txt
-JOB_DIR=${JOB_FOLDER}/${JOB_NAME}-${SLURM_JOB_ID} # make a specific directory for this particular job
-mkdir -p $JOB_DIR
-
-# define a unique log file in the right place
-logfile="${JOB_DIR}/slurm-${SLURM_JOB_ID}.out"
-echo "Writing to "$logfile
-
-export MASTER_PORT=12355
-export WORLD_SIZE=$(($SLURM_NNODES * $SLURM_NTASKS_PER_NODE))
-echo "WORLD_SIZE="$WORLD_SIZE >> $logfile
-echo "MASTER_PORT="$MASTER_PORT >> $logfile
-
-master_addr=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
-export MASTER_ADDR=$master_addr
-echo "MASTER_ADDR="$MASTER_ADDR >> $logfile
-
-module purge
-module load python
-conda activate networkAlignmentAnalysis
-
-start_time=$(date +%s)
-
-# (nvidia-smi --query-gpu=utilization.gpu --format=csv --loop=1 --filename=gpu_utilization.log)
-
-srun python cluster/imagenet-example/imagenet.py --epochs=10 --use_wandb --log-interval=100 --job-folder=${JOB_DIR} --job-name=${JOB_NAME} --save-model >> $logfile
-
-end_time=$(date +%s)
-
-total_time=$((end_time-start_time))
-echo "Total Time= "$total_time" seconds" >> $logfile
-echo "Total Time= "$total_time" seconds"
diff --git a/_arxiv/archive/cluster/lightning-example/lightning_train.sbatch b/_arxiv/archive/cluster/lightning-example/lightning_train.sbatch
deleted file mode 100644
index 86a127f2..00000000
--- a/_arxiv/archive/cluster/lightning-example/lightning_train.sbatch
+++ /dev/null
@@ -1,35 +0,0 @@
-#! /bin/bash
-#SBATCH --job-name=vast_benchmark
-#SBATCH --time=1-00:00:00
-#SBATCH --partition=kempner_requeue
-#SBATCH --nodes=1
-#SBATCH --ntasks-per-node=1
-#SBATCH --cpus-per-task=4
-#SBATCH --mem=1000G
-#SBATCH --output=benchmark_lightning.out
-#SBATCH --error=benchmark_lightning.err
-#SBATCH --account=kempner_dev
-#SBATCH --dependency=singleton
-#SBATCH --gres=gpu:1
-
-if [[ -z $1 ]] ; then
- model="alexnet"
-else
- model=$1
-fi
-
-
-VAST_CREDENTIALS_PATH=/n/holylabs/LABS/kempner_dev/Users/$USER/vast-benchmarking/.load_credentials.sh
-METRIC_GATHERING_SCRIPT_PATH=/n/holylabs/LABS/kempner_dev/Users/$USER/vast-benchmarking/get_vast_stats.sh
-DATADIR=/n/fas-vast/kempner_benchmark/preproccessed_imagenet/imagenet21k_resized
-CONTAINER_PATH=/n/holylabs/LABS/kempner_dev/Lab/containers/pytorch_2.1.2-cuda12.1-cudnn8-runtime-lightning.sif
-
-
-source $VAST_CREDENTIALS_PATH
-metrics_file=vast_metrics/vast_metrics_${SLURM_NNODES}nodes_${SLURM_GPUS_ON_NODE}gpus_${model}_${SLURM_JOB_ID}.csv
-
-$METRIC_GATHERING_SCRIPT_PATH $metrics_file 30 &
-
-time srun singularity exec --bind $DATADIR --nv $CONTAINER_PATH python3 train_model_lightning.py --dataset $DATADIR/imagenet21k_train \
- --epochs 1 --batch_size 64 --num_workers $SLURM_CPUS_PER_TASK --num_gpus $SLURM_GPUS_ON_NODE --num_nodes $SLURM_NNODES \
- --model $model
\ No newline at end of file
diff --git a/_arxiv/archive/cluster/lightning-example/submit_lightning.sh b/_arxiv/archive/cluster/lightning-example/submit_lightning.sh
deleted file mode 100644
index 3e246f8f..00000000
--- a/_arxiv/archive/cluster/lightning-example/submit_lightning.sh
+++ /dev/null
@@ -1,26 +0,0 @@
-#! /bin/bash
-# Script used to submit a matrix of jobs to the cluster
-
-mkdir -p logs
-mkdir -p vast_metrics
-
-for model in alexnet resnet50 ; do
- for nodes in 1 2 4 8 16 32 ; do
- if [[ $nodes -eq 1 ]]; then
- for gpus in 1 2 4 ; do
- sbatch -N $nodes --ntasks-per-node $gpus --gres gpu:$gpus \
- -o logs/vast_benchmark_${nodes}nodes_${gpus}gpus_${model}.out \
- -e logs/vast_benchmark_${nodes}nodes_${gpus}gpus_${model}.err \
- lightning_train.sbatch $model
- sleep 3
- done
- else
- gpus=4
- sbatch -N $nodes --ntasks-per-node $gpus --gres gpu:$gpus \
- -o logs/vast_benchmark_${nodes}nodes_${gpus}gpus_${model}.out \
- -e logs/vast_benchmark_${nodes}nodes_${gpus}gpus_${model}.err \
- lightning_train.sbatch $model
- sleep 3
- fi
- done
-done
\ No newline at end of file
diff --git a/_arxiv/archive/cluster/lightning-example/train_model_lightning.py b/_arxiv/archive/cluster/lightning-example/train_model_lightning.py
deleted file mode 100644
index 3589f688..00000000
--- a/_arxiv/archive/cluster/lightning-example/train_model_lightning.py
+++ /dev/null
@@ -1,110 +0,0 @@
-from typing import Any
-from lightning.pytorch.utilities.types import STEP_OUTPUT
-import torch
-import torchvision
-from torchvision import datasets, transforms
-from torch.utils.data import DataLoader
-
-import lightning as L
-
-import argparse
-import logging
-
-
-class LightningModel(L.LightningModule):
- def __init__(self, model, optimizer, lr) -> None:
- super().__init__()
- self.model = model
- self.optimizer = optimizer
- self.lr = lr
-
- def training_step(self, batch, batch_idx):
- images, labels = batch
- output = self.model(images)
- loss = torch.nn.functional.cross_entropy(output, labels)
- return loss
-
- def configure_optimizers(self):
- optimizer = self.optimizer(self.parameters(), lr=self.lr)
- return optimizer
-
-
-# def train_one_epoch(model, loader, optimizer, device):
-# model.train()
-# batch_num = 0
-# for images, labels in loader:
-# images, labels = images.to(device), labels.to(device)
-# optimizer.zero_grad()
-# output = model(images)
-# loss = torch.nn.functional.cross_entropy(output, labels)
-# loss.backward()
-# optimizer.step()
-# if batch_num % 100 == 0:
-# logging.info(f"Batch Num: {batch_num} Loss: {loss.item()}")
-# batch_num += 1
-
-
-def parse_args() -> argparse.Namespace:
- parser = argparse.ArgumentParser(description="Train a model")
- parser.add_argument("--dataset", type=str, default="imagenet21k_train", help="dataset to use")
- parser.add_argument("--batch_size", type=int, default=64, help="batch size")
- parser.add_argument("--epochs", type=int, default=10, help="number of epochs")
- parser.add_argument("--lr", type=float, default=0.01, help="learning rate")
- parser.add_argument("--model", type=str, default="alexnet", help="model to use")
- parser.add_argument("--num_workers", type=int, default=4, help="number of workers")
- parser.add_argument("--num_gpus", type=int, default=1, help="number of gpus per node")
- parser.add_argument("--num_nodes", type=int, default=1, help="number of nodes")
- parser.add_argument(
- "--checkpoint_steps",
- type=int,
- default=0,
- help="number of training steps to save checkpoint after",
- )
- parser.add_argument("--checkpoint_dir", type=str, default="checkpoints", help="directory to save checkpoints")
-
- return parser.parse_args()
-
-
-def main():
- args = parse_args()
- preprocess = transforms.Compose(
- [
- transforms.Resize(256),
- transforms.CenterCrop(224),
- transforms.ToTensor(),
- transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
- ]
- )
- dataset = datasets.ImageFolder(args.dataset, transform=preprocess)
- logging.info(f"Data loaded, Found {len(dataset.classes)} classes, {len(dataset)} images")
- num_classes = len(dataset.classes)
- loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers)
- model = torchvision.models.get_model(args.model, num_classes=num_classes)
- optimizer = torch.optim.SGD
- lighting_model = LightningModel(model, optimizer, args.lr)
-
- callbacks = None
- if args.checkpoint_steps > 0:
- callbacks = [
- L.callbacks.ModelCheckpoint(
- dirpath=args.checkpoint_dir,
- save_top_k=-1,
- save_last=True,
- every_n_train_steps=args.checkpoint_steps,
- )
- ]
-
- trainer = L.Trainer(
- max_epochs=args.epochs,
- min_epochs=args.epochs,
- strategy="ddp",
- devices=args.num_gpus,
- num_nodes=args.num_nodes,
- callbacks=callbacks,
- )
- trainer.fit(lighting_model, train_dataloaders=loader)
-
-
-if __name__ == "__main__":
- logging.basicConfig(level=logging.INFO)
- main()
diff --git a/_arxiv/archive/cluster/slurm_settings.txt b/_arxiv/archive/cluster/slurm_settings.txt
deleted file mode 100644
index 229c3711..00000000
--- a/_arxiv/archive/cluster/slurm_settings.txt
+++ /dev/null
@@ -1,3 +0,0 @@
-# Job Folder
-JOB_FOLDER=./jobs
-
diff --git a/_arxiv/archive/cluster/testing-examples/test_job.slurm b/_arxiv/archive/cluster/testing-examples/test_job.slurm
deleted file mode 100644
index 7ab81e4d..00000000
--- a/_arxiv/archive/cluster/testing-examples/test_job.slurm
+++ /dev/null
@@ -1,32 +0,0 @@
-#!/bin/bash
-#SBATCH --job-name=test_files # create a short name for your job
-#SBATCH --nodes=1 # node count
-#SBATCH --ntasks-per-node=1 # total number of tasks per node
-#SBATCH --cpus-per-task=1 # cpu-cores per task (>1 if multi-threaded tasks)
-#SBATCH --mem=100M # total memory per node (4 GB per cpu-core is default)
-#SBATCH --time=00:20:00 # total run time limit (HH:MM:SS)
-#SBATCH --mail-type=begin # send email when job begins
-#SBATCH --mail-type=end # send email when job ends
-#SBATCH --output=/dev/null # no output message
-
-source cluster/slurm_settings.txt
-mkdir -p $JOB_FOLDER
-logfile="${JOB_FOLDER}/slurm-${SLURM_JOB_ID}.out"
-
-echo "Writing to ${logfile}"
-echo "Writing in ${logfile}" >> $logfile
-
-module purge
-module load python
-
-start_time=$(date +%s)
-
-python -c "print('hello world')" >> $logfile
-python test_print.py >> $logfile
-
-end_time=$(date +%s)
-
-total_time=$((end_time-start_time))
-echo "Total Time= "$total_time" seconds" >> $logfile
-
-echo "Finished."
diff --git a/_arxiv/archive/refactoring_docs/GAUSSIAN_MI_SUMMARY.md b/_arxiv/archive/refactoring_docs/GAUSSIAN_MI_SUMMARY.md
deleted file mode 100644
index 58e65f3b..00000000
--- a/_arxiv/archive/refactoring_docs/GAUSSIAN_MI_SUMMARY.md
+++ /dev/null
@@ -1,84 +0,0 @@
-# Gaussian Mutual Information with Edgeworth Expansion
-
-## Overview
-I've implemented a new metric `gaussian_mi_analytic` that computes mutual information between inputs and outputs of neural network nodes, assuming approximately Gaussian distributions with analytic expansions for non-Gaussian corrections.
-
-## Features
-
-### 1. **Analytic Gaussian MI Calculation**
-For linear transformations Y = WX + ε where X and ε are Gaussian:
-- Exact formula: I(X;Y) = 1/2 * log(det(Σ_X) * det(Σ_Y) / det(Σ_joint))
-- Efficient computation using covariance matrices
-- Numerical stability with regularization
-
-### 2. **Edgeworth Expansion Corrections**
-The metric includes corrections for non-Gaussian distributions up to order 3:
-
-- **Order 0**: Pure Gaussian assumption
-- **Order 1**: First-order correction using third cumulants (skewness)
-- **Order 2**: Second-order correction using fourth cumulants (kurtosis) and mixed terms
-- **Order 3**: Third-order correction with cross-terms between skewness and kurtosis
-
-### 3. **Key Parameters**
-- `expansion_order`: Controls the order of Edgeworth expansion (0-3)
-- `noise_std`: Assumed noise level in the system
-- `regularization`: Numerical stability parameter
-- `per_neuron`: Compute MI for each neuron separately or jointly
-
-## Mathematical Foundation
-
-### Cumulants Used
-- κ₂ = variance
-- κ₃ = E[X³] (related to skewness)
-- κ₄ = E[X⁴] - 3σ⁴ (excess kurtosis)
-
-### Edgeworth Corrections
-The corrections capture deviations from Gaussianity:
-
-1. **First-order**: Involves normalized third cumulants (γ₁ = κ₃/σ³)
-2. **Second-order**: Involves normalized fourth cumulants (γ₂ = κ₄/σ⁴)
-3. **Third-order**: Cross-terms between different order cumulants
-
-## Test Results
-
-The metric was tested on:
-1. **Pure Gaussian data**: Corrections are minimal as expected
-2. **Skewed data (Chi-squared)**: Small corrections observed
-3. **Heavy-tailed data (Student-t)**: Significant corrections, especially at orders 2 and 3
-4. **Different noise levels**: MI decreases with increasing noise as expected
-
-### Example Results
-- Pure Gaussian (Order 0): MI ≈ 2.99
-- Heavy-tailed (Order 0): MI ≈ 2.99
-- Heavy-tailed (Order 2): MI ≈ 4.16 (showing significant correction)
-
-## Usage Example
-
-```python
-from src.alignment.core.registry import METRIC_REGISTRY
-
-# Create metric with second-order corrections
-metric = METRIC_REGISTRY.get("gaussian_mi_analytic")(
- expansion_order=2,
- noise_std=0.1,
- per_neuron=True
-)
-
-# Compute MI scores for each neuron
-mi_scores = metric.compute(inputs=inputs, weights=weights)
-```
-
-## Benefits
-
-1. **Analytic computation**: Fast and exact for Gaussian case
-2. **Non-Gaussian handling**: Edgeworth expansion captures deviations
-3. **Flexible**: Can adjust expansion order based on data characteristics
-4. **Per-neuron analysis**: Can analyze individual neuron information transfer
-
-## Integration
-
-The metric is fully integrated into the alignment framework:
-- Registered as `gaussian_mi_analytic`
-- Available through the metric registry
-- Follows the standard BaseMetric interface
-- Properly handles different tensor dimensions
\ No newline at end of file
diff --git a/_arxiv/archive/scripts/README.md b/_arxiv/archive/scripts/README.md
deleted file mode 100644
index f5b9f704..00000000
--- a/_arxiv/archive/scripts/README.md
+++ /dev/null
@@ -1,35 +0,0 @@
-# Alignment Scripts
-
-This directory contains utility scripts for the Alignment library. These scripts provide various functionality for running experiments, demonstrations, and utilities.
-
-## Contents
-
-- **direct_pruning_test.py**: Tests pruning functionality directly without relying on the experiment infrastructure.
-- **run_multi_strategy_experiment.py**: Runs experiments with multiple pruning strategies.
-- **run_fixed_experiment.py**: Runs experiments with fixed configurations.
-- **run_cascading_with_plots.py**: Runs cascading pruning experiments and generates plots.
-- **run_cascading_test.py**: Tests cascading pruning methodology.
-- **continual_mnist.py**: Implementation of continual learning on MNIST.
-- **teacher_student.py**: Implementation of teacher-student neural network model.
-
-## Shell Scripts
-
-- **run_benchmark.sh**: Convenience script for running benchmarks.
-- **run_cascading_test.sh**: Shell script for testing cascading pruning.
-
-## Usage
-
-Most scripts have command-line arguments for configuration:
-
-```bash
-python scripts/run_fixed_experiment.py --config configs/config_alignment_experiment.yaml
-```
-
-## Adding New Scripts
-
-When adding new scripts, please follow these guidelines:
-
-1. Include clear documentation at the top of the script about its purpose and usage
-2. Add command-line arguments for configuration using argparse
-3. Include appropriate error handling and logging
-4. Update this README with information about the new script
\ No newline at end of file
diff --git a/_arxiv/archive/scripts/__init__.py b/_arxiv/archive/scripts/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/_arxiv/archive/scripts/continual_mnist.py b/_arxiv/archive/scripts/continual_mnist.py
deleted file mode 100644
index 021b5a55..00000000
--- a/_arxiv/archive/scripts/continual_mnist.py
+++ /dev/null
@@ -1,144 +0,0 @@
-# Some Code for continual learning with permuted MNIST
-import time
-from functools import partial
-import numpy as np
-import scipy as sp
-import sklearn
-import torch
-import torch.nn.functional as F
-from torch import nn
-from torchvision.transforms import v2 as transforms
-
-from tqdm import tqdm
-import matplotlib as mpl
-from matplotlib import pyplot as plt
-
-import os
-import sys
-
-mainPath = os.path.dirname(os.path.abspath(__file__)) + "/../.."
-sys.path.append(mainPath)
-
-from alignment.models.registry import get_model
-from alignment.datasets import get_dataset
-from alignment.experiments.registry import get_experiment
-from alignment import utils
-from alignment import files
-from alignment import train
-
-
-def permute(batch, batch_dim=True, shuffle_idx=None):
- if shuffle_idx is not None:
- original_size = batch[0].shape
- if batch_dim:
- batch[0] = batch[0][:, shuffle_idx]
- else:
- batch[0] = batch[0][shuffle_idx]
- batch[0].reshape(original_size)
- return batch
-
-
-def add_permutation(dataset, num_pixels=784):
- perm = partial(permute, batch_dim=True, shuffle_idx=torch.randperm(num_pixels))
- if dataset.extra_transform is None:
- dataset.extra_transform = []
- dataset.extra_transform.append(perm)
- return dataset
-
-
-def update_permutation(dataset, num_pixels=784):
- perm = partial(permute, batch_dim=True, shuffle_idx=torch.randperm(num_pixels))
- dataset.extra_transform[-1] = perm
- return dataset
-
-
-def do_permuted_round(nets, optimizers, dataset, train_epochs=1, verbose=False):
- parameters = dict(
- verbose=verbose,
- num_epochs=train_epochs,
- alignment=False,
- )
- dataset = update_permutation(dataset)
-
- # train and test
- train_results = train.train(nets, optimizers, dataset, **parameters)
- test_results = train.test(nets, dataset, **parameters)
-
- return train_results, test_results
-
-
-def main():
- DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
- print("using device: ", DEVICE)
-
- model_name = "MLP"
- dataset_name = "MNIST"
-
- hidden_widths = [200, 200, 200]
-
- lrs = [1e-1, 3e-2, 1e-2, 3e-3]
- num_replicates = 3
-
- nets = []
- optimizers = []
- net_lr = []
- for lr in lrs:
- for _ in range(num_replicates):
- net = get_model(model_name, build=True, dataset=dataset_name, hidden_widths=hidden_widths, dropout=0.0, ignore_flag=False)
- net.to(DEVICE)
-
- optimizer = torch.optim.SGD(net.parameters(), lr=lr)
-
- nets.append(net)
- optimizers.append(optimizer)
- net_lr.append(lr)
-
- loader_parameters = dict(
- shuffle=True,
- batch_size=5,
- )
- dataset = get_dataset(dataset_name, build=True, transform_parameters=net, loader_parameters=loader_parameters, device=DEVICE)
- dataset = add_permutation(dataset)
-
- num_rounds = 200
-
- train_loss = []
- train_accuracy = []
- test_loss = []
- test_accuracy = []
- for round in tqdm(range(num_rounds)):
- c_train_res, c_test_res = do_permuted_round(nets, optimizers, dataset, verbose=False)
- train_loss.append(c_train_res["loss"])
- train_accuracy.append(c_train_res["accuracy"])
- test_loss.append(c_test_res["loss"])
- test_accuracy.append(c_test_res["accuracy"])
-
- loss = torch.stack([torch.tensor(l) for l in test_loss])
- accuracy = torch.stack([torch.tensor(a) for a in test_accuracy])
-
- type_loss = utils.compute_stats_by_type(loss, len(lrs), 1)[0]
- type_accuracy = utils.compute_stats_by_type(accuracy, len(lrs), 1)[0]
-
- # print(loss.shape, accuracy.shape, type_loss.shape, type_accuracy.shape)
-
- cols = mpl.colormaps["Set1"].resampled(len(lrs))
- names = [f"lr={lr}" for lr in lrs]
-
- fig, ax = plt.subplots(1, 2, figsize=(6, 3), layout="constrained")
- for ii in range(len(lrs)):
- ax[0].plot(range(num_rounds), type_loss[:, ii], c=cols(ii), label=names[ii])
- ax[1].plot(range(num_rounds), type_accuracy[:, ii], c=cols(ii), label=names[ii])
-
- ax[0].set_xlabel("Rounds of Permuted MNIST")
- ax[1].set_xlabel("Rounds of Permuted MNIST")
- ax[0].set_ylabel("loss")
- ax[1].set_ylabel("accuracy")
- ax[0].legend(loc="best")
- ax[1].legend(loc="best")
- plt.show()
-
- print("Hello, add a debugger here to evaluate while testing")
-
-
-if __name__ == "__main__":
- main()
diff --git a/_arxiv/archive/scripts/direct_pruning_test.py b/_arxiv/archive/scripts/direct_pruning_test.py
deleted file mode 100644
index efb97975..00000000
--- a/_arxiv/archive/scripts/direct_pruning_test.py
+++ /dev/null
@@ -1,257 +0,0 @@
-#!/usr/bin/env python
-"""
-Direct Pruning Test
-
-This script tests pruning functionality directly, without relying on the complex
-experiment infrastructure. It loads a model, applies pruning, and logs the results
-including before/after weights and accuracies.
-"""
-
-import os
-import sys
-import logging
-import argparse
-import copy
-import numpy as np
-import torch
-import torch.nn as nn
-
-# Set up logging
-logging.basicConfig(
- level=logging.INFO,
- format='%(asctime)s [%(levelname)s] %(message)s',
- handlers=[
- logging.StreamHandler(),
- logging.FileHandler('direct_pruning_test.log', mode='w')
- ]
-)
-logger = logging.getLogger(__name__)
-
-# Add src to path if not already there
-src_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'src')
-if src_path not in sys.path:
- sys.path.insert(0, src_path)
-
-# Import alignment modules
-from alignment.models.registry import create_model
-from alignment.datasets import load_dataset
-from alignment.metrics import get_metric
-from alignment.dropout import progressive_dropout
-from alignment.config import ExperimentConfig
-
-def count_zero_weights(model):
- """Count zero weights in the model."""
- zero_weights = 0
- total_weights = 0
-
- for name, param in model.named_parameters():
- if 'weight' in name:
- layer_zeros = (param.data == 0).sum().item()
- layer_total = param.data.numel()
- zero_weights += layer_zeros
- total_weights += layer_total
-
- if hasattr(param, 'shape'):
- logger.info(f"Layer {name}: {layer_zeros}/{layer_total} zeros "
- f"({100.0*layer_zeros/layer_total:.2f}% pruned), shape {param.shape}")
-
- if total_weights > 0:
- logger.info(f"Total: {zero_weights}/{total_weights} zeros "
- f"({100.0*zero_weights/total_weights:.2f}% pruned)")
-
- return zero_weights, total_weights
-
-def verify_pruning(config_path=None, model_name='mlp', dataset_name='mnist', strategy='low_rq',
- pruning_mode='layer_wise', dropout_mode='unscaled', pruning_percent=0.5,
- device=None):
- """Test pruning directly."""
- # Use GPU if available
- if device is None:
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- logger.info(f"Using device: {device}")
-
- # Fix dataset name casing - use uppercase for standard datasets
- if dataset_name.lower() == 'mnist':
- dataset_name = 'MNIST'
- elif dataset_name.lower() == 'cifar10':
- dataset_name = 'CIFAR10'
- elif dataset_name.lower() == 'cifar100':
- dataset_name = 'CIFAR100'
-
- # Define data path
- data_root = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'data')
-
- # Create dataset config
- dataset_config = {
- 'dataset_name': dataset_name,
- 'batch_size': 128,
- 'data_dir': data_root,
- 'download': True
- }
-
- # Make sure data directory exists
- os.makedirs(data_root, exist_ok=True)
-
- dataset = load_dataset(dataset_config)
-
- # Get input dimension from dataset
- if dataset_name.lower() in ['mnist', 'fashion_mnist']:
- input_dim = 784 # 28*28
- else:
- # Try to infer from the dataset
- sample_data = next(iter(dataset.train_loader))[0]
- input_dim = sample_data[0].flatten().shape[0]
- logger.info(f"Inferred input dimension: {input_dim}")
-
- # Create a simple model
- class MLP(nn.Module):
- def __init__(self, input_dim=784, hidden_dim=100, output_dim=10):
- super().__init__()
- self.fc1 = nn.Linear(input_dim, hidden_dim)
- self.fc2 = nn.Linear(hidden_dim, hidden_dim)
- self.fc3 = nn.Linear(hidden_dim, output_dim)
- self.relu = nn.ReLU()
-
- # Define alignment layers for pruning
- self.alignment_layers = [self.fc1, self.fc2, self.fc3]
- self.alignment_names = ['layer_0', 'layer_1', 'layer_2']
- self.hidden = {}
-
- def forward(self, x):
- # Store inputs for hooking
- batch_size = x.size(0)
- x = x.view(batch_size, -1) # Flatten
-
- # Forward with storing activations
- self.hidden['layer_0'] = x
- x = self.relu(self.fc1(x))
-
- self.hidden['layer_1'] = x
- x = self.relu(self.fc2(x))
-
- self.hidden['layer_2'] = x
- x = self.fc3(x)
-
- return x
-
- # Create model based on dataset
- model = MLP(input_dim=input_dim, hidden_dim=100, output_dim=10)
- model.to(device)
-
- # Check for alignment layers
- if not hasattr(model, 'alignment_layers'):
- logger.error("Model doesn't have alignment_layers attribute")
- return
-
- logger.info(f"Created model with {len(model.alignment_layers)} alignment layers")
-
- # Make a copy for pruning
- pruned_model = copy.deepcopy(model)
- pruned_model.to(device)
-
- # Make sure parameters are on device
- for param in pruned_model.parameters():
- param.data = param.data.to(device)
-
- # Evaluate original model
- model.eval()
- orig_accuracy, orig_loss = dataset.evaluate(model, device)
- logger.info(f"Original model: accuracy={orig_accuracy:.2f}%, loss={orig_loss:.4f}")
-
- # Check zero weights before pruning
- logger.info("Weights before pruning:")
- orig_zeros, orig_total = count_zero_weights(pruned_model)
-
- # Setup metric
- metric = get_metric('rq')
-
- # Apply pruning
- logger.info(f"Applying pruning: strategy={strategy}, mode={pruning_mode}, "
- f"dropout_mode={dropout_mode}, percent={pruning_percent*100:.1f}%")
-
- # Call progressive_dropout directly
- network_accuracies, network_losses = progressive_dropout(
- [pruned_model],
- dataset,
- [0.0, pruning_percent], # Use just two fractions: 0% and the target %
- metric,
- device,
- pruning_mode=pruning_mode,
- dropout_mode=dropout_mode,
- strategy=strategy,
- show_progress=False
- )
-
- # Check if pruning was applied
- logger.info("Weights after pruning:")
- pruned_zeros, pruned_total = count_zero_weights(pruned_model)
-
- # Calculate how many weights were pruned
- weights_pruned = pruned_zeros - orig_zeros
- pruned_percent = 100.0 * weights_pruned / pruned_total if pruned_total > 0 else 0
- logger.info(f"Weights pruned: {weights_pruned}/{pruned_total} ({pruned_percent:.2f}%)")
-
- # Evaluate pruned model
- pruned_model.eval()
- pruned_accuracy, pruned_loss = dataset.evaluate(pruned_model, device)
- logger.info(f"Pruned model: accuracy={pruned_accuracy:.2f}%, loss={pruned_loss:.4f}")
-
- # Calculate accuracy change
- acc_change = pruned_accuracy - orig_accuracy
- logger.info(f"Accuracy change: {acc_change:.2f}% points")
-
- # Summary
- logger.info("\nSUMMARY:")
- logger.info(f"Strategy: {strategy}")
- logger.info(f"Pruning mode: {pruning_mode}")
- logger.info(f"Dropout mode: {dropout_mode}")
- logger.info(f"Pruning percentage: {pruning_percent*100:.1f}%")
- logger.info(f"Original accuracy: {orig_accuracy:.2f}%")
- logger.info(f"Pruned accuracy: {pruned_accuracy:.2f}%")
- logger.info(f"Accuracy change: {acc_change:.2f}% points")
- logger.info(f"Weights pruned: {weights_pruned}/{pruned_total} ({pruned_percent:.2f}%)")
-
- return {
- 'orig_accuracy': orig_accuracy,
- 'pruned_accuracy': pruned_accuracy,
- 'acc_change': acc_change,
- 'weights_pruned': weights_pruned,
- 'total_weights': pruned_total,
- 'pruned_percent': pruned_percent
- }
-
-def main():
- """Main function."""
- parser = argparse.ArgumentParser(description='Test pruning directly')
- parser.add_argument('--model', type=str, default='mlp', help='Model name')
- parser.add_argument('--dataset', type=str, default='mnist', help='Dataset name')
- parser.add_argument('--strategy', type=str, default='low_rq',
- choices=['high_rq', 'low_rq', 'random'], help='Pruning strategy')
- parser.add_argument('--pruning-mode', type=str, default='layer_wise',
- choices=['global_joint', 'layer_wise', 'layer_isolated'],
- help='Pruning mode')
- parser.add_argument('--dropout-mode', type=str, default='unscaled',
- choices=['scaled', 'unscaled'], help='Dropout mode')
- parser.add_argument('--pruning-percent', type=float, default=0.5,
- help='Pruning percentage (0.0-1.0)')
- parser.add_argument('--device', type=str, default=None, help='Device (cuda/cpu)')
-
- args = parser.parse_args()
-
- # Convert percent to fraction if needed
- if args.pruning_percent > 1.0:
- args.pruning_percent /= 100.0
-
- # Run test
- verify_pruning(
- model_name=args.model,
- dataset_name=args.dataset,
- strategy=args.strategy,
- pruning_mode=args.pruning_mode,
- dropout_mode=args.dropout_mode,
- pruning_percent=args.pruning_percent,
- device=args.device
- )
-
-if __name__ == '__main__':
- main()
\ No newline at end of file
diff --git a/_arxiv/archive/scripts/run_benchmark.sh b/_arxiv/archive/scripts/run_benchmark.sh
deleted file mode 100755
index db589054..00000000
--- a/_arxiv/archive/scripts/run_benchmark.sh
+++ /dev/null
@@ -1,42 +0,0 @@
-#!/bin/bash
-
-# Benchmark script for testing tensorized network training
-# This script runs multiple benchmark configurations to compare
-# training speeds for different numbers of networks
-
-ROOT_DIR="/n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment"
-cd $ROOT_DIR
-
-# Set up common parameters
-EPOCHS=1
-DEVICE="cuda"
-BATCH_SIZE=64
-HIDDEN_SIZES="32,16"
-
-echo "----------------------------------------------"
-echo "Network Training Method Benchmark"
-echo "----------------------------------------------"
-echo "Date: $(date)"
-echo "Device: $DEVICE"
-echo "Hidden sizes: $HIDDEN_SIZES"
-echo "Epochs: $EPOCHS"
-echo "Batch size: $BATCH_SIZE"
-echo "----------------------------------------------"
-echo ""
-
-# Run benchmarks with different network counts
-for NUM_NETWORKS in 1 3 5 10 20
-do
- echo "=== Running benchmark with $NUM_NETWORKS networks ==="
- python benchmark_network_training.py \
- --num_networks $NUM_NETWORKS \
- --hidden_sizes $HIDDEN_SIZES \
- --epochs $EPOCHS \
- --device $DEVICE \
- --batch_size $BATCH_SIZE
- echo ""
-done
-
-echo "----------------------------------------------"
-echo "Benchmark completed"
-echo "----------------------------------------------"
\ No newline at end of file
diff --git a/_arxiv/archive/scripts/run_cascading_test.py b/_arxiv/archive/scripts/run_cascading_test.py
deleted file mode 100644
index 5e49e42f..00000000
--- a/_arxiv/archive/scripts/run_cascading_test.py
+++ /dev/null
@@ -1,73 +0,0 @@
-#!/usr/bin/env python
-"""
-Test script for running alignment experiments with cascading layer pruning
-"""
-
-import os
-import sys
-import logging
-import time
-
-# Add the src directory to the Python path
-script_dir = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(script_dir, "src"))
-
-try:
- from alignment.config import ExperimentConfig
- from alignment.experiments.alignment_experiments import AlignmentExperiment
- from alignment.experiments.alignment_experiments import set_logging_level
- print("Successfully imported alignment modules")
-except ImportError as e:
- print(f"Error importing alignment modules: {e}")
- sys.exit(1)
-
-def main():
- # Define experiment parameters
- config_file = "configs/config_alignment_experiment.yaml"
- experiment_type = "progressive_dropout"
- pruning_mode = "cascading_layer"
- run_name = "cascading_test"
-
- print(f"Starting alignment experiment with cascading layer pruning at {time.ctime()}")
- start_time = time.time()
-
- # Set up logging
- set_logging_level(logging.INFO)
-
- # Load configuration
- config_path = os.path.join(script_dir, config_file)
- config = ExperimentConfig.load(config_path)
-
- # Set experiment parameters
- config.experiment_type = experiment_type
-
- # Make sure the extra attribute exists
- if not hasattr(config, 'extra'):
- config.extra = type('ExtraConfig', (), {})()
-
- config.extra.dropout_pruning_mode = pruning_mode
-
- # Set experiment name
- config.experiment_name = run_name
-
- # Modify configuration for faster testing
- config.training.replicates = 3 # Use fewer networks
- config.alignment.dropout_steps = 3 # Use fewer dropout steps
-
- # Print configuration for debugging
- print(f"Using pruning mode: {pruning_mode}")
- print(f"Number of replicates: {config.training.replicates}")
- print(f"Number of dropout steps: {config.alignment.dropout_steps}")
-
- # Create and run experiment
- experiment = AlignmentExperiment(config)
- results, networks = experiment.run()
-
- end_time = time.time()
- print(f"Alignment experiment finished at {time.ctime()}")
- print(f"Total duration: {end_time - start_time:.2f} seconds")
-
- return 0
-
-if __name__ == "__main__":
- sys.exit(main())
\ No newline at end of file
diff --git a/_arxiv/archive/scripts/run_cascading_test.sh b/_arxiv/archive/scripts/run_cascading_test.sh
deleted file mode 100755
index a233ad25..00000000
--- a/_arxiv/archive/scripts/run_cascading_test.sh
+++ /dev/null
@@ -1,50 +0,0 @@
-#!/bin/bash
-
-# Load necessary modules (if needed)
-# module purge
-# module load python/3.12.5-fasrc01
-# module load cuda/12.4.1-fasrc01
-# module load cudnn/8.9.2.26_cuda12-fasrc01
-
-# Define the experiment parameters
-CONFIG_FILE="configs/config_alignment_experiment.yaml"
-EXPERIMENT_TYPE="progressive_dropout"
-PRUNING_MODE="cascading_layer"
-RUN_NAME="cascading_test"
-
-echo "Starting alignment experiment with cascading layer pruning at $(date)"
-start_time=$(date +%s)
-
-# Add the alignment package to the Python path
-export PYTHONPATH=/n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment
-
-# Execute the alignment experiment
-cd /n/holylabs/LABS/kempner_dev/Users/hsafaai/Code/alignment
-python -c "
-import logging
-import sys
-from alignment.config import ExperimentConfig
-from alignment.experiments.alignment_experiments import AlignmentExperiment, set_logging_level
-
-# Set up logging
-set_logging_level(logging.INFO)
-
-# Load configuration
-config_path = '${CONFIG_FILE}'
-config = ExperimentConfig.load(config_path)
-
-# Set experiment parameters
-config.experiment_type = '${EXPERIMENT_TYPE}'
-config.extra.dropout_pruning_mode = '${PRUNING_MODE}'
-config.run.name = '${RUN_NAME}'
-
-# Create and run experiment
-experiment = AlignmentExperiment(config)
-results, networks = experiment.run()
-
-print('Experiment completed successfully')
-"
-
-end_time=$(date +%s)
-echo "Alignment experiment finished at $(date)"
-echo "Total duration: $((end_time - start_time)) seconds."
\ No newline at end of file
diff --git a/_arxiv/archive/scripts/run_cascading_with_plots.py b/_arxiv/archive/scripts/run_cascading_with_plots.py
deleted file mode 100755
index e4cba3b2..00000000
--- a/_arxiv/archive/scripts/run_cascading_with_plots.py
+++ /dev/null
@@ -1,179 +0,0 @@
-#!/usr/bin/env python
-"""
-Run alignment experiment with cascading layer pruning and enhanced visualization
-"""
-
-import os
-import sys
-import logging
-import time
-import matplotlib.pyplot as plt
-import numpy as np
-import torch
-
-# Add the src directory to the Python path
-script_dir = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, os.path.join(script_dir, "src"))
-
-# Import alignment modules
-from alignment.config import ExperimentConfig
-from alignment.experiments.alignment_experiments import AlignmentExperiment, set_logging_level
-from alignment.utils.plotting import plot_dropout_results
-
-# Set style to match the alignment_preref version
-plt.style.use('seaborn-v0_8-whitegrid')
-plt.rcParams['figure.figsize'] = (10, 6)
-plt.rcParams['lines.linewidth'] = 2.5
-plt.rcParams['axes.grid'] = True
-plt.rcParams['grid.alpha'] = 0.3
-plt.rcParams['axes.labelsize'] = 14
-plt.rcParams['xtick.labelsize'] = 12
-plt.rcParams['ytick.labelsize'] = 12
-plt.rcParams['legend.fontsize'] = 12
-plt.rcParams['figure.titlesize'] = 16
-
-# Ensure wandb is available
-try:
- import wandb
-except ImportError:
- print("Warning: wandb not installed. Install with 'pip install wandb' for full functionality.")
- wandb = None
-
-def enhanced_plotting(results, figure_path, pruning_mode, dropout_mode, experiment_name):
- """Enhanced plotting function with consistent colors and styles"""
-
- # Define consistent colors and styles to match alignment_preref
- saved_figures = []
-
- # Extract data
- dropout_fractions = results["dropout_fractions"]
- accuracies = np.array(results["accuracies"])
-
- # Plot mean accuracy vs. dropout fraction with enhanced styling
- plt.figure(figsize=(12, 8))
-
- # Use a richer color palette and add markers for better visibility
- plt.plot(dropout_fractions, np.mean(accuracies, axis=0), 'o-',
- color='#1f77b4', linewidth=2.5, markersize=8,
- label="Mean Accuracy")
-
- # Add error bands with semi-transparency
- plt.fill_between(
- dropout_fractions,
- np.mean(accuracies, axis=0) - np.std(accuracies, axis=0),
- np.mean(accuracies, axis=0) + np.std(accuracies, axis=0),
- alpha=0.25, color='#1f77b4'
- )
-
- # Enhance plot formatting
- plt.xlabel("Dropout Fraction", fontsize=14)
- plt.ylabel("Accuracy (%)", fontsize=14)
- plt.title(f"{experiment_name}: {pruning_mode} Pruning", fontsize=16)
- plt.grid(True, alpha=0.3)
- plt.ylim(0, 100) # Set consistent y-axis range
- plt.xlim(0, max(dropout_fractions) * 1.05) # Add a small margin
- plt.legend(loc="upper right", fontsize=12)
-
- # Add more information to the plot
- plt.text(0.02, 0.02, f"Dropout Mode: {dropout_mode}", transform=plt.gca().transAxes,
- fontsize=10, bbox=dict(facecolor='white', alpha=0.8))
-
- # Save the enhanced figure
- if figure_path is not None:
- os.makedirs(figure_path, exist_ok=True)
- filename = os.path.join(
- figure_path,
- f"enhanced_{pruning_mode}_{dropout_mode}.png"
- )
- plt.savefig(filename, dpi=300, bbox_inches="tight")
- saved_figures.append(filename)
-
- # Log to wandb if available
- if wandb is not None and wandb.run is not None:
- wandb.log({"dropout_plot": wandb.Image(filename)})
-
- plt.close()
- else:
- plt.show()
-
- # Also generate the standard plots for comparison
- standard_plots = plot_dropout_results(
- results,
- figure_path,
- pruning_mode=pruning_mode,
- dropout_mode=dropout_mode,
- title_prefix=experiment_name
- )
-
- saved_figures.extend(standard_plots)
- return saved_figures
-
-def run_experiment(config_path):
- """Run the experiment with enhanced visualization"""
- print(f"Starting cascading layer pruning experiment at {time.ctime()}")
- start_time = time.time()
-
- # Set up logging
- set_logging_level(logging.INFO)
-
- # Load configuration
- config = ExperimentConfig.load(config_path)
-
- # Ensure we're using progressive dropout with cascading layer
- config.experiment_type = "progressive_dropout"
-
- # Make sure the extra attribute exists
- if not hasattr(config, 'extra'):
- config.extra = type('ExtraConfig', (), {})()
-
- # Set the pruning mode to cascading layer
- config.extra.dropout_pruning_mode = "cascading_layer"
-
- # Initialize wandb if available
- if wandb is not None and getattr(config.checkpointing, "use_wandb", False):
- wandb.init(
- project=getattr(config, "wandb_project", "neural_alignment"),
- entity=getattr(config, "wandb_entity", None),
- name=getattr(config, "experiment_name", "cascading_layer_test"),
- config=config.to_dict()
- )
- wandb.run.log_code(".", include_fn=lambda path: path.endswith(".py"))
-
- # Create and run experiment
- experiment = AlignmentExperiment(config)
- results, networks = experiment.run()
-
- # Enhance the default plots
- if "progressive_dropout" in results:
- dropout_results = results["progressive_dropout"]
- pruning_mode = config.extra.dropout_pruning_mode
- dropout_mode = getattr(config.extra, "dropout_mode", "scaled")
-
- # Create enhanced plots
- saved_figures = enhanced_plotting(
- dropout_results,
- experiment.figure_path,
- pruning_mode,
- dropout_mode,
- getattr(config, "experiment_name", "Progressive Dropout")
- )
-
- print(f"Generated {len(saved_figures)} plot files")
-
- end_time = time.time()
- print(f"Experiment finished at {time.ctime()}")
- print(f"Total duration: {end_time - start_time:.2f} seconds")
-
- # Finish wandb run
- if wandb is not None and wandb.run is not None:
- wandb.finish()
-
- return results, networks
-
-if __name__ == "__main__":
- if len(sys.argv) > 1:
- config_path = sys.argv[1]
- else:
- config_path = "configs/config_alignment_experiment.yaml"
-
- results, _ = run_experiment(config_path)
\ No newline at end of file
diff --git a/_arxiv/archive/scripts/run_fixed_experiment.py b/_arxiv/archive/scripts/run_fixed_experiment.py
deleted file mode 100644
index 2935332c..00000000
--- a/_arxiv/archive/scripts/run_fixed_experiment.py
+++ /dev/null
@@ -1,402 +0,0 @@
-#!/usr/bin/env python
-"""
-Run a fixed version of the alignment experiment.
-"""
-
-import os
-import sys
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.optim as optim
-import numpy as np
-import logging
-from tqdm import tqdm
-
-# Configure basic logging
-logging.basicConfig(level=logging.INFO,
- format='%(asctime)s [%(levelname)s] %(message)s',
- handlers=[logging.StreamHandler()])
-logger = logging.getLogger(__name__)
-
-# Add the project to path
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".")))
-
-# Import from the alignment package
-from alignment.config import ExperimentConfig
-from alignment.models.registry import create_model
-from alignment.datasets import load_dataset
-from alignment.metrics import get_metric
-from alignment.utils.plotting import plot_dropout_results
-
-def train_networks(networks, dataset, device, num_epochs=5, learning_rate=0.001):
- """Train multiple networks on the given dataset."""
- logger.info(f"Training {len(networks)} networks for {num_epochs} epochs")
-
- # Track training history for plotting
- training_history = {
- 'train_loss': [],
- 'train_acc': [],
- 'test_loss': [],
- 'test_acc': []
- }
-
- # Train each network
- trained_networks = []
- for i, network in enumerate(networks):
- logger.info(f"Training network {i+1}/{len(networks)}")
-
- # Move network to the device
- network = network.to(device)
-
- # Create optimizer
- optimizer = optim.Adam(network.parameters(), lr=learning_rate)
-
- # Training loop
- network.train()
- history = {
- 'train_loss': [],
- 'train_acc': [],
- 'test_loss': [],
- 'test_acc': []
- }
-
- for epoch in range(num_epochs):
- running_loss = 0.0
- correct = 0
- total = 0
-
- # Train on each batch
- for inputs, targets in tqdm(dataset.train_loader, desc=f"Network {i+1}, Epoch {epoch+1}/{num_epochs}"):
- inputs, targets = inputs.to(device), targets.to(device)
-
- # Zero the parameter gradients
- optimizer.zero_grad()
-
- # Forward pass
- outputs = network(inputs)
- loss = F.cross_entropy(outputs, targets)
-
- # Backward pass and optimize
- loss.backward()
- optimizer.step()
-
- # Track statistics
- running_loss += loss.item()
- _, predicted = outputs.max(1)
- total += targets.size(0)
- correct += predicted.eq(targets).sum().item()
-
- # Calculate epoch statistics
- train_loss = running_loss / len(dataset.train_loader)
- train_acc = 100.0 * correct / total
-
- # Evaluate on test set
- network.eval()
- test_correct = 0
- test_total = 0
- test_loss_sum = 0.0
-
- with torch.no_grad():
- for inputs, targets in dataset.test_loader:
- inputs, targets = inputs.to(device), targets.to(device)
- outputs = network(inputs)
-
- # Calculate loss
- loss = F.cross_entropy(outputs, targets, reduction='sum')
- test_loss_sum += loss.item()
-
- # Calculate accuracy
- _, predicted = outputs.max(1)
- test_total += targets.size(0)
- test_correct += predicted.eq(targets).sum().item()
-
- test_acc = 100.0 * test_correct / test_total
- test_loss = test_loss_sum / test_total
-
- # Store history
- history['train_loss'].append(train_loss)
- history['train_acc'].append(train_acc)
- history['test_loss'].append(test_loss)
- history['test_acc'].append(test_acc)
-
- # Log progress
- logger.info(f"Network {i+1}, Epoch {epoch+1}/{num_epochs}: "
- f"Train Loss={train_loss:.4f}, Train Acc={train_acc:.2f}%, "
- f"Test Loss={test_loss:.4f}, Test Acc={test_acc:.2f}%")
-
- # Add to trained networks
- trained_networks.append(network)
-
- # Accumulate training history (average across networks)
- if i == 0:
- # First network, initialize history
- training_history = history
- else:
- # Average with previous networks
- for key in training_history:
- if key in history:
- # Calculate running average
- for epoch_idx in range(len(history[key])):
- if epoch_idx < len(training_history[key]):
- training_history[key][epoch_idx] = (training_history[key][epoch_idx] * i + history[key][epoch_idx]) / (i + 1)
-
- logger.info(f"Completed training {len(networks)} networks")
- return trained_networks, training_history
-
-def test_pruning_strategies(networks, dataset, dropout_fractions, device):
- """Test different pruning strategies on trained networks."""
- logger.info(f"Testing pruning strategies on {len(networks)} networks")
-
- # Initialize results
- results = {
- 'dropout_fractions': dropout_fractions,
- 'accuracies': {'high_rq': [], 'low_rq': [], 'random': []},
- 'stds': {'high_rq': [], 'low_rq': [], 'random': []},
- 'losses': {'high_rq': [], 'low_rq': [], 'random': []}
- }
-
- # For each pruning percentage, test all networks
- for prune_idx, prune_percent in enumerate(dropout_fractions):
- logger.info(f"Testing pruning percentage: {prune_percent*100:.1f}%")
-
- # Store results for each network at this pruning percentage
- strategy_results = {
- 'high_rq': {'acc': [], 'loss': []},
- 'low_rq': {'acc': [], 'loss': []},
- 'random': {'acc': [], 'loss': []}
- }
-
- # Process each network
- for net_idx, network in enumerate(networks):
- logger.info(f"Processing network {net_idx+1}/{len(networks)}")
-
- # Save original weights
- original_weights = {}
- original_biases = {}
-
- for i, layer in enumerate(network.alignment_layers):
- if hasattr(layer, "weight") and layer.weight is not None:
- original_weights[i] = layer.weight.data.clone()
- if hasattr(layer, "bias") and layer.bias is not None:
- original_biases[i] = layer.bias.data.clone()
-
- # Get original accuracy
- if prune_idx == 0 or net_idx == 0:
- network.eval()
- correct = 0
- total = 0
- total_loss = 0.0
-
- with torch.no_grad():
- for inputs, targets in dataset.test_loader:
- inputs, targets = inputs.to(device), targets.to(device)
- outputs = network(inputs)
-
- # Compute loss
- loss = F.cross_entropy(outputs, targets, reduction='sum')
- total_loss += loss.item()
-
- # Compute accuracy
- _, predicted = outputs.max(1)
- total += targets.size(0)
- correct += predicted.eq(targets).sum().item()
-
- orig_accuracy = 100.0 * correct / total
- orig_loss = total_loss / total
-
- logger.info(f"Original accuracy: {orig_accuracy:.2f}%, loss: {orig_loss:.4f}")
-
- # Test each strategy
- for strategy in ["high_rq", "low_rq", "random"]:
- # Restore original weights
- for i, layer in enumerate(network.alignment_layers):
- if i in original_weights:
- layer.weight.data = original_weights[i].clone()
- if i in original_biases and hasattr(layer, "bias") and layer.bias is not None:
- layer.bias.data = original_biases[i].clone()
-
- # Skip pruning for 0% case
- if prune_percent == 0.0:
- acc = orig_accuracy
- loss = orig_loss
- else:
- # Apply pruning based on strategy
- total_neurons = 0
- total_pruned = 0
-
- # Apply pruning to each layer
- for i, layer in enumerate(network.alignment_layers):
- if i not in original_weights:
- continue
-
- # Compute neuron importance scores (using weight magnitude as proxy)
- weights = layer.weight.data
- input_dim = weights.shape[1]
- total_neurons += input_dim
-
- # Calculate importance scores (weight magnitude)
- neuron_scores = [torch.norm(weights[:, j]).item() for j in range(input_dim)]
-
- # Calculate how many neurons to prune
- num_to_drop = max(1, int(input_dim * prune_percent)) if prune_percent > 0 else 0
- total_pruned += num_to_drop
-
- if num_to_drop > 0:
- # Get indices to drop based on strategy
- if strategy == "high_rq": # Drop highest alignment neurons
- sorted_indices = np.argsort(neuron_scores)[::-1] # Sort descending
- to_drop = sorted_indices[:num_to_drop]
- elif strategy == "low_rq": # Drop lowest alignment neurons
- sorted_indices = np.argsort(neuron_scores) # Sort ascending
- to_drop = sorted_indices[:num_to_drop]
- else: # Random pruning
- all_indices = list(range(input_dim))
- np.random.shuffle(all_indices)
- to_drop = all_indices[:num_to_drop]
-
- # Zero out weights for these neurons
- for idx in to_drop:
- if idx < weights.shape[1]:
- layer.weight.data[:, idx] = 0.0
- if hasattr(layer, "bias") and layer.bias is not None and idx < layer.bias.data.shape[0]:
- layer.bias.data[idx] = 0.0
-
- logger.info(f"Pruned {total_pruned}/{total_neurons} neurons with strategy: {strategy}")
-
- # Evaluate pruned network
- network.eval()
- correct = 0
- total = 0
- total_loss = 0.0
-
- with torch.no_grad():
- for inputs, targets in dataset.test_loader:
- inputs, targets = inputs.to(device), targets.to(device)
- outputs = network(inputs)
-
- # Compute loss
- loss = F.cross_entropy(outputs, targets, reduction='sum')
- total_loss += loss.item()
-
- # Compute accuracy
- _, predicted = outputs.max(1)
- total += targets.size(0)
- correct += predicted.eq(targets).sum().item()
-
- acc = 100.0 * correct / total
- loss = total_loss / total
-
- # Store results for this network and strategy
- strategy_results[strategy]['acc'].append(acc)
- strategy_results[strategy]['loss'].append(loss)
-
- logger.info(f"Strategy {strategy}, Accuracy: {acc:.2f}%, Loss: {loss:.4f}")
-
- # Calculate mean and std for each strategy at this pruning percentage
- for strategy in ["high_rq", "low_rq", "random"]:
- # Get all networks' results for this strategy and pruning percentage
- accs = strategy_results[strategy]['acc']
- losses = strategy_results[strategy]['loss']
-
- # Calculate statistics
- mean_acc = np.mean(accs)
- std_acc = np.std(accs)
- mean_loss = np.mean(losses)
-
- # Add to results
- results['accuracies'][strategy].append(mean_acc)
- results['stds'][strategy].append(std_acc)
- results['losses'][strategy].append(mean_loss)
-
- logger.info(f"Pruning {prune_percent*100:.1f}%, Strategy {strategy}: "
- f"Mean acc: {mean_acc:.2f}%, Std: {std_acc:.2f}%")
-
- return results
-
-def main():
- """Main function to run the experiment."""
- # Set random seed for reproducibility
- torch.manual_seed(42)
- np.random.seed(42)
-
- # Load configuration
- config_path = "configs/config_alignment_experiment.yaml"
- config = ExperimentConfig.load(config_path)
-
- # Set device
- device = "cuda" if torch.cuda.is_available() else "cpu"
- logger.info(f"Using device: {device}")
-
- # Load dataset
- logger.info(f"Loading dataset: {config.dataset.dataset_name}")
- dataset = load_dataset(config.dataset)
-
- # Create multiple networks (replicates)
- num_networks = min(getattr(config.training, "replicates", 5), 2) # Limit to 2 networks for faster testing
- logger.info(f"Creating {num_networks} networks")
-
- networks = []
- for i in range(num_networks):
- logger.info(f"Creating model {i+1}/{num_networks}: {config.model.model_name}")
- network = create_model(config.model)
- networks.append(network)
-
- # Get dropout fractions from config
- dropin_min = config.alignment.dropout_min
- dropin_max = config.alignment.dropout_max
- num_dropout_fractions = config.alignment.dropout_steps
- dropout_fractions = np.linspace(dropin_min, dropin_max, num_dropout_fractions).tolist()
-
- # Train all networks
- trained_networks, training_history = train_networks(
- networks,
- dataset,
- device,
- num_epochs=getattr(config.training, "epochs", 5),
- learning_rate=getattr(config.training, "learning_rate", 0.001)
- )
-
- # Test pruning strategies
- pruning_results = test_pruning_strategies(
- trained_networks,
- dataset,
- dropout_fractions,
- device
- )
-
- # Add training history to results
- pruning_results['training_history'] = training_history
-
- # Create output directory
- results_dir = "debug_output"
- os.makedirs(results_dir, exist_ok=True)
-
- # Generate plots
- try:
- logger.info("Generating pruning plots with error bars")
-
- # Get pruning mode and dropout mode from config
- pruning_mode = getattr(config.extra, "dropout_pruning_mode", "global_joint")
- dropout_mode = getattr(config.extra, "dropout_mode", "scaled")
-
- saved_plots = plot_dropout_results(
- pruning_results,
- results_dir,
- title_prefix="Fixed Pruning Test",
- pruning_mode=pruning_mode,
- dropout_mode=dropout_mode
- )
-
- if saved_plots:
- logger.info(f"Generated {len(saved_plots)} plots: {saved_plots}")
-
- except Exception as e:
- logger.error(f"Error generating dropout plots: {str(e)}")
- import traceback
- logger.error(traceback.format_exc())
-
- logger.info("Experiment completed successfully!")
-
-if __name__ == "__main__":
- main()
\ No newline at end of file
diff --git a/_arxiv/archive/scripts/run_multi_strategy_experiment.py b/_arxiv/archive/scripts/run_multi_strategy_experiment.py
deleted file mode 100644
index 95d2a8cc..00000000
--- a/_arxiv/archive/scripts/run_multi_strategy_experiment.py
+++ /dev/null
@@ -1,614 +0,0 @@
-#!/usr/bin/env python
-"""
-Run an experiment with multiple pruning strategies.
-
-This script executes neural network pruning using three different strategies:
-1. high_rq: Prune neurons with highest alignment scores (weight magnitudes)
-2. low_rq: Prune neurons with lowest alignment scores (weight magnitudes)
-3. random: Prune neurons randomly
-
-The results are plotted for comparison.
-"""
-
-import os
-import sys
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.optim as optim
-import numpy as np
-import logging
-from tqdm import tqdm
-import matplotlib.pyplot as plt
-import copy
-
-# Configure basic logging
-logging.basicConfig(level=logging.INFO,
- format='%(asctime)s [%(levelname)s] %(message)s',
- handlers=[logging.StreamHandler()])
-logger = logging.getLogger(__name__)
-
-# Add the project to path
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".")))
-
-# Import from the alignment package
-from alignment.config import ExperimentConfig
-from alignment.models.registry import create_model
-from alignment.datasets import load_dataset
-from alignment.metrics import get_metric
-from alignment.utils.plotting import plot_dropout_results
-
-def train_networks(networks, dataset, device, num_epochs=5, learning_rate=0.001):
- """Train multiple networks efficiently in one training loop."""
- logger.info(f"Training {len(networks)} networks for {num_epochs} epochs")
-
- # Track training history for plotting
- training_history = {
- 'train_loss': [],
- 'train_acc': [],
- 'test_loss': [],
- 'test_acc': []
- }
-
- # Set up optimizer for each network
- optimizers = []
- for network in networks:
- network.to(device)
- optimizers.append(torch.optim.Adam(network.parameters(), lr=learning_rate))
-
- # Train all networks for each epoch
- epoch_pbar = tqdm(range(num_epochs), desc="Training epochs", position=0)
- for epoch in epoch_pbar:
- # Initialize epoch stats
- epoch_train_loss = 0.0
- epoch_train_acc = 0.0
- epoch_test_loss = 0.0
- epoch_test_acc = 0.0
-
- # Training phase
- net_pbar = tqdm(enumerate(zip(networks, optimizers)),
- desc=f"Epoch {epoch+1}/{num_epochs} networks",
- total=len(networks),
- position=1,
- leave=False)
-
- for network_idx, (network, optimizer) in net_pbar:
- network.train()
- running_loss = 0.0
- correct = 0
- total = 0
-
- # Train on each batch
- batch_pbar = tqdm(dataset.train_loader,
- desc=f"Network {network_idx+1}/{len(networks)}",
- position=2,
- leave=False)
-
- for inputs, targets in batch_pbar:
- inputs, targets = inputs.to(device), targets.to(device)
-
- # Zero the parameter gradients
- optimizer.zero_grad()
-
- # Forward pass
- outputs = network(inputs)
- loss = F.cross_entropy(outputs, targets)
-
- # Backward pass and optimize
- loss.backward()
- optimizer.step()
-
- # Track statistics
- running_loss += loss.item()
- _, predicted = outputs.max(1)
- total += targets.size(0)
- correct += predicted.eq(targets).sum().item()
-
- # Update batch progress bar
- if total > 0:
- batch_pbar.set_postfix({
- 'loss': running_loss / len(batch_pbar),
- 'acc': 100.0 * correct / total
- })
-
- # Calculate network training statistics
- if total > 0:
- network_train_loss = running_loss / len(dataset.train_loader)
- network_train_acc = 100.0 * correct / total
-
- # Accumulate for epoch average
- epoch_train_loss += network_train_loss
- epoch_train_acc += network_train_acc
-
- # Update network progress bar
- net_pbar.set_postfix({
- 'train_loss': network_train_loss,
- 'train_acc': network_train_acc
- })
-
- # Evaluation phase
- network.eval()
- test_correct = 0
- test_total = 0
- test_loss_sum = 0.0
-
- # Evaluation progress bar
- eval_pbar = tqdm(dataset.test_loader,
- desc=f"Evaluating network {network_idx+1}",
- position=2,
- leave=False)
-
- with torch.no_grad():
- for inputs, targets in eval_pbar:
- inputs, targets = inputs.to(device), targets.to(device)
- outputs = network(inputs)
-
- # Calculate loss
- loss = F.cross_entropy(outputs, targets, reduction='sum')
- test_loss_sum += loss.item()
-
- # Calculate accuracy
- _, predicted = outputs.max(1)
- test_total += targets.size(0)
- test_correct += predicted.eq(targets).sum().item()
-
- # Update eval progress bar
- if test_total > 0:
- eval_pbar.set_postfix({
- 'test_loss': test_loss_sum / test_total,
- 'test_acc': 100.0 * test_correct / test_total
- })
-
- # Calculate network testing statistics
- if test_total > 0:
- network_test_loss = test_loss_sum / test_total
- network_test_acc = 100.0 * test_correct / test_total
-
- # Accumulate for epoch average
- epoch_test_loss += network_test_loss
- epoch_test_acc += network_test_acc
-
- # Calculate and store epoch averages
- epoch_train_loss /= len(networks)
- epoch_train_acc /= len(networks)
- epoch_test_loss /= len(networks)
- epoch_test_acc /= len(networks)
-
- training_history['train_loss'].append(epoch_train_loss)
- training_history['train_acc'].append(epoch_train_acc)
- training_history['test_loss'].append(epoch_test_loss)
- training_history['test_acc'].append(epoch_test_acc)
-
- # Update epoch progress bar
- epoch_pbar.set_postfix({
- 'train_loss': epoch_train_loss,
- 'train_acc': f"{epoch_train_acc:.2f}%",
- 'test_loss': epoch_test_loss,
- 'test_acc': f"{epoch_test_acc:.2f}%"
- })
-
- # Log progress
- logger.info(f"Epoch {epoch+1}/{num_epochs}: "
- f"Train Loss={epoch_train_loss:.4f}, Train Acc={epoch_train_acc:.2f}%, "
- f"Test Loss={epoch_test_loss:.4f}, Test Acc={epoch_test_acc:.2f}%")
-
- logger.info(f"Completed training {len(networks)} networks.")
- return networks, training_history
-
-def run_pruning(networks, dataset, dropout_fractions, device, pruning_mode="layer_wise", dropout_mode="scaled"):
- """Run pruning with three different strategies and collect results in a more efficient way."""
- # Define the strategies to evaluate
- strategies = ["high_rq", "low_rq", "random"]
-
- # Initialize results structure
- results = {
- 'dropout_fractions': dropout_fractions,
- 'accuracies': {strategy: [] for strategy in strategies},
- 'stds': {strategy: [] for strategy in strategies},
- 'losses': {strategy: [] for strategy in strategies}
- }
-
- # Store original weights and biases for all networks
- logger.info("Saving original network weights")
- original_weights = {}
- original_biases = {}
-
- for net_idx, network in enumerate(networks):
- original_weights[net_idx] = {}
- original_biases[net_idx] = {}
-
- for layer_idx, layer in enumerate(network.alignment_layers):
- if hasattr(layer, "weight") and layer.weight is not None:
- original_weights[net_idx][layer_idx] = layer.weight.data.clone()
- if hasattr(layer, "bias") and layer.bias is not None:
- original_biases[net_idx][layer_idx] = layer.bias.data.clone()
-
- # First, get original accuracy for all networks (0% pruning)
- logger.info("Evaluating original network performance (0% pruning)")
- orig_accs = []
- orig_losses = []
-
- # Move all networks to device
- for network in networks:
- network.to(device)
-
- # Evaluate all networks at baseline
- for network in networks:
- network.eval()
- correct = 0
- total = 0
- total_loss = 0.0
-
- with torch.no_grad():
- for inputs, targets in dataset.test_loader:
- inputs, targets = inputs.to(device), targets.to(device)
- outputs = network(inputs)
-
- # Compute loss
- loss = F.cross_entropy(outputs, targets, reduction='sum')
- total_loss += loss.item()
-
- # Compute accuracy
- _, predicted = outputs.max(1)
- total += targets.size(0)
- correct += predicted.eq(targets).sum().item()
-
- if total > 0:
- orig_accs.append(100.0 * correct / total)
- orig_losses.append(total_loss / total)
-
- # Calculate baseline statistics
- orig_mean_acc = np.mean(orig_accs)
- orig_std_acc = np.std(orig_accs)
- orig_mean_loss = np.mean(orig_losses)
-
- # Add baseline values to all strategies (they start from the same point)
- for strategy in strategies:
- results['accuracies'][strategy].append(orig_mean_acc)
- results['stds'][strategy].append(orig_std_acc)
- results['losses'][strategy].append(orig_mean_loss)
-
- logger.info(f"Original accuracy: {orig_mean_acc:.2f}% ± {orig_std_acc:.2f}%")
-
- # Iterate through pruning fractions (skip the first one which is 0%)
- fraction_pbar = tqdm(enumerate(dropout_fractions),
- total=len(dropout_fractions),
- desc="Pruning fractions",
- position=0)
-
- for fraction_idx, fraction in fraction_pbar:
- if fraction_idx == 0 and fraction == 0.0:
- continue
-
- fraction_pbar.set_description(f"Pruning fraction: {fraction:.2f}")
-
- # For each strategy, process all networks at once
- strategy_pbar = tqdm(strategies, desc=f"Pruning strategies", position=1, leave=False)
-
- for strategy in strategy_pbar:
- strategy_pbar.set_description(f"Strategy: {strategy}")
-
- # Restore original weights for all networks at once
- for net_idx, network in enumerate(networks):
- for layer_idx in original_weights[net_idx]:
- layer = network.alignment_layers[layer_idx]
- layer.weight.data = original_weights[net_idx][layer_idx].clone()
- if layer_idx in original_biases[net_idx] and hasattr(layer, "bias") and layer.bias is not None:
- layer.bias.data = original_biases[net_idx][layer_idx].clone()
-
- # Apply pruning to all networks in parallel
- for net_idx, network in enumerate(networks):
- # Apply pruning based on the mode
- if pruning_mode == "global_joint":
- # Collect all neurons and their scores across all layers
- all_neurons = []
-
- for layer_idx, layer in enumerate(network.alignment_layers):
- if layer_idx not in original_weights[net_idx]:
- continue
-
- weights = layer.weight.data
- input_dim = weights.shape[1]
-
- # Calculate importance scores
- neuron_scores = torch.norm(weights, dim=0).cpu().numpy()
-
- # Store (layer_idx, neuron_idx, score) tuples
- for j, score in enumerate(neuron_scores):
- all_neurons.append((layer_idx, j, score))
-
- # Sort neurons by score based on strategy
- if strategy == "high_rq":
- # Sort by highest scores first
- all_neurons.sort(key=lambda x: x[2], reverse=True)
- elif strategy == "low_rq":
- # Sort by lowest scores first
- all_neurons.sort(key=lambda x: x[2])
- elif strategy == "random":
- # Shuffle neurons randomly
- import random
- random.shuffle(all_neurons)
-
- # Calculate how many neurons to prune
- total_neurons = len(all_neurons)
- num_to_drop = int(total_neurons * fraction)
-
- if num_to_drop > 0:
- # Get indices to drop
- to_drop = all_neurons[:num_to_drop]
-
- # Apply pruning
- for layer_idx, neuron_idx, _ in to_drop:
- layer = network.alignment_layers[layer_idx]
- # Zero out weights for this neuron
- if neuron_idx < layer.weight.data.shape[1]:
- layer.weight.data[:, neuron_idx] = 0.0
- if hasattr(layer, "bias") and layer.bias is not None and neuron_idx < layer.bias.data.shape[0]:
- layer.bias.data[neuron_idx] = 0.0
-
- elif pruning_mode == "layer_wise":
- # Apply pruning to each layer individually - can optimize with tensor operations
- for layer_idx, layer in enumerate(network.alignment_layers):
- if layer_idx not in original_weights[net_idx]:
- continue
-
- weights = layer.weight.data
- input_dim = weights.shape[1]
-
- # Calculate importance scores - vectorized version
- neuron_scores = torch.norm(weights, dim=0).cpu().numpy()
-
- # Calculate how many neurons to prune in this layer
- num_to_drop = int(input_dim * fraction)
-
- if num_to_drop > 0:
- # Get neurons to drop based on strategy
- if strategy == "high_rq":
- # Sort by highest scores first (descending)
- sorted_indices = np.argsort(neuron_scores)[::-1]
- to_drop = sorted_indices[:num_to_drop]
- elif strategy == "low_rq":
- # Sort by lowest scores first (ascending)
- sorted_indices = np.argsort(neuron_scores)
- to_drop = sorted_indices[:num_to_drop]
- elif strategy == "random":
- # Choose random neurons
- all_indices = np.arange(input_dim)
- np.random.shuffle(all_indices)
- to_drop = all_indices[:num_to_drop]
-
- # Apply pruning
- if to_drop.size > 0:
- # Create a mask tensor
- mask = torch.ones_like(weights)
-
- # Convert numpy array to tensor and ensure it's properly formatted for indexing
- to_drop_tensor = torch.tensor(to_drop.tolist(), device=weights.device, dtype=torch.long)
-
- # Use tensor indexing
- mask[:, to_drop_tensor] = 0.0
-
- # Apply mask to weights
- layer.weight.data = weights * mask
-
- # Apply to bias if it exists
- if hasattr(layer, "bias") and layer.bias is not None:
- bias_mask = torch.ones_like(layer.bias.data)
- valid_indices = [idx for idx in to_drop if idx < layer.bias.data.shape[0]]
- if valid_indices:
- bias_indices = torch.tensor(valid_indices, device=layer.bias.device, dtype=torch.long)
- bias_mask[bias_indices] = 0.0
- layer.bias.data = layer.bias.data * bias_mask
-
- # Evaluate all pruned networks efficiently
- network_accs = []
- network_losses = []
-
- for network in networks:
- # Evaluate
- network.eval()
- correct = 0
- total = 0
- total_loss = 0.0
-
- with torch.no_grad():
- for inputs, targets in dataset.test_loader:
- inputs, targets = inputs.to(device), targets.to(device)
- outputs = network(inputs)
-
- # Compute loss
- loss = F.cross_entropy(outputs, targets, reduction='sum')
- total_loss += loss.item()
-
- # Compute accuracy
- _, predicted = outputs.max(1)
- total += targets.size(0)
- correct += predicted.eq(targets).sum().item()
-
- if total > 0:
- accuracy = 100.0 * correct / total
- loss = total_loss / total
-
- network_accs.append(accuracy)
- network_losses.append(loss)
-
- # Calculate statistics for this fraction and strategy
- if network_accs:
- mean_acc = np.mean(network_accs)
- std_acc = np.std(network_accs)
- mean_loss = np.mean(network_losses)
-
- # Add to results
- results['accuracies'][strategy].append(mean_acc)
- results['stds'][strategy].append(std_acc)
- results['losses'][strategy].append(mean_loss)
-
- logger.info(f" {strategy}: Accuracy {mean_acc:.2f}% ± {std_acc:.2f}%")
-
- # Restore original weights for all networks
- logger.info("Restoring original weights")
- for net_idx, network in enumerate(networks):
- for layer_idx in original_weights[net_idx]:
- layer = network.alignment_layers[layer_idx]
- layer.weight.data = original_weights[net_idx][layer_idx].clone()
- if layer_idx in original_biases[net_idx] and hasattr(layer, "bias") and layer.bias is not None:
- layer.bias.data = original_biases[net_idx][layer_idx].clone()
-
- return results
-
-def plot_results(results, pruning_mode, dropout_mode, output_dir="multi_strategy_output"):
- """Create plots showing the results of different pruning strategies."""
- # Ensure output directory exists
- os.makedirs(output_dir, exist_ok=True)
-
- # Extract data for plotting
- fractions = results['dropout_fractions']
- strategies = list(results['accuracies'].keys())
-
- # Set up plot colors and markers
- colors = {'high_rq': 'red', 'low_rq': 'green', 'random': 'blue'}
- markers = {'high_rq': 'o', 'low_rq': 's', 'random': '^'}
- labels = {
- 'high_rq': 'Prune Highest Magnitude',
- 'low_rq': 'Prune Lowest Magnitude',
- 'random': 'Prune Random'
- }
-
- # Create accuracy plot
- plt.figure(figsize=(10, 6))
-
- for strategy in strategies:
- accs = results['accuracies'][strategy]
- stds = results['stds'][strategy]
-
- plt.errorbar(
- fractions,
- accs,
- yerr=stds,
- label=labels[strategy],
- color=colors[strategy],
- marker=markers[strategy],
- capsize=4,
- markersize=8,
- linewidth=2
- )
-
- plt.title(f'Accuracy vs. Pruning Fraction\n({pruning_mode} pruning, {dropout_mode} mode)', fontsize=14)
- plt.xlabel('Pruning Fraction', fontsize=12)
- plt.ylabel('Accuracy (%)', fontsize=12)
- plt.grid(True, linestyle='--', alpha=0.7)
- plt.legend(fontsize=12)
- plt.tight_layout()
-
- # Save plot
- accuracy_plot_file = os.path.join(output_dir, f'accuracy_plot_{pruning_mode}.png')
- plt.savefig(accuracy_plot_file, dpi=300)
- plt.close()
-
- # Also use the alignment plotting utility if available
- try:
- from alignment.utils.plotting import plot_dropout_results
-
- saved_plots = plot_dropout_results(
- results,
- output_dir,
- title_prefix=f"Multi-Strategy Pruning Test",
- pruning_mode=pruning_mode,
- dropout_mode=dropout_mode
- )
-
- if saved_plots:
- logger.info(f"Generated {len(saved_plots)} plots using alignment plotting utility")
- return saved_plots
- except Exception as e:
- logger.warning(f"Could not use alignment plotting utility: {str(e)}")
-
- return [accuracy_plot_file]
-
-def main():
- """Main function to run the multi-strategy pruning experiment."""
- # Set random seed for reproducibility
- torch.manual_seed(42)
- np.random.seed(42)
-
- # Load configuration
- config_path = "configs/config_alignment_experiment.yaml"
- config = ExperimentConfig.load(config_path)
-
- # Set device
- device = "cuda" if torch.cuda.is_available() else "cpu"
- logger.info(f"Using device: {device}")
-
- # Load dataset
- logger.info(f"Loading dataset: {config.dataset.dataset_name}")
- dataset = load_dataset(config.dataset)
-
- # Create multiple networks (replicates)
- num_networks = min(getattr(config.training, "replicates", 5), 3) # Limit to 3 networks for faster testing
- logger.info(f"Creating {num_networks} networks")
-
- networks = []
- for i in range(num_networks):
- logger.info(f"Creating model {i+1}/{num_networks}: {config.model.model_name}")
- network = create_model(config.model)
- networks.append(network)
-
- # Get dropout fractions from config
- dropin_min = config.alignment.dropout_min
- dropin_max = config.alignment.dropout_max
- num_dropout_fractions = config.alignment.dropout_steps
- dropout_fractions = np.linspace(dropin_min, dropin_max, num_dropout_fractions).tolist()
-
- # Get pruning and dropout modes
- pruning_mode = getattr(config.extra, "dropout_pruning_mode", "layer_wise")
- dropout_mode = getattr(config.extra, "dropout_mode", "scaled")
-
- # Train all networks
- trained_networks, training_history = train_networks(
- networks,
- dataset,
- device,
- num_epochs=getattr(config.training, "epochs", 10),
- learning_rate=getattr(config.training, "learning_rate", 0.001)
- )
-
- # Run pruning with different strategies
- pruning_results = run_pruning(
- trained_networks,
- dataset,
- dropout_fractions,
- device,
- pruning_mode=pruning_mode,
- dropout_mode=dropout_mode
- )
-
- # Add training history to results
- pruning_results['training_history'] = training_history
-
- # Create output directory
- results_dir = "multi_strategy_output"
- os.makedirs(results_dir, exist_ok=True)
-
- # Generate plots
- try:
- logger.info("Generating multi-strategy pruning plots")
-
- saved_plots = plot_results(
- pruning_results,
- pruning_mode,
- dropout_mode,
- output_dir=results_dir
- )
-
- if saved_plots:
- logger.info(f"Generated {len(saved_plots)} plots: {saved_plots}")
-
- except Exception as e:
- logger.error(f"Error generating plots: {str(e)}")
- import traceback
- logger.error(traceback.format_exc())
-
- logger.info("Multi-strategy experiment completed successfully!")
-
-if __name__ == "__main__":
- main()
\ No newline at end of file
diff --git a/_arxiv/archive/scripts/teacher_student.py b/_arxiv/archive/scripts/teacher_student.py
deleted file mode 100644
index 3109341d..00000000
--- a/_arxiv/archive/scripts/teacher_student.py
+++ /dev/null
@@ -1,93 +0,0 @@
-from tqdm import tqdm
-import numpy as np
-import torch
-
-from matplotlib import pyplot as plt
-
-import os
-import sys
-
-mainPath = os.path.dirname(os.path.abspath(__file__)) + "/.."
-sys.path.append(mainPath)
-
-from networkAlignmentAnalysis.models.registry import get_model
-from networkAlignmentAnalysis import utils
-from networkAlignmentAnalysis import train
-
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-
-def get_data(modes, signal_dist, noise_amplitude, batch_size):
- signal = signal_dist.sample((batch_size,))
- data = signal @ modes + torch.normal(0, noise_amplitude, (batch_size, modes.size(1)))
- return signal, data
-
-
-def train(teacher, student, modes, signal_dist, noise_amplitude, batch_size=256, num_epochs=1000):
- optimizer = torch.optim.SGD(student.parameters(), lr=1e-2)
- lossfn = torch.nn.MSELoss()
- track_loss = []
- alignment = []
- for epoch in tqdm(range(num_epochs)):
- signal, data = get_data(modes, signal_dist, noise_amplitude, batch_size)
- signal, data = signal.to(DEVICE), data.to(DEVICE)
-
- optimizer.zero_grad()
- target = teacher(signal)
- output = student(data, store_hidden=True)
- loss = lossfn(output, target)
- loss.backward()
- optimizer.step()
- track_loss.append(loss.item())
- alignment.append(student.measure_alignment(data, precomputed=True))
-
- alignment = [torch.stack(a) for a in utils.transpose_list(alignment)] # list across layers, (epochs x dim)
- return student, track_loss, alignment
-
-
-if __name__ == "__main__":
- print("using device: ", DEVICE)
-
- N = 100
- num_modes = 50
- shape = 0.7 * torch.ones((num_modes,))
- modes = torch.normal(0, 1, (num_modes, N))
- amplitude = torch.rand((num_modes,)) * 0.9 + 0.1
- signal_dist = torch.distributions.gamma.Gamma(shape, amplitude)
- noise_amplitude = 0.2
-
- model_name = "MLP"
- teacher = get_model(model_name, build=True, input_dim=num_modes, dropout=0.0, ignore_flag=False, linear=True).to(DEVICE).eval()
- student = get_model(model_name, build=True, input_dim=N, dropout=0.0, ignore_flag=False).to(DEVICE)
-
- student, track_loss, alignment = train(teacher, student, modes, signal_dist, noise_amplitude)
- mean_alignment = torch.stack([torch.mean(align, dim=1) for align in alignment])
-
- fig, ax = plt.subplots(1, 3, figsize=(9, 3), layout="constrained")
- ax[0].plot(track_loss)
- ax[0].set_ylim(0, 0.5)
- for layer in range(len(alignment)):
- ax[1].plot(mean_alignment[layer], label=f"layer{layer}")
- ax[1].legend()
-
- signal, data = get_data(modes, signal_dist, noise_amplitude, 1)
- ax[2].scatter(teacher(signal).detach(), student(data).detach())
- plt.show()
-
- N, S = 2, 10000
- num_modes = 2
- nonnormality = 1.5
- a = b = 1 / nonnormality
- modes = np.random.normal(0, 1, (num_modes, N)) * np.random.normal(0, 1, num_modes).reshape(-1, 1)
- signal = np.random.gamma(0.5, 1, (S, num_modes)) * np.sign(np.random.random((S, num_modes)) - 0.5)
- data = signal @ modes
- w, v = [np.array(o) for o in utils.smart_pca(torch.tensor(data).T)]
- print(w.shape, v.shape)
-
- plt.scatter(data[:, 0], data[:, 1], s=1, c=("b", 0.4))
- for mode in modes:
- plt.plot([0, mode[0]], [0, mode[1]], c="r")
- for evec in v.T:
- plt.plot([0, evec[0]], [0, evec[1]], c="k")
-
- plt.show()
diff --git a/_arxiv/archive/scripts/timetests/test_batched_alignment.py b/_arxiv/archive/scripts/timetests/test_batched_alignment.py
deleted file mode 100644
index 738ce96c..00000000
--- a/_arxiv/archive/scripts/timetests/test_batched_alignment.py
+++ /dev/null
@@ -1,41 +0,0 @@
-import time
-import torch
-from alignment import utils
-
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-# B, D, S, C = 1024, 25, 784, 32
-B, D, S, C = 1024, 800, 196, 64
-input = torch.normal(0, 1, (B, D, S)).to(DEVICE)
-weight = torch.normal(0, 1, (C, D)).to(DEVICE)
-
-
-def get_align_usual(input, weight):
- var_stride = torch.mean(torch.var(input, dim=1), dim=0)
- align_stride = torch.stack([utils.alignment(input[:, :, i], weight) for i in range(S)], dim=1)
- return utils.weighted_average(align_stride, var_stride.view(1, -1), 1, ignore_nan=True)
-
-
-def get_align_new(input, weight):
- var_stride = torch.mean(torch.var(input, dim=1), dim=0)
- cc = utils.batch_cov(input.transpose(0, 2))
- rq = torch.sum(torch.matmul(weight, cc) * weight, axis=2) / torch.sum(weight * weight, axis=1)
- prq = rq / torch.diagonal(cc, dim1=1, dim2=2).sum(1, keepdim=True)
- return utils.weighted_average(prq, var_stride.view(-1, 1), 0, ignore_nan=True)
-
-
-au = get_align_usual(input, weight)
-an = get_align_new(input, weight)
-print("align_usual and align_new produce same result?", torch.allclose(au, an))
-
-num_tests = 100
-
-t = time.time()
-for _ in range(num_tests):
- _ = get_align_usual(input, weight)
-print("get_align_usual : time per test=", (time.time() - t) / num_tests)
-
-t = time.time()
-for _ in range(num_tests):
- _ = get_align_new(input, weight)
-print("get_align_new : time per test=", (time.time() - t) / num_tests)
diff --git a/_arxiv/archive/scripts/timetests/test_dataloader.py b/_arxiv/archive/scripts/timetests/test_dataloader.py
deleted file mode 100644
index 130c2070..00000000
--- a/_arxiv/archive/scripts/timetests/test_dataloader.py
+++ /dev/null
@@ -1,106 +0,0 @@
-import os
-import sys
-
-mainPath = os.path.dirname(os.path.abspath(__file__)) + "/.."
-sys.path.append(mainPath)
-
-import time
-from tqdm import tqdm
-
-import torch
-import torchvision
-from torchvision.transforms import v2 as transforms
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.optim as optim
-
-from alignment import files
-
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-
-class Net(nn.Module):
- def __init__(self):
- super().__init__()
- self.conv1 = nn.Conv2d(3, 6, 5)
- self.pool = nn.MaxPool2d(2, 2)
- self.conv2 = nn.Conv2d(6, 16, 5)
- self.fc1 = nn.Linear(16 * 5 * 5, 120)
- self.fc2 = nn.Linear(120, 84)
- self.fc3 = nn.Linear(84, 10)
-
- def forward(self, x):
- x = self.pool(F.relu(self.conv1(x)))
- x = self.pool(F.relu(self.conv2(x)))
- x = torch.flatten(x, 1) # flatten all dimensions except batch
- x = F.relu(self.fc1(x))
- x = F.relu(self.fc2(x))
- x = self.fc3(x)
- return x
-
-
-def check_time(num_workers, batch_size, pin_memory, fast_loader):
- transform = transforms.Compose(
- transforms.ToImage(),
- transforms.ToDtype(torch.float32, scale=True),
- transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
- )
-
- trainset = torchvision.datasets.CIFAR10(root=files.data_path(), train=True, download=False, transform=transform)
-
- loader_kwargs = dict(batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=pin_memory)
-
- if fast_loader:
- trainloader = torch.utils.data.DataLoader(trainset, **loader_kwargs, persistent_workers=fast_loader)
- else:
- trainloader = torch.utils.data.DataLoader(trainset, **loader_kwargs, persistent_workers=fast_loader)
-
- net = Net()
- net.to(DEVICE)
-
- criterion = nn.CrossEntropyLoss()
- optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
-
- time_per_batch = torch.zeros((2, len(trainloader)))
- batch_time = time.time()
- for cycle in range(2):
- for idx, (images, labels) in enumerate(trainloader):
- images, labels = images.to(DEVICE), labels.to(DEVICE)
-
- optimizer.zero_grad()
-
- outputs = net(images)
- loss = criterion(outputs, labels)
- loss.backward()
- optimizer.step()
-
- time_per_batch[cycle, idx] = time.time() - batch_time
- batch_time = time.time()
-
- avg_batch = torch.mean(time_per_batch[:, 1:])
- total_time = torch.sum(time_per_batch)
- init_time = time_per_batch[0, 0] - avg_batch # estimate of prepare time
- second_init = time_per_batch[1, 0] - avg_batch
- return init_time, avg_batch, total_time, second_init
-
-
-if __name__ == "__main__":
- print("using device: ", DEVICE)
-
- # tests
- num_workers = [0, 2, 4]
- batch_size = [8, 1024]
- pin_memory = [True, False]
- use_fast_loader = [True, False]
-
- for nw in num_workers:
- for bs in batch_size:
- for pin in pin_memory:
- for fast in use_fast_loader:
- if nw == 0 and fast:
- continue # persistent memory is irrelevant when num_workers=0
-
- init, avg, total, second = check_time(nw, bs, pin, fast)
- print(f"NW={nw}, BS={bs}, Pin={pin}, Persistent={fast}")
- print(f"Dur={total:.3f}, PerBatch={avg:.2f}, Prep={init:.2f}, SecondPrep={second:.2f}")
- print("")
diff --git a/configs/README.md b/configs/README.md
index 26b6f47d..2144ae9e 100644
--- a/configs/README.md
+++ b/configs/README.md
@@ -1,41 +1,31 @@
# Configuration Guide
-All experiments are configured via YAML files.
-
----
+Experiments are configured using YAML files.
## Template
-`template.yaml` - Complete template with all available parameters documented inline.
-
----
-
-## Examples
-
-Compact, ready-to-use configurations:
-
-### LLaMA-3
+`template.yaml` - Complete template with all available parameters.
-- `examples/llama3_scoring.yaml` - Compute per-neuron importance scores
-- `examples/llama3_pruning.yaml` - Redundancy-aware pruning
+## Example Configurations
-### Vision Models
+Ready-to-use configurations in `examples/`:
-- `examples/resnet_pruning.yaml` - ResNet-18 pruning with dependency handling
-- `examples/mnist_basic.yaml` - Simple MNIST analysis
+**Vision Models**
+- `mnist_basic.yaml` - MLP analysis on MNIST
+- `resnet_pruning.yaml` - ResNet-18 pruning on CIFAR-10
----
+**LLaMA Models**
+- `llama3_scoring.yaml` - Compute importance scores
+- `llama3_pruning.yaml` - Apply pruning
## Usage
-Run any experiment:
-
+Run experiment:
```bash
python scripts/run_experiment.py --config configs/examples/resnet_pruning.yaml
```
Override parameters:
-
```bash
python scripts/run_experiment.py \
--config configs/examples/resnet_pruning.yaml \
@@ -44,27 +34,23 @@ python scripts/run_experiment.py \
--target-sparsity 0.5
```
----
-
-## Creating Custom Configs
+## Creating Configurations
1. Copy template: `cp configs/template.yaml configs/my_config.yaml`
-2. Modify parameters (all options documented inline)
+2. Edit parameters as needed
3. Run: `python scripts/run_experiment.py --config configs/my_config.yaml`
----
-
-## Parameter Categories
+## Configuration Sections
-- `experiment`: Name, seed, device, output directory
-- `model`: Architecture, pretrained, layers to track
-- `dataset`: Data source, batch size, preprocessing
-- `metrics`: Which metrics to compute and their parameters
-- `training`: Training parameters (if training from scratch)
-- `pruning`: Pruning strategy, distribution, scoring method
-- `layer_config`: Architecture-specific settings (CNN, transformer)
-- `analysis`: Analysis options (class-conditioned, save options)
-- `visualization`: Plot generation settings
-- `advanced`: Backend, parallelization, optimization options
+- `experiment` - Name, seed, device, output directory
+- `model` - Architecture, pretrained weights, layers to track
+- `dataset` - Data source, batch size, preprocessing
+- `metrics` - Metrics to compute and their parameters
+- `training` - Training parameters
+- `pruning` - Pruning strategy, distribution, scoring
+- `layer_config` - Architecture-specific settings
+- `analysis` - Analysis options
+- `visualization` - Plot settings
+- `advanced` - Backend and optimization options
-See `template.yaml` for complete parameter documentation.
+See `template.yaml` for detailed parameter documentation.
diff --git a/docs/README.md b/docs/README.md
index 3925a8db..fc33b8fd 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -1,45 +1,34 @@
# Alignment Framework Documentation
-Neural Network Alignment Analysis & Intelligent Pruning
-
----
+Neural network alignment analysis and intelligent pruning.
## Getting Started
- [Installation Guide](installation.md) - Setup and dependencies
-- [Usage Guide](usage.md) - How to run experiments with YAML configs
-- [User Guide](user_guide.md) - Complete usage documentation
-
----
+- [Usage Guide](usage.md) - Running experiments with configuration files
+- [User Guide](user_guide.md) - Comprehensive framework documentation
## Reference
-- [API Reference](api_reference.md) - Class and method documentation
-- [Quick Reference](quick_reference.md) - Code examples and patterns
-- [Configuration Reference](../configs/template.yaml) - All parameters with options
-
----
-
-## Version Information
-
-- [Changelog](changelog.md) - Release notes and version history
-
----
+- [API Reference](api_reference.md) - API documentation
+- [Quick Reference](quick_reference.md) - Code examples
+- [Configuration Template](../configs/template.yaml) - All configuration parameters
## Examples
-Working example configurations:
-
-- `configs/examples/mnist_basic.yaml` - Simple MNIST analysis
-- `configs/examples/resnet_pruning.yaml` - ResNet pruning on CIFAR-10
-- `configs/examples/llama3_scoring.yaml` - LLaMA-3 per-neuron scoring
-- `configs/examples/llama3_pruning.yaml` - LLaMA-3 pruning
+Example configuration files in `configs/examples/`:
-Run: `python scripts/run_experiment.py --config [path]`
+- `mnist_basic.yaml` - MLP analysis on MNIST
+- `resnet_pruning.yaml` - ResNet pruning on CIFAR-10
+- `llama3_scoring.yaml` - LLaMA scoring
+- `llama3_pruning.yaml` - LLaMA pruning
----
+Run experiments:
+```bash
+python scripts/run_experiment.py --config configs/examples/mnist_basic.yaml
+```
-## Quick Example
+## Quick Start
```python
from alignment import ModelWrapper, get_metric
@@ -53,9 +42,7 @@ weights = wrapper.get_layer_weights()
scores = rq.compute(acts['layer_input'], weights['layer'])
```
----
-
-## Building HTML Documentation
+## Building Documentation
```bash
cd docs
diff --git a/docs/api_reference.md b/docs/api_reference.md
index ed32cedc..22e4d492 100644
--- a/docs/api_reference.md
+++ b/docs/api_reference.md
@@ -1,8 +1,6 @@
# API Reference
-**Quick reference for alignment framework APIs**
-
----
+Quick reference for framework APIs.
## Core Classes
@@ -35,8 +33,6 @@ metric.requires_outputs # bool: needs layer outputs
metric.compute(inputs, weights, outputs, **kwargs) # Returns scores
```
----
-
## Metrics
### Rayleigh Quotient
@@ -102,8 +98,6 @@ alignment = grad_align.compute(
# Returns correlation between local signal and backprop
```
----
-
## Services
### ActivationCaptureService
@@ -176,8 +170,6 @@ masks = MaskOperations.global_threshold_mask(
)
```
----
-
## Pruning
### Quick Pruning
@@ -229,8 +221,6 @@ results = optimizer.compare_strategies_parallel(
)
```
----
-
## Training
### Training Callback
@@ -253,8 +243,6 @@ callback.on_batch_end(wrapper, inputs, targets, step)
history = callback.get_history()
```
----
-
## Model Wrappers
### Generic Wrapper
@@ -301,8 +289,6 @@ wrapper.ffn_layers # {'expansion': [...], 'contraction': [...]}
wrapper.attention_layers # {'q': [...], 'k': [...], 'v': [...]}
```
----
-
## Utility Functions
### Layer Detection
@@ -329,8 +315,6 @@ cov = estimate_covariance(
)
```
----
-
## Configuration Parameters
### Metric Parameters
@@ -356,8 +340,6 @@ cov = estimate_covariance(
**Direction:** 'low' (prune unimportant), 'high' (ablation), 'random' (baseline)
----
-
## Common Patterns
### Analyze Pretrained Model
@@ -408,15 +390,11 @@ results = optimizer.compare_strategies_parallel(
)
```
----
-
## Examples
See `examples/` directory for complete workflows:
- `07_mnist_intelligent_pruning.py` - End-to-end pruning
-- `08_llama_ffn_pruning.py` - LLM FFN analysis
-- `09_attention_neuron_vs_head_pruning.py` - Attention analysis
-
-Run: `python examples/07_mnist_intelligent_pruning.py`
+- `08_llama_ffn_pruning.py` - LLM feed-forward analysis
+- `09_attention_neuron_vs_head_pruning.py` - Attention layer analysis
diff --git a/docs/quick_reference.md b/docs/quick_reference.md
index a2bd8137..e25f0cac 100644
--- a/docs/quick_reference.md
+++ b/docs/quick_reference.md
@@ -1,8 +1,6 @@
# Quick Reference
-**Common code patterns and examples**
-
----
+Common code patterns and examples.
## Basic Usage
@@ -66,8 +64,6 @@ result = prune_with_all_options(
)
```
----
-
## Metrics
### Rayleigh Quotient
@@ -124,30 +120,33 @@ scorer = NodeScoringService(
scores = scorer.compute_composite_scores(inputs, weights, targets)
```
----
-
## Pruning
-### Pruning Strategies
+### Available Strategies
```python
-# Available strategies:
strategies = [
- 'magnitude', # L1/L2 norm
- 'random', # Baseline
- 'gradient', # Gradient-based
- 'alignment', # RQ-based
- 'composite', # Redundancy-aware
- 'movement', # Training-aware
- 'adaptive', # Adaptive per-layer
- 'ultimate' # Multi-stage (best)
+ 'magnitude',
+ 'random',
+ 'gradient',
+ 'alignment',
+ 'composite',
+ 'movement',
+ 'adaptive',
+ 'ultimate'
]
```
### Distribution Methods
```python
-# How to allocate sparsity across layers:
+distributions = [
+ 'uniform',
+ 'global_threshold',
+ 'adaptive_sensitivity',
+ 'importance_weighted',
+ 'size_proportional'
+]
distributions = [
'uniform', # Same % per layer
'global_threshold', # Global score threshold
@@ -159,16 +158,14 @@ distributions = [
]
```
-### Direction Options
+### Pruning Direction
```python
-direction='low' # Prune low-scoring (production default)
-direction='high' # Prune high-scoring (ablation study)
-direction='random' # Random (baseline)
+direction='low' # Prune low-scoring neurons
+direction='high' # Prune high-scoring neurons (ablation)
+direction='random' # Random pruning (baseline)
```
----
-
## Training Integration
### Track Metrics During Training
@@ -206,8 +203,6 @@ alignment = grad_align.compute(
# High alignment = Hebbian rule works for this neuron
```
----
-
## Architecture-Specific
### CNNs
@@ -256,8 +251,6 @@ scores = redundancy.compute(outputs=outputs)
# [11008] - one per neuron
```
----
-
## Configuration
### Minimal Config
@@ -290,8 +283,6 @@ pruning:
Run: `python scripts/run_experiment.py --config my_config.yaml`
----
-
## Performance Tips
### For Large Models (LLMs)
@@ -320,8 +311,6 @@ optimizer = ParallelPruningOptimizer()
results = optimizer.prune_ensemble_parallel(networks, ...)
```
----
-
## Common Patterns
### Complete Pruning Workflow
@@ -361,8 +350,6 @@ train(model, train_loader, epochs=20)
accuracy = evaluate(model, test_loader)
```
----
-
## Troubleshooting
### Singular Covariance
@@ -387,7 +374,5 @@ from alignment.pruning.dependency_aware import prune_model_with_dependencies
result = prune_model_with_dependencies(model, scores, amount)
```
----
-
-See [USER_GUIDE.md](USER_GUIDE.md) for detailed documentation.
+See [user_guide.md](user_guide.md) for detailed documentation.
diff --git a/docs/usage.md b/docs/usage.md
index b531e2a7..e86eae76 100644
--- a/docs/usage.md
+++ b/docs/usage.md
@@ -1,61 +1,48 @@
-# How to Run Experiments with YAML Configs
+# Usage Guide
-## Quick Start
+## Running Experiments
-### Step 1: Activate Environment
+The framework uses YAML configuration files to specify experiments. This approach allows reproducible experiments and easy parameter management.
+
+### Basic Usage
```bash
conda activate alignment
cd /path/to/alignment
-```
-
-### Step 2: Run Experiment
-
-```bash
python scripts/run_experiment.py --config configs/examples/mnist_basic.yaml
```
----
-
-## Available Example Configs
+### Example Configurations
-### 1. MNIST Basic Analysis
+The `configs/examples/` directory contains pre-configured experiments:
+**MNIST Basic Analysis**
```bash
python scripts/run_experiment.py --config configs/examples/mnist_basic.yaml
```
+Trains an MLP on MNIST and computes alignment scores.
-Computes: RQ scores for simple MLP on MNIST
-
-### 2. ResNet Pruning
-
+**ResNet Pruning**
```bash
python scripts/run_experiment.py --config configs/examples/resnet_pruning.yaml
```
+Applies pruning to ResNet-18 on CIFAR-10 using alignment-based importance scores.
-Performs: Redundancy-aware pruning on ResNet-18 with CIFAR-10
-
-### 3. LLaMA-3 Scoring
-
+**LLaMA-3 Scoring**
```bash
python scripts/run_experiment.py --config configs/examples/llama3_scoring.yaml
```
+Computes per-neuron importance scores for LLaMA model feed-forward layers.
-Computes: Per-neuron importance scores for LLaMA-3 FFN
-
-### 4. LLaMA-3 Pruning
-
+**LLaMA-3 Pruning**
```bash
python scripts/run_experiment.py --config configs/examples/llama3_pruning.yaml
```
-
-Performs: Redundancy-aware pruning of LLaMA-3 model
-
----
+Prunes LLaMA model using information-theoretic importance scores.
## Command-Line Overrides
-Override any parameter:
+Override configuration parameters from the command line:
```bash
python scripts/run_experiment.py \
@@ -65,62 +52,149 @@ python scripts/run_experiment.py \
--target-sparsity 0.5
```
-Common overrides:
-- `--device cuda:0` - GPU selection
-- `--batch-size 64` - Batch size
-- `--target-sparsity 0.7` - Pruning amount
-- `--epochs 50` - Training epochs
-- `--output-dir ./my_results` - Output directory
+Common override options:
+- `--device cuda:0` - Select GPU device
+- `--batch-size 64` - Set batch size
+- `--target-sparsity 0.7` - Set pruning target
+- `--epochs 50` - Set training epochs
+- `--output-dir ./results` - Set output directory
----
+## Creating Custom Configurations
-## Creating Custom Configs
+### From Template
-### Method 1: Copy and Modify Template
+Copy the template and modify for your needs:
```bash
cp configs/template.yaml configs/my_experiment.yaml
-# Edit my_experiment.yaml
+# Edit my_experiment.yaml with desired parameters
python scripts/run_experiment.py --config configs/my_experiment.yaml
```
-### Method 2: Copy Existing Example
+### From Existing Example
+
+Start with an example configuration:
```bash
cp configs/examples/resnet_pruning.yaml configs/my_resnet.yaml
-# Modify my_resnet.yaml
+# Modify specific parameters
python scripts/run_experiment.py --config configs/my_resnet.yaml
```
----
-
-## Config Structure
+## Configuration Structure
-All configs have the same structure:
+All configuration files follow the same structure:
```yaml
-experiment: # Experiment settings
- name: "..."
- seed: 42
+```yaml
+experiment:
+ name: "my_experiment"
device: "cuda"
-
-model: # Model architecture
- name: "..." # 'resnet18', 'mlp', 'meta-llama/...'
+
+model:
+ name: "resnet18"
pretrained: true
-
-dataset: # Data source
- name: "..." # 'mnist', 'cifar10', 'wikitext'
+
+dataset:
+ name: "cifar10"
batch_size: 128
-metrics: # Metrics to compute
+metrics:
enabled: ['rayleigh_quotient']
-pruning: # Pruning settings (optional)
+pruning:
+ enabled: false
+```
+
+See `configs/template.yaml` for all available parameters.
+
+## Experiment Types
+
+### Computing Metrics
+
+Compute alignment and information-theoretic scores:
+
+```yaml
+metrics:
+ enabled: ['rayleigh_quotient', 'pairwise_redundancy_gaussian', 'synergy_gaussian_mmi']
+
+ rayleigh_quotient:
+ relative: true
+ regularization: 1.0e-6
+
+ pairwise_redundancy_gaussian:
+ mode: 'output_based'
+ num_pairs: 10
+
+training:
+ enabled: false
+pruning:
+ enabled: false
+```
+
+### Training Networks
+
+Train from scratch with optional metric tracking:
+
+```yaml
+training:
enabled: true
- strategy: '...'
+ epochs: 100
+ learning_rate: 0.001
+ optimizer: 'adam'
+ compute_metrics_during_training: false
+```
+
+### Pruning Networks
+
+Apply pruning with specified strategy:
+
+```yaml
+pruning:
+ enabled: true
+ strategy: 'composite'
target_sparsity: 0.7
+ distribution: 'adaptive_sensitivity'
+ scoring: 'rayleigh_quotient'
+ structured: true
+
+ fine_tune:
+ enabled: true
+ epochs: 20
+ learning_rate: 0.0001
+```
+
+### Multi-Level Pruning
+
+Test multiple sparsity levels:
+
+```yaml
+pruning:
+ enabled: true
+ sparsity_levels: [0.3, 0.5, 0.7, 0.9]
+ strategy: 'magnitude'
+```
-# ... other sections as needed
+## Output Structure
+
+Results are saved to the specified output directory:
+
+```
+results/[experiment_name]/
+├── config.yaml # Configuration used
+├── results.json # Numerical results
+├── scores/ # Per-layer importance scores
+├── plots/ # Visualizations
+└── checkpoints/ # Model checkpoints
+```
+
+## Workflow
+
+1. Create or select configuration file
+2. Activate environment: `conda activate alignment`
+3. Run experiment: `python scripts/run_experiment.py --config [path]`
+4. Results saved to output directory
+5. Analyze results and visualizations
```
See `configs/template.yaml` for complete parameter reference.
diff --git a/docs/user_guide.md b/docs/user_guide.md
index ebc8b3b4..217406c4 100644
--- a/docs/user_guide.md
+++ b/docs/user_guide.md
@@ -1,63 +1,128 @@
# Alignment Framework User Guide
-**Comprehensive guide to using the alignment framework for neural network analysis and pruning**
-
----
+Comprehensive guide to neural network alignment analysis and pruning.
## Table of Contents
1. [Core Concepts](#core-concepts)
-2. [Computing Metrics](#computing-metrics)
-3. [Pruning Strategies](#pruning-strategies)
-4. [Architecture Support](#architecture-support)
-5. [Configuration](#configuration)
-6. [Advanced Features](#advanced-features)
-
----
+2. [Available Metrics](#available-metrics)
+3. [Computing Metrics](#computing-metrics)
+4. [Pruning Strategies](#pruning-strategies)
+5. [Architecture Support](#architecture-support)
+6. [Configuration](#configuration)
## Core Concepts
-### Rayleigh Quotient (RQ)
+### Alignment
-Measures how well neuron weights align with input principal components:
+Alignment measures how well neuron weights align with the structure of their input activations. The Rayleigh Quotient quantifies the proportion of input variance captured by each neuron's weight vector.
+For weight vector w and input covariance Σ:
```
RQ(w) = (w^T Σ w) / (w^T w · tr(Σ))
```
-Higher RQ indicates better alignment with dominant input variance.
+Higher RQ values indicate the neuron aligns with dominant input variance directions.
-### Class-Conditioned RQ (ΔRQ)
+### Class-Conditioned Alignment
-Measures task-relevant alignment:
+Class-conditioned analysis compares alignment across different classes versus overall alignment. The difference (ΔRQ) indicates task-relevant alignment:
```
ΔRQ = RQ(overall) - E[RQ(class-conditioned)]
```
-Positive ΔRQ indicates the neuron captures discriminative features.
+Positive ΔRQ means the neuron captures discriminative features between classes.
-### Redundancy
+### Information Theory
-Measures overlap between neuron pairs:
+The framework uses information-theoretic measures to quantify relationships between neurons:
-```
-R(i,j) = I(Y_i; Y_j) = -0.5 · log(1 - ρ²)
-```
+- **Mutual Information**: Shared information between two variables
+- **Redundancy**: Information overlap between neuron pairs
+- **Synergy**: Information that emerges only from joint neuron outputs
+- **Partial Information Decomposition (PID)**: Decomposes information into unique, redundant, and synergistic components
-High redundancy means neurons capture similar information.
+## Available Metrics
-### Synergy
+### Alignment Metrics
-Measures complementary information:
+**Rayleigh Quotient (RQ)**
+- Measures alignment with input covariance structure
+- Returns per-neuron scores indicating variance captured
+- Supports relative normalization and regularization
+- Works with linear, convolutional, and transformer layers
-```
-S(Z; Y_i, Y_j) = I(Z; Y_i, Y_j) - I(Z; Y_i) - I(Z; Y_j) + min(I(Z; Y_i), I(Z; Y_j))
-```
+**Class-Conditioned RQ**
+- Computes RQ separately for each class
+- Returns ΔRQ showing task-relevant alignment
+- Requires labeled data
+
+**Spectral Alignment**
+- Analyzes eigenvalue decomposition of weight-covariance product
+- Provides detailed spectral properties
+
+### Information-Theoretic Metrics
+
+**Mutual Information (MI)**
+- Gaussian MI using analytic formula: MI = -0.5 log(1 - ρ²)
+- Can compute MI between neuron outputs and targets
+- Supports both continuous and categorical targets
+
+**Pairwise Redundancy**
+- Computes redundancy between neuron pairs
+- Output-based mode for efficiency with large models
+- Covariance-based mode for precise estimates
+- Sampling strategies: random, nearest, or all pairs
+- Returns per-neuron average redundancy scores
+
+**Synergy (MMI)**
+- Measures complementary information from neuron pairs
+- Uses Minimum Mutual Information redundancy definition
+- Returns per-neuron synergy scores
+- Requires target labels
+
+**Partial Information Decomposition (PID)**
+- Decomposes information into unique, redundant, and synergistic parts
+- Shared information: overlap between two neurons
+- Unique information: contribution from single neurons
+- Synergistic information: emergent from joint outputs
+
+**Higher-Order Metrics**
+- Total correlation: generalization of MI to multiple variables
+- Interaction information: higher-order dependencies
+- Connected information: network-level measures
+
+### Gradient-Based Metrics
-High synergy means neurons provide unique joint information.
+**Gradient Alignment**
+- Compares gradient-based learning with local learning rules
+- Analyzes alignment between backpropagation and Hebbian-style updates
+- Useful for biologically-plausible learning analysis
----
+**Local Learning Rule Search**
+- Finds optimal local learning rules per neuron
+- Searches over combinations of pre/post-synaptic activity
+- Returns best-fit local rule for each neuron
+
+### Task-Specific Metrics
+
+**Classification Alignment**
+- Measures discriminative power for classification tasks
+- Analyzes separation of class representations
+
+**Vision Task Alignment**
+- Specialized for vision architectures
+- Analyzes feature selectivity patterns
+
+### Similarity Metrics
+
+**Centered Kernel Alignment (CKA)**
+- Compares representations between layers or models
+- Invariant to isotropic scaling
+
+**Canonical Correlation Analysis (CCA)**
+- Finds maximally correlated directions between representations
## Computing Metrics
@@ -66,13 +131,13 @@ High synergy means neurons provide unique joint information.
```python
from alignment import ModelWrapper, get_metric
-# Wrap model
+# Wrap model to track layers
wrapper = ModelWrapper(model, tracked_layers=['conv1', 'fc1'])
-# Get metric
+# Get metric instance
rq_metric = get_metric('rayleigh_quotient')
-# Capture activations
+# Forward pass with activation capture
outputs, acts = wrapper.forward_with_activations(inputs)
weights = wrapper.get_layer_weights()
@@ -83,26 +148,6 @@ scores = rq_metric.compute(
)
```
-### Available Metrics
-
-**Alignment:**
-- `rayleigh_quotient` - Standard RQ
-- `delta_alignment` - Change in alignment over training
-
-**Information-Theoretic:**
-- `gaussian_mi_analytic` - Mutual information (Gaussian approximation)
-- `pairwise_redundancy_gaussian` - Per-neuron redundancy
-- `synergy_gaussian_mmi` - Per-neuron synergy
-- `conditional_mutual_information` - Conditional MI
-
-**Gradient-Based:**
-- `gradient_alignment` - Alignment between local signal and backprop
-- `local_learning_rule_search` - Find optimal local rules per neuron
-
-**Task-Specific:**
-- `classification_alignment` - Classification-specific importance
-- `vision_task_alignment` - Vision-specific patterns
-
### Composite Scoring
Combine multiple metrics for robust importance estimation:
@@ -113,7 +158,7 @@ from alignment.services import NodeScoringService
scorer = NodeScoringService(
metrics={
'rq': get_metric('rayleigh_quotient'),
- 'redundancy': get_metric('pairwise_redundancy_gaussian', mode='output_based'),
+ 'redundancy': get_metric('pairwise_redundancy_gaussian'),
'synergy': get_metric('synergy_gaussian_mmi')
},
alpha_mi=0.0,
@@ -123,290 +168,375 @@ scorer = NodeScoringService(
)
scores = scorer.compute_composite_scores(inputs, weights, targets)
-# score = β·synergy - γ·redundancy + δ·log(RQ)
```
----
+The composite score combines: synergy (positive), redundancy (negative), and alignment.
## Pruning Strategies
-### Available Strategies
+### Overview
-**Basic:**
-- `magnitude` - L1/L2 norm
-- `random` - Random baseline
-- `gradient` - Gradient-based
-- `fisher` - Fisher information
+Pruning removes less important network parameters to reduce model size and computational cost. The framework provides multiple strategies based on different importance criteria.
-**Alignment-Based:**
-- `alignment` - RQ-based structured pruning
-- `global_alignment` - Global threshold across layers
-- `hybrid` - Magnitude + Alignment combination
+### Magnitude-Based Strategies
-**Advanced:**
-- `movement` - Prune weights moving toward zero (training-aware)
-- `adaptive` - Adaptive per-layer amounts based on sensitivity
-- `ultimate` - Multi-stage combining all best practices
+**MagnitudePruning**
+- Prunes weights with smallest absolute values
+- Fast and simple baseline
+- Works for unstructured and structured pruning
+- Applies L1 or L2 norm for importance scores
-**Novel:**
-- `composite` - Redundancy-aware (preserves synergistic neurons)
+**IterativeMagnitudePruning**
+- Applies magnitude pruning in multiple rounds
+- Allows network to recover between pruning steps
+- Typically achieves better accuracy than one-shot pruning
-### Distribution Across Layers
+**GlobalMagnitudePruning**
+- Applies single threshold across all layers
+- Natural adaptation to layer sensitivity
+- Some layers may be pruned more than others
-**Uniform:** Same percentage per layer
-```python
-# 70% removed from each layer
-```
+### Gradient-Based Strategies
-**Global Threshold:** Single threshold across all layers
-```python
-# Naturally varying amounts based on score distribution
-```
+**GradientPruning**
+- Uses gradient magnitude as importance signal
+- Can use gradient alone or gradient-weight product (Taylor approximation)
+- Requires backward pass to compute gradients
-**Adaptive Sensitivity:** Per-layer amounts based on importance
-```python
-# Sensitive layers: 50%
-# Robust layers: 85%
-# Overall: 70% average
-```
+**FisherPruning**
+- Uses Fisher information for importance
+- Accounts for both gradient and Hessian information
+- More accurate but computationally expensive
-### Using the Orchestrator
+**MomentumPruning**
+- Incorporates momentum in importance computation
+- Smooths importance over training iterations
-```python
-from alignment.pruning.orchestrator import prune_with_all_options
-
-result = prune_with_all_options(
- model,
- target_sparsity=0.7,
-
- # Distribution: how to allocate across layers
- # Options: 'uniform', 'global_threshold', 'adaptive_sensitivity',
- # 'importance_weighted', 'cascading', 'size_proportional', 'hybrid'
- distribution='adaptive_sensitivity',
-
- # Scoring: what importance metric
- # Options: 'magnitude', 'rayleigh_quotient', 'composite', 'movement'
- scoring='composite',
-
- # Direction: which neurons to remove
- # Options: 'low' (remove unimportant), 'high' (ablation), 'random' (baseline)
- direction='low',
-
- # Data
- train_loader=train_loader, # For fine-tuning
- val_loader=val_loader, # For evaluation
- trainer_fn=train_function, # Training function
- eval_fn=eval_function, # Evaluation function
-
- # Options
- fine_tune_epochs=20
-)
-```
+### Alignment-Based Strategies
----
+**AlignmentPruning**
+- Uses Rayleigh Quotient or other alignment metrics
+- Structured pruning (removes entire neurons/channels)
+- Considers input-weight relationships
-## Architecture Support
+**GlobalAlignmentPruning**
+- Global threshold across layers based on alignment
+- Automatically handles layer-wise sensitivity
-### MLPs
+**HybridPruning**
+- Combines magnitude and alignment scores
+- Balances simple magnitude with structural alignment
-```python
-layer = model.fc1 # Linear(784, 256)
-scores = rq.compute(inputs, layer.weight)
-# [256] - one score per neuron
-```
+**CascadingAlignmentPruning**
+- Prunes layers sequentially
+- Recomputes alignment after each layer
+- Accounts for pruning effects on downstream layers
-### CNNs
+### Random Strategies
-```python
-conv = model.conv1 # Conv2d(64, 128, 3, 3)
-scores = rq.compute(inputs, conv.weight)
-# [128] - one score per channel
+**RandomPruning**
+- Prunes weights randomly
+- Baseline for comparison
+- Useful for ablation studies
-# Dependency-aware pruning automatically handles channel propagation
-```
+**LayerwiseRandomPruning**
+- Random pruning with per-layer amounts
-### Transformers & LLMs
+**BernoulliPruning**
+- Stochastic pruning with Bernoulli sampling
-```python
-from alignment.models.transformer_enhanced import LLaMAWrapper
+### Parallel Strategies
-wrapper = LLaMAWrapper(llama_model, track_ffn=True)
+**ParallelModePruning**
+- Applies multiple pruning modes simultaneously (low/high/random)
+- Returns separate masks for each mode
+- Useful for comparing ablations
-# FFN neurons: up_proj has 11,008 neurons
-ffn_up = model.model.layers[0].mlp.up_proj
-scores = rq.compute(inputs, ffn_up.weight)
-# [11008] - one score per FFN neuron
+**TensorizedPruning**
+- GPU-optimized parallel execution
+- Efficient for multiple strategies
-# Attention: Can analyze per-head or per-neuron
-```
+**AsyncParallelPruning**
+- Asynchronous parallel execution
+- For distributed pruning experiments
----
+### Advanced Strategies
-## Configuration
+**MovementPruning**
+- Tracks weight movement during training
+- Prunes weights moving toward zero
+- Training-aware strategy
-### Quick Start Config
+**AdaptiveMovementPruning**
+- Adapts pruning amount per layer based on movement patterns
-```yaml
-# configs/quickstart.yaml
-experiment:
- name: "quickstart"
+### Pruning Distribution
-model:
- name: "resnet18"
- pretrained: true
+Distribution strategies determine how pruning is allocated across layers.
-dataset:
- name: "cifar10"
- batch_size: 128
+**Uniform Distribution**
+- Same sparsity percentage for all layers
+- Simple and predictable
+- May not respect layer-specific sensitivity
-metrics:
- enabled: ['rayleigh_quotient']
-```
+**Global Threshold**
+- Single importance threshold across all layers
+- Natural adaptation based on score distributions
+- Different layers pruned by different amounts
-### Complete Options
+**Adaptive Sensitivity**
+- Per-layer pruning amounts based on sensitivity analysis
+- Sensitive layers pruned less
+- Robust layers pruned more
+- Maintains target overall sparsity
-See `configs/template_master_v2.yaml` for all parameters with documentation.
+**Importance Weighted**
+- Allocates pruning inversely to layer importance
+- More important layers retain more parameters
-Key sections:
-- `experiment`: Name, seed, device, output
-- `model`: Architecture, pretrained, layers to track
-- `dataset`: Data source, batch size, augmentation
-- `metrics`: Which metrics to compute and their parameters
-- `pruning`: Strategy, distribution, scoring, fine-tuning
-- `training`: Training parameters, metric tracking
-- `advanced`: Backend selection, parallelization
+**Size Proportional**
+- Pruning amount scales with layer size
+- Larger layers can tolerate more pruning
----
+**Cascading**
+- Sequential layer-by-layer pruning
+- Recomputes scores after each layer
-## Advanced Features
+### Pruning Modes
-### Training-Time Metrics
+**Low Mode** (default)
+- Prunes weights/neurons with lowest importance scores
+- Standard pruning approach
-Track alignment evolution during training with zero overhead:
+**High Mode**
+- Prunes weights/neurons with highest importance scores
+- Used for ablation studies to test importance hypotheses
-```python
-from alignment.training.callbacks import AlignmentMetricsCallback
+**Random Mode**
+- Random pruning regardless of scores
+- Baseline for comparison
-callback = AlignmentMetricsCallback(
- metrics={'rq': get_metric('rayleigh_quotient')},
- layers=['conv1', 'fc1'],
- frequency=100 # Every 100 steps
-)
+### Structured vs Unstructured
+
+**Unstructured Pruning**
+- Prunes individual weights
+- Arbitrary sparsity patterns
+- Requires sparse tensor support for speedup
-# In training loop:
-for inputs, targets in train_loader:
- outputs = model(inputs)
- loss.backward()
- optimizer.step()
-
- callback.on_batch_end(wrapper, inputs, targets, global_step)
+**Structured Pruning**
+- Prunes entire neurons or channels
+- Maintains dense tensor operations
+- Immediate speedup without specialized hardware
+- Handles dependencies automatically
-# Analyze evolution
-history = callback.get_history()
+### Using Configuration Files
+
+Specify pruning in YAML configuration:
+
+```yaml
+pruning:
+ enabled: true
+ strategy: 'composite'
+ target_sparsity: 0.7
+ distribution: 'adaptive_sensitivity'
+ scoring: 'rayleigh_quotient'
+ direction: 'low'
+ structured: true
+ dependency_aware: true
+
+ fine_tune:
+ enabled: true
+ epochs: 20
+ learning_rate: 0.0001
+
+ composite_weights:
+ beta_synergy: 0.3
+ gamma_redundancy: 0.4
+ delta_rq: 0.3
```
-### Gradient-Based Local Learning
+Then run:
+```bash
+python scripts/run_experiment.py --config configs/my_pruning.yaml
+```
+
+## Architecture Support
+
+### Linear Layers (MLPs)
-Design bio-plausible learning rules:
+For fully-connected layers, metrics compute per-neuron scores:
```python
-from alignment.metrics.gradient_based import LocalLearningRuleSearch
+layer = model.fc1 # Linear(in_features=784, out_features=256)
+scores = rq.compute(inputs, layer.weight) # Returns [256]
+```
-searcher = LocalLearningRuleSearch()
+Each neuron receives an importance score. Structured pruning removes entire neurons.
-# After backward pass:
-best_rules = searcher.compute(
- inputs, outputs,
- gradients=layer.weight.grad
-)
-# Returns best local rule per neuron
+### Convolutional Layers (CNNs)
+
+For convolutional layers, metrics operate on channels:
+
+```python
+conv = model.conv1 # Conv2d(in_channels=64, out_channels=128, kernel_size=3)
+scores = rq.compute(inputs, conv.weight) # Returns [128]
```
-### Pairwise Metrics
+Each output channel receives a score. The framework handles:
+- Spatial weight dimensions automatically
+- Different conv modes: unfold, patchwise, channel variance
+- Dependency propagation between layers
-Any pairwise metric can aggregate to single-neuron scores:
+**CNN-Specific Modes:**
+- `unfold`: Unfolds spatial dimensions into feature vectors
+- `patchwise`: Treats each spatial location separately
+- `channel_variance`: Computes variance across spatial dimensions
+- `batch_patch_combined`: Combines batch and spatial statistics
+### Transformer Layers
+
+For transformer architectures, the framework supports:
+
+**Feed-Forward Networks (FFN)**
```python
-redundancy = get_metric('pairwise_redundancy_gaussian',
- mode='output_based', # Fast!
- num_pairs=10, # Sample 10 partners
- aggregation='mean') # How to aggregate
+ffn_layer = model.layers[0].mlp.up_proj
+scores = rq.compute(inputs, ffn_layer.weight)
+```
-scores = redundancy.compute(outputs=layer_outputs)
-# [N] - per-neuron redundancy
+Each FFN neuron receives a score. For LLaMA models, typical FFN dimensions are 11,008 neurons.
-# Or get full matrix:
-matrix = redundancy.compute(outputs, return_matrix=True)
-# [N, N] - all pairwise relationships
+**Attention Layers**
+```python
+# Per-head analysis
+wrapper = TransformerWrapper(model, track_per_head=True)
+
+# Per-neuron within attention projections
+q_proj = model.layers[0].self_attn.q_proj
+scores = rq.compute(inputs, q_proj.weight)
```
-### Dependency-Aware Pruning
+Can analyze:
+- Query, Key, Value projections separately
+- Per-head importance
+- Per-neuron importance within projections
+- Attention output projections
+
+**Aggregation Options:**
+- `sequence_mean`: Average over sequence length
+- `token_level`: Per-token analysis
-Automatically handles inter-layer dependencies:
+### Large Language Models
+
+The framework provides specialized wrappers for LLMs:
```python
-from alignment.pruning.dependency_aware import prune_model_with_dependencies
+from alignment.models.transformer_enhanced import LLaMAWrapper
-result = prune_model_with_dependencies(
- model,
- layer_scores={'conv1': scores1, 'conv2': scores2},
- amount=0.5,
- verbose=True
+wrapper = LLaMAWrapper(
+ llama_model,
+ track_ffn=True,
+ track_attention=True
)
+```
-# Automatically propagates:
-# conv1.out_channels → conv2.in_channels
-# Maintains shape compatibility
+Supports analysis and pruning of:
+- FFN intermediate layers (up_proj, down_proj)
+- Attention projections (q_proj, k_proj, v_proj, o_proj)
+- Layer-wise or model-wide analysis
+
+## Configuration
+
+Experiments are configured using YAML files. The framework provides a template with all available options.
+
+### Basic Configuration
+
+```yaml
+experiment:
+ name: "my_experiment"
+ seed: 42
+ device: "cuda"
+ output_dir: "./results"
+
+model:
+ name: "resnet18"
+ pretrained: true
+ tracked_layers: null # Auto-detect
+
+dataset:
+ name: "cifar10"
+ data_path: "./data"
+ batch_size: 128
+
+metrics:
+ enabled: ['rayleigh_quotient']
```
----
+### Configuration Sections
+
+**experiment**
+- Experiment name, random seed, device selection
+- Output directory for results
+
+**model**
+- Model architecture name or path
+- Whether to use pretrained weights
+- Checkpoint path for custom models
+- Which layers to track for analysis
-## Performance
+**dataset**
+- Dataset name and path
+- Batch size and data loading workers
+- Data augmentation settings
-### Computation Time
+**metrics**
+- List of metrics to compute
+- Per-metric configuration parameters
+- Regularization and sampling settings
-| Model Scale | Per-Layer Time | Throughput |
-|-------------|----------------|------------|
-| Small (N=256) | ~28ms | 35 layers/sec |
-| Medium (N=1024) | ~105ms | 10 layers/sec |
-| Large (N=4096) | ~440ms | 2.3 layers/sec |
+**pruning**
+- Pruning strategy and target sparsity
+- Distribution method across layers
+- Scoring method for importance
+- Structured vs unstructured
+- Fine-tuning configuration
-### Speedup Techniques
+**training**
+- Training parameters if training from scratch
+- Learning rate, optimizer, scheduler
+- Metric tracking during training
-- **Output-based redundancy**: 30x faster for large models
-- **Shared computation**: 2-3x when computing multiple metrics
-- **Parallel strategies**: M strategies in ~1.3x time
+**advanced**
+- Covariance computation methods
+- Hook management settings
+- Layer-specific configurations
----
+See `configs/template.yaml` for complete documentation of all parameters.
## Testing
+Run tests to verify installation and functionality:
+
```bash
-# Run all tests
+# All tests
pytest tests/
-# Scientific correctness validation
-python tests/unit/metrics/test_scientific_correctness.py
+# Specific test modules
+pytest tests/unit/metrics/
+pytest tests/unit/pruning/
-# Specific module
-pytest tests/unit/services/
+# Scientific correctness validation
+pytest tests/unit/metrics/test_scientific_correctness.py
```
----
-
-## Contributing
-
-1. Fork the repository
-2. Create a feature branch
-3. Add tests for new features
-4. Submit pull request
-
----
+## Examples
-## Support
+The `examples/` directory contains scripts demonstrating framework capabilities:
-- **Documentation**: See guides in repository root
-- **Examples**: `examples/` directory
-- **Issues**: GitHub issues
-- **API Reference**: Run `cd docs && make html`
+- `01_quick_start.py` - Basic usage
+- `02_complete_experiment.py` - Full experiment workflow
+- `03_pruning_strategies.py` - Pruning strategy comparison
+- `06_redundancy_aware_pruning.py` - Information-theoretic pruning
+- `07_mnist_intelligent_pruning.py` - MNIST pruning example
+- `08_llama_ffn_pruning.py` - LLM feed-forward pruning
+- `09_attention_neuron_vs_head_pruning.py` - Attention layer analysis
diff --git a/drafts/alignment_notes b/drafts/alignment_notes
deleted file mode 160000
index fea87596..00000000
--- a/drafts/alignment_notes
+++ /dev/null
@@ -1 +0,0 @@
-Subproject commit fea87596a5b84f587955c3af61fde1acc4359d5d
diff --git a/examples/01_quick_start.py b/examples/01_quick_start.py
index 462167e4..0459145c 100644
--- a/examples/01_quick_start.py
+++ b/examples/01_quick_start.py
@@ -6,13 +6,13 @@
Usage:
python quick_demo.py
-
+
No configuration needed - this script runs with default settings and creates its own model.
Requirements:
- PyTorch
- alignment package installed
-
+
Output:
- Console output showing:
* Model structure
@@ -24,8 +24,9 @@
import torch
import torch.nn as nn
-from alignment.models import ModelWrapper
+
from alignment.metrics import get_metric
+from alignment.models import ModelWrapper
def main():
@@ -37,55 +38,55 @@ def main():
nn.ReLU(),
nn.Linear(128, 10)
)
-
+
print("Model created with 3 linear layers")
-
+
# 2. Wrap the model to track activations
wrapped_model = ModelWrapper(model)
print(f"\nTracked layers: {wrapped_model.tracked_layers}")
-
+
# 3. Generate some dummy data
batch_size = 32
inputs = torch.randn(batch_size, 784)
-
+
# 4. Forward pass with activation tracking
outputs, activations = wrapped_model.forward_with_activations(inputs)
print(f"\nCollected activations from {len(activations)} points")
print(f"Activation keys: {list(activations.keys())}")
-
+
# 5. Get layer weights
weights = wrapped_model.get_layer_weights()
print(f"Extracted weights from {len(weights)} layers")
-
+
# 6. Compute alignment metrics
print("\n" + "="*60)
print("Computing Alignment Metrics")
print("="*60)
-
+
# Rayleigh Quotient
RQMetric = get_metric('rayleigh_quotient')
rq_metric = RQMetric() # Instantiate the metric
-
+
for layer_name in wrapped_model.tracked_layers:
# Get inputs and weights for this layer
layer_inputs = activations[f"{layer_name}_input"]
layer_weights = weights[layer_name]
-
+
# Compute RQ scores
scores = rq_metric.compute(inputs=layer_inputs, weights=layer_weights)
-
+
print(f"\nLayer: {layer_name}")
print(f" Input shape: {layer_inputs.shape}")
print(f" Weight shape: {layer_weights.shape}")
print(f" RQ scores: mean={scores.mean():.4f}, std={scores.std():.4f}")
print(f" Min neuron score: {scores.min():.4f}")
print(f" Max neuron score: {scores.max():.4f}")
-
+
# 7. Try other metrics
print("\n" + "="*60)
print("Trying Other Metrics")
print("="*60)
-
+
# Weight cosine similarity
WeightSimMetric = get_metric('weight_cosine_similarity')
weight_sim = WeightSimMetric() # Instantiate
@@ -93,35 +94,35 @@ def main():
layer_weights = weights[layer_name]
sim_scores = weight_sim.compute(weights=layer_weights)
print(f"\nWeight Cosine Similarity ({layer_name}): mean={sim_scores.mean():.4f}")
-
+
# 8. Demonstrate pruning
print("\n" + "="*60)
print("Pruning Demo")
print("="*60)
-
- from alignment.pruning import get_pruning_strategy, PruningConfig
-
+
+ from alignment.pruning import PruningConfig, get_pruning_strategy
+
# Use magnitude-based pruning
config = PruningConfig(amount=0.5, pruning_mode='low')
strategy = get_pruning_strategy('magnitude', config=config)
-
+
# Apply pruning to first layer
first_layer = model[0] # First linear layer
mask = strategy.prune(first_layer)
-
- print(f"\nPruning 50% of weights in first layer")
+
+ print("\nPruning 50% of weights in first layer")
print(f"Sparsity achieved: {(mask == 0).float().mean():.2%}")
-
+
# Apply mask
first_layer.weight.data *= mask
-
+
# Test forward pass after pruning
pruned_outputs = model(inputs)
print(f"\nOutput shape after pruning: {pruned_outputs.shape}")
print("Pruning applied successfully!")
-
+
print("\nDemo completed successfully!")
if __name__ == "__main__":
- main()
\ No newline at end of file
+ main()
diff --git a/examples/02_complete_experiment.py b/examples/02_complete_experiment.py
index 945133fb..d3b8e19a 100644
--- a/examples/02_complete_experiment.py
+++ b/examples/02_complete_experiment.py
@@ -12,7 +12,7 @@
Usage:
python standard_alignment_experiment.py
-
+
No configuration needed - this script runs with default settings.
The script will:
- Download MNIST dataset automatically (if not present)
@@ -37,21 +37,21 @@
- comparison_grid.png: Comprehensive analysis grid
"""
+import json
+from pathlib import Path
+
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
-import numpy as np
-from pathlib import Path
-import json
-from typing import Dict, List, Any
+
+from alignment.analysis.visualization import PruningVisualizer
+from alignment.metrics import get_metric
# Alignment framework imports
from alignment.models import ModelWrapper
-from alignment.metrics import get_metric
-from alignment.pruning import get_pruning_strategy, PruningConfig
-from alignment.analysis.visualization import PruningVisualizer
+from alignment.pruning import PruningConfig, get_pruning_strategy
class SimpleNet(nn.Module):
@@ -64,7 +64,7 @@ def __init__(self, hidden_sizes=[784, 256, 128, 10]):
if i < len(hidden_sizes) - 2: # No ReLU after last layer
layers.append(nn.ReLU())
self.model = nn.Sequential(*layers)
-
+
def forward(self, x):
x = x.view(x.size(0), -1)
return self.model(x)
@@ -75,9 +75,9 @@ def train_model(model, train_loader, val_loader, epochs=10, device='cuda'):
model = model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
-
+
history = {'train_loss': [], 'val_acc': []}
-
+
for epoch in range(epochs):
# Training
model.train()
@@ -90,7 +90,7 @@ def train_model(model, train_loader, val_loader, epochs=10, device='cuda'):
loss.backward()
optimizer.step()
train_loss += loss.item()
-
+
# Validation
model.eval()
correct = 0
@@ -100,15 +100,15 @@ def train_model(model, train_loader, val_loader, epochs=10, device='cuda'):
output = model(data)
pred = output.argmax(dim=1)
correct += pred.eq(target).sum().item()
-
+
train_loss /= len(train_loader)
val_acc = 100. * correct / len(val_loader.dataset)
-
+
history['train_loss'].append(train_loss)
history['val_acc'].append(val_acc)
-
+
print(f'Epoch {epoch+1}/{epochs}: Loss: {train_loss:.4f}, Acc: {val_acc:.2f}%')
-
+
return history
@@ -116,36 +116,36 @@ def compute_alignment_metrics(model, data_loader, device='cuda'):
"""Compute various alignment metrics for the model."""
model = model.to(device)
wrapped_model = ModelWrapper(model.model) # Wrap the sequential model
-
+
# Get a batch of data for metric computation
data_iter = iter(data_loader)
inputs, _ = next(data_iter)
inputs = inputs.to(device)
inputs = inputs.view(inputs.size(0), -1) # Flatten for the model
-
+
# Forward pass to collect activations
_, activations = wrapped_model.forward_with_activations(inputs)
weights = wrapped_model.get_layer_weights()
-
+
# Initialize metrics
metrics = {
'rayleigh_quotient': get_metric('rayleigh_quotient')(),
'weight_cosine_similarity': get_metric('weight_cosine_similarity')(),
}
-
+
results = {}
-
+
# Compute metrics for each layer
for layer_name in wrapped_model.tracked_layers:
results[layer_name] = {}
-
+
# Rayleigh Quotient
if f"{layer_name}_input" in activations:
layer_inputs = activations[f"{layer_name}_input"]
layer_weights = weights[layer_name]
-
+
rq_scores = metrics['rayleigh_quotient'].compute(
- inputs=layer_inputs,
+ inputs=layer_inputs,
weights=layer_weights
)
results[layer_name]['rayleigh_quotient'] = {
@@ -154,7 +154,7 @@ def compute_alignment_metrics(model, data_loader, device='cuda'):
'min': float(rq_scores.min()),
'max': float(rq_scores.max())
}
-
+
# Weight Cosine Similarity
layer_weights = weights[layer_name]
sim_scores = metrics['weight_cosine_similarity'].compute(weights=layer_weights)
@@ -162,40 +162,40 @@ def compute_alignment_metrics(model, data_loader, device='cuda'):
'mean': float(sim_scores.mean()),
'std': float(sim_scores.std())
}
-
+
return results
-def apply_pruning_experiment(model, val_loader, strategies=['magnitude', 'random'],
+def apply_pruning_experiment(model, val_loader, strategies=['magnitude', 'random'],
sparsities=[0.1, 0.3, 0.5, 0.7, 0.9], device='cuda'):
"""Apply different pruning strategies and evaluate."""
results = {}
-
+
for strategy_name in strategies:
results[strategy_name] = {}
-
+
for sparsity in sparsities:
# Clone model for this experiment
pruned_model = SimpleNet()
pruned_model.load_state_dict(model.state_dict())
pruned_model = pruned_model.to(device)
-
+
# Apply pruning
config = PruningConfig(amount=sparsity, pruning_mode='low')
strategy = get_pruning_strategy(strategy_name, config=config)
-
+
# Prune each linear layer
for name, module in pruned_model.named_modules():
if isinstance(module, nn.Linear):
mask = strategy.prune(module)
module.weight.data *= mask
-
+
# Evaluate
pruned_model.eval()
correct = 0
total_loss = 0
criterion = nn.CrossEntropyLoss()
-
+
with torch.no_grad():
for data, target in val_loader:
data, target = data.to(device), target.to(device)
@@ -204,18 +204,18 @@ def apply_pruning_experiment(model, val_loader, strategies=['magnitude', 'random
total_loss += loss.item()
pred = output.argmax(dim=1)
correct += pred.eq(target).sum().item()
-
+
accuracy = 100. * correct / len(val_loader.dataset)
avg_loss = total_loss / len(val_loader)
-
+
results[strategy_name][sparsity] = {
'accuracy': accuracy,
'loss': avg_loss
}
-
+
print(f"{strategy_name} - Sparsity {sparsity:.1%}: "
f"Acc: {accuracy:.2f}%, Loss: {avg_loss:.4f}")
-
+
return results
@@ -223,24 +223,24 @@ def visualize_results(pruning_results, output_dir='results/standard_experiment')
"""Generate visualizations of the results."""
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
-
+
visualizer = PruningVisualizer()
-
+
# Create performance comparison plot
- fig1 = visualizer.plot_pruning_performance(
+ visualizer.plot_pruning_performance(
pruning_results,
metrics=['accuracy', 'loss'],
save_path=output_dir / 'pruning_performance.png',
title='Pruning Strategy Comparison on MNIST',
show_confidence=False
)
-
+
# Create comprehensive comparison grid
- fig2 = visualizer.plot_pruning_comparison_grid(
+ visualizer.plot_pruning_comparison_grid(
pruning_results,
save_path=output_dir / 'comparison_grid.png'
)
-
+
print(f"\nVisualizations saved to {output_dir}")
@@ -249,75 +249,75 @@ def main():
print("=" * 60)
print("Standard Alignment Experiment")
print("=" * 60)
-
+
# Setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"\nUsing device: {device}")
-
+
# Data loading
print("\nLoading MNIST dataset...")
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
-
+
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
val_dataset = datasets.MNIST('./data', train=False, transform=transform)
-
+
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=256, shuffle=False)
-
+
# Create model
print("\nCreating model...")
model = SimpleNet()
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
-
+
# Train model
print("\nTraining model...")
history = train_model(model, train_loader, val_loader, epochs=5, device=device)
-
+
# Save training history
output_dir = Path('results/standard_experiment')
output_dir.mkdir(parents=True, exist_ok=True)
-
+
with open(output_dir / 'training_history.json', 'w') as f:
json.dump(history, f, indent=2)
-
+
# Compute alignment metrics
print("\nComputing alignment metrics...")
metrics = compute_alignment_metrics(model, val_loader, device=device)
-
+
print("\nAlignment Metrics:")
for layer, layer_metrics in metrics.items():
print(f"\nLayer {layer}:")
for metric_name, values in layer_metrics.items():
if 'mean' in values:
print(f" {metric_name}: mean={values['mean']:.4f}, std={values['std']:.4f}")
-
+
# Save metrics
with open(output_dir / 'alignment_metrics.json', 'w') as f:
json.dump(metrics, f, indent=2)
-
+
# Pruning experiments
print("\n" + "=" * 60)
print("Pruning Experiments")
print("=" * 60)
-
+
pruning_results = apply_pruning_experiment(
model, val_loader,
strategies=['magnitude', 'random'],
sparsities=[0.1, 0.3, 0.5, 0.7, 0.9],
device=device
)
-
+
# Save pruning results
with open(output_dir / 'pruning_results.json', 'w') as f:
json.dump(pruning_results, f, indent=2)
-
+
# Generate visualizations
print("\nGenerating visualizations...")
visualize_results(pruning_results, output_dir)
-
+
# Summary
print("\n" + "=" * 60)
print("Experiment Complete!")
@@ -329,12 +329,12 @@ def main():
print(" - pruning_results.json: Pruning experiment results")
print(" - pruning_performance.png: Performance comparison plot")
print(" - comparison_grid.png: Comprehensive analysis grid")
-
+
# Print key findings
print("\nKey Findings:")
baseline_acc = history['val_acc'][-1]
print(f" - Baseline accuracy: {baseline_acc:.2f}%")
-
+
for strategy in pruning_results:
acc_50 = pruning_results[strategy][0.5]['accuracy']
acc_90 = pruning_results[strategy][0.9]['accuracy']
@@ -344,4 +344,4 @@ def main():
if __name__ == "__main__":
- main()
\ No newline at end of file
+ main()
diff --git a/examples/03_pruning_strategies.py b/examples/03_pruning_strategies.py
index f3849f26..53baae99 100644
--- a/examples/03_pruning_strategies.py
+++ b/examples/03_pruning_strategies.py
@@ -8,7 +8,7 @@
Usage:
python pruning_strategies_demo.py
-
+
No configuration needed - this script runs standalone demonstrations.
This comprehensive demo showcases:
@@ -35,12 +35,13 @@
- Key takeaways and insights
"""
-import torch
-import torch.nn as nn
-import numpy as np
-from pathlib import Path
import sys
import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent / 'src'))
@@ -48,10 +49,10 @@
# Direct imports
from alignment.pruning.base import PruningConfig
from alignment.pruning.strategies import (
- MagnitudePruning,
GradientPruning,
- RandomPruning,
+ MagnitudePruning,
ParallelModePruning,
+ RandomPruning,
TensorizedPruning,
)
@@ -65,7 +66,7 @@ def __init__(self, input_size=784, hidden_size=256, num_classes=10):
self.fc2 = nn.Linear(hidden_size, 128)
self.relu2 = nn.ReLU()
self.fc3 = nn.Linear(128, num_classes)
-
+
def forward(self, x):
x = x.view(x.size(0), -1)
x = self.relu1(self.fc1(x))
@@ -77,65 +78,65 @@ def forward(self, x):
def demo_basic_pruning():
"""Demonstrate basic pruning with different modes."""
print("\n=== Basic Pruning Demo ===\n")
-
+
# Create model
model = SimpleNet(hidden_size=512)
-
+
# Initialize weights
for m in model.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_normal_(m.weight)
-
+
# Test different pruning modes on first layer
layer = model.fc1
sparsity = 0.5
-
+
print(f"Layer shape: {layer.weight.shape}")
print(f"Target sparsity: {sparsity*100}%\n")
-
+
# 1. Magnitude pruning - low mode (prune small weights)
print("1. Magnitude Pruning - Low Mode (prune small weights):")
config = PruningConfig(amount=sparsity, pruning_mode='low')
pruner = MagnitudePruning(config)
-
+
scores = pruner.compute_importance_scores(layer)
mask = pruner.create_pruning_mask(scores)
-
+
weights = layer.weight.data.flatten()
kept_weights = weights[mask.flatten().bool()]
pruned_weights = weights[~mask.flatten().bool()]
-
+
print(f" Sparsity achieved: {(mask == 0).float().mean():.2%}")
print(f" Avg magnitude of kept weights: {kept_weights.abs().mean():.4f}")
print(f" Avg magnitude of pruned weights: {pruned_weights.abs().mean():.4f}")
print(f" Max pruned weight magnitude: {pruned_weights.abs().max():.4f}")
-
+
# 2. Magnitude pruning - high mode (prune large weights)
print("\n2. Magnitude Pruning - High Mode (prune large weights):")
config = PruningConfig(amount=sparsity, pruning_mode='high')
pruner = MagnitudePruning(config)
-
+
mask = pruner.create_pruning_mask(scores)
-
+
kept_weights = weights[mask.flatten().bool()]
pruned_weights = weights[~mask.flatten().bool()]
-
+
print(f" Sparsity achieved: {(mask == 0).float().mean():.2%}")
print(f" Avg magnitude of kept weights: {kept_weights.abs().mean():.4f}")
print(f" Avg magnitude of pruned weights: {pruned_weights.abs().mean():.4f}")
print(f" Min pruned weight magnitude: {pruned_weights.abs().min():.4f}")
-
+
# 3. Random pruning
print("\n3. Random Pruning:")
config = PruningConfig(amount=sparsity)
pruner = RandomPruning(config)
-
+
scores = pruner.compute_importance_scores(layer)
mask = pruner.create_pruning_mask(scores)
-
+
kept_weights = weights[mask.flatten().bool()]
pruned_weights = weights[~mask.flatten().bool()]
-
+
print(f" Sparsity achieved: {(mask == 0).float().mean():.2%}")
print(f" Avg magnitude of kept weights: {kept_weights.abs().mean():.4f}")
print(f" Avg magnitude of pruned weights: {pruned_weights.abs().mean():.4f}")
@@ -145,62 +146,62 @@ def demo_basic_pruning():
def demo_parallel_pruning():
"""Demonstrate parallel pruning strategies."""
print("\n\n=== Parallel Pruning Demo ===\n")
-
+
# Create model
model = SimpleNet(hidden_size=256)
-
+
# Initialize weights
for m in model.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_normal_(m.weight)
-
+
layer = model.fc1
sparsity = 0.5
-
+
# 1. Parallel mode pruning
print("1. Parallel Mode Pruning (all modes simultaneously):")
parallel_pruner = ParallelModePruning(
modes=['low', 'high', 'random'],
base_strategy='magnitude'
)
-
+
start_time = time.time()
result = parallel_pruner.prune_parallel(layer, amount=sparsity)
parallel_time = time.time() - start_time
-
+
print(f"\n Execution time: {parallel_time:.4f}s")
- print(f"\n Sparsity by mode:")
+ print("\n Sparsity by mode:")
for mode, sparse in result.sparsities.items():
print(f" {mode}: {sparse:.2%}")
-
+
# Analyze overlaps
- print(f"\n Mask overlaps:")
+ print("\n Mask overlaps:")
low_high_overlap = (result.masks['low'] * result.masks['high']).sum().item()
low_random_overlap = (result.masks['low'] * result.masks['random']).sum().item()
high_random_overlap = (result.masks['high'] * result.masks['random']).sum().item()
total = result.masks['low'].numel()
-
+
print(f" low ∩ high: {low_high_overlap/total:.2%}")
print(f" low ∩ random: {low_random_overlap/total:.2%}")
print(f" high ∩ random: {high_random_overlap/total:.2%}")
-
+
# 2. Sequential comparison
print("\n2. Sequential execution (for comparison):")
-
+
start_time = time.time()
masks_sequential = {}
-
+
for mode in ['low', 'high', 'random']:
config = PruningConfig(amount=sparsity, pruning_mode=mode)
pruner = MagnitudePruning(config)
scores = pruner.compute_importance_scores(layer)
masks_sequential[mode] = pruner.create_pruning_mask(scores)
-
+
sequential_time = time.time() - start_time
-
+
print(f"\n Execution time: {sequential_time:.4f}s")
print(f" Speedup from parallel: {sequential_time/parallel_time:.2f}x")
-
+
# Verify results match
print(f"\n Results match: {all(torch.allclose(result.masks[mode], masks_sequential[mode]) for mode in ['low', 'high', 'random'])}")
@@ -208,7 +209,7 @@ def demo_parallel_pruning():
def demo_tensorized_pruning():
"""Demonstrate tensorized pruning for GPU efficiency."""
print("\n\n=== Tensorized Pruning Demo ===\n")
-
+
# Create larger model for better timing comparison
model = nn.Sequential(
nn.Linear(1024, 2048),
@@ -217,29 +218,29 @@ def demo_tensorized_pruning():
nn.ReLU(),
nn.Linear(1024, 512),
)
-
+
# Move to GPU if available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
-
+
print(f"Device: {device}")
print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
-
+
# Initialize weights
for m in model.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_normal_(m.weight)
-
+
# 1. Tensorized pruning - compute multiple sparsity levels at once
print("\n1. Tensorized Pruning (multiple sparsity levels):")
-
+
tensorized_pruner = TensorizedPruning()
-
+
# Test on first layer
layer = model[0]
modes = ['low', 'high', 'random']
amounts = [0.1, 0.3, 0.5, 0.7, 0.9]
-
+
start_time = time.time()
pruning_tensor = tensorized_pruner.compute_pruning_tensor(
layer,
@@ -247,29 +248,29 @@ def demo_tensorized_pruning():
amounts=amounts
)
tensorized_time = time.time() - start_time
-
+
print(f"\n Execution time: {tensorized_time:.4f}s")
print(f" Pruning tensor shape: {pruning_tensor.shape}")
print(f" (modes={len(modes)}, amounts={len(amounts)}, *weight_shape)")
-
+
# Analyze patterns
print("\n Analyzing pruning patterns...")
analysis = tensorized_pruner.analyze_pruning_patterns(pruning_tensor)
-
- print(f"\n Sparsity progression (avg across weight dims):")
+
+ print("\n Sparsity progression (avg across weight dims):")
sparsity_prog = analysis['sparsity_progression']
for i, mode in enumerate(modes):
print(f" {mode}: ", end="")
for j, amount in enumerate(amounts):
print(f"{sparsity_prog[i, j]:.2f} ", end="")
print()
-
+
# 2. Compare with sequential approach
print("\n2. Sequential approach (for comparison):")
-
+
start_time = time.time()
sequential_masks = {}
-
+
for mode in modes:
sequential_masks[mode] = {}
for amount in amounts:
@@ -278,32 +279,32 @@ def demo_tensorized_pruning():
pruner = RandomPruning(config)
else:
pruner = MagnitudePruning(config)
-
+
scores = pruner.compute_importance_scores(layer)
mask = pruner.create_pruning_mask(scores)
sequential_masks[mode][amount] = mask
-
+
sequential_time = time.time() - start_time
-
+
print(f"\n Execution time: {sequential_time:.4f}s")
print(f" Speedup from tensorized: {sequential_time/tensorized_time:.2f}x")
-
+
# 3. Memory efficiency analysis
print("\n3. Memory efficiency:")
-
+
# Tensorized approach stores everything in one tensor
tensorized_memory = pruning_tensor.element_size() * pruning_tensor.nelement()
-
+
# Sequential approach stores individual masks
sequential_memory = 0
for mode_masks in sequential_masks.values():
for mask in mode_masks.values():
sequential_memory += mask.element_size() * mask.nelement()
-
+
print(f" Tensorized memory: {tensorized_memory / (1024**2):.2f} MB")
print(f" Sequential memory: {sequential_memory / (1024**2):.2f} MB")
print(f" Memory efficiency: {sequential_memory/tensorized_memory:.2f}x")
-
+
# 4. Show overlap analysis
if 'mode_overlap' in analysis:
print("\n4. Mode overlap analysis (low ∩ high):")
@@ -320,57 +321,57 @@ def demo_tensorized_pruning():
def demo_gradient_based_pruning():
"""Demonstrate gradient-based pruning."""
print("\n\n=== Gradient-Based Pruning Demo ===\n")
-
+
# Create model and dummy data
model = SimpleNet(input_size=784, hidden_size=256)
criterion = nn.CrossEntropyLoss()
-
+
# Create dummy batch
batch_size = 32
x = torch.randn(batch_size, 784)
y = torch.randint(0, 10, (batch_size,))
-
+
# Forward and backward pass to get gradients
output = model(x)
loss = criterion(output, y)
loss.backward()
-
+
layer = model.fc1
sparsity = 0.5
-
+
print(f"Layer shape: {layer.weight.shape}")
print(f"Target sparsity: {sparsity*100}%")
-
+
# 1. Gradient magnitude pruning - low mode
print("\n1. Gradient Pruning - Low Mode (prune small gradients):")
config = PruningConfig(amount=sparsity, pruning_mode='low')
pruner = GradientPruning(config)
-
+
scores = pruner.compute_importance_scores(layer)
mask = pruner.create_pruning_mask(scores)
-
+
grad_magnitudes = layer.weight.grad.abs().flatten()
kept_grads = grad_magnitudes[mask.flatten().bool()]
pruned_grads = grad_magnitudes[~mask.flatten().bool()]
-
+
print(f" Sparsity achieved: {(mask == 0).float().mean():.2%}")
print(f" Avg gradient magnitude of kept: {kept_grads.mean():.6f}")
print(f" Avg gradient magnitude of pruned: {pruned_grads.mean():.6f}")
-
+
# 2. Gradient magnitude pruning - high mode
print("\n2. Gradient Pruning - High Mode (prune large gradients):")
config = PruningConfig(amount=sparsity, pruning_mode='high')
pruner = GradientPruning(config)
-
+
mask = pruner.create_pruning_mask(scores)
-
+
kept_grads = grad_magnitudes[mask.flatten().bool()]
pruned_grads = grad_magnitudes[~mask.flatten().bool()]
-
+
print(f" Sparsity achieved: {(mask == 0).float().mean():.2%}")
print(f" Avg gradient magnitude of kept: {kept_grads.mean():.6f}")
print(f" Avg gradient magnitude of pruned: {pruned_grads.mean():.6f}")
-
+
# Compare with weight magnitude
print("\n3. Gradient vs Weight magnitude correlation:")
weight_mags = layer.weight.abs().flatten()
@@ -383,29 +384,29 @@ def main():
print("=" * 60)
print("Pruning Strategies Demo")
print("=" * 60)
-
+
# Set random seed
torch.manual_seed(42)
np.random.seed(42)
-
+
try:
# Run demos
demo_basic_pruning()
demo_parallel_pruning()
demo_tensorized_pruning()
demo_gradient_based_pruning()
-
+
print("\n" + "=" * 60)
print("Demo Complete!")
print("=" * 60)
-
+
print("\nKey takeaways:")
print("1. Low mode pruning: Removes small magnitude weights (common approach)")
print("2. High mode pruning: Removes large magnitude weights (adversarial)")
print("3. Parallel pruning: Compute multiple modes simultaneously for efficiency")
print("4. Tensorized pruning: GPU-optimized for multiple sparsity levels")
print("5. Gradient pruning: Uses gradient information for importance")
-
+
except Exception as e:
print(f"\nError: {e}")
import traceback
@@ -413,4 +414,4 @@ def main():
if __name__ == "__main__":
- main()
\ No newline at end of file
+ main()
diff --git a/examples/04_visualization_gallery.py b/examples/04_visualization_gallery.py
index 6fb3dc2c..c1d732cd 100644
--- a/examples/04_visualization_gallery.py
+++ b/examples/04_visualization_gallery.py
@@ -6,7 +6,7 @@
Usage:
python visualization_demo.py
-
+
This demo showcases:
- Enhanced accuracy vs sparsity plots with relative performance
- Weight distribution comparison before/after pruning
@@ -19,15 +19,17 @@
Results are saved to: results/enhanced_visualizations/
"""
-import torch
-import numpy as np
from pathlib import Path
+
+import numpy as np
+import torch
+
from alignment.analysis.visualization import PruningVisualizer
def create_demo_results():
"""Create comprehensive demo results for visualization."""
-
+
# Progressive dropout results
dropout_results = {
'dropout_rates': [0.0, 0.2, 0.4, 0.6, 0.8, 0.9],
@@ -53,7 +55,7 @@ def create_demo_results():
'random': [46, 55, 62, 67, 71, 74, 76, 77, 78, 79]
}
}
-
+
# Standard pruning results (sparsity-based)
pruning_results = {}
for strategy in ['magnitude', 'gradient', 'fisher', 'random']:
@@ -68,12 +70,12 @@ def create_demo_results():
}
accuracy = base_acc * (1 - sparsity * decay_factors[strategy])
accuracy += np.random.normal(0, 1.0)
-
+
pruning_results[strategy][sparsity] = {
'accuracy': max(10, min(100, accuracy)),
'loss': 0.1 + sparsity * 0.8 * (1 + decay_factors[strategy])
}
-
+
# Multi-seed results
multi_seed_results = {}
for strategy in ['magnitude', 'gradient', 'random']:
@@ -88,36 +90,36 @@ def create_demo_results():
'loss': base['loss'] + np.random.normal(0, 0.1)
}
multi_seed_results[strategy].append(seed_data)
-
+
# Layer-wise sparsity patterns
layer_sparsities = {
'magnitude': {'conv1': 0.3, 'conv2': 0.5, 'fc1': 0.7, 'fc2': 0.9},
'gradient': {'conv1': 0.4, 'conv2': 0.5, 'fc1': 0.6, 'fc2': 0.8},
'random': {'conv1': 0.5, 'conv2': 0.5, 'fc1': 0.5, 'fc2': 0.5}
}
-
+
model_accuracy = {
'magnitude': 85.3,
'gradient': 82.1,
'random': 68.5
}
-
+
# Weight distributions (simulated)
np.random.seed(42)
weights_before = {
'conv1': torch.randn(64, 3, 3, 3),
'fc1': torch.randn(128, 512)
}
-
+
weights_after = {}
for layer_name, weights in weights_before.items():
# Simulate pruning effects
mask_magnitude = torch.abs(weights) > torch.quantile(torch.abs(weights), 0.5)
weights_after[f"{layer_name}_magnitude"] = weights * mask_magnitude
-
+
mask_random = torch.rand_like(weights) > 0.5
weights_after[f"{layer_name}_random"] = weights * mask_random
-
+
# Multi-metric comparison data
strategy_metrics = {
'magnitude_low': {
@@ -149,7 +151,7 @@ def create_demo_results():
'Memory': 0.90
}
}
-
+
# Efficiency curve data
efficiency_data = {}
for strategy in ['magnitude_low', 'magnitude_high', 'gradient_low', 'random']:
@@ -163,9 +165,9 @@ def create_demo_results():
factor = efficiency_factors[strategy]
accuracies = 100 * np.exp(-0.3 * (compression_ratios - 1) / factor)
accuracies = np.clip(accuracies, 10, 100)
-
+
efficiency_data[strategy] = list(zip(compression_ratios, accuracies))
-
+
return {
'dropout_results': dropout_results,
'pruning_results': pruning_results,
@@ -184,87 +186,87 @@ def main():
print("=" * 60)
print("Enhanced Pruning Visualization Demo")
print("=" * 60)
-
+
# Create output directory
output_dir = Path('results/enhanced_visualizations')
output_dir.mkdir(parents=True, exist_ok=True)
-
+
# Create demo data
print("\nGenerating demo data...")
demo_data = create_demo_results()
-
+
# Initialize visualizer
visualizer = PruningVisualizer()
-
+
# 1. Enhanced accuracy vs sparsity plot
print("\n1. Creating enhanced accuracy vs sparsity plot...")
- fig1 = visualizer.plot_accuracy_vs_sparsity_enhanced(
+ visualizer.plot_accuracy_vs_sparsity_enhanced(
demo_data['dropout_results'],
save_path=output_dir / 'enhanced_accuracy_vs_sparsity.png'
)
-
+
# 2. Standard pruning performance
print("2. Creating standard pruning performance plot...")
- fig2 = visualizer.plot_pruning_performance(
+ visualizer.plot_pruning_performance(
demo_data['pruning_results'],
metrics=['accuracy', 'loss'],
save_path=output_dir / 'pruning_performance.png',
title='Pruning Strategy Performance Comparison'
)
-
+
# 3. Weight distribution comparison
print("3. Creating weight distribution comparison...")
- fig3 = visualizer.plot_weight_distribution_comparison(
+ visualizer.plot_weight_distribution_comparison(
demo_data['weights_before'],
demo_data['weights_after'],
strategies=['magnitude', 'random'],
save_path=output_dir / 'weight_distributions.png'
)
-
+
# 4. Multi-metric radar chart
print("4. Creating multi-metric radar chart...")
- fig4 = visualizer.plot_multi_metric_radar(
+ visualizer.plot_multi_metric_radar(
demo_data['strategy_metrics'],
save_path=output_dir / 'multi_metric_radar.png'
)
-
+
# 5. Pruning efficiency curves
print("5. Creating pruning efficiency curves...")
- fig5 = visualizer.plot_pruning_efficiency_curve(
+ visualizer.plot_pruning_efficiency_curve(
demo_data['efficiency_data'],
save_path=output_dir / 'efficiency_curves.png'
)
-
+
# 6. Comprehensive dashboard
print("6. Creating comprehensive dashboard...")
- fig6 = visualizer.plot_comprehensive_dashboard(
+ visualizer.plot_comprehensive_dashboard(
demo_data['dropout_results'],
save_path=output_dir / 'comprehensive_dashboard.png'
)
-
+
# 7. Comparison grid
print("7. Creating comparison grid...")
- fig7 = visualizer.plot_pruning_comparison_grid(
+ visualizer.plot_pruning_comparison_grid(
demo_data['pruning_results'],
save_path=output_dir / 'comparison_grid.png'
)
-
+
# 8. Multi-seed analysis
print("8. Creating multi-seed analysis...")
- fig8 = visualizer.plot_multi_seed_results(
+ visualizer.plot_multi_seed_results(
demo_data['multi_seed_results'],
metric='accuracy',
save_path=output_dir / 'multi_seed_analysis.png'
)
-
+
# 9. Layer-wise pruning visualization
print("9. Creating layer-wise pruning visualization...")
- fig9 = visualizer.plot_layer_wise_pruning(
+ visualizer.plot_layer_wise_pruning(
demo_data['layer_sparsities'],
demo_data['model_accuracy'],
save_path=output_dir / 'layer_wise_pruning.png'
)
-
+
# Summary
print("\n" + "=" * 60)
print("Enhanced Visualization Demo Complete!")
@@ -280,7 +282,7 @@ def main():
print(" 7. comparison_grid.png - 6-panel strategy comparison")
print(" 8. multi_seed_analysis.png - Statistical analysis across seeds")
print(" 9. layer_wise_pruning.png - Layer-specific sparsity patterns")
-
+
print("\nVisualization Features:")
print(" - Consistent color scheme across all plots")
print(" - High-resolution output (300 DPI)")
@@ -290,4 +292,4 @@ def main():
if __name__ == "__main__":
- main()
\ No newline at end of file
+ main()
diff --git a/examples/06_redundancy_aware_pruning.py b/examples/06_redundancy_aware_pruning.py
index 5c02fcca..f74703a0 100644
--- a/examples/06_redundancy_aware_pruning.py
+++ b/examples/06_redundancy_aware_pruning.py
@@ -19,17 +19,17 @@
import torch
import torch.nn as nn
-from torchvision import datasets, transforms, models
-import numpy as np
+from torchvision import datasets, transforms
+
+from alignment.metrics import get_metric
# Alignment framework imports
from alignment.models import BaseModelWrapper
from alignment.services import (
ActivationCaptureService,
+ MaskOperations,
NodeScoringService,
- MaskOperations
)
-from alignment.metrics import get_metric
def create_simple_cnn():
@@ -54,7 +54,7 @@ def evaluate_model(model, dataloader, device='cpu'):
model.eval()
correct = 0
total = 0
-
+
with torch.no_grad():
for inputs, targets in dataloader:
inputs, targets = inputs.to(device), targets.to(device)
@@ -62,7 +62,7 @@ def evaluate_model(model, dataloader, device='cpu'):
_, predicted = torch.max(outputs, 1)
total += targets.size(0)
correct += (predicted == targets).sum().item()
-
+
return 100 * correct / total
@@ -71,58 +71,58 @@ def main():
print("=" * 80)
print("Redundancy-Aware Pruning Demonstration")
print("=" * 80)
-
+
# Setup
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"\nUsing device: {device}")
-
+
# Load MNIST
print("\n1. Loading MNIST dataset...")
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
-
+
test_dataset = datasets.MNIST(
root='./data',
train=False,
download=True,
transform=transform
)
-
+
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=128,
shuffle=False
)
-
+
# Create and wrap model
print("\n2. Creating model...")
model = create_simple_cnn().to(device)
-
+
# Simple training (or load pretrained)
# For demo, we'll use random weights
print(" (Using random initialization for demo)")
-
+
wrapper = BaseModelWrapper(
model,
tracked_layers=['3', '6'], # Conv2 and Linear1
track_inputs=True,
track_outputs=True
)
-
+
print(f" Tracking layers: {wrapper.tracked_layers}")
-
+
# Initialize services
print("\n3. Initializing services...")
capture_service = ActivationCaptureService(wrapper)
-
+
# Initialize metrics
print("\n4. Initializing metrics...")
rq_metric = get_metric('rayleigh_quotient', relative=True, regularization=1e-6)
redundancy_metric = get_metric('pairwise_redundancy_gaussian', num_pairs=10)
synergy_metric = get_metric('synergy_gaussian_mmi', num_pairs=10)
-
+
scoring_service = NodeScoringService(
metrics={
'rq': rq_metric,
@@ -134,21 +134,21 @@ def main():
gamma_redundancy=0.4, # Redundancy weight (negative)
delta_rq=0.3 # RQ weight
)
-
+
# Capture activations on a subset
print("\n5. Capturing activations...")
batch_inputs, batch_targets = next(iter(test_loader))
batch_inputs = batch_inputs.to(device)
batch_targets = batch_targets.to(device)
-
+
activation_data = capture_service.capture(
batch_inputs,
layers=wrapper.tracked_layers,
include_weights=True
)
-
+
print(f" Captured data from {len(activation_data.layer_names)} layers")
-
+
# Compute composite scores
print("\n6. Computing composite scores...")
layerwise_scores = scoring_service.compute_layerwise_scores(
@@ -157,13 +157,13 @@ def main():
include_redundancy=True,
include_synergy=True
)
-
+
# Display results for first layer
layer_name = wrapper.tracked_layers[0]
if layer_name in layerwise_scores:
scores = layerwise_scores[layer_name]
print(f"\n Results for layer '{layer_name}':")
-
+
if scores.rq is not None:
print(f" - RQ: mean={scores.rq.mean():.4f}, std={scores.rq.std():.4f}")
if scores.redundancy is not None:
@@ -172,7 +172,7 @@ def main():
print(f" - Synergy: mean={scores.synergy.mean():.4f}, std={scores.synergy.std():.4f}")
if scores.composite is not None:
print(f" - Composite: mean={scores.composite.mean():.4f}, std={scores.composite.std():.4f}")
-
+
# Demonstrate class-conditioned RQ
print("\n7. Computing class-conditioned RQ (ΔRQ)...")
layer_name = wrapper.tracked_layers[-1] # Linear layer
@@ -183,73 +183,73 @@ def main():
targets=batch_targets,
return_delta_rq=True
)
-
+
print(f" Layer '{layer_name}':")
print(f" - RQ (unconditional): {delta_rq_results['rq_uncond'].mean():.4f}")
print(f" - RQ (class-cond): {delta_rq_results['rq_cond'].mean():.4f}")
print(f" - ΔRQ: {delta_rq_results['delta_rq'].mean():.4f}")
- print(f" ΔRQ measures task-relevant alignment (larger = more discriminative)")
-
+ print(" ΔRQ measures task-relevant alignment (larger = more discriminative)")
+
# Create pruning masks
print("\n8. Creating pruning masks...")
pruning_amount = 0.3 # Prune 30%
-
+
for layer_name, scores in layerwise_scores.items():
if scores.composite is None:
continue
-
+
# Create mask using composite scores
mask = MaskOperations.create_structured_mask(
scores.composite,
amount=pruning_amount,
mode='low' # Prune low-importance neurons
)
-
+
stats = MaskOperations.get_mask_statistics(mask)
print(f" Layer '{layer_name}': {stats['kept_elements']}/{stats['total_elements']} neurons kept")
-
+
# Compare methods
print("\n9. Comparing pruning methods...")
print(" Method comparison (on same layer):")
layer_name = wrapper.tracked_layers[-1]
if layer_name in layerwise_scores:
scores = layerwise_scores[layer_name]
-
+
# Method 1: Random
mask_random = MaskOperations.create_structured_mask(
torch.rand_like(scores.composite),
amount=pruning_amount,
mode='random'
)
-
+
# Method 2: RQ only
mask_rq = MaskOperations.create_structured_mask(
scores.rq,
amount=pruning_amount,
mode='low'
)
-
+
# Method 3: Composite (redundancy-aware)
mask_composite = MaskOperations.create_structured_mask(
scores.composite,
amount=pruning_amount,
mode='low'
)
-
+
# Check overlap
overlap_rq_composite = (mask_rq & mask_composite).sum().item() / mask_rq.sum().item()
overlap_random_composite = (mask_random & mask_composite).sum().item() / mask_random.sum().item()
-
+
print(f" - Overlap (RQ vs Composite): {overlap_rq_composite:.2%}")
print(f" - Overlap (Random vs Composite): {overlap_random_composite:.2%}")
print(f" → Redundancy-awareness changes {(1-overlap_rq_composite)*100:.1f}% of pruning decisions")
-
+
# Evaluate baseline
print("\n10. Evaluating model...")
baseline_acc = evaluate_model(model, test_loader, device)
print(f" Baseline accuracy: {baseline_acc:.2f}%")
print(" (Would apply masks and fine-tune for full experiment)")
-
+
print("\n" + "=" * 80)
print("Summary")
print("=" * 80)
@@ -257,7 +257,7 @@ def main():
Key Features Demonstrated:
✓ ActivationCaptureService - Clean API for activation capture
✓ PairwiseRedundancyGaussian - Identifies redundant neurons
-✓ SynergyGaussianMMI - Identifies complementary neurons
+✓ SynergyGaussianMMI - Identifies complementary neurons
✓ Class-conditioned RQ - Measures task-relevant alignment (ΔRQ)
✓ NodeScoringService - Composite scoring with configurable weights
✓ MaskOperations - Flexible mask creation and analysis
diff --git a/examples/07_mnist_intelligent_pruning.py b/examples/07_mnist_intelligent_pruning.py
index 78e85d84..e19826cc 100644
--- a/examples/07_mnist_intelligent_pruning.py
+++ b/examples/07_mnist_intelligent_pruning.py
@@ -1,36 +1,34 @@
"""
-Complete Example: Intelligent Pruning on MNIST
+MNIST Intelligent Pruning Example
-This script demonstrates the full workflow:
-1. Train a simple model on MNIST
-2. Compute redundancy-aware composite scores
-3. Prune using multiple strategies
+Demonstrates complete workflow:
+1. Train MLP on MNIST
+2. Compute composite importance scores
+3. Apply pruning with different strategies
4. Compare results
-
-This is a practical, runnable example showing real improvements.
"""
+from pathlib import Path
+
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
-import matplotlib.pyplot as plt
-from pathlib import Path
+
+from alignment.metrics import get_metric
# Alignment framework
from alignment.models import BaseModelWrapper
from alignment.services import (
ActivationCaptureService,
+ MaskOperations,
NodeScoringService,
- MaskOperations
)
-from alignment.metrics import get_metric
-# Simple MLP for MNIST
class SimpleMLP(nn.Module):
- """Simple MLP: 784 -> 128 -> 64 -> 10"""
+ """MLP: 784 -> 128 -> 64 -> 10"""
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(784, 128)
@@ -38,7 +36,7 @@ def __init__(self):
self.fc2 = nn.Linear(128, 64)
self.relu2 = nn.ReLU()
self.fc3 = nn.Linear(64, 10)
-
+
def forward(self, x):
x = x.view(x.size(0), -1)
x = self.relu1(self.fc1(x))
@@ -48,34 +46,34 @@ def forward(self, x):
def train_model(model, train_loader, epochs=5, device='cpu'):
- """Quick training."""
+ """Train model."""
model.train()
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
-
+
print(f"Training for {epochs} epochs...")
for epoch in range(epochs):
total_loss = 0
correct = 0
total = 0
-
+
for batch_idx, (inputs, targets) in enumerate(train_loader):
inputs, targets = inputs.to(device), targets.to(device)
-
+
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
-
+
total_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
-
+
if batch_idx % 100 == 0:
print(f' Epoch {epoch+1}, Batch {batch_idx}: Loss={loss.item():.4f}')
-
+
acc = 100. * correct / total
print(f'Epoch {epoch+1}: Loss={total_loss/len(train_loader):.4f}, Acc={acc:.2f}%')
@@ -85,7 +83,7 @@ def evaluate(model, test_loader, device='cpu'):
model.eval()
correct = 0
total = 0
-
+
with torch.no_grad():
for inputs, targets in test_loader:
inputs, targets = inputs.to(device), targets.to(device)
@@ -93,24 +91,19 @@ def evaluate(model, test_loader, device='cpu'):
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
-
+
return 100. * correct / total
def prune_model(model, wrapper, pruning_method, pruning_amount, val_loader, device='cpu'):
- """
- Prune model using specified method.
-
- Args:
- pruning_method: 'random', 'magnitude', 'rq', or 'composite'
- """
+ """Prune model using specified method."""
print(f"\nPruning with {pruning_method} (amount={pruning_amount:.1%})...")
-
+
# Get a batch for metric computation
inputs_batch, targets_batch = next(iter(val_loader))
inputs_batch = inputs_batch.to(device)
targets_batch = targets_batch.to(device)
-
+
# Capture activations
capture_service = ActivationCaptureService(wrapper)
data = capture_service.capture(
@@ -118,30 +111,30 @@ def prune_model(model, wrapper, pruning_method, pruning_amount, val_loader, devi
layers=wrapper.tracked_layers,
include_weights=True
)
-
+
# Compute scores based on method
masks = {}
-
+
for layer_name in wrapper.tracked_layers:
if layer_name not in data.inputs or layer_name not in data.weights:
continue
-
+
inputs = data.inputs[layer_name]
weights = data.weights[layer_name]
-
+
if pruning_method == 'random':
# Random scores
scores = torch.rand(weights.shape[0])
-
+
elif pruning_method == 'magnitude':
# L2 norm of weights
scores = torch.norm(weights, p=2, dim=1)
-
+
elif pruning_method == 'rq':
# RQ only
rq_metric = get_metric('rayleigh_quotient')
scores = rq_metric.compute(inputs, weights)
-
+
elif pruning_method == 'composite':
# Redundancy-aware composite
scoring_service = NodeScoringService(
@@ -155,15 +148,15 @@ def prune_model(model, wrapper, pruning_method, pruning_amount, val_loader, devi
gamma_redundancy=0.4,
delta_rq=0.3
)
-
+
layer_scores = scoring_service.compute_composite_scores(
inputs, weights, targets_batch
)
scores = layer_scores.composite
-
+
else:
raise ValueError(f"Unknown method: {pruning_method}")
-
+
# Create mask
mask = MaskOperations.create_structured_mask(
scores,
@@ -171,13 +164,13 @@ def prune_model(model, wrapper, pruning_method, pruning_amount, val_loader, devi
mode='low'
)
masks[layer_name] = mask
-
+
stats = MaskOperations.get_mask_statistics(mask)
print(f" {layer_name}: {stats['kept_elements']}/{stats['total_elements']} kept")
-
+
# Apply masks
apply_masks_to_model(model, masks, wrapper.tracked_layers)
-
+
return masks
@@ -186,13 +179,13 @@ def apply_masks_to_model(model, masks, layer_names):
for name, module in model.named_modules():
if name in masks and hasattr(module, 'weight'):
mask = masks[name]
-
+
# Expand mask to weight dimensions
if isinstance(module, nn.Linear):
# Zero out entire rows (neurons)
weight_mask = mask.unsqueeze(1).expand_as(module.weight)
module.weight.data *= weight_mask.float()
-
+
if module.bias is not None:
module.bias.data *= mask.float()
@@ -202,55 +195,55 @@ def main():
print("=" * 80)
print("MNIST Intelligent Pruning - Complete Workflow")
print("=" * 80)
-
+
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"\nDevice: {device}")
-
+
# Load MNIST
print("\n1. Loading MNIST...")
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
-
+
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('./data', train=False, transform=transform)
-
+
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)
-
+
# Validation loader (smaller for metric computation)
val_loader = DataLoader(test_dataset, batch_size=512, shuffle=False)
-
+
# Train baseline model
print("\n2. Training baseline model...")
model = SimpleMLP().to(device)
train_model(model, train_loader, epochs=3, device=device)
-
+
baseline_acc = evaluate(model, test_loader, device)
print(f"\n✓ Baseline accuracy: {baseline_acc:.2f}%")
-
+
# Wrap model
- wrapper = BaseModelWrapper(
+ BaseModelWrapper(
model,
tracked_layers=['fc1', 'fc2'], # Track hidden layers
track_inputs=True,
track_outputs=True
)
-
+
# Pruning experiments
print("\n3. Pruning experiments...")
print("-" * 80)
-
+
pruning_amount = 0.5 # Prune 50%
methods = ['random', 'magnitude', 'rq', 'composite']
results = {}
-
+
for method in methods:
# Create fresh copy of model
model_copy = SimpleMLP().to(device)
model_copy.load_state_dict(model.state_dict())
-
+
# Wrap
wrapper_copy = BaseModelWrapper(
model_copy,
@@ -258,75 +251,57 @@ def main():
track_inputs=True,
track_outputs=True
)
-
+
# Prune
- masks = prune_model(
- model_copy, wrapper_copy, method,
+ prune_model(
+ model_copy, wrapper_copy, method,
pruning_amount, val_loader, device
)
-
+
# Evaluate immediately after pruning (no fine-tuning)
acc_pruned = evaluate(model_copy, test_loader, device)
-
+
# Fine-tune
print(f"\n Fine-tuning {method}...")
train_model(model_copy, train_loader, epochs=2, device=device)
-
+
# Final evaluation
acc_final = evaluate(model_copy, test_loader, device)
-
+
results[method] = {
'acc_before': baseline_acc,
'acc_pruned': acc_pruned,
'acc_final': acc_final,
'drop': baseline_acc - acc_final
}
-
+
print(f" {method}: {baseline_acc:.2f}% → {acc_pruned:.2f}% → {acc_final:.2f}% (drop: {baseline_acc - acc_final:.2f}%)")
-
+
# Summary
print("\n" + "=" * 80)
print("RESULTS SUMMARY")
print("=" * 80)
print(f"\nPruning amount: {pruning_amount:.0%} of neurons")
print(f"Baseline accuracy: {baseline_acc:.2f}%\n")
-
+
print("Method | After Pruning | After Fine-tune | Accuracy Drop")
print("-" * 70)
for method, res in results.items():
print(f"{method:15s} | {res['acc_pruned']:13.2f}% | {res['acc_final']:15.2f}% | {res['drop']:13.2f}%")
-
+
# Find best method
best_method = min(results.keys(), key=lambda m: results[m]['drop'])
print(f"\n✓ Best method: {best_method} (smallest accuracy drop: {results[best_method]['drop']:.2f}%)")
-
- # Expected outcome
- print("\n" + "=" * 80)
- print("EXPECTED OUTCOME")
- print("=" * 80)
- print("""
-Redundancy-aware pruning (composite) should outperform others because:
-• Preserves high-synergy neuron pairs (complementary information)
-• Removes high-redundancy neurons (overlapping information)
-• Uses task-relevance when targets available (ΔRQ)
-
-Expected ranking (best to worst):
-1. Composite (redundancy-aware) - Smallest drop
-2. RQ (alignment-aware)
-3. Magnitude
-4. Random - Largest drop
-
-If composite matches or beats magnitude with same drop, that's a success!
-At higher sparsity (70-90%), the gap should be even larger (+3-5% better).
- """)
-
+
+ print("\nComposite pruning considers redundancy and synergy in addition to alignment.")
+
# Save results
output_dir = Path('results/mnist_intelligent_pruning')
output_dir.mkdir(parents=True, exist_ok=True)
-
+
torch.save(results, output_dir / 'pruning_results.pt')
print(f"\n✓ Results saved to {output_dir}")
-
+
return results
diff --git a/examples/08_llama_ffn_pruning.py b/examples/08_llama_ffn_pruning.py
index 67d9698e..05ae45e2 100644
--- a/examples/08_llama_ffn_pruning.py
+++ b/examples/08_llama_ffn_pruning.py
@@ -12,59 +12,59 @@
up_proj: Linear(4096 → 11,008) ← 11,008 neurons we can analyze!
gate_proj: Linear(4096 → 11,008)
down_proj: Linear(11,008 → 4096)
-
+
Forward: down_proj(SiLU(gate_proj(x)) * up_proj(x))
"""
+from typing import List
+
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
-from typing import Dict, Optional
+
+from alignment.metrics import get_metric
# Alignment framework
from alignment.models.transformer_enhanced import LLaMAWrapper
from alignment.services import (
ActivationCaptureService,
- NodeScoringService,
- MaskOperations
+ MaskOperations,
)
-from alignment.metrics import get_metric
-from alignment.pruning.dependency_aware import DependencyAwarePruning
def load_llama_model(model_name: str = 'gpt2', use_small_model: bool = True):
"""
Load LLaMA or similar model.
-
+
Args:
model_name: HuggingFace model name
use_small_model: If True, use GPT-2 as proxy (faster for demo)
-
+
Returns:
model, tokenizer
"""
# For demo, use GPT-2 (similar architecture, much smaller)
# Replace with 'meta-llama/Meta-Llama-3-8B' for actual LLaMA-3
-
+
if use_small_model:
model_name = 'gpt2' # 12 layers, 768 hidden, much faster to load
- print(f"Using GPT-2 as demo (same architecture principles as LLaMA)")
-
+ print("Using GPT-2 as demo (same architecture principles as LLaMA)")
+
print(f"Loading {model_name}...")
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
-
+
# Set padding token
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
-
+
model.eval()
-
+
print(f"✓ Loaded {model_name}")
print(f" - Layers: {model.config.num_hidden_layers}")
print(f" - Hidden size: {model.config.hidden_size}")
print(f" - Attention heads: {model.config.num_attention_heads}")
-
+
return model, tokenizer
@@ -77,32 +77,32 @@ def analyze_ffn_layer(
):
"""
Analyze a single FFN layer - compute per-neuron scores.
-
+
Args:
model: LLaMA/GPT model
wrapper: LLaMAWrapper instance
layer_idx: Which transformer layer to analyze (0-indexed)
input_text: List of input strings
tokenizer: HuggingFace tokenizer
-
+
Returns:
Dict with per-neuron scores
"""
print(f"\n{'='*80}")
print(f"Analyzing Layer {layer_idx} FFN")
print(f"{'='*80}")
-
+
# Tokenize inputs
inputs = tokenizer(input_text, padding=True, truncation=True, return_tensors='pt')
-
+
# Get FFN layer names (HF naming convention)
# GPT-2: transformer.h.{i}.mlp.c_fc (up), c_proj (down)
# LLaMA: model.layers.{i}.mlp.up_proj, down_proj, gate_proj
-
+
# Try to infer FFN layer names
ffn_up = None
ffn_down = None
-
+
for name, module in model.named_modules():
if f'.{layer_idx}.' in name or f'.h.{layer_idx}.' in name:
if 'mlp' in name:
@@ -111,51 +111,51 @@ def analyze_ffn_layer(
ffn_up = name
elif 'down_proj' in name or 'c_proj' in name:
ffn_down = name
-
+
if ffn_up is None:
print(f" ⚠️ Could not find FFN layers for layer {layer_idx}")
return {}
-
- print(f" FFN layers found:")
+
+ print(" FFN layers found:")
print(f" - Up projection: {ffn_up}")
print(f" - Down projection: {ffn_down}")
-
+
# Get module to check dimensions
ffn_up_module = dict(model.named_modules())[ffn_up]
num_neurons = ffn_up_module.out_features
hidden_dim = ffn_up_module.in_features
-
- print(f" Dimensions:")
+
+ print(" Dimensions:")
print(f" - Hidden: {hidden_dim}")
print(f" - FFN neurons: {num_neurons} ← Can analyze each one!")
-
+
# Wrap model to track this layer
wrapper_ffn = LLaMAWrapper(model, tracked_layers=[ffn_up])
-
+
# Capture activations
- capture = ActivationCaptureService(wrapper_ffn)
-
+ ActivationCaptureService(wrapper_ffn)
+
with torch.no_grad():
# Use model's forward to get activations
- outputs = model(**inputs)
-
+ model(**inputs)
+
# Get cached activations
acts = wrapper_ffn._activation_cache
-
+
# Get layer activations and weights
ffn_input = acts.get(f'{ffn_up}_input')
ffn_output = acts.get(f'{ffn_up}_output')
ffn_weights = ffn_up_module.weight.detach() # [num_neurons, hidden_dim]
-
+
if ffn_input is None or ffn_output is None:
print(" ⚠️ Could not capture activations")
return {}
-
- print(f"\n Activation shapes:")
+
+ print("\n Activation shapes:")
print(f" - Input: {ffn_input.shape}") # [B, T, hidden_dim]
print(f" - Output: {ffn_output.shape}") # [B, T, num_neurons]
print(f" - Weights: {ffn_weights.shape}") # [num_neurons, hidden_dim]
-
+
# Preprocess for metrics (average over sequence)
if ffn_input.ndim == 3:
ffn_input_2d = ffn_input.mean(dim=1) # [B, hidden_dim]
@@ -163,37 +163,37 @@ def analyze_ffn_layer(
else:
ffn_input_2d = ffn_input
ffn_output_2d = ffn_output
-
- print(f"\n Computing per-neuron metrics...")
-
+
+ print("\n Computing per-neuron metrics...")
+
# Compute multiple metrics per neuron
metrics_results = {}
-
+
# 1. Rayleigh Quotient (alignment)
- print(f" - RQ (alignment with input PCs)...")
+ print(" - RQ (alignment with input PCs)...")
rq_metric = get_metric('rayleigh_quotient')
rq_scores = rq_metric.compute(inputs=ffn_input_2d, weights=ffn_weights)
metrics_results['rq'] = rq_scores
print(f" Range: [{rq_scores.min():.4f}, {rq_scores.max():.4f}], Mean: {rq_scores.mean():.4f}")
-
+
# 2. Redundancy (overlap with other neurons)
- print(f" - Redundancy (overlap with other neurons)...")
- redundancy_metric = get_metric('pairwise_redundancy_gaussian',
+ print(" - Redundancy (overlap with other neurons)...")
+ redundancy_metric = get_metric('pairwise_redundancy_gaussian',
mode='output_based', # FAST!
num_pairs=20)
redundancy_scores = redundancy_metric.compute(outputs=ffn_output_2d)
metrics_results['redundancy'] = redundancy_scores
print(f" Range: [{redundancy_scores.min():.4f}, {redundancy_scores.max():.4f}], Mean: {redundancy_scores.mean():.4f}")
-
+
# 3. Show per-neuron breakdown
- print(f"\n Per-neuron analysis (first 10 neurons):")
+ print("\n Per-neuron analysis (first 10 neurons):")
print(f" {'Neuron':<8} {'RQ':<10} {'Redundancy':<12} {'Assessment'}")
print(f" {'-'*60}")
-
+
for i in range(min(10, num_neurons)):
rq_val = rq_scores[i].item()
red_val = redundancy_scores[i].item()
-
+
# Assess importance
if red_val > 0.5:
assessment = "Redundant (candidate for pruning)"
@@ -201,18 +201,18 @@ def analyze_ffn_layer(
assessment = "Important (high alignment)"
else:
assessment = "Low importance"
-
+
print(f" {i:<8} {rq_val:<10.4f} {red_val:<12.4f} {assessment}")
-
+
# Show statistics
high_redundancy = (redundancy_scores > 0.5).sum().item()
low_rq = (rq_scores < 0.005).sum().item()
-
- print(f"\n Summary:")
+
+ print("\n Summary:")
print(f" - High redundancy neurons (>0.5): {high_redundancy}/{num_neurons} ({100*high_redundancy/num_neurons:.1f}%)")
print(f" - Low RQ neurons (<0.005): {low_rq}/{num_neurons} ({100*low_rq/num_neurons:.1f}%)")
print(f" - Potential pruning candidates: {max(high_redundancy, low_rq)} neurons")
-
+
return metrics_results
@@ -224,7 +224,7 @@ def prune_ffn_neurons(
):
"""
Prune individual neurons in FFN layer with dependency awareness.
-
+
Args:
model: LLaMA model
layer_idx: Layer index
@@ -234,11 +234,11 @@ def prune_ffn_neurons(
print(f"\n{'='*80}")
print(f"Pruning Layer {layer_idx} FFN (amount={amount:.0%})")
print(f"{'='*80}")
-
+
# Find FFN layers
ffn_up = None
ffn_down = None
-
+
for name, module in model.named_modules():
if f'.{layer_idx}.' in name and 'mlp' in name:
if isinstance(module, nn.Linear):
@@ -246,46 +246,46 @@ def prune_ffn_neurons(
ffn_up = (name, module)
elif 'down_proj' in name or 'c_proj' in name:
ffn_down = (name, module)
-
+
if ffn_up is None or ffn_down is None:
print(" ⚠️ Could not find FFN layers")
return
-
+
up_name, up_module = ffn_up
down_name, down_module = ffn_down
-
+
# Create mask
mask = MaskOperations.create_structured_mask(scores, amount=amount, mode='low')
num_kept = mask.sum().item()
num_total = len(mask)
-
+
print(f"\n Mask: {num_kept}/{num_total} neurons kept ({100*num_kept/num_total:.1f}%)")
-
+
# Apply to up_proj (output dimension)
print(f"\n Applying to {up_name}:")
print(f" - Before: {up_module.weight.shape}")
-
+
with torch.no_grad():
# Mask output neurons (rows of up_proj)
up_module.weight.data *= mask.unsqueeze(1).float()
if up_module.bias is not None:
up_module.bias.data *= mask.float()
-
+
print(f" - Pruned: {(~mask).sum().item()} neurons zeroed out")
-
+
# Apply to down_proj (input dimension) - DEPENDENCY!
print(f"\n Applying to {down_name} (dependency propagation):")
print(f" - Before: {down_module.weight.shape}")
-
+
with torch.no_grad():
# Mask input neurons (columns of down_proj)
down_module.weight.data *= mask.unsqueeze(0).float()
-
+
print(f" - Pruned: {(~mask).sum().item()} inputs zeroed out (matches up_proj)")
-
+
# Verify dependency handled
print(f"\n ✓ Dependency handled: up_proj outputs ({num_kept}) = down_proj inputs ({num_kept})")
-
+
return mask
@@ -294,85 +294,85 @@ def main():
print("=" * 80)
print("LLaMA FFN Per-Neuron Analysis & Pruning")
print("=" * 80)
-
+
# Setup
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"\nDevice: {device}")
-
+
# Load model (using GPT-2 for demo, same principles apply to LLaMA-3)
model, tokenizer = load_llama_model(use_small_model=True)
model = model.to(device)
-
+
# Wrap model
wrapper = LLaMAWrapper(
model,
track_ffn=True,
track_attention=True
)
-
+
# Sample inputs
input_texts = [
"The capital of France is",
"Artificial intelligence is",
"The meaning of life is"
]
-
+
# Analyze first layer's FFN
layer_idx = 0
scores_dict = analyze_ffn_layer(
model, wrapper, layer_idx, input_texts, tokenizer
)
-
+
if not scores_dict:
print("\n⚠️ Could not analyze FFN layer")
return
-
+
# Create composite score
print(f"\n{'='*80}")
print("Creating Composite Importance Score")
print(f"{'='*80}")
-
+
# Combine RQ and redundancy
rq = scores_dict['rq']
redundancy = scores_dict['redundancy']
-
+
# Composite: High RQ + Low Redundancy = Important
# Score = RQ - 0.5*Redundancy
composite = torch.log(rq + 1e-8) - 0.5 * redundancy
-
- print(f" Composite = log(RQ) - 0.5*Redundancy")
+
+ print(" Composite = log(RQ) - 0.5*Redundancy")
print(f" Range: [{composite.min():.4f}, {composite.max():.4f}]")
-
+
# Identify most/least important neurons
top_important = torch.argsort(composite, descending=True)[:5]
top_redundant = torch.argsort(redundancy, descending=True)[:5]
-
+
print(f"\n Most important neurons (high composite): {top_important.tolist()}")
print(f" Most redundant neurons (high redundancy): {top_redundant.tolist()}")
-
+
# Prune using composite scores
print(f"\n{'='*80}")
print("Pruning FFN Neurons (Dependency-Aware)")
print(f"{'='*80}")
-
+
pruning_amount = 0.3 # Prune 30% of FFN neurons
-
+
mask = prune_ffn_neurons(
model,
layer_idx=layer_idx,
scores=composite,
amount=pruning_amount
)
-
+
# Verify model still runs
print(f"\n{'='*80}")
print("Verifying Pruned Model")
print(f"{'='*80}")
-
+
with torch.no_grad():
# Tokenize
test_input = tokenizer("After pruning, the model", return_tensors='pt').to(device)
-
+
# Forward pass
try:
output = model(**test_input)
@@ -380,7 +380,7 @@ def main():
print(f" ✓ Output shape: {output.logits.shape}")
except Exception as e:
print(f" ✗ Model failed after pruning: {e}")
-
+
# Summary
print(f"\n{'='*80}")
print("Summary")
diff --git a/examples/09_attention_neuron_vs_head_pruning.py b/examples/09_attention_neuron_vs_head_pruning.py
index f178fdf2..a781e6dd 100644
--- a/examples/09_attention_neuron_vs_head_pruning.py
+++ b/examples/09_attention_neuron_vs_head_pruning.py
@@ -16,11 +16,10 @@
import torch
import torch.nn as nn
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM
-from alignment.models.transformer_enhanced import TransformerWrapperEnhanced
-from alignment.services import NodeScoringService, MaskOperations
from alignment.metrics import get_metric
+from alignment.services import MaskOperations
def analyze_attention_structure(model):
@@ -28,10 +27,10 @@ def analyze_attention_structure(model):
print("=" * 80)
print("Attention Layer Structure")
print("=" * 80)
-
+
# Get first attention layer
attn_layer = model.transformer.h[0].attn if hasattr(model, 'transformer') else model.model.layers[0].self_attn
-
+
# Find projection layers
for name, module in attn_layer.named_modules():
if isinstance(module, nn.Linear):
@@ -39,7 +38,7 @@ def analyze_attention_structure(model):
print(f" Weight shape: {module.weight.shape}")
print(f" Output neurons: {module.weight.shape[0]}")
print(f" Input dimension: {module.weight.shape[1]}")
-
+
if module.weight.shape[0] % 32 == 0: # Assuming 32 heads
neurons_per_head = module.weight.shape[0] // 32
print(f" → Organized as: 32 heads × {neurons_per_head} dims/head")
@@ -49,13 +48,13 @@ def analyze_attention_structure(model):
def neuron_level_pruning_demo(model, layer_idx=0):
"""
Demonstrate NEURON-LEVEL pruning in attention.
-
+
Treats Q/K/V/O projections as Linear layers and prunes individual neurons.
"""
print("\n" + "=" * 80)
print("NEURON-LEVEL Attention Pruning")
print("=" * 80)
-
+
# Get attention layer
try:
attn = model.transformer.h[layer_idx].attn # GPT-2 style
@@ -65,12 +64,9 @@ def neuron_level_pruning_demo(model, layer_idx=0):
except:
attn = model.model.layers[layer_idx].self_attn # LLaMA style
q_proj = attn.q_proj
- k_proj = attn.k_proj
- v_proj = attn.v_proj
- o_proj = attn.o_proj
print(" Model type: LLaMA (separate Q/K/V/O projections)")
separate_projections = True
-
+
if not separate_projections:
print(" Note: GPT-2 combines Q/K/V, will demonstrate on combined projection")
target_proj = q_proj
@@ -78,82 +74,82 @@ def neuron_level_pruning_demo(model, layer_idx=0):
else:
target_proj = q_proj
proj_name = "Q projection"
-
+
num_neurons = target_proj.weight.shape[0]
hidden_dim = target_proj.weight.shape[1]
-
+
print(f"\n{proj_name}:")
print(f" Shape: {target_proj.weight.shape}")
print(f" Total neurons: {num_neurons}")
print(f" Input dimension: {hidden_dim}")
-
+
if num_neurons % 32 == 0:
print(f" Organized as: {num_neurons // 128} heads × 128 dims")
-
+
# Create dummy inputs for scoring
- print(f"\nComputing per-neuron importance...")
+ print("\nComputing per-neuron importance...")
inputs = torch.randn(16, hidden_dim) # [B, hidden_dim]
outputs = target_proj(inputs) # [B, num_neurons]
-
+
# Compute per-neuron scores
rq = get_metric('rayleigh_quotient')
rq_scores = rq.compute(inputs, target_proj.weight)
-
+
redundancy = get_metric('pairwise_redundancy_gaussian', mode='output_based', num_pairs=20)
redundancy_scores = redundancy.compute(outputs=outputs)
-
+
# Composite score
composite = torch.log(rq_scores + 1e-8) - 0.4 * redundancy_scores
-
+
# Show individual neuron analysis
print(f"\n Individual neuron breakdown (first 10 of {num_neurons}):")
print(f" {'Neuron':<8} {'RQ':<10} {'Redundancy':<12} {'Composite':<12}")
print(f" {'-'*50}")
-
+
for i in range(min(10, num_neurons)):
print(f" {i:<8} {rq_scores[i]:.4f} {redundancy_scores[i]:.4f} {composite[i]:.4f}")
-
+
# Prune neurons
pruning_amount = 0.3 # 30%
mask = MaskOperations.create_structured_mask(composite, amount=pruning_amount, mode='low')
-
+
num_pruned = (~mask).sum().item()
num_kept = mask.sum().item()
-
- print(f"\n Pruning plan:")
+
+ print("\n Pruning plan:")
print(f" Amount: {pruning_amount:.0%}")
print(f" Neurons pruned: {num_pruned}")
print(f" Neurons kept: {num_kept}")
-
+
# Show which neurons are pruned
pruned_indices = torch.where(~mask)[0]
if len(pruned_indices) <= 20:
print(f" Pruned neuron indices: {pruned_indices.tolist()}")
else:
print(f" Pruned neuron indices: {pruned_indices[:10].tolist()} ... (showing first 10)")
-
+
# Apply pruning
- print(f"\n Applying neuron-level mask...")
+ print("\n Applying neuron-level mask...")
target_proj.weight.data *= mask.unsqueeze(1).float()
if target_proj.bias is not None:
target_proj.bias.data *= mask.float()
-
+
print(f" ✓ {num_pruned} neurons zeroed out")
- print(f" ✓ Can prune ANY subset of neurons (not constrained to heads)")
-
+ print(" ✓ Can prune ANY subset of neurons (not constrained to heads)")
+
return mask, composite
def head_level_pruning_demo(model, layer_idx=0):
"""
Demonstrate HEAD-LEVEL pruning in attention.
-
+
Aggregates neurons into heads and prunes entire heads.
"""
print("\n" + "=" * 80)
print("HEAD-LEVEL Attention Pruning")
print("=" * 80)
-
+
try:
attn = model.model.layers[layer_idx].self_attn
q_proj = attn.q_proj
@@ -162,69 +158,69 @@ def head_level_pruning_demo(model, layer_idx=0):
except:
print(" Using GPT-2 (combined QKV, skipping head demo)")
return None, None
-
- print(f"\n Attention configuration:")
+
+ print("\n Attention configuration:")
print(f" Number of heads: {num_heads}")
print(f" Dimension per head: {head_dim}")
print(f" Total dimensions: {num_heads * head_dim}")
-
+
# Compute neuron-level scores first
inputs = torch.randn(16, 4096)
outputs = q_proj(inputs)
-
+
redundancy = get_metric('pairwise_redundancy_gaussian', mode='output_based', num_pairs=20)
neuron_scores = redundancy.compute(outputs=outputs) # [4096]
-
+
# Aggregate to head-level
print(f"\n Aggregating {neuron_scores.shape[0]} neurons into {num_heads} heads...")
-
+
head_scores = []
for head_idx in range(num_heads):
start_idx = head_idx * head_dim
end_idx = start_idx + head_dim
-
+
# Average neuron scores within this head
head_score = neuron_scores[start_idx:end_idx].mean()
head_scores.append(head_score)
-
+
head_scores = torch.tensor(head_scores) # [32]
-
+
# Show per-head scores
- print(f"\n Per-head redundancy:")
+ print("\n Per-head redundancy:")
print(f" {'Head':<6} {'Redundancy':<12} {'Assessment'}")
print(f" {'-'*40}")
-
+
for h in range(min(10, num_heads)):
assessment = "Redundant (prune)" if head_scores[h] > 0.5 else "Unique (keep)"
print(f" {h:<6} {head_scores[h]:.4f} {assessment}")
-
+
# Prune heads
pruning_amount = 0.25 # 25% = 8 heads
head_mask = MaskOperations.create_structured_mask(head_scores, amount=pruning_amount, mode='low')
-
+
num_heads_pruned = (~head_mask).sum().item()
-
- print(f"\n Pruning plan:")
+
+ print("\n Pruning plan:")
print(f" Amount: {pruning_amount:.0%} of heads")
print(f" Heads pruned: {num_heads_pruned}/{num_heads}")
print(f" Pruned head indices: {torch.where(~head_mask)[0].tolist()}")
-
+
# Expand to neuron-level mask
neuron_mask = head_mask.repeat_interleave(head_dim) # [4096]
-
- print(f"\n Expanding to neuron-level:")
+
+ print("\n Expanding to neuron-level:")
print(f" Neuron mask shape: {neuron_mask.shape}")
print(f" Neurons pruned: {(~neuron_mask).sum().item()} (= {num_heads_pruned} heads × {head_dim} dims)")
-
+
# Apply to Q/K/V/O (all get same mask for consistency)
- print(f"\n Applying to all projections (Q/K/V/O)...")
+ print("\n Applying to all projections (Q/K/V/O)...")
q_proj.weight.data *= neuron_mask.unsqueeze(1).float()
attn.k_proj.weight.data *= neuron_mask.unsqueeze(1).float()
attn.v_proj.weight.data *= neuron_mask.unsqueeze(1).float()
attn.o_proj.weight.data *= neuron_mask.unsqueeze(0).float() # Input dim
-
+
print(f" ✓ {num_heads_pruned} heads pruned consistently across Q/K/V/O")
-
+
return head_mask, head_scores
@@ -233,32 +229,32 @@ def compare_approaches(neuron_mask, head_mask, num_heads=32, head_dim=128):
print("\n" + "=" * 80)
print("Comparison: Neuron-Level vs Head-Level")
print("=" * 80)
-
+
# Convert head mask to neuron mask for comparison
if head_mask is not None:
head_neuron_mask = head_mask.repeat_interleave(head_dim)
else:
print(" Head-level demo skipped (GPT-2 model)")
return
-
+
# Statistics
neuron_pruned = (~neuron_mask).sum().item()
head_equivalent_pruned = (~head_neuron_mask).sum().item()
-
- print(f"\nNeuron-Level Pruning:")
+
+ print("\nNeuron-Level Pruning:")
print(f" Neurons pruned: {neuron_pruned}")
- print(f" Flexibility: Can prune ANY neurons")
- print(f" Granularity: Individual neurons")
-
- print(f"\nHead-Level Pruning:")
+ print(" Flexibility: Can prune ANY neurons")
+ print(" Granularity: Individual neurons")
+
+ print("\nHead-Level Pruning:")
print(f" Neurons pruned: {head_equivalent_pruned} (= {head_equivalent_pruned//head_dim} heads × {head_dim})")
print(f" Flexibility: Must prune in multiples of {head_dim}")
- print(f" Granularity: Entire heads")
-
- print(f"\nKey Differences:")
- print(f" ✓ Neuron-level: More flexible (any %, any neurons)")
- print(f" ✓ Head-level: Cleaner (maintains head structure)")
- print(f" ✓ Both supported by your framework!")
+ print(" Granularity: Entire heads")
+
+ print("\nKey Differences:")
+ print(" ✓ Neuron-level: More flexible (any %, any neurons)")
+ print(" ✓ Head-level: Cleaner (maintains head structure)")
+ print(" ✓ Both supported by your framework!")
def main():
@@ -266,35 +262,35 @@ def main():
print("=" * 80)
print("Attention Layer Pruning: Neuron vs Head Level")
print("=" * 80)
-
+
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"\nDevice: {device}")
-
+
# Load model (GPT-2 for demo, same principles for LLaMA)
print("\nLoading model (GPT-2 for demo)...")
model = AutoModelForCausalLM.from_pretrained('gpt2')
model = model.to(device)
model.eval()
-
- print(f"✓ Loaded GPT-2")
+
+ print("✓ Loaded GPT-2")
print(f" - Layers: {model.config.n_layer}")
print(f" - Hidden size: {model.config.n_embd}")
print(f" - Attention heads: {model.config.n_head}")
-
+
# Analyze structure
analyze_attention_structure(model)
-
+
# Demonstrate neuron-level pruning
neuron_mask, neuron_scores = neuron_level_pruning_demo(model, layer_idx=0)
-
+
# Demonstrate head-level pruning
# (Skip for GPT-2 since it has combined QKV, but would work for LLaMA)
head_mask, head_scores = head_level_pruning_demo(model, layer_idx=0)
-
+
# Compare
if head_mask is not None:
compare_approaches(neuron_mask, head_mask)
-
+
# Summary
print("\n" + "=" * 80)
print("Summary: You Can Prune Attention at BOTH Levels!")
diff --git a/examples/README.md b/examples/README.md
index 918bc121..080d2ff9 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -1,30 +1,37 @@
-## Examples
+# Examples
-This directory contains example scripts demonstrating the alignment framework functionality.
+Example scripts demonstrating framework capabilities.
-### Available Examples
+## Available Examples
-1. **01_quick_start.py** - Minimal example showing basic model wrapping and metrics
-2. **02_complete_experiment.py** - Full workflow with training, pruning, and visualization
-3. **03_pruning_strategies.py** - Comparison of different pruning methods
-4. **04_visualization_gallery.py** - Comprehensive visualization capabilities
+**Basic Usage**
+- `01_quick_start.py` - Model wrapping and basic metrics
+- `02_complete_experiment.py` - Full workflow with training and pruning
+- `03_pruning_strategies.py` - Pruning strategy comparison
+- `04_visualization_gallery.py` - Visualization examples
-### Running Examples
+**Advanced Usage**
+- `06_redundancy_aware_pruning.py` - Information-theoretic pruning
+- `07_mnist_intelligent_pruning.py` - Complete MNIST pruning workflow
+- `08_llama_ffn_pruning.py` - LLaMA feed-forward pruning
+- `09_attention_neuron_vs_head_pruning.py` - Attention analysis
-Each example can be run independently:
+## Running Examples
+
+Each example is self-contained:
```bash
python examples/01_quick_start.py
-python examples/02_complete_experiment.py
-python examples/03_pruning_strategies.py
-python examples/04_visualization_gallery.py
+python examples/07_mnist_intelligent_pruning.py
+python examples/08_llama_ffn_pruning.py
```
-### Configuration-Based Experiments
+## Configuration-Based Experiments
-For experiments using YAML configurations, see:
+For YAML-based experiments:
```bash
-python scripts/run_experiment.py --config configs/examples/resnet18_analysis.yaml
+python scripts/run_experiment.py --config configs/examples/mnist_basic.yaml
+python scripts/run_experiment.py --config configs/examples/resnet_pruning.yaml
```
-See the `configs/examples/` directory for configuration templates.
\ No newline at end of file
+See `configs/examples/` for available configurations.
\ No newline at end of file
diff --git a/examples/quick_demo.py b/examples/quick_demo.py
index 462167e4..0459145c 100644
--- a/examples/quick_demo.py
+++ b/examples/quick_demo.py
@@ -6,13 +6,13 @@
Usage:
python quick_demo.py
-
+
No configuration needed - this script runs with default settings and creates its own model.
Requirements:
- PyTorch
- alignment package installed
-
+
Output:
- Console output showing:
* Model structure
@@ -24,8 +24,9 @@
import torch
import torch.nn as nn
-from alignment.models import ModelWrapper
+
from alignment.metrics import get_metric
+from alignment.models import ModelWrapper
def main():
@@ -37,55 +38,55 @@ def main():
nn.ReLU(),
nn.Linear(128, 10)
)
-
+
print("Model created with 3 linear layers")
-
+
# 2. Wrap the model to track activations
wrapped_model = ModelWrapper(model)
print(f"\nTracked layers: {wrapped_model.tracked_layers}")
-
+
# 3. Generate some dummy data
batch_size = 32
inputs = torch.randn(batch_size, 784)
-
+
# 4. Forward pass with activation tracking
outputs, activations = wrapped_model.forward_with_activations(inputs)
print(f"\nCollected activations from {len(activations)} points")
print(f"Activation keys: {list(activations.keys())}")
-
+
# 5. Get layer weights
weights = wrapped_model.get_layer_weights()
print(f"Extracted weights from {len(weights)} layers")
-
+
# 6. Compute alignment metrics
print("\n" + "="*60)
print("Computing Alignment Metrics")
print("="*60)
-
+
# Rayleigh Quotient
RQMetric = get_metric('rayleigh_quotient')
rq_metric = RQMetric() # Instantiate the metric
-
+
for layer_name in wrapped_model.tracked_layers:
# Get inputs and weights for this layer
layer_inputs = activations[f"{layer_name}_input"]
layer_weights = weights[layer_name]
-
+
# Compute RQ scores
scores = rq_metric.compute(inputs=layer_inputs, weights=layer_weights)
-
+
print(f"\nLayer: {layer_name}")
print(f" Input shape: {layer_inputs.shape}")
print(f" Weight shape: {layer_weights.shape}")
print(f" RQ scores: mean={scores.mean():.4f}, std={scores.std():.4f}")
print(f" Min neuron score: {scores.min():.4f}")
print(f" Max neuron score: {scores.max():.4f}")
-
+
# 7. Try other metrics
print("\n" + "="*60)
print("Trying Other Metrics")
print("="*60)
-
+
# Weight cosine similarity
WeightSimMetric = get_metric('weight_cosine_similarity')
weight_sim = WeightSimMetric() # Instantiate
@@ -93,35 +94,35 @@ def main():
layer_weights = weights[layer_name]
sim_scores = weight_sim.compute(weights=layer_weights)
print(f"\nWeight Cosine Similarity ({layer_name}): mean={sim_scores.mean():.4f}")
-
+
# 8. Demonstrate pruning
print("\n" + "="*60)
print("Pruning Demo")
print("="*60)
-
- from alignment.pruning import get_pruning_strategy, PruningConfig
-
+
+ from alignment.pruning import PruningConfig, get_pruning_strategy
+
# Use magnitude-based pruning
config = PruningConfig(amount=0.5, pruning_mode='low')
strategy = get_pruning_strategy('magnitude', config=config)
-
+
# Apply pruning to first layer
first_layer = model[0] # First linear layer
mask = strategy.prune(first_layer)
-
- print(f"\nPruning 50% of weights in first layer")
+
+ print("\nPruning 50% of weights in first layer")
print(f"Sparsity achieved: {(mask == 0).float().mean():.2%}")
-
+
# Apply mask
first_layer.weight.data *= mask
-
+
# Test forward pass after pruning
pruned_outputs = model(inputs)
print(f"\nOutput shape after pruning: {pruned_outputs.shape}")
print("Pruning applied successfully!")
-
+
print("\nDemo completed successfully!")
if __name__ == "__main__":
- main()
\ No newline at end of file
+ main()
diff --git a/scripts/README.md b/scripts/README.md
index 4c0e4ae0..9deacea9 100644
--- a/scripts/README.md
+++ b/scripts/README.md
@@ -1,21 +1,27 @@
-# Production Scripts
+# Scripts
-This directory contains production-ready scripts for running alignment experiments.
+Scripts for running experiments.
-## Scripts
+## Main Script
-### `run_experiment.py`
-The main experiment runner for the alignment framework. This is the primary tool for running research experiments with full configuration support.
+`run_experiment.py` - Experiment runner with YAML configuration support.
-**Usage:**
+Usage:
```bash
-# From repository root
-python scripts/run_experiment.py --config configs/unified_config.yaml
+python scripts/run_experiment.py --config configs/examples/mnist_basic.yaml
```
-## Difference from Examples
+The script handles:
+- Model loading (pretrained or training from scratch)
+- Metric computation
+- Pruning with various strategies
+- Evaluation and visualization
-- **`examples/`**: Educational scripts for learning the framework
-- **`scripts/`**: Production tools for research experiments
+## Configuration
-The examples demonstrate concepts with hardcoded parameters for clarity, while the scripts here are fully configurable tools designed for actual research work.
\ No newline at end of file
+Experiments are configured via YAML files. See `configs/template.yaml` for all parameters and `configs/examples/` for working examples.
+
+## Examples vs Scripts
+
+- `examples/` - Standalone demonstration scripts
+- `scripts/` - Configuration-based experiment runner
\ No newline at end of file
diff --git a/scripts/run_experiment.py b/scripts/run_experiment.py
index ccb6bc41..9a78ea71 100644
--- a/scripts/run_experiment.py
+++ b/scripts/run_experiment.py
@@ -14,19 +14,14 @@
"""
import argparse
+import json
import logging
import os
import sys
-import json
-import yaml
-from pathlib import Path
from datetime import datetime
-from typing import Dict, List, Any, Optional, Tuple
+from pathlib import Path
-import torch
-import torch.nn as nn
-import numpy as np
-from tqdm import tqdm
+import yaml
# Add the src directory to Python path
current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -34,8 +29,10 @@
# Import from the alignment package
from alignment.experiments.general_alignment import GeneralAlignmentExperiment
+from alignment.pruning.experiments.cascading_layer import (
+ CascadingLayerPruningExperiment,
+)
from alignment.pruning.experiments.layer_wise import LayerIsolatedPruningExperiment
-from alignment.pruning.experiments.cascading_layer import CascadingLayerPruningExperiment
logger = logging.getLogger(__name__)
@@ -44,7 +41,7 @@ def load_config(config_path, overrides=None):
"""Load and merge configuration."""
with open(config_path, 'r') as f:
config = yaml.safe_load(f)
-
+
# Apply overrides
if overrides:
for key, value in overrides.items():
@@ -55,7 +52,7 @@ def load_config(config_path, overrides=None):
current[k] = {}
current = current[k]
current[keys[-1]] = value
-
+
return config
@@ -63,11 +60,11 @@ def create_experiment_config(unified_config):
"""Convert unified config to experiment config object."""
from alignment.experiments.base import ExperimentConfig
from alignment.experiments.general_alignment import GeneralAlignmentConfig
-
+
# Extract model config and handle different naming conventions
model_config = unified_config.get('model', {})
model_name = model_config.get('name', model_config.get('architecture', 'mlp'))
-
+
# Extract dataset config
dataset_config = unified_config.get('dataset', {})
if isinstance(dataset_config, str):
@@ -76,7 +73,7 @@ def create_experiment_config(unified_config):
dataset_config = {'name': dataset_name}
else:
dataset_name = dataset_config.get('name', dataset_config.get('dataset', 'mnist'))
-
+
# Build base parameters common to all experiment types
base_params = {
'name': unified_config.get('experiment_name', 'unified_experiment'),
@@ -88,7 +85,7 @@ def create_experiment_config(unified_config):
'num_workers': dataset_config.get('num_workers', unified_config.get('data', {}).get('num_workers', 4)),
'metrics': unified_config.get('alignment', {}).get('metrics', unified_config.get('analysis', {}).get('metrics', ['rayleigh_quotient'])),
}
-
+
# Build model config with proper parameter names
model_kwargs = {}
if model_name == 'mlp':
@@ -109,9 +106,9 @@ def create_experiment_config(unified_config):
else:
model_kwargs['input_dim'] = model_config.get('input_dim', 784)
model_kwargs['output_dim'] = model_config.get('output_dim', 10)
-
+
base_params['model_config'] = model_kwargs
-
+
# Add training config if present
training_config = unified_config.get('training', {})
if training_config:
@@ -123,15 +120,15 @@ def create_experiment_config(unified_config):
else:
base_params['learning_rate'] = training_config.get('learning_rate', 0.001)
base_params['optimizer'] = training_config.get('optimizer', 'adam')
-
+
# Infer experiment type based on enabled blocks
pruning_cfg = unified_config.get('pruning', {})
dropout_cfg = unified_config.get('dropout', {})
-
+
pruning_enabled = pruning_cfg.get('enabled', False)
cascading_scope = pruning_cfg.get('scope', 'layer') == 'cascading'
dropout_enabled = dropout_cfg.get('enabled', False)
-
+
# Determine if we need specialized experiment types or GeneralAlignmentConfig
if cascading_scope and pruning_enabled:
# For cascading layer pruning, use base ExperimentConfig
@@ -149,14 +146,14 @@ def create_experiment_config(unified_config):
general_params['aggregate_metrics'] = unified_config.get('aggregate_metrics', True)
general_params['save_individual_networks'] = unified_config.get('save_individual_networks', False)
general_params['save_checkpoints'] = unified_config.get('save_checkpoints', False)
-
+
config = GeneralAlignmentConfig(**general_params)
-
+
# Set flags based on enabled blocks
config.do_train = training_config.get('epochs', training_config.get('do_train', 0)) > 0
config.do_dropout_analysis = dropout_enabled
config.do_pruning_experiments = pruning_enabled
-
+
# Determine inferred type for logging
if pruning_enabled:
inferred_type = 'standard_pruning'
@@ -164,12 +161,12 @@ def create_experiment_config(unified_config):
inferred_type = 'progressive_dropout'
else:
inferred_type = 'alignment_analysis'
-
+
# Handle plot generation
visualization_config = unified_config.get('visualization', {})
output_config = unified_config.get('output', {})
analysis_config = unified_config.get('analysis', {})
-
+
# Default to True if not explicitly set to False
generate_plots = True
if 'generate_plots' in visualization_config:
@@ -178,66 +175,66 @@ def create_experiment_config(unified_config):
generate_plots = output_config['generate_plots']
elif 'generate_plots' in analysis_config:
generate_plots = analysis_config['generate_plots']
-
+
config.generate_plots = generate_plots
logger.info(f"Plot generation enabled: {config.generate_plots}")
-
+
# Pruning specific configuration
if pruning_enabled:
# Get algorithms
algorithms = pruning_cfg.get('algorithms', ['magnitude'])
config.pruning_strategies = algorithms if isinstance(algorithms, list) else [algorithms]
-
+
# Get sparsity levels
sparsity_levels = pruning_cfg.get('sparsity_levels', [0.0, 0.1, 0.3, 0.5, 0.7, 0.9])
config.pruning_amounts = sparsity_levels
-
+
# Get selection modes
selection_modes = pruning_cfg.get('selection_modes', pruning_cfg.get('selection_mode', ['low']))
config.pruning_selection_mode = selection_modes if isinstance(selection_modes, list) else [selection_modes]
-
+
# Fine-tuning settings
config.fine_tune_after_pruning = pruning_cfg.get('fine_tune_after_pruning', True)
config.fine_tune_epochs = pruning_cfg.get('fine_tune_epochs', 5)
config.fine_tune_learning_rate = pruning_cfg.get('fine_tune_learning_rate', 0.0001)
-
+
# Scope settings
scope = pruning_cfg.get('scope', 'layer')
config.pruning_scope = scope
-
+
# Alignment metric settings
config.pruning_alignment_metric = pruning_cfg.get('alignment_metric', 'rayleigh_quotient')
config.pruning_hybrid_alpha = pruning_cfg.get('hybrid_alpha', 0.5)
-
+
# Ultra-parallel evaluation settings
config.use_ultra_parallel_eval = pruning_cfg.get('use_ultra_parallel_eval', False)
config.eval_batches = pruning_cfg.get('eval_batches', None)
-
+
logger.info(f"Pruning enabled: algorithms={config.pruning_strategies}, levels={config.pruning_amounts}, modes={config.pruning_selection_mode}")
logger.info(f"Ultra-parallel eval: {config.use_ultra_parallel_eval}, eval_batches: {config.eval_batches}")
-
+
# Dropout specific configuration
if dropout_enabled:
config.dropout_rates = dropout_cfg.get('rates', [0.0, 0.1, 0.3, 0.5, 0.7, 0.9])
logger.info(f"Dropout enabled: rates={config.dropout_rates}")
-
+
# Common configuration for all experiment types
if not isinstance(config, GeneralAlignmentConfig):
# Ensure plot generation is enabled for other experiment types too
visualization_config = unified_config.get('visualization', {})
output_config = unified_config.get('output', {})
config.generate_plots = visualization_config.get('generate_plots', output_config.get('generate_plots', True))
-
+
# Add additional attributes for compatibility
config.training_config = training_config
config.train_model = config.training_config.get('epochs', 0) > 0
config.alignment_metrics = config.metrics
config.apply_pruning = pruning_enabled
-
+
# Legacy support for pruning_analysis and network_compression blocks
pruning_analysis = unified_config.get('pruning_analysis', {})
network_compression = unified_config.get('network_compression', {})
-
+
# Use the appropriate config based on what's enabled
if pruning_analysis.get('enabled', False):
active_pruning_config = pruning_analysis
@@ -249,12 +246,12 @@ def create_experiment_config(unified_config):
# Fallback to pruning config or empty
active_pruning_config = pruning_cfg if pruning_enabled else {}
config.pruning_strategy = pruning_cfg.get('algorithms', ['magnitude'])[0] if pruning_enabled else 'magnitude'
-
+
config.pruning_config = active_pruning_config
config.analysis_config = unified_config.get('visualization', unified_config.get('analysis', {}))
config.eval_model = True
config.cnn_mode = model_config.get('cnn_mode', 'unfold')
-
+
# Set dropout rates (handle multiple sources)
if hasattr(config, 'dropout_rates'):
# Already set for GeneralAlignmentConfig
@@ -263,7 +260,7 @@ def create_experiment_config(unified_config):
config.dropout_rates = pruning_analysis.get('dropout_rates', [0.0, 0.1, 0.3, 0.5, 0.7, 0.9])
else:
config.dropout_rates = dropout_cfg.get('rates', [0.0, 0.1, 0.3, 0.5, 0.7, 0.9])
-
+
# Handle selection modes for compatibility
if pruning_analysis.get('enabled', False):
selection_mode = pruning_analysis.get('selection_strategies', ['low'])
@@ -271,19 +268,19 @@ def create_experiment_config(unified_config):
selection_mode = [network_compression.get('selection_strategy', 'low')]
else:
selection_mode = pruning_cfg.get('selection_modes', ['low']) if pruning_enabled else ['low']
-
+
config.pruning_modes = selection_mode if isinstance(selection_mode, list) else [selection_mode]
-
+
config.cascade_direction = unified_config.get('experiment_specific', {}).get('cascade_direction', 'forward')
config.recompute_scores = True
-
+
# Note: Output directories are now set in main() after creating timestamped folders
config.plot_dpi = unified_config.get('visualization', {}).get('plot_dpi', unified_config.get('output', {}).get('plot_dpi', 300))
-
+
# Log the inferred experiment type
logger.info(f"Inferred experiment type: {inferred_type} based on config blocks")
config._inferred_experiment_type = inferred_type # Store for later use
-
+
return config
@@ -294,57 +291,57 @@ def main():
parser.add_argument('--device', type=str, help='Override device')
parser.add_argument('--seed', type=int, help='Override seed')
parser.add_argument('--output-dir', type=str, help='Override output directory')
-
+
args, unknown = parser.parse_known_args()
-
+
# Parse additional overrides
overrides = {}
if args.device:
overrides['device'] = args.device
if args.seed:
overrides['seed'] = args.seed
-
+
# Load config using the proper config loader
from alignment.configs.config_loader import load_config as proper_load_config
config = proper_load_config(args.config)
-
+
# Apply overrides to the loaded config
if overrides:
for key, value in overrides.items():
if hasattr(config, key):
setattr(config, key, value)
-
+
# Create timestamped output directory
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
experiment_name = getattr(config, 'name', 'experiment')
-
+
if args.output_dir:
output_dir = Path(args.output_dir)
else:
# Create a unique directory with experiment name and timestamp
output_dir = Path(f"results/{experiment_name}_{timestamp}")
-
+
output_dir.mkdir(parents=True, exist_ok=True)
-
+
# Save the configuration used
config_save_path = output_dir / 'experiment_config.yaml'
config.save(config_save_path)
-
+
# Update config with timestamped directories
config.checkpoint_dir = str(output_dir / 'checkpoints')
config.log_dir = str(output_dir / 'logs')
config.experiment_dir = str(output_dir) # Add experiment_dir for compatibility
-
+
# Create plots directory in results folder (not in logs)
plots_dir = output_dir / 'plots'
config.plots_dir = str(plots_dir) # Add plots_dir for visualization
-
+
# Ensure directories exist
Path(config.checkpoint_dir).mkdir(parents=True, exist_ok=True)
Path(config.log_dir).mkdir(parents=True, exist_ok=True)
plots_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Created plots directory: {plots_dir}")
-
+
# Setup logging to both file and console
log_file = output_dir / 'experiment.log'
logging.basicConfig(
@@ -355,10 +352,10 @@ def main():
logging.StreamHandler()
]
)
-
+
# Print experiment info
print(f"\n{'='*60}")
- print(f"Running Alignment Experiment")
+ print("Running Alignment Experiment")
print(f"{'='*60}")
print(f"Configuration: {args.config}")
print(f"Output directory: {output_dir}")
@@ -366,13 +363,13 @@ def main():
print(f"Plot generation: {getattr(config, 'generate_plots', True)}")
print(f"Plots directory: {plots_dir}")
print(f"{'='*60}\n")
-
+
# Use the inferred experiment type from config
experiment_type = getattr(config, '_inferred_experiment_type', 'alignment_analysis')
-
+
logger.info(f"Running {experiment_type} experiment")
-
-
+
+
# Create experiment based on inferred type
if experiment_type in ['standard_pruning', 'progressive_dropout', 'alignment_analysis']:
experiment = GeneralAlignmentExperiment(config)
@@ -382,14 +379,14 @@ def main():
experiment = CascadingLayerPruningExperiment(config)
else:
raise ValueError(f"Unknown experiment type: {experiment_type}")
-
-
+
+
# Run experiment
results = experiment.run()
-
+
# Save results with timestamp
results_file = output_dir / f'results_{timestamp}.json'
-
+
# Convert numpy arrays to lists for JSON serialization
def convert_to_serializable(obj):
if hasattr(obj, 'tolist'):
@@ -399,12 +396,12 @@ def convert_to_serializable(obj):
elif isinstance(obj, list):
return [convert_to_serializable(i) for i in obj]
return obj
-
+
serializable_results = convert_to_serializable(results)
-
+
with open(results_file, 'w') as f:
json.dump(serializable_results, f, indent=2)
-
+
# Create experiment summary
summary_file = output_dir / 'experiment_summary.txt'
with open(summary_file, 'w') as f:
@@ -414,19 +411,19 @@ def convert_to_serializable(obj):
f.write(f"Experiment Type: {experiment_type} (inferred from config blocks)\n")
f.write(f"Plot Generation: {getattr(config, 'generate_plots', True)}\n")
f.write("=" * 50 + "\n\n")
-
+
# Add results summary
if 'test_results' in results:
f.write("Final Model Performance:\n")
f.write(f" - Accuracy: {results['test_results'].get('final_accuracy', 'N/A'):.2f}%\n")
f.write(f" - Loss: {results['test_results'].get('final_loss', 'N/A'):.4f}\n\n")
-
+
if 'pruning_results' in results and results['pruning_results']:
f.write("Pruning Experiments:\n")
strategies = results['pruning_results'].get('strategies', {})
f.write(f" - Strategies tested: {list(strategies.keys())}\n")
f.write(f" - Plots saved in: {config.log_dir}/plots/\n")
-
+
# List plots created
plots_created = list(plots_dir.glob('*'))
if plots_created:
@@ -436,28 +433,28 @@ def convert_to_serializable(obj):
f.write(f" - {plot_file.name}\n")
else:
f.write("\nNo plots were generated.\n")
-
+
f.write("\nGenerated Files:\n")
for file_path in sorted(output_dir.rglob('*')):
if file_path.is_file():
relative_path = file_path.relative_to(output_dir)
f.write(f" - {relative_path}\n")
-
+
# Print completion message
print(f"\n{'='*60}")
print("Experiment Complete!")
print(f"{'='*60}")
-
+
if 'test_results' in results:
print(f"Final model accuracy: {results['test_results'].get('final_accuracy', 'N/A'):.2f}%")
print(f"Final model loss: {results['test_results'].get('final_loss', 'N/A'):.4f}")
-
+
print(f"\nAll results saved in: {output_dir}")
print(f" - Configuration: {config_save_path}")
print(f" - Results: {results_file}")
print(f" - Summary: {summary_file}")
print(f" - Logs: {log_file}")
-
+
# Check and report on plots
plots_created = list(plots_dir.glob('*'))
if plots_created:
@@ -466,13 +463,13 @@ def convert_to_serializable(obj):
if plot_file.is_file():
print(f" * {plot_file.name}")
else:
- print(f" - No plots generated (check generate_plots setting and experiment configuration)")
+ print(" - No plots generated (check generate_plots setting and experiment configuration)")
print(f" * Pruning enabled: {getattr(config, 'do_pruning_experiments', False)}")
print(f" * Plot generation: {getattr(config, 'generate_plots', False)}")
print(f" * Pruning results in output: {'pruning_results' in results}")
-
+
print(f"{'='*60}\n")
if __name__ == '__main__':
- main()
\ No newline at end of file
+ main()
diff --git a/src/alignment/__init__.py b/src/alignment/__init__.py
index 205b0035..dea0bd96 100644
--- a/src/alignment/__init__.py
+++ b/src/alignment/__init__.py
@@ -10,27 +10,27 @@
# Core functionality
from .core.base import BaseMetric
from .core.registry import METRIC_REGISTRY
-from .models import ModelWrapper
-
-# Metrics
-from .metrics import get_metric, list_metrics
# Data processing
from .data.processing import BatchMetricProcessor
-# Pruning
-from .pruning import PruningConfig, get_pruning_strategy
-
# Experiment tracking
from .experiments.tracking import create_tracker
+# Metrics
+from .metrics import get_metric, list_metrics
+from .models import ModelWrapper
+
+# Pruning
+from .pruning import PruningConfig, get_pruning_strategy
+
# Services (NEW in v0.2.0)
from .services import (
ActivationCaptureService,
ActivationData,
- NodeScoringService,
CompositeScores,
MaskOperations,
+ NodeScoringService,
)
# Visualization
@@ -43,33 +43,33 @@
__all__ = [
# Core
"ModelWrapper",
- "BaseMetric",
+ "BaseMetric",
"METRIC_REGISTRY",
-
+
# Metrics
"get_metric",
"list_metrics",
-
+
# Data processing
"BatchMetricProcessor",
-
+
# Services (NEW in v0.2.0)
"ActivationCaptureService",
"ActivationData",
"NodeScoringService",
"CompositeScores",
"MaskOperations",
-
+
# Pruning
"PruningConfig",
"get_pruning_strategy",
-
+
# Experiment tracking
"create_tracker",
-
+
# Visualization
"AlignmentVisualizer",
-
+
# Version
"__version__"
-]
\ No newline at end of file
+]
diff --git a/src/alignment/analysis/__init__.py b/src/alignment/analysis/__init__.py
index 5a3401e0..051e90fc 100644
--- a/src/alignment/analysis/__init__.py
+++ b/src/alignment/analysis/__init__.py
@@ -10,23 +10,13 @@
"""
# Aggregation
-from .aggregation import (
- ResultAggregator,
- MetricAggregator,
- LayerAggregator,
-)
+from .aggregation import LayerAggregator, MetricAggregator, ResultAggregator
# Unified Reporting
-from .unified_reporter import (
- UnifiedReporter,
- generate_quick_report,
-)
+from .unified_reporter import UnifiedReporter, generate_quick_report
# Unified Visualization
-from .visualization.unified_visualizer import (
- UnifiedVisualizer,
- plot_quick_summary,
-)
+from .visualization.unified_visualizer import UnifiedVisualizer, plot_quick_summary
__all__ = [
# Aggregation
@@ -39,4 +29,4 @@
# Visualization
'UnifiedVisualizer',
'plot_quick_summary',
-]
\ No newline at end of file
+]
diff --git a/src/alignment/analysis/aggregation/__init__.py b/src/alignment/analysis/aggregation/__init__.py
index 083196f7..84e56b12 100644
--- a/src/alignment/analysis/aggregation/__init__.py
+++ b/src/alignment/analysis/aggregation/__init__.py
@@ -2,12 +2,12 @@
Aggregation utilities for experiment analysis.
"""
-from .results import ResultAggregator
-from .metrics import MetricAggregator
from .layers import LayerAggregator
+from .metrics import MetricAggregator
+from .results import ResultAggregator
__all__ = [
'ResultAggregator',
'MetricAggregator',
'LayerAggregator',
-]
\ No newline at end of file
+]
diff --git a/src/alignment/analysis/aggregation/layers.py b/src/alignment/analysis/aggregation/layers.py
index a8f74623..321b8e40 100644
--- a/src/alignment/analysis/aggregation/layers.py
+++ b/src/alignment/analysis/aggregation/layers.py
@@ -2,10 +2,11 @@
Layer-wise metric aggregation for analyzing patterns across network layers.
"""
-from typing import Dict, List, Tuple
+import logging
from collections import defaultdict
+from typing import Dict, List, Tuple
+
import numpy as np
-import logging
logger = logging.getLogger(__name__)
@@ -14,15 +15,15 @@ class LayerAggregator:
"""
Aggregates metrics by layer to analyze layer-wise patterns.
"""
-
+
def __init__(self):
"""Initialize layer aggregator."""
self.layer_metrics = defaultdict(lambda: defaultdict(list))
-
+
def add_metrics(self, metrics: Dict[str, Dict[str, float]]):
"""
Add metrics from a single evaluation.
-
+
Args:
metrics: Dictionary of metrics by name and layer
"""
@@ -30,20 +31,20 @@ def add_metrics(self, metrics: Dict[str, Dict[str, float]]):
if isinstance(layer_values, dict):
for layer_name, value in layer_values.items():
self.layer_metrics[layer_name][metric_name].append(value)
-
+
def get_layer_summary(self, layer_name: str) -> Dict[str, Dict[str, float]]:
"""
Get summary statistics for a specific layer.
-
+
Args:
layer_name: Name of the layer
-
+
Returns:
Dictionary mapping metric names to statistics
"""
if layer_name not in self.layer_metrics:
return {}
-
+
summary = {}
for metric_name, values in self.layer_metrics[layer_name].items():
if values:
@@ -55,9 +56,9 @@ def get_layer_summary(self, layer_name: str) -> Dict[str, Dict[str, float]]:
'max': float(np.max(values_array)),
'count': len(values)
}
-
+
return summary
-
+
def rank_layers(
self,
metric_name: str,
@@ -66,21 +67,21 @@ def rank_layers(
) -> List[Tuple[str, float]]:
"""
Rank layers by a specific metric.
-
+
Args:
metric_name: Metric to rank by
criterion: Statistic to use ('mean', 'max', 'min', 'std')
ascending: Whether to sort in ascending order
-
+
Returns:
List of (layer_name, value) tuples
"""
layer_values = []
-
+
for layer_name, metrics in self.layer_metrics.items():
if metric_name in metrics and metrics[metric_name]:
values_array = np.array(metrics[metric_name])
-
+
if criterion == 'mean':
value = np.mean(values_array)
elif criterion == 'max':
@@ -91,14 +92,14 @@ def rank_layers(
value = np.std(values_array)
else:
raise ValueError(f"Unknown criterion: {criterion}")
-
+
layer_values.append((layer_name, float(value)))
-
+
# Sort by value
layer_values.sort(key=lambda x: x[1], reverse=not ascending)
-
+
return layer_values
-
+
def find_anomalous_layers(
self,
metric_name: str,
@@ -106,35 +107,35 @@ def find_anomalous_layers(
) -> List[str]:
"""
Find layers with anomalous metric values.
-
+
Args:
metric_name: Metric to analyze
threshold_std: Number of standard deviations for anomaly detection
-
+
Returns:
List of anomalous layer names
"""
# Collect all values
all_values = []
layer_means = {}
-
+
for layer_name, metrics in self.layer_metrics.items():
if metric_name in metrics and metrics[metric_name]:
mean_value = np.mean(metrics[metric_name])
layer_means[layer_name] = mean_value
all_values.append(mean_value)
-
+
if len(all_values) < 3:
return []
-
+
# Compute global statistics
global_mean = np.mean(all_values)
global_std = np.std(all_values)
-
+
# Find anomalous layers
anomalous = []
for layer_name, mean_value in layer_means.items():
if abs(mean_value - global_mean) > threshold_std * global_std:
anomalous.append(layer_name)
-
- return anomalous
\ No newline at end of file
+
+ return anomalous
diff --git a/src/alignment/analysis/aggregation/metrics.py b/src/alignment/analysis/aggregation/metrics.py
index ac9b2941..a3d8c571 100644
--- a/src/alignment/analysis/aggregation/metrics.py
+++ b/src/alignment/analysis/aggregation/metrics.py
@@ -2,10 +2,11 @@
Metric aggregation utilities for tracking metrics over time.
"""
-from typing import Dict, List, Tuple, Any
+import logging
from collections import defaultdict
+from typing import Any, Dict, List, Tuple
+
import numpy as np
-import logging
logger = logging.getLogger(__name__)
@@ -13,32 +14,32 @@
class MetricAggregator:
"""
Aggregates metrics across time steps or iterations.
-
+
Useful for analyzing metric evolution during training or experiments.
"""
-
+
def __init__(self):
"""Initialize metric aggregator."""
self.metrics_over_time = defaultdict(lambda: defaultdict(list))
self.steps = []
-
+
def add_step(self, step: int, metrics: Dict[str, Dict[str, float]]):
"""
Add metrics from a single step.
-
+
Args:
step: Step/iteration number
metrics: Dictionary of metrics by name and layer
"""
self.steps.append(step)
-
+
for metric_name, layer_values in metrics.items():
if isinstance(layer_values, dict):
for layer_name, value in layer_values.items():
self.metrics_over_time[metric_name][layer_name].append(value)
else:
self.metrics_over_time[metric_name]['value'].append(layer_values)
-
+
def get_metric_evolution(
self,
metric_name: str,
@@ -46,23 +47,23 @@ def get_metric_evolution(
) -> Tuple[List[int], List[float]]:
"""
Get the evolution of a metric over time.
-
+
Args:
metric_name: Name of the metric
layer_name: Name of the layer
-
+
Returns:
Tuple of (steps, values)
"""
if metric_name not in self.metrics_over_time:
return [], []
-
+
if layer_name not in self.metrics_over_time[metric_name]:
return [], []
-
+
values = self.metrics_over_time[metric_name][layer_name]
return self.steps[:len(values)], values
-
+
def compute_trends(
self,
metric_name: str,
@@ -71,41 +72,41 @@ def compute_trends(
) -> Dict[str, Any]:
"""
Compute trend statistics for a metric.
-
+
Args:
metric_name: Name of the metric
layer_name: Name of the layer
window_size: Window size for moving average
-
+
Returns:
Dictionary with trend statistics
"""
steps, values = self.get_metric_evolution(metric_name, layer_name)
-
+
if len(values) < 2:
return {}
-
+
values_array = np.array(values)
-
+
# Compute moving average
if len(values) >= window_size:
- moving_avg = np.convolve(values_array,
- np.ones(window_size) / window_size,
+ moving_avg = np.convolve(values_array,
+ np.ones(window_size) / window_size,
mode='valid')
else:
moving_avg = values_array
-
+
# Compute linear trend
coeffs = np.polyfit(range(len(values)), values_array, 1)
slope = coeffs[0]
-
+
# Find change points
if len(values) > 2:
diffs = np.diff(values_array)
change_points = np.where(np.abs(diffs) > 2 * np.std(diffs))[0]
else:
change_points = []
-
+
return {
'initial_value': float(values[0]),
'final_value': float(values[-1]),
@@ -115,4 +116,4 @@ def compute_trends(
'percent_change': float((values[-1] - values[0]) / (values[0] + 1e-8) * 100),
'moving_average': moving_avg.tolist() if len(moving_avg) > 0 else [],
'change_points': change_points.tolist()
- }
\ No newline at end of file
+ }
diff --git a/src/alignment/analysis/aggregation/results.py b/src/alignment/analysis/aggregation/results.py
index 1b6c8d25..e282643d 100644
--- a/src/alignment/analysis/aggregation/results.py
+++ b/src/alignment/analysis/aggregation/results.py
@@ -2,12 +2,13 @@
Result aggregation utilities for analyzing experiment outputs.
"""
-from typing import Dict, List, Optional, Any, Union, Tuple
+import json
+import logging
from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
import numpy as np
import pandas as pd
-import json
-import logging
logger = logging.getLogger(__name__)
@@ -15,19 +16,19 @@
class ResultAggregator:
"""
Aggregates results from multiple experiments or runs.
-
+
This class provides utilities for:
- Loading results from multiple sources
- Computing statistics across runs
- Extracting specific metrics
- Comparing experiments
"""
-
+
def __init__(self):
"""Initialize result aggregator."""
self.results = {}
self.metadata = {}
-
+
def add_results(
self,
name: str,
@@ -36,7 +37,7 @@ def add_results(
):
"""
Add results from an experiment.
-
+
Args:
name: Experiment name/identifier
results: Experiment results dictionary
@@ -46,11 +47,11 @@ def add_results(
if metadata:
self.metadata[name] = metadata
logger.info(f"Added results for experiment: {name}")
-
+
def load_from_file(self, path: Union[str, Path], name: Optional[str] = None):
"""
Load results from a JSON file.
-
+
Args:
path: Path to results file
name: Name to use (defaults to filename)
@@ -58,16 +59,16 @@ def load_from_file(self, path: Union[str, Path], name: Optional[str] = None):
path = Path(path)
if not name:
name = path.stem
-
+
with open(path, 'r') as f:
results = json.load(f)
-
+
self.add_results(name, results)
-
+
def load_from_directory(self, directory: Union[str, Path], pattern: str = "*_results.json"):
"""
Load all matching result files from a directory.
-
+
Args:
directory: Directory containing result files
pattern: Glob pattern for result files
@@ -75,9 +76,9 @@ def load_from_directory(self, directory: Union[str, Path], pattern: str = "*_res
directory = Path(directory)
for result_file in directory.glob(pattern):
self.load_from_file(result_file)
-
+
logger.info(f"Loaded {len(self.results)} result files from {directory}")
-
+
def get_metric_values(
self,
metric_name: str,
@@ -86,35 +87,35 @@ def get_metric_values(
) -> Dict[str, Union[float, Dict[str, float]]]:
"""
Extract specific metric values across experiments.
-
+
Args:
metric_name: Name of the metric
layer_name: Specific layer (None for all layers)
experiment_names: Experiments to include (None for all)
-
+
Returns:
Dictionary mapping experiment names to metric values
"""
if experiment_names is None:
experiment_names = list(self.results.keys())
-
+
metric_values = {}
-
+
for exp_name in experiment_names:
if exp_name not in self.results:
continue
-
+
exp_results = self.results[exp_name]
-
+
# Navigate to metrics
if 'metrics' in exp_results:
metrics = exp_results['metrics']
-
+
# Get final metrics (last step)
if metrics:
last_step = max(int(k) for k in metrics.keys())
step_metrics = metrics[str(last_step)]
-
+
if metric_name in step_metrics:
if layer_name:
value = step_metrics[metric_name].get(layer_name)
@@ -122,9 +123,9 @@ def get_metric_values(
metric_values[exp_name] = value
else:
metric_values[exp_name] = step_metrics[metric_name]
-
+
return metric_values
-
+
def compute_statistics(
self,
metric_name: str,
@@ -133,31 +134,31 @@ def compute_statistics(
) -> Dict[str, float]:
"""
Compute statistics for a metric across experiments.
-
+
Args:
metric_name: Name of the metric
layer_name: Layer to analyze
experiment_pattern: Pattern to filter experiments
-
+
Returns:
Dictionary with statistics (mean, std, min, max, etc.)
"""
# Filter experiments
if experiment_pattern:
- exp_names = [name for name in self.results.keys()
+ exp_names = [name for name in self.results.keys()
if experiment_pattern in name]
else:
exp_names = list(self.results.keys())
-
+
# Get metric values
values_dict = self.get_metric_values(metric_name, layer_name, exp_names)
values = list(values_dict.values())
-
+
if not values:
return {}
-
+
values_array = np.array(values)
-
+
return {
'mean': float(np.mean(values_array)),
'std': float(np.std(values_array)),
@@ -168,7 +169,7 @@ def compute_statistics(
'q3': float(np.percentile(values_array, 75)),
'count': len(values)
}
-
+
def to_dataframe(
self,
metrics: Optional[List[str]] = None,
@@ -176,43 +177,43 @@ def to_dataframe(
) -> pd.DataFrame:
"""
Convert results to a pandas DataFrame.
-
+
Args:
metrics: Metrics to include (None for all)
layers: Layers to include (None for all)
-
+
Returns:
DataFrame with experiments as rows and metric/layer combinations as columns
"""
data = []
-
+
for exp_name, results in self.results.items():
row = {'experiment': exp_name}
-
+
# Add metadata
if exp_name in self.metadata:
row.update(self.metadata[exp_name])
-
+
# Add metrics
if 'metrics' in results and results['metrics']:
# Get final metrics
last_step = max(int(k) for k in results['metrics'].keys())
step_metrics = results['metrics'][str(last_step)]
-
+
for metric_name, layer_values in step_metrics.items():
if metrics and metric_name not in metrics:
continue
-
+
if isinstance(layer_values, dict):
for layer_name, value in layer_values.items():
if layers and layer_name not in layers:
continue
-
+
col_name = f"{metric_name}_{layer_name}"
row[col_name] = value
else:
row[metric_name] = layer_values
-
+
data.append(row)
-
- return pd.DataFrame(data)
\ No newline at end of file
+
+ return pd.DataFrame(data)
diff --git a/src/alignment/analysis/dynamic_scoring.py b/src/alignment/analysis/dynamic_scoring.py
index 456664ac..9c79f091 100644
--- a/src/alignment/analysis/dynamic_scoring.py
+++ b/src/alignment/analysis/dynamic_scoring.py
@@ -5,10 +5,10 @@
with loss changes to identify truly important neurons.
"""
-import torch
-import numpy as np
-from typing import Dict, List, Optional, Tuple
import logging
+from typing import Dict, List, Optional
+
+import torch
logger = logging.getLogger(__name__)
@@ -16,13 +16,13 @@
class DynamicScoreAggregator:
"""
Aggregate scores from training history with loss correlation.
-
+
Combines:
- Final scores (current state)
- Score evolution (trend over training)
- Loss correlation (impact on optimization)
- Stability (consistency)
-
+
Example:
>>> aggregator = DynamicScoreAggregator()
>>> dynamic_scores = aggregator.aggregate(
@@ -30,7 +30,7 @@ class DynamicScoreAggregator:
... loss_history=training_losses
... )
"""
-
+
def __init__(
self,
weight_final: float = 0.4,
@@ -40,7 +40,7 @@ def __init__(
):
"""
Initialize dynamic score aggregator.
-
+
Args:
weight_final: Weight for final score value
weight_trend: Weight for trend (increase/decrease)
@@ -51,14 +51,14 @@ def __init__(
self.weight_trend = weight_trend
self.weight_loss_corr = weight_loss_corr
self.weight_stability = weight_stability
-
+
# Normalize weights
total = sum([weight_final, weight_trend, weight_loss_corr, weight_stability])
self.weight_final /= total
self.weight_trend /= total
self.weight_loss_corr /= total
self.weight_stability /= total
-
+
def aggregate(
self,
score_history: Dict[str, Dict[str, List[float]]],
@@ -68,40 +68,40 @@ def aggregate(
) -> torch.Tensor:
"""
Aggregate scores for a layer using training dynamics.
-
+
Args:
score_history: From AlignmentMetricsCallback.get_history()
loss_history: Training loss at each step
layer_name: Layer to process
metric_name: Which metric to aggregate
-
+
Returns:
Dynamic scores per neuron [num_neurons]
"""
if layer_name not in score_history['history']:
raise ValueError(f"No history for layer {layer_name}")
-
+
if metric_name not in score_history['history'][layer_name]:
raise ValueError(f"No {metric_name} history for {layer_name}")
-
+
# Get score evolution
scores_over_time = score_history['history'][layer_name][metric_name]
# This is list of scalar means - need per-neuron history
# For now, work with what we have
-
+
# If we have per-neuron history (not yet implemented in callback):
- # scores_over_time = [step1_scores, step2_scores, ...]
+ # scores_over_time = [step1_scores, step2_scores, ...]
# where each is [num_neurons]
-
+
# For now, provide framework for when per-neuron tracking is added
logger.warning(
"Current callback tracks scalar means. "
"For per-neuron dynamic scoring, need to track full tensors."
)
-
+
# Return placeholder
return torch.tensor(scores_over_time[-1]) # Final value
-
+
def compute_loss_correlation(
self,
score_evolution: torch.Tensor, # [num_steps, num_neurons]
@@ -109,100 +109,100 @@ def compute_loss_correlation(
) -> torch.Tensor:
"""
Compute correlation between each neuron's score and training loss.
-
+
High positive correlation: Neuron's importance grew as loss decreased
→ Neuron is important for learning
-
+
Negative/low correlation: Neuron's importance didn't track loss
→ Neuron might be less critical
-
+
Args:
score_evolution: Score over time per neuron
loss_evolution: Loss over time
-
+
Returns:
Correlation per neuron [num_neurons]
"""
num_steps, num_neurons = score_evolution.shape
-
+
# Convert loss to tensor
loss_tensor = torch.tensor(loss_evolution, dtype=torch.float32)
-
+
# Compute correlation for each neuron
correlations = torch.zeros(num_neurons)
-
+
for neuron_idx in range(num_neurons):
neuron_scores = score_evolution[:, neuron_idx]
-
+
# Pearson correlation
corr = torch.corrcoef(torch.stack([neuron_scores, loss_tensor]))[0, 1]
-
+
correlations[neuron_idx] = corr.abs() # Use absolute value
-
+
return correlations
-
+
def compute_trend(
self,
score_evolution: torch.Tensor # [num_steps, num_neurons]
) -> torch.Tensor:
"""
Compute trend (increasing/decreasing) for each neuron.
-
+
Positive trend: Importance increased → likely important
Negative trend: Importance decreased → less critical
-
+
Returns:
Trend per neuron [num_neurons]
"""
# Simple: final - initial
- trend = score_evolution[-1] - score_evolution[0]
-
+ score_evolution[-1] - score_evolution[0]
+
# More sophisticated: linear regression slope
num_steps, num_neurons = score_evolution.shape
time_steps = torch.arange(num_steps, dtype=torch.float32)
-
+
trends_fitted = torch.zeros(num_neurons)
-
+
for neuron_idx in range(num_neurons):
y = score_evolution[:, neuron_idx]
-
+
# Fit line: y = a + b*t
# b = cov(t, y) / var(t)
mean_t = time_steps.mean()
mean_y = y.mean()
-
+
cov_ty = ((time_steps - mean_t) * (y - mean_y)).sum()
var_t = ((time_steps - mean_t) ** 2).sum()
-
+
slope = cov_ty / (var_t + 1e-8)
trends_fitted[neuron_idx] = slope
-
+
return trends_fitted
-
+
def compute_stability(
self,
score_evolution: torch.Tensor # [num_steps, num_neurons]
) -> torch.Tensor:
"""
Compute stability (inverse variance) for each neuron.
-
+
Low variance: Consistently important → reliable signal
High variance: Fluctuating → less reliable
-
+
Returns:
Stability per neuron [num_neurons]
"""
# Variance over time
variance = score_evolution.var(dim=0) # [num_neurons]
-
+
# Stability = inverse variance
stability = 1.0 / (variance + 1e-4)
-
+
# Normalize to [0, 1]
stability = (stability - stability.min()) / (stability.max() - stability.min() + 1e-8)
-
+
return stability
-
+
def aggregate_full(
self,
score_evolution: torch.Tensor, # [num_steps, num_neurons]
@@ -210,11 +210,11 @@ def aggregate_full(
) -> torch.Tensor:
"""
Full aggregation using all components.
-
+
Args:
score_evolution: Score history per neuron
loss_evolution: Loss history
-
+
Returns:
Aggregated dynamic scores [num_neurons]
"""
@@ -223,16 +223,16 @@ def aggregate_full(
trends = self.compute_trend(score_evolution)
loss_corr = self.compute_loss_correlation(score_evolution, loss_evolution)
stability = self.compute_stability(score_evolution)
-
+
# Normalize each component to [0, 1]
def normalize(x):
return (x - x.min()) / (x.max() - x.min() + 1e-8)
-
+
final_norm = normalize(final_scores)
trend_norm = normalize(trends + trends.abs().max()) # Make positive
loss_corr_norm = normalize(loss_corr)
stability_norm = stability # Already normalized
-
+
# Weighted combination
dynamic_scores = (
self.weight_final * final_norm +
@@ -240,39 +240,39 @@ def normalize(x):
self.weight_loss_corr * loss_corr_norm +
self.weight_stability * stability_norm
)
-
+
return dynamic_scores
class TrainingAwareScoring:
"""
Enhanced scoring using full training history.
-
+
Requires per-neuron tracking during training (not just scalar means).
"""
-
+
@staticmethod
def enhance_callback_for_per_neuron_tracking():
"""
Instructions for enhancing callback to track per-neuron evolution.
-
+
Current callback tracks: scalar mean per step
Enhanced version should track: full tensor per step (memory intensive!)
-
+
Modification needed in AlignmentMetricsCallback:
-
+
```python
# Instead of:
score_value = scores.mean().item()
self.history[layer][metric].append(score_value)
-
+
# Do:
if self.track_per_neuron:
self.history[layer][metric].append(scores.cpu()) # Full tensor
else:
self.history[layer][metric].append(scores.mean().item())
```
-
+
Then dynamic scoring becomes very powerful!
"""
pass
@@ -287,18 +287,18 @@ def compute_dynamic_importance(
) -> torch.Tensor:
"""
Convenience function for dynamic importance computation.
-
+
Args:
score_history: Training history from callback
loss_history: Loss values during training
layer_name: Layer to process
metric_name: Metric to use
aggregation_weights: Optional custom weights
-
+
Returns:
Dynamic importance scores
"""
aggregator = DynamicScoreAggregator(**(aggregation_weights or {}))
-
+
return aggregator.aggregate(score_history, loss_history, layer_name, metric_name)
diff --git a/src/alignment/analysis/unified_reporter.py b/src/alignment/analysis/unified_reporter.py
index f589e224..686e86f9 100644
--- a/src/alignment/analysis/unified_reporter.py
+++ b/src/alignment/analysis/unified_reporter.py
@@ -5,12 +5,13 @@
providing a single interface for generating reports in various formats.
"""
-from typing import Dict, List, Optional, Any, Union
-from pathlib import Path
-import pandas as pd
import json
import logging
from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+import pandas as pd
logger = logging.getLogger(__name__)
@@ -22,14 +23,14 @@ class UnifiedReporter:
- Markdown
- JSON
- LaTeX (optional)
-
+
This consolidates functionality from HTMLReporter, MarkdownReporter, and JSONReporter.
"""
-
+
def __init__(self, title: str = "Alignment Analysis Report"):
"""
Initialize unified reporter.
-
+
Args:
title: Report title
"""
@@ -43,11 +44,11 @@ def __init__(self, title: str = "Alignment Analysis Report"):
"created": datetime.now().isoformat(),
"version": "1.0"
}
-
+
def add_section(self, name: str, content: Any, section_type: str = "text"):
"""
Add a section to the report.
-
+
Args:
name: Section name
content: Section content (text, DataFrame, dict, etc.)
@@ -58,11 +59,11 @@ def add_section(self, name: str, content: Any, section_type: str = "text"):
"content": content,
"type": section_type
})
-
+
def add_figure(self, figure_path: Union[str, Path], caption: str = "", width: Optional[str] = None):
"""
Add a figure to the report.
-
+
Args:
figure_path: Path to the figure
caption: Figure caption
@@ -73,11 +74,11 @@ def add_figure(self, figure_path: Union[str, Path], caption: str = "", width: Op
"caption": caption,
"width": width
})
-
+
def add_table(self, name: str, data: Union[pd.DataFrame, Dict, List]):
"""
Add a table to the report.
-
+
Args:
name: Table name
data: Table data (DataFrame, dict, or list of dicts)
@@ -88,36 +89,36 @@ def add_table(self, name: str, data: Union[pd.DataFrame, Dict, List]):
df = pd.DataFrame(data)
else:
df = data
-
+
self.tables.append({
"name": name,
"data": df
})
-
+
def add_metadata(self, key: str, value: Any):
"""Add metadata to the report."""
self.metadata[key] = value
-
+
# ========== HTML Generation ==========
-
+
def generate_html(self, output_path: Union[str, Path], include_toc: bool = True):
"""
Generate HTML report.
-
+
Args:
output_path: Path to save the HTML file
include_toc: Whether to include table of contents
"""
html = self._generate_html_header(include_toc)
-
+
# Add table of contents
if include_toc:
html += self._generate_html_toc()
-
+
# Add sections
for section in self.sections:
html += self._generate_html_section(section)
-
+
# Add tables
if self.tables:
html += 'Tables
'
@@ -125,26 +126,26 @@ def generate_html(self, output_path: Union[str, Path], include_toc: bool = True)
html += f'{table["name"]}
'
html += table["data"].to_html(classes='data-table', index=False)
html += ''
-
+
# Add figures
if self.figures:
html += 'Figures
'
for fig in self.figures:
html += self._generate_html_figure(fig)
html += ''
-
+
html += self._generate_html_footer()
-
+
# Save
output_path = Path(output_path)
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text(html)
-
+
logger.info(f"Generated HTML report: {output_path}")
-
+
def _generate_html_header(self, include_toc: bool) -> str:
"""Generate HTML header with styles."""
- return f"""
+ html = f"""
{self.title}
@@ -244,42 +245,41 @@ def _generate_html_header(self, include_toc: bool) -> str:
{self.title}
Generated: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
- Version: {self.metadata.get('version', '1.0')}
-"""
-
+ Version: {self.metadata.get('version', '1.0')}"""
+
# Add custom metadata
for key, value in self.metadata.items():
if key not in ['title', 'created', 'version']:
html += f"
{key.title()}: {value}"
-
+
html += "
"
return html
-
+
def _generate_html_toc(self) -> str:
"""Generate table of contents."""
toc = 'Table of Contents
'
-
+
for i, section in enumerate(self.sections):
section_id = f"section-{i}"
toc += f'- {section["name"]}
'
-
+
if self.tables:
toc += '- Tables
'
-
+
if self.figures:
toc += '- Figures
'
-
+
toc += '
{section["name"]}
'
-
+
content = section["content"]
section_type = section["type"]
-
+
if section_type == "text":
html += f"
{content}
"
elif section_type == "code":
@@ -294,42 +294,42 @@ def _generate_html_section(self, section: Dict) -> str:
elif section_type == "table":
if isinstance(content, pd.DataFrame):
html += content.to_html(classes='data-table', index=False)
-
+
html += '