"""vi Analyze STDP weight changes between checkpoints to evaluate learning progress. Usage: python -m STDP_Communicator.analyze_stdp_weights --early 5 --final 20 --visualize This script loads checkpoint files from different epochs, compares synaptic weights, and provides statistics and optional visualizations of the weight changes. """ import os import sys import torch import argparse import logging import numpy as np import matplotlib.pyplot as plt from pathlib import Path from typing import Dict, List, Tuple, Optional, Union # Configure logging logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) def load_checkpoint(checkpoint_path: str) -> Dict: """ Load a checkpoint file and return its contents. Args: checkpoint_path: Path to the checkpoint file Returns: Dictionary containing checkpoint data """ if not os.path.exists(checkpoint_path): raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") try: logger.info(f"Loading checkpoint: {checkpoint_path}") checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu')) return checkpoint except Exception as e: logger.error(f"Error loading checkpoint: {e}") raise def extract_synaptic_weights(checkpoint: Dict) -> torch.Tensor: """ Extract synaptic weights from a checkpoint. Args: checkpoint: Loaded checkpoint data Returns: Tensor containing synaptic weights """ if "synaptic_weights" in checkpoint: logger.info("Found direct synaptic_weights key") return checkpoint["synaptic_weights"] # Try to find weights in model state dict if "model_state_dict" in checkpoint: state_dict = checkpoint["model_state_dict"] weight_keys = [k for k in state_dict.keys() if "weight" in k.lower()] if weight_keys: logger.info(f"Using weights from key: {weight_keys[0]}") return state_dict[weight_keys[0]] # Look for weights in different formats for key in checkpoint.keys(): if isinstance(checkpoint[key], dict) and "weights" in checkpoint[key]: logger.info(f"Found weights in nested dictionary: {key}.weights") return checkpoint[key]["weights"] raise ValueError("Could not find synaptic weights in checkpoint") def analyze_weight_changes(early_weights: torch.Tensor, final_weights: torch.Tensor) -> Dict: """ Analyze changes between two sets of weights. Args: early_weights: Weights from earlier epoch final_weights: Weights from later epoch Returns: Dictionary of statistics about weight changes """ # Ensure tensors are on the same device if early_weights.device != final_weights.device: final_weights = final_weights.to(early_weights.device) # Calculate absolute differences diff = torch.abs(final_weights - early_weights) # Calculate statistics stats = { "mean_change": diff.mean().item(), "max_change": diff.max().item(), "std_change": diff.std().item(), "percent_changed": (diff > 0.0001).float().mean().item() * 100, # Percentage of weights changed by >0.0001 "early_weights_mean": early_weights.mean().item(), "early_weights_std": early_weights.std().item(), "final_weights_mean": final_weights.mean().item(), "final_weights_std": final_weights.std().item() } return stats def visualize_weights( early_weights: torch.Tensor, final_weights: torch.Tensor, output_dir: str = "weight_analysis" ) -> None: """ Create visualizations of weight changes and save to files. Args: early_weights: Weights from earlier epoch final_weights: Weights from later epoch output_dir: Directory to save visualizations """ os.makedirs(output_dir, exist_ok=True) # Convert to numpy for easier plotting early_np = early_weights.flatten().detach().numpy() final_np = final_weights.flatten().detach().numpy() diff_np = np.abs(final_np - early_np) # Create plots plt.figure(figsize=(15, 10)) # 1. Weight distribution histograms plt.subplot(2, 2, 1) plt.hist(early_np, bins=50, alpha=0.5, label="Early Weights") plt.hist(final_np, bins=50, alpha=0.5, label="Final Weights") plt.title("Weight Distribution Comparison") plt.xlabel("Weight Value") plt.ylabel("Count") plt.legend() # 2. Changes histogram plt.subplot(2, 2, 2) plt.hist(diff_np, bins=50, color='green') plt.title("Weight Changes") plt.xlabel("Absolute Change") plt.ylabel("Count") # 3. Scatter plot of early vs final weights plt.subplot(2, 2, 3) # If there are too many weights, sample a subset max_points = 10000 if len(early_np) > max_points: indices = np.random.choice(len(early_np), max_points, replace=False) early_sample = early_np[indices] final_sample = final_np[indices] else: early_sample = early_np final_sample = final_np plt.scatter(early_sample, final_sample, alpha=0.1) plt.plot([early_np.min(), early_np.max()], [early_np.min(), early_np.max()], 'r--') # y=x line plt.title("Early vs Final Weights") plt.xlabel("Early Weights") plt.ylabel("Final Weights") # 4. Weight change heatmap (if weights are 2D) plt.subplot(2, 2, 4) if len(early_weights.shape) == 2: diff_2d = torch.abs(final_weights - early_weights).detach().numpy() plt.imshow(diff_2d, cmap='hot', interpolation='nearest') plt.colorbar() plt.title("Weight Change Heatmap") else: # If weights aren't 2D, show a different visualization sorted_early = np.sort(early_np) sorted_final = np.sort(final_np) plt.plot(sorted_early, label="Early Weights (Sorted)") plt.plot(sorted_final, label="Final Weights (Sorted)") plt.title("Sorted Weight Comparison") plt.legend() # Save and show the figure plt.tight_layout() plt.savefig(os.path.join(output_dir, "weight_analysis.png"), dpi=300) plt.close() logger.info(f"Visualizations saved to {output_dir}/weight_analysis.png") def parse_args(): parser = argparse.ArgumentParser(description="Analyze STDP weights between checkpoints") parser.add_argument("--early", type=int, default=5, help="Early epoch number") parser.add_argument("--final", type=int, default=20, help="Final epoch number") parser.add_argument("--checkpoint-dir", type=str, default="checkpoints", help="Directory containing checkpoint files") parser.add_argument("--output-dir", type=str, default="weight_analysis", help="Directory to save analysis outputs") parser.add_argument("--visualize", action="store_true", help="Create visualizations") return parser.parse_args() def main(): args = parse_args() try: # Resolve paths checkpoint_dir = Path(args.checkpoint_dir) if not checkpoint_dir.is_absolute(): # Get project root - two directories up from this script project_root = Path(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) checkpoint_dir = project_root / checkpoint_dir # Construct checkpoint paths early_path = checkpoint_dir / f"stdp_model_epoch_{args.early}.pt" final_path = checkpoint_dir / f"stdp_model_epoch_{args.final}.pt" # Load checkpoints early_checkpoint = load_checkpoint(str(early_path)) final_checkpoint = load_checkpoint(str(final_path)) # Extract weights early_weights = extract_synaptic_weights(early_checkpoint) final_weights = extract_synaptic_weights(final_checkpoint) # Analyze changes stats = analyze_weight_changes(early_weights, final_weights) # Print results logger.info("Weight Change Analysis:") logger.info(f"Epochs {args.early} → {args.final}") logger.info(f"Average weight change: {stats['mean_change']:.6f}") logger.info(f"Maximum weight change: {stats['max_change']:.6f}") logger.info(f"Standard deviation of changes: {stats['std_change']:.6f}") logger.info(f"Percentage of weights changed: {stats['percent_changed']:.2f}%") logger.info(f"Early weights - mean: {stats['early_weights_mean']:.6f}, std: {stats['early_weights_std']:.6f}") logger.info(f"Final weights - mean: {stats['final_weights_mean']:.6f}, std: {stats['final_weights_std']:.6f}") # Create visualizations if requested if args.visualize: output_dir = Path(args.output_dir) if not output_dir.is_absolute(): output_dir = project_root / output_dir visualize_weights(early_weights, final_weights, str(output_dir)) except Exception as e: logger.error(f"Error during analysis: {e}", exc_info=True) return 1 return 0 if __name__ == "__main__": sys.exit(main())