Update README.md

README.md

@@ -2,4 +2,305 @@
 license: mit
 tags:
 - braindecode
----
+---
+
+Shallow conversion from the original weight for braindecode.
+
+```python
+
+#!/usr/bin/env python3
+"""
+Complete LaBraM Weight Transfer Script
+
+Combines explicit weight mapping with full backbone transfer.
+Uses precise key renaming to transfer all compatible parameters.
+
+Transfers weights from LaBraM checkpoint to Braindecode Labram model.
+"""
+
+import torch
+import argparse
+from braindecode.models import Labram
+
+
+def create_weight_mapping():
+    """
+    Create comprehensive weight mapping from LaBraM to Braindecode.
+    
+    Includes:
+    - Temporal convolution layers (patch_embed)
+    - All transformer blocks
+    - Position embeddings
+    - Other backbone components
+    """
+    return {
+        # Temporal Convolution Layers
+        'student.patch_embed.conv1.weight': 'patch_embed.temporal_conv.conv1.weight',
+        'student.patch_embed.conv1.bias': 'patch_embed.temporal_conv.conv1.bias',
+        'student.patch_embed.norm1.weight': 'patch_embed.temporal_conv.norm1.weight',
+        'student.patch_embed.norm1.bias': 'patch_embed.temporal_conv.norm1.bias',
+        'student.patch_embed.conv2.weight': 'patch_embed.temporal_conv.conv2.weight',
+        'student.patch_embed.conv2.bias': 'patch_embed.temporal_conv.conv2.bias',
+        'student.patch_embed.norm2.weight': 'patch_embed.temporal_conv.norm2.weight',
+        'student.patch_embed.norm2.bias': 'patch_embed.temporal_conv.norm2.bias',
+        'student.patch_embed.conv3.weight': 'patch_embed.temporal_conv.conv3.weight',
+        'student.patch_embed.conv3.bias': 'patch_embed.temporal_conv.conv3.bias',
+        'student.patch_embed.norm3.weight': 'patch_embed.temporal_conv.norm3.weight',
+        'student.patch_embed.norm3.bias': 'patch_embed.temporal_conv.norm3.bias',
+        # Note: Other backbone layers (blocks, embeddings, norm, fc_norm) are handled
+        # by removing 'student.' prefix in process_state_dict()
+    }
+
+
+def process_state_dict(state_dict, weight_mapping):
+    """
+    Process checkpoint state dict with explicit mapping.
+    
+    Parameters:
+    -----------
+    state_dict : dict
+        Original checkpoint state dictionary
+    weight_mapping : dict
+        Explicit mapping for special layers (patch_embed)
+        
+    Returns:
+    --------
+    dict : Processed state dict ready for Braindecode model
+    """
+    new_state = {}
+    mapped_keys = []
+    skipped_keys = []
+    
+    for key, value in state_dict.items():
+        # Skip classification head (task-specific)
+        if 'head' in key:
+            skipped_keys.append((key, 'head layer'))
+            continue
+        
+        # Use explicit mapping for patch_embed temporal_conv
+        if key in weight_mapping:
+            new_key = weight_mapping[key]
+            new_state[new_key] = value
+            mapped_keys.append((key, new_key))
+            continue
+        
+        # Skip original patch_embed if not in mapping (SegmentPatch)
+        if 'patch_embed' in key and 'temporal_conv' not in key:
+            skipped_keys.append((key, 'patch_embed (non-temporal)'))
+            continue
+        
+        # For backbone layers, remove 'student.' prefix
+        if key.startswith('student.'):
+            new_key = key.replace('student.', '')
+            new_state[new_key] = value
+            mapped_keys.append((key, new_key))
+            continue
+        
+        # Keep other keys as-is
+        new_state[key] = value
+        mapped_keys.append((key, key))
+    
+    return new_state, mapped_keys, skipped_keys
+
+
+def transfer_labram_weights(
+    checkpoint_path,
+    n_times=1600,
+    n_chans=64,
+    n_outputs=4,
+    output_path=None,
+    verbose=True
+):
+    """
+    Transfer LaBraM weights to Braindecode Labram using explicit mapping.
+    
+    Parameters:
+    -----------
+    checkpoint_path : str
+        Path to LaBraM checkpoint
+    n_times : int
+        Number of time samples
+    n_chans : int
+        Number of channels
+    n_outputs : int
+        Number of output classes
+    output_path : str
+        Where to save the model
+    verbose : bool
+        Print transfer details
+        
+    Returns:
+    --------
+    model : Labram
+        Model with transferred weights
+    stats : dict
+        Transfer statistics
+    """
+    
+    print("\n" + "="*70)
+    print("LaBraM → Braindecode Weight Transfer")
+    print("="*70)
+    
+    # Load checkpoint
+    print(f"\nLoading checkpoint: {checkpoint_path}")
+    checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
+    
+    # Extract model state
+    if isinstance(checkpoint, dict) and 'model' in checkpoint:
+        state = checkpoint['model']
+    else:
+        state = checkpoint
+    
+    original_params = len(state)
+    print(f"Original checkpoint: {original_params} parameters")
+    
+    # Create weight mapping
+    weight_mapping = create_weight_mapping()
+    
+    # Process state dict
+    print("\nProcessing checkpoint...")
+    new_state, mapped_keys, skipped_keys = process_state_dict(state, weight_mapping)
+    
+    transferred_params = len(mapped_keys)
+    print(f"Mapped keys: {transferred_params} ({transferred_params/original_params*100:.1f}%)")
+    print(f"Skipped keys: {len(skipped_keys)}")
+    
+    if verbose and skipped_keys:
+        print(f"\nSkipped layers:")
+        for key, reason in skipped_keys[:5]:  # Show first 5
+            print(f"  - {key:50s} ({reason})")
+        if len(skipped_keys) > 5:
+            print(f"  ... and {len(skipped_keys) - 5} more")
+    
+    # Create model
+    print(f"\nCreating Labram model:")
+    print(f"  n_times: {n_times}")
+    print(f"  n_chans: {n_chans}")
+    print(f"  n_outputs: {n_outputs}")
+    model = Labram(
+        n_times=n_times,
+        n_chans=n_chans,
+        n_outputs=n_outputs,
+        neural_tokenizer=True,
+    )
+    
+    # Load weights
+    print("\nLoading weights into model...")
+    incompatible = model.load_state_dict(new_state, strict=False)
+    
+    missing_count = len(incompatible.missing_keys) if incompatible.missing_keys else 0
+    unexpected_count = len(incompatible.unexpected_keys) if incompatible.unexpected_keys else 0
+    
+    if missing_count > 0:
+        print(f"  Missing keys: {missing_count} (expected - will be initialized)")
+    if unexpected_count > 0:
+        print(f"  Unexpected keys: {unexpected_count}")
+    
+    # Test forward pass
+    if verbose:
+        print("\nTesting forward pass...")
+        x = torch.randn(2, n_chans, n_times)
+        with torch.no_grad():
+            output = model(x)
+        print(f"  Input shape:  {x.shape}")
+        print(f"  Output shape: {output.shape}")
+        print("  ✅ Forward pass successful!")
+    
+    # Save model if output_path provided
+    if output_path:
+        print(f"\nSaving model to: {output_path}")
+        torch.save(model.state_dict(), output_path)
+        print(f"  ✅ Model saved")
+    
+    stats = {
+        'original': original_params,
+        'transferred': transferred_params,
+        'skipped': len(skipped_keys),
+        'transfer_rate': f"{transferred_params/original_params*100:.1f}%"
+    }
+    
+    return model, stats
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description='Transfer LaBraM weights to Braindecode Labram',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  # Default transfer (backbone parameters)
+  python labram_complete_transfer.py
+  
+  # Transfer and save model
+  python labram_complete_transfer.py --output labram_weights.pt
+  
+  # Custom EEG parameters
+  python labram_complete_transfer.py --n-times 2000 --n-chans 62 --n-outputs 2
+  
+  # Custom checkpoint path
+  python labram_complete_transfer.py --checkpoint path/to/checkpoint.pth
+        """
+    )
+    
+    parser.add_argument(
+        '--checkpoint',
+        type=str,
+        default='LaBraM/checkpoints/labram-base.pth',
+        help='Path to LaBraM checkpoint (default: LaBraM/checkpoints/labram-base.pth)'
+    )
+    parser.add_argument(
+        '--n-times',
+        type=int,
+        default=1600,
+        help='Number of time samples (default: 1600)'
+    )
+    parser.add_argument(
+        '--n-chans',
+        type=int,
+        default=64,
+        help='Number of channels (default: 64)'
+    )
+    parser.add_argument(
+        '--n-outputs',
+        type=int,
+        default=4,
+        help='Number of output classes (default: 4)'
+    )
+    parser.add_argument(
+        '--output',
+        type=str,
+        default=None,
+        help='Output file path to save model weights'
+    )
+    parser.add_argument(
+        '--device',
+        type=str,
+        default='cpu',
+        help='Device to use (default: cpu)'
+    )
+    
+    args = parser.parse_args()
+    
+    print("="*70)
+    print("LaBraM → Braindecode Weight Transfer")
+    print("="*70)
+    
+    # Transfer weights
+    model, stats = transfer_labram_weights(
+        checkpoint_path=args.checkpoint,
+        n_times=args.n_times,
+        n_chans=args.n_chans,
+        n_outputs=args.n_outputs,
+        output_path=args.output,
+        verbose=True
+    )
+    
+    print("\n" + "="*70)
+    print("✅ TRANSFER COMPLETE")
+    print("="*70)
+    print(f"Original parameters:   {stats['original']}")
+    print(f"Transferred:           {stats['transferred']} ({stats['transfer_rate']})")
+    print(f"Skipped:               {stats['skipped']}")
+    print("="*70)
+
+```