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