| | --- |
| | 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) |
| | |
| | ``` |
