Upload 12 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+pipeline.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,87 @@
+# SLIM-BRAIN: A DATA- AND TRAINING-EFFICIENT FOUNDATION MODEL FOR FMRI DATA ANALYSIS
+
+<div align="center">
+  
+[![arXiv](https://img.shields.io/badge/arXiv-2512.21881-b31b1b.svg?style=flat-square)](https://www.arxiv.org/abs/2512.21881)
+[![GitHub](https://img.shields.io/badge/GitHub-Repository-181717?style=flat-square&logo=github)](https://github.com/OneMore1/SLIM-Brain2026)
+[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Models-blue)](https://huggingface.co/OneMore1/Slim-Brain)
+
+</div>
+
+This repository contains the official implementation of SLIM-Brain. SLIM-Brain is a two-stage, selective-compute pipeline for voxel-level fMRI representation learning. A lightweight global branch ranks informative temporal windows; a high-capacity 4D Hiera–JEPA encoder processes only those windows, focusing compute on brain voxels and drastically reducing memory.
+
+
+<p align="center">
+  <img src="pipeline.png" width="800" alt="framework">
+</p>
+
+---
+
+## Installation
+
+Setting up the environment requires Python 3.13 and CUDA-compatible PyTorch for GPU acceleration:
+
+```bash
+conda create -n hiera-jepa python=3.13.5
+conda activate hiera-jepa
+
+# Install dependencies
+pip install -r requirements.txt
+```
+
+## Project Structure
+
+The codebase is organized into modular components for easy navigation and extension:
+
+```
+hiera-jepa/
+├── configs/                    # YAML configuration files for training and model parameters
+├── checkpoints/                # Saved model weights and training checkpoints
+├── hiera/                      # Hierarchical Vision Transformer backbone implementation
+├── scripts/                   # Bash....
+├── finetune.py               # Downstream task training and feature extraction script
+└── requirements.txt            # Python package dependencies
+```
+
+## Downstream evaluation
+
+1. Ensure your pre-train data structure as follow:
+
+```
+data_root/
+├── ABIDE_train/                
+├── ABIDE_val/                  
+├── HCP_val/              
+└── HCP_train/              
+    ├── 0010001/                # Subject ID
+    └── 0010002/                
+        ├── 0010002_run-1_0000-0199_1.npz  # Data chunk 1 
+        ├── 0010002_run-1_0000-0199_2.npz  # Data chunk 2
+```
+
+2. Loading downstream datasets as following data structure:
+
+```yaml
+task:
+  csv: "/path/to/data_csv"
+
+data:
+  data_root: /path/to/data_root
+  datasets: ["HCP"]
+  mode: "directory"
+```
+
+3. Start downstream training:
+
+```bash
+# running downstream training
+sh scripts/finetune.sh
+```
+
+#### Model Checkpoints
+
+Our pre-trained model weights can be found in the checkpoints directory:  `./checkpoints/best_model.pth`
+
+
+
+

checkpoints/__init__.py

@@ -0,0 +1 @@
+

configs/finetune_config.yaml

@@ -0,0 +1,106 @@
+experiment:
+  name: "finetune_classification"
+  output_dir: "./output/hiera_finetune"
+  seed: 44
+  resume: null  # Path to checkpoint to resume from 
+  pretrained_checkpoint: "/vePFS-0x0d/home/yewh/Hiera_MAE/checkpoint/checkpoint_epoch_39.pth"
+
+# --- Task Settings ---
+task:
+  task_type: "regression"  # "classification" or "regression"
+  num_classes: 1  # Number of classes for classification (e.g., 2 for binary classification)
+  mean: 33.9289
+  std: 21.5580
+  csv: "/vePFS-0x0d/home/yewh/data_csv/RBC-NKI.csv"  # CSV with columns: Subject, DX_GROUP (for classification) or age (for regression)
+
+# --- Data Settings ---
+data:
+  data_root: "/vePFS-0x0d/fmri-data/WAR_NPYZ"
+  datasets: ["HCP"]
+  train_split_suffixes: ["train_40"]
+  val_split_suffixes: ["val_40"]
+  test_split_suffixes: ["test_40"]
+  input_seq_len: 40  # Temporal length to crop (T dimension)
+
+  # Data dimensions (D, H, W, T) -> will be permuted to (T, D, H, W)
+  spatial_dims: [96, 96, 96]  # D, H, W
+
+  # DataLoader settings
+  batch_size: 2  # Per GPU batch size (can be larger than pretraining)
+  num_workers: 4
+  pin_memory: true
+  prefetch_factor: 2
+
+# --- Model Settings ---
+model:
+  # Input configuration
+  input_size: [40, 96, 96, 96]  # [T, D, H, W]
+  in_chans: 1
+
+  # Patch embedding configuration
+  patch_kernel: [1, 4, 4, 4]  # [T, D, H, W]
+  patch_stride: [1, 4, 4, 4]  # [T, D, H, W]
+  patch_padding: [0, 0, 0, 0]  # [T, D, H, W]
+
+  # Hiera architecture
+  embed_dim: 64  
+  num_heads: 1   
+  stages: [2, 3, 16, 3]  
+  q_pool: 2      
+  q_stride: [2, 2, 2, 2]  # Stride for q_pool [T, D, H, W]
+  mask_unit_size: [8, 8, 8, 8]  # Mask unit size [T, D, H, W]
+  mlp_ratio: 4.0  
+  mask_unit_attn: [true, true, False, False]
+
+# --- Training Settings ---
+training:
+  # Optimization
+  optimizer: "adamw"
+  learning_rate: 1.0e-4  # Lower learning rate for fine-tuning
+  head_lr: 5.0e-4  # Higher learning rate for classification head
+  layer_decay: 0.75
+  weight_decay: 0.05
+  betas: [0.9, 0.99]
+
+  # Learning rate schedule
+  lr_scheduler: "cosine"
+  warmup_epochs: 2  # Warmup epochs
+
+  # Weight freezing
+  freeze_encoder: true  # Set to true to freeze the entire encoder and only train the head
+  min_lr: 1.0e-6  # Minimum learning rate at the end of schedule
+
+  # Training duration
+  epochs: 200
+
+  # Gradient settings
+  clip_grad: 1.0  # Gradient clipping value, null to disable
+  accum_iter: 8  # Gradient accumulation steps (usually 1 for fine-tuning)
+
+  # Mixed precision
+  use_amp: true  # Use automatic mixed precision
+
+# --- Distributed Training Settings ---
+distributed:
+  backend: "nccl"
+  init_method: "env://"
+  world_size: -1  # Will be set automatically
+  rank: -1  # Will be set automatically
+  dist_url: "env://"
+
+# --- Logging Settings ---
+logging:
+  print_freq: 40  # Print frequency (iterations)
+  log_freq: 40  # Log frequency (iterations)
+  save_freq: 5  # Checkpoint save frequency (epochs)
+
+  # Weights & Biases
+  use_wandb: false
+  wandb_project: "hiera_fmri_finetune"
+  wandb_entity: null  # Your wandb username/team
+
+# --- Validation Settings ---
+validation:
+  val_freq: 1  # Validation frequency (epochs)
+  save_best: true  # Save best model based on validation metric
+

data/downstream_dataset.py

@@ -0,0 +1,150 @@
+import os
+import glob
+import re 
+import numpy as np
+import pandas as pd 
+import torch
+from torch.utils.data import Dataset
+from typing import List, Tuple, Union, Literal
+import torch.nn.functional as F
+from .pretrain_dataset import fMRIDataset
+import io  
+import nibabel as nib
+
+class fMRITaskDataset(fMRIDataset):
+
+    def __init__(
+        self,
+        data_root: str,
+        datasets: List[str],
+        split_suffixes: List[str],
+        crop_length: int,
+        label_csv_path: str,
+        task_type: Literal['classification', 'regression'] = 'classification',
+        downstream=True,
+    ):
+        super().__init__(data_root, datasets, split_suffixes, crop_length, downstream)
+        
+        self.task_type = task_type
+        self.labels_map = self._load_and_process_labels(label_csv_path)
+
+        initial_file_count = len(self.file_paths)
+        self.file_paths = [
+            path for path in self.file_paths 
+            if self._extract_subject_id(path) in self.labels_map
+        ]
+        
+        if len(self.file_paths) < initial_file_count:
+            print(f"Warning: Dropped {initial_file_count - len(self.file_paths)} files due to missing labels in CSV.")
+        
+        print(f"Task Dataset ready for {self.task_type}. Usable files: {len(self.file_paths)}")
+
+
+    def _extract_subject_id(self, file_path: str) -> str:
+
+
+            # folder_name = os.path.basename(os.path.dirname(file_path))
+            # match = re.search(r'(\d{7})', folder_name)
+
+            match = re.search(r'(\d{6})', os.path.basename(file_path))
+            
+            if match:
+                subject_id_with_zeros = match.group(1)
+                subject_id = subject_id_with_zeros.lstrip('0') 
+                
+                return subject_id
+                
+            return "" 
+
+    def _load_and_process_labels(self, csv_path: str) -> dict:
+
+        if not os.path.exists(csv_path):
+            raise FileNotFoundError(f"Label CSV file not found at: {csv_path}")
+            
+        print(f"Loading labels from {csv_path}...")
+        df = pd.read_csv(csv_path)
+        
+        df['Subject'] = df['Subject'].astype(str)
+        df.dropna(subset=['Subject'], inplace=True) 
+
+        labels_map = {}
+        
+        if self.task_type == 'classification':
+            label_col = None
+            if 'Gender' in df.columns:
+                label_col = 'Gender'
+            elif 'gender' in df.columns:
+                label_col = 'gender'
+            elif 'age_group' in df.columns: 
+                label_col = 'age_group'
+            
+            if label_col is None:
+                raise ValueError("CSV must contain 'sex', 'gender' or 'age_group' column for classification.")
+
+            print(f"Using column '{label_col}' as label.")
+              
+            # unique_vals = df[label_col].unique() 
+
+            sex_mapping = {'F': 0, 'M': 1, 'f': 0, 'm': 1}
+            
+            if df[label_col].dtype == object and df[label_col].astype(str).iloc[0].upper() in ['F', 'M']:
+                print(f"Encoding {label_col} (F/M) to Integers (0/1)...")
+                df = df[df[label_col].isin(sex_mapping.keys())]
+                df[label_col] = df[label_col].map(sex_mapping)
+            else:
+                df[label_col] = pd.to_numeric(df[label_col], errors='coerce').astype(int)
+            
+            for _, row in df.iterrows():
+                subject_id = row['Subject']
+                labels_map[subject_id] = torch.tensor(row[label_col], dtype=torch.long)
+
+        elif self.task_type == 'regression':
+            label_col = 'age'
+            if label_col not in df.columns:
+                 raise ValueError(f"Regression task requires '{label_col}' column.")
+            df[label_col] = pd.to_numeric(df[label_col], errors='coerce')
+            df.dropna(subset=[label_col], inplace=True)
+            
+            for _, row in df.iterrows():
+                subject_id = row['Subject']
+                labels_map[subject_id] = torch.tensor(row[label_col], dtype=torch.float32).view(1)
+
+        else:
+            raise ValueError(f"Unsupported task_type: {self.task_type}")
+
+        print(f"Successfully loaded {len(labels_map)} subjects' labels.")
+        return labels_map
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        retries = 0
+        max_retries = 100 
+        while retries < max_retries:
+            try:
+                data_tensor = super().__getitem__(idx)
+
+                if data_tensor is None:
+                    raise ValueError(f"Failed to load data at index {idx} (super returned None)")
+
+                file_path = self.file_paths[idx]
+                
+                subject_id = self._extract_subject_id(file_path)
+
+                data_tensor = data_tensor.unsqueeze(0)
+                
+                if subject_id in self.labels_map:
+                    label_tensor = self.labels_map[subject_id]
+
+                    return data_tensor, label_tensor
+                else:
+                    raise KeyError(f"Label not found for subject ID: {subject_id}")
+
+            except Exception as e:
+                # print(f"Warning: Error loading index {idx}: {e}. Retrying...")
+                
+                idx = np.random.randint(0, len(self))
+                retries += 1
+        
+        raise RuntimeError(f"Failed to load any valid data after {max_retries} retries.")
+            
+        return data_tensor, label_tensor

data/pretrain_dataset.py

@@ -0,0 +1,61 @@
+import os
+import glob
+import numpy as np
+from typing import Any, Callable, Dict, Optional, Set, Tuple
+import torch
+from torch.utils.data import Dataset
+import random
+
+class fMRIDataset(Dataset):
+    def __init__(self, 
+                 data_root, datasets, split_suffixes, crop_length=40, downstream=False):
+
+        self.file_paths = []
+        self.crop_length = crop_length
+        self.downstream = downstream
+        for dataset_name in datasets:
+            for suffix in split_suffixes:
+                folder_name = f"{dataset_name}_{suffix}"
+                folder_path = os.path.join(data_root, folder_name)
+                if not os.path.exists(folder_path):
+                    print(f"Warning: Folder not found: {folder_path}")
+                    continue
+
+                for root, dirs, files in os.walk(folder_path):
+                    npz_files = glob.glob(os.path.join(root, "*.npz"))
+                    if len(npz_files) > 1:
+                        # sample_size = max(1, int(len(npz_files) * 0.5)) 
+                        # npz_files = random.sample(npz_files, sample_size)
+                        npz_files = sorted(npz_files)[:1]
+                    self.file_paths.extend(npz_files)
+
+        print(f"Dataset loaded. Total files found: {len(self.file_paths)}")
+
+    def __len__(self):
+        return len(self.file_paths)
+
+    def __getitem__(self, idx):
+
+        file_path = self.file_paths[idx]
+        try:
+            with np.load(file_path) as data_file:
+                key = list(data_file.keys())[0]
+                fmri_data = data_file[key] 
+                fmri_data = fmri_data.astype(np.float32)
+        except Exception as e:
+            print(f"Error loading file {file_path}: {e}")
+            return None
+
+        total_time_frames = fmri_data.shape[-1]
+        if total_time_frames > self.crop_length:
+            start_idx = np.random.randint(0, total_time_frames - self.crop_length + 1)
+            end_idx = start_idx + self.crop_length
+            cropped_data = fmri_data[..., start_idx:end_idx]
+        else:
+            cropped_data = fmri_data[..., :self.crop_length]
+
+        data_tensor = torch.from_numpy(cropped_data)
+
+        data_tensor = data_tensor.permute(3, 0, 1, 2)
+
+        return data_tensor

finetune.py

@@ -0,0 +1,586 @@
+import os
+import sys
+import argparse
+import yaml
+import datetime
+import numpy as np
+from pathlib import Path
+from sklearn.metrics import f1_score
+
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data import DataLoader, DistributedSampler
+from torch.cuda.amp import GradScaler, autocast
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'hiera'))
+
+from hiera.hiera_mae import HieraClassifier
+from data.downstream_dataset import fMRITaskDataset, fMRITaskDataset1, EmoFMRIDataset, HCPtaskDataset
+from data.adni_dataset import ADNIDataset
+
+from utils.utils import MetricLogger, load_config, log_to_file, count_parameters, save_checkpoint, load_checkpoint, LabelScaler
+from utils.optim import create_optimizer, create_lr_scheduler
+from utils.ddp import setup_distributed, set_seed, cleanup_distributed
+
+
+def create_model(config):
+    """Create Hiera Classifier model from config"""
+    task_config = config['task']
+    exp_config = config['experiment']
+
+    model_config = config['model'] 
+    pretrained_checkpoint_path = exp_config.get('pretrained_checkpoint', None)
+
+    if pretrained_checkpoint_path:
+        pretrain_config_path = Path(pretrained_checkpoint_path).parent.parent / 'config.yaml'
+        if os.path.exists(pretrain_config_path):
+            print(f"Loading model architecture from pretrained config: {pretrain_config_path}")
+            pretrain_config = load_config(pretrain_config_path)
+            model_config = pretrain_config['model']
+        else:
+            print(f"Warning: Pretrained config not found at {pretrain_config_path}. Using finetune config for model architecture.")
+
+    model = HieraClassifier(
+        num_classes=task_config['num_classes'],
+        task_type=task_config['task_type'],
+        input_size=tuple(model_config['input_size']),
+        in_chans=model_config['in_chans'],
+        patch_kernel=tuple(model_config['patch_kernel']),
+        patch_stride=tuple(model_config['patch_stride']),
+        patch_padding=tuple(model_config['patch_padding']),
+        q_stride=tuple(model_config['q_stride']),
+        mask_unit_size=tuple(model_config['mask_unit_size']),
+        embed_dim=model_config['embed_dim'],
+        num_heads=model_config['num_heads'],
+        stages=tuple(model_config['stages']),
+        q_pool=model_config['q_pool'],
+        mlp_ratio=model_config['mlp_ratio'],
+    )
+
+    # Load pretrained weights if specified
+    if pretrained_checkpoint_path:
+        if os.path.exists(pretrained_checkpoint_path):
+            model.load_pretrained_mae(pretrained_checkpoint_path)
+        else:
+            print(f"Warning: Pretrained checkpoint not found at {pretrained_checkpoint_path}. Model is randomly initialized.")
+    else:
+        print("Warning: No pretrained checkpoint specified. Model is randomly initialized.")
+
+    return model
+
+
+
+def create_dataloaders(config, is_distributed, rank, world_size):
+    """Create train, validation, and test dataloaders"""
+    data_config = config['data']
+    task_config = config['task']
+
+    train_dataset = fMRITaskDataset(
+        data_root=data_config['data_root'],
+        datasets=data_config['datasets'],
+        split_suffixes=data_config['train_split_suffixes'],
+        crop_length=data_config['input_seq_len'],
+        label_csv_path=task_config['csv'],
+        task_type=task_config['task_type']
+    )
+
+    val_dataset = fMRITaskDataset(
+        data_root=data_config['data_root'],
+        datasets=data_config['datasets'],
+        split_suffixes=data_config['val_split_suffixes'],
+        crop_length=data_config['input_seq_len'],
+        label_csv_path=task_config['csv'],
+        task_type=task_config['task_type']
+    )
+
+
+    test_dataset = fMRITaskDataset(
+        data_root=data_config['data_root'],
+        datasets=data_config['datasets'],
+        split_suffixes=data_config.get('test_split_suffixes', ['test']),
+        crop_length=data_config['input_seq_len'],
+        label_csv_path=task_config['csv'],
+        task_type=task_config['task_type']
+    )
+    
+
+
+    # Create samplers
+    if is_distributed:
+        train_sampler = DistributedSampler(
+            train_dataset,
+            num_replicas=world_size,
+            rank=rank,
+            shuffle=True,
+            seed=config['experiment']['seed']
+        )
+        val_sampler = DistributedSampler(val_dataset, num_replicas=world_size, rank=rank, shuffle=False)
+        test_sampler = DistributedSampler(test_dataset, num_replicas=world_size, rank=rank, shuffle=False)
+    else:
+        train_sampler = None
+        val_sampler = None
+        test_sampler = None
+
+    # Create dataloaders
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=data_config['batch_size'],
+        sampler=train_sampler,
+        shuffle=(train_sampler is None),
+        num_workers=data_config['num_workers'],
+        pin_memory=data_config['pin_memory'],
+        prefetch_factor=data_config.get('prefetch_factor', 2),
+        drop_last=True
+    )
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=data_config['batch_size'],
+        sampler=val_sampler,
+        shuffle=False,
+        num_workers=data_config['num_workers'],
+        pin_memory=data_config['pin_memory'],
+        prefetch_factor=data_config.get('prefetch_factor', 2),
+        drop_last=False
+    )
+
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=data_config['batch_size'],
+        sampler=test_sampler,
+        shuffle=False,
+        num_workers=data_config['num_workers'],
+        pin_memory=data_config['pin_memory'],
+        prefetch_factor=data_config.get('prefetch_factor', 2),
+        drop_last=False
+    )
+
+    return train_loader, val_loader, test_loader, train_sampler
+
+
+def train_one_epoch(model, train_loader, criterion, optimizer, scheduler, scaler, epoch, config,
+                    rank, world_size, label_scaler=None,log_file=None):
+    """Train for one epoch"""
+    model.train()
+
+    metric_logger = MetricLogger(delimiter="  ")
+    header = f'Epoch: [{epoch}]'
+
+    train_config = config['training']
+    log_config = config['logging']
+    task_config = config['task']
+
+    accum_iter = train_config['accum_iter']
+    use_amp = train_config['use_amp']
+    clip_grad = train_config.get('clip_grad', None)
+
+    optimizer.zero_grad()
+
+    for data_iter_step, (samples, labels) in enumerate(metric_logger.log_every(train_loader, log_config['print_freq'], header)):
+        # Adjust learning rate per iteration
+        if data_iter_step % accum_iter == 0:
+            scheduler.step()
+
+        # Move data to GPU
+        samples = samples.cuda(rank, non_blocking=True)
+        labels = labels.cuda(rank, non_blocking=True)
+
+
+        # Forward pass with mixed precision
+        with autocast(enabled=use_amp):
+            outputs = model(samples)
+
+            # Calculate loss based on task type
+            if task_config['task_type'] == 'classification':
+                if labels.dim() > 1:
+                    labels = labels.squeeze()
+
+                loss = criterion(outputs, labels)
+                # Calculate accuracy
+                _, predicted = outputs.max(1)
+                correct = predicted.eq(labels).sum().item()
+                accuracy = correct / labels.size(0)
+            else:  # regression
+                if label_scaler is not None:
+                    target_for_loss = label_scaler.transform(labels)
+                else:
+                    target_for_loss = labels
+                loss = criterion(outputs.squeeze(), target_for_loss.squeeze())
+                accuracy = 0.0  # Not applicable for regression
+
+            loss = loss / accum_iter
+
+        # Backward pass
+        if use_amp:
+            scaler.scale(loss).backward()
+
+            if (data_iter_step + 1) % accum_iter == 0:
+                if clip_grad is not None:
+                    scaler.unscale_(optimizer)
+                    nn.utils.clip_grad_norm_(model.parameters(), clip_grad)
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad()
+        else:
+            loss.backward()
+
+            if (data_iter_step + 1) % accum_iter == 0:
+                if clip_grad is not None:
+                    nn.utils.clip_grad_norm_(model.parameters(), clip_grad)
+                optimizer.step()
+                optimizer.zero_grad()
+
+        # Synchronize loss across GPUs
+        loss_value = loss.item() * accum_iter
+        if not np.isfinite(loss_value):
+            print(f"Loss is {loss_value}, stopping training")
+            sys.exit(1)
+
+        metric_logger.update(loss=loss_value)
+        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+        if task_config['task_type'] == 'classification':
+            metric_logger.update(acc=accuracy)
+
+    # Gather stats from all processes
+    metric_logger.synchronize_between_processes()
+    print(f"Averaged stats: {metric_logger}")
+
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+@torch.no_grad()
+def evaluate(model, data_loader, criterion, config, rank, epoch=None, label_scaler=None, mode='val'):
+
+    model.eval()
+    metric_logger = MetricLogger(delimiter="  ")
+    header = f'{mode.capitalize()} Epoch: [{epoch}]' if epoch is not None else f'{mode.capitalize()}:'
+    
+    task_type = config['task']['task_type']
+
+    all_preds, all_targets = [], []
+
+    for samples, labels in metric_logger.log_every(data_loader, 50, header):
+        samples = samples.cuda(rank, non_blocking=True)
+        labels = labels.cuda(rank, non_blocking=True)
+
+        outputs = model(samples)
+
+        if task_type == 'classification':
+            labels = labels.squeeze().long() if labels.dim() > 1 else labels.long()
+            loss = criterion(outputs, labels)
+            
+            preds = outputs.argmax(1)
+            acc = (preds == labels).float().mean().item()
+            metric_logger.update(loss=loss.item(), acc=acc)
+            
+            all_preds.append(preds.cpu())
+            all_targets.append(labels.cpu())
+            
+        else:  
+            if label_scaler is not None:
+                target_norm = label_scaler.transform(labels)
+            loss = criterion(outputs.view(-1), target_norm.view(-1))
+            
+            metric_logger.update(loss=loss.item())
+            all_preds.append(outputs.detach().cpu().view(-1))
+            all_targets.append(target_norm.detach().cpu().view(-1))
+
+    if len(all_preds) > 0:
+        all_preds = torch.cat(all_preds)
+        all_targets = torch.cat(all_targets)
+
+        if task_type == 'classification':
+            f1 = f1_score(all_targets.numpy(), all_preds.numpy(), average='weighted')
+            metric_logger.update(f1=f1)
+        else:
+            mse = torch.mean((all_preds - all_targets) ** 2).item()
+            mae = torch.mean(torch.abs(all_preds - all_targets)).item()
+            
+            ss_res = torch.sum((all_targets - all_preds) ** 2)
+            ss_tot = torch.sum((all_targets - all_targets.mean()) ** 2)
+            r2 = (1 - ss_res / (ss_tot + 1e-8)).item()
+            
+            vx = all_preds - all_preds.mean()
+            vy = all_targets - all_targets.mean()
+            corr = (torch.sum(vx * vy) / (torch.sqrt(torch.sum(vx**2)) * torch.sqrt(torch.sum(vy**2)) + 1e-8)).item()
+            
+            metric_logger.update(mse=mse, mae=mae, r2=r2, corr=corr)
+
+    metric_logger.synchronize_between_processes()
+    
+    if rank == 0:
+        print(f"[{mode.upper()}] Global stats: {metric_logger}")
+
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}
+
+
+def main():
+    """Main fine-tuning function"""
+    # Parse arguments
+    parser = argparse.ArgumentParser(description='Hiera MAE 4D fMRI Downstream Fine-tuning')
+    parser.add_argument('--config', type=str, default='configs/finetune_config.yaml',
+                        help='Path to config file')
+    parser.add_argument('--resume', type=str, default=None,
+                        help='Path to checkpoint to resume from')
+    parser.add_argument('--output_dir', type=str, default=None,
+                        help='Output directory (overrides config)')
+    args = parser.parse_args()
+
+    # Load config
+    config = load_config(args.config)
+
+    # Override config with command line arguments
+    if args.resume is not None:
+        config['experiment']['resume'] = args.resume
+    if args.output_dir is not None:
+        config['experiment']['output_dir'] = args.output_dir
+
+    # Setup distributed training
+    is_distributed, rank, world_size, gpu = setup_distributed()
+
+    # Set random seed
+    set_seed(config['experiment']['seed'], rank)
+
+    # Create output directories
+    if rank == 0:
+        output_dir = Path(config['experiment']['output_dir'])
+        checkpoint_dir = output_dir / 'checkpoints'
+        log_dir = output_dir / 'logs'
+
+        output_dir.mkdir(parents=True, exist_ok=True)
+        checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        log_dir.mkdir(parents=True, exist_ok=True)
+
+        # Save config
+        with open(output_dir / 'config.yaml', 'w') as f:
+            yaml.dump(config, f, default_flow_style=False)
+
+        # Setup text log file
+        log_file = output_dir / 'training_log.txt'
+        with open(log_file, 'w') as f:
+            f.write(f"Fine-tuning started at {datetime.datetime.now()}\n")
+            f.write("="*80 + "\n")
+            f.write(f"Config: {args.config}\n")
+            f.write(f"Output directory: {config['experiment']['output_dir']}\n")
+            f.write(f"Task type: {config['task']['task_type']}\n")
+            f.write("="*80 + "\n\n")
+    else:
+        checkpoint_dir = None
+        log_file = None
+
+    if is_distributed:
+        dist.barrier()
+
+    model = create_model(config)
+    model = model.cuda(gpu)
+
+    if rank == 0:
+        print("\nAnalyzing model architecture...")
+        count_parameters(model, verbose=True)
+
+    if is_distributed:
+        model = DDP(model, device_ids=[gpu], find_unused_parameters=True)
+
+    model_without_ddp = model.module if is_distributed else model
+
+    if rank == 0:
+        print("Creating dataloaders...")
+    train_loader, val_loader, test_loader, train_sampler = create_dataloaders(
+        config, is_distributed, rank, world_size
+    )
+
+    label_scaler = None
+    if config['task']['task_type'] == 'regression':
+        if rank == 0:
+            mean_val = config['task']['mean']
+            scale_val = config['task']['std']
+            print(f"StandardScaler fit complete. Mean: {mean_val:.4f}, Std: {scale_val:.4f}")
+
+        norm_mean = torch.tensor(mean_val, device=gpu, dtype=torch.float32)
+        norm_std = torch.tensor(scale_val, device=gpu, dtype=torch.float32)
+
+        if is_distributed:
+            dist.broadcast(norm_mean, src=0)
+            dist.broadcast(norm_std, src=0)
+
+        label_scaler = LabelScaler(norm_mean, norm_std)
+
+    if rank == 0:
+        print(f"Training samples: {len(train_loader.dataset)}")
+        print(f"Validation samples: {len(val_loader.dataset)}")
+        print(f"Test samples: {len(test_loader.dataset)}")
+        print(f"Batches per epoch: {len(train_loader)}")
+
+    # Create loss criterion
+    task_config = config['task']
+    if task_config['task_type'] == 'classification':
+        criterion = nn.CrossEntropyLoss(label_smoothing=0.0)
+    else:  # regression
+        criterion = nn.MSELoss()
+
+    # Optionally freeze the encoder
+    if config['training'].get('freeze_encoder', False):
+        if rank == 0:
+            print("Freezing encoder weights. Only the head will be trained.")
+        for name, param in model_without_ddp.named_parameters():
+            if 'head' not in name:
+                param.requires_grad = False
+
+        # Log which parameters are trainable
+        if rank == 0:
+            print("Trainable parameters:")
+            for name, param in model_without_ddp.named_parameters():
+                if param.requires_grad:
+                    print(name)
+
+    # Create optimizer and scheduler
+    optimizer = create_optimizer(model_without_ddp, config)
+    scheduler = create_lr_scheduler(optimizer, config, len(train_loader))
+
+    # Create gradient scaler for mixed precision
+    scaler = GradScaler() if config['training']['use_amp'] else None
+
+    # Load checkpoint if resuming
+    start_epoch = 0
+    best_metric = 0.0  # For classification: accuracy
+    best_loss = float('inf') # For regression: loss
+
+    if config['experiment'].get('resume', None) is not None:
+        start_epoch, best_metric, best_loss = load_checkpoint(
+            config['experiment']['resume'],
+            model_without_ddp,
+            optimizer,
+            scheduler,
+            scaler
+        )
+        print(f"Resumed from epoch {start_epoch}. Best metric: {best_metric:.4f}, Best loss: {best_loss:.4f}")
+    else:
+        # Initialize best_metric for new run based on task
+        if config['task']['task_type'] == 'classification':
+            best_metric = 0.0  # Accuracy starts at 0
+        else: # regression
+            best_metric = float('inf') 
+
+    # Training loop
+    if rank == 0:
+        print("Starting fine-tuning...")
+        print(f"Training from epoch {start_epoch} to {config['training']['epochs']}")
+
+    for epoch in range(start_epoch, config['training']['epochs']):
+        if is_distributed and train_sampler is not None:
+            train_sampler.set_epoch(epoch)
+
+        # Train for one epoch
+        train_stats = train_one_epoch(
+            model, train_loader, criterion, optimizer, scheduler, scaler,
+            epoch, config, rank, world_size, label_scaler, log_file
+        )
+
+        # Log training stats
+        if rank == 0:
+            log_msg = f"Epoch {epoch} Training - "
+            log_msg += " | ".join([f"{k}: {v:.4f}" for k, v in train_stats.items()])
+            print(log_msg)
+            log_to_file(log_file, log_msg)
+
+        # Validate
+        if epoch % config['validation']['val_freq'] == 0 or epoch == config['training']['epochs'] - 1:
+            print(f"DEBUG: label_scaler type is {type(label_scaler)}, value is {label_scaler}")
+            val_stats = evaluate(
+                model, val_loader, criterion, config, rank, epoch, label_scaler, 'val'
+            )
+            test_stats = evaluate(model, test_loader, criterion, config, rank, epoch, label_scaler, 'test' )
+
+            # Log validation stats
+            if rank == 0:
+                log_msg = f"Epoch {epoch} Validation - "
+                log_msg += " | ".join([f"{k}: {v:.4f}" for k, v in val_stats.items()])
+                print(log_msg)
+                log_to_file(log_file, log_msg)
+
+                log_msg = f"Epoch {epoch} Test - "
+                log_msg += " | ".join([f"{k}: {v:.4f}" for k, v in test_stats.items()])
+                print(log_msg)
+                log_to_file(log_file, log_msg)
+
+            # Determine best model based on task type
+            if rank == 0:
+                if task_config['task_type'] == 'classification':
+                    # For classification, higher accuracy is better
+                    current_metric = val_stats.get('acc', 0.0)
+                    is_best = current_metric > best_metric
+                    if is_best:
+                        best_metric = current_metric
+                        best_loss = val_stats['loss']
+                else:
+                    # For regression, lower loss is better
+                    is_best = val_stats['loss'] < best_loss
+                    if is_best:
+                        best_loss = val_stats['loss']
+                        best_metric = -best_loss  # Store negative loss as metric
+
+                checkpoint_state = {
+                    'epoch': epoch + 1,
+                    'model_state_dict': model_without_ddp.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'scheduler_state_dict': scheduler.state_dict(),
+                    'best_metric': best_metric,
+                    'best_loss': best_loss,
+                    'config': config,
+                    'train_stats': train_stats,
+                    'val_stats': val_stats,
+                }
+
+                if scaler is not None:
+                    checkpoint_state['scaler_state_dict'] = scaler.state_dict()
+
+                save_checkpoint(
+                    checkpoint_state,
+                    is_best,
+                    checkpoint_dir,
+                    filename=f'checkpoint_epoch_{epoch}.pth'
+                )
+
+                checkpoint_msg = f"Checkpoint saved at epoch {epoch}"
+                print(checkpoint_msg)
+                log_to_file(log_file, checkpoint_msg)
+
+                if is_best:
+                    if task_config['task_type'] == 'classification':
+                        best_msg = f"New best validation accuracy: {best_metric:.4f}"
+                    else:
+                        best_msg = f"New best validation loss: {best_loss:.4f}"
+                    print(best_msg)
+                    log_to_file(log_file, best_msg)
+
+        # Save periodic checkpoint
+        if rank == 0 and (epoch + 1) % config['logging']['save_freq'] == 0:
+            checkpoint_state = {
+                'epoch': epoch + 1,
+                'model_state_dict': model_without_ddp.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'scheduler_state_dict': scheduler.state_dict(),
+                'best_metric': best_metric,
+                'best_loss': best_loss,
+                'config': config,
+            }
+
+            if scaler is not None:
+                checkpoint_state['scaler_state_dict'] = scaler.state_dict()
+
+            save_checkpoint(
+                checkpoint_state,
+                False,
+                checkpoint_dir,
+                filename=f'checkpoint_epoch_{epoch}.pth'
+            )
+
+
+    # Cleanup
+    cleanup_distributed()
+
+
+if __name__ == '__main__':
+    main()

requirements.txt

@@ -0,0 +1,84 @@
+attrs==25.3.0
+certifi==2025.8.3
+charset-normalizer==3.4.3
+click==8.2.1
+cloudpickle==3.1.1
+contourpy==1.3.3
+cycler==0.12.1
+einops==0.8.1
+et_xmlfile==2.0.0
+filelock==3.18.0
+fonttools==4.59.0
+fsspec==2025.7.0
+future==1.0.0
+h5py==3.14.0
+hf-xet==1.1.8
+huggingface-hub==0.34.4
+hyperopt==0.2.7
+idna==3.10
+Jinja2==3.1.6
+joblib==1.5.1
+jsonschema==4.25.0
+jsonschema-specifications==2025.4.1
+kiwisolver==1.4.8
+lightning-utilities==0.15.0
+MarkupSafe==3.0.2
+matplotlib==3.10.3
+mpmath==1.3.0
+msgpack==1.1.1
+networkx==3.5
+nibabel==5.3.2
+numpy==2.3.2
+nvidia-cublas-cu12==12.6.4.1
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cudnn-cu12==9.5.1.17
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cufile-cu12==1.11.1.6
+nvidia-curand-cu12==10.3.7.77
+nvidia-cusolver-cu12==11.7.1.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cusparselt-cu12==0.6.3
+nvidia-nccl-cu12==2.26.2
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-nvtx-cu12==12.6.77
+openpyxl==3.1.5
+packaging==25.0
+pandas==2.3.1
+pillow==11.3.0
+protobuf==6.32.0
+psutil==7.0.0
+py4j==0.10.9.9
+pyaml==25.7.0
+pyarrow==21.0.0
+pyparsing==3.2.3
+python-dateutil==2.9.0.post0
+pytz==2025.2
+PyYAML @ file:///croot/pyyaml_1731006091482/work
+pyzstd==0.17.0
+ray==2.48.0
+referencing==0.36.2
+requests==2.32.4
+rpds-py==0.27.0
+safetensors==0.6.2
+scikit-learn==1.7.1
+scikit-optimize==0.10.2
+scipy==1.16.1
+seaborn==0.13.2
+setuptools==78.1.1
+six==1.17.0
+sympy==1.14.0
+threadpoolctl==3.6.0
+timm==1.0.19
+torch==2.7.1
+torchaudio==2.7.1
+torchmetrics==1.8.0
+torchsummary==1.5.1
+torchvision==0.22.1
+tqdm==4.67.1
+triton==3.3.1
+typing_extensions==4.14.1
+tzdata==2025.2
+urllib3==2.5.0
+wheel==0.45.1

scripts/finetune.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+
+# Set environment variables
+export CUDA_VISIBLE_DEVICES=3  
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+
+# Configuration
+CONFIG_FILE="/vePFS-0x0d/home/yewh/Hiera_MAE/configs/finetune_config.yaml"
+NUM_GPUS=1  # Fixed: Changed from 0 to 2 (number of available GPUs)
+MASTER_PORT=29503
+
+# Optional: Output directory
+OUTPUT_DIR="/vePFS-0x0d/home/yewh/Hiera_MAE/output/downstream/nki/age-lp3"
+
+# Optional: Resume from checkpoint
+# RESUME_CHECKPOINT="output/hiera_finetune/checkpoints/checkpoint_epoch_10.pth"
+
+echo "Starting DDP fine-tuning with $NUM_GPUS GPUs..."
+echo "Config: $CONFIG_FILE"
+echo "Output directory: $OUTPUT_DIR"
+
+# Launch training with torchrun (recommended for PyTorch >= 1.10)
+if [ -z "$RESUME_CHECKPOINT" ]; then
+    # Start from scratch (or from pretrained MAE)
+    torchrun \
+        --standalone \
+        --nnodes=1 \
+        --nproc_per_node=$NUM_GPUS \
+        --master_port=$MASTER_PORT \
+        /vePFS-0x0d/home/yewh/Hiera_MAE/finetune.py \
+        --config $CONFIG_FILE \
+        --output_dir $OUTPUT_DIR
+else
+    # Resume from checkpoint
+    torchrun \
+        --standalone \
+        --nnodes=1 \
+        --nproc_per_node=$NUM_GPUS \
+        --master_port=$MASTER_PORT \
+        /vePFS-0x0d/home/yewh/Hiera_MAE/finetune.py \
+        --config $CONFIG_FILE \
+        --output_dir $OUTPUT_DIR \
+        --resume $RESUME_CHECKPOINT
+fi
+
+echo "Fine-tuning completed!"

utils/ddp.py

@@ -0,0 +1,46 @@
+import os
+import torch
+import datetime
+import numpy as np
+import torch.distributed as dist
+
+def setup_distributed():
+    """Initialize distributed training"""
+    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
+        rank = int(os.environ["RANK"])
+        world_size = int(os.environ['WORLD_SIZE'])
+        gpu = int(os.environ['LOCAL_RANK'])
+    elif 'SLURM_PROCID' in os.environ:
+        rank = int(os.environ['SLURM_PROCID'])
+        gpu = rank % torch.cuda.device_count()
+        world_size = int(os.environ['SLURM_NTASKS'])
+    else:
+        print('Not using distributed mode')
+        return False, 0, 1, 0
+
+    torch.cuda.set_device(gpu)
+    dist.init_process_group(
+        backend='nccl',
+        init_method='env://',
+        world_size=world_size,
+        rank=rank,
+        timeout=datetime.timedelta(minutes=30)
+    )
+    dist.barrier()
+    return True, rank, world_size, gpu
+
+
+def cleanup_distributed():
+    """Cleanup distributed training"""
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
+
+def set_seed(seed, rank=0):
+    """Set random seed for reproducibility"""
+    seed = seed + rank
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)

utils/optim.py

@@ -0,0 +1,90 @@
+import torch
+import numpy as np
+
+def create_optimizer(model, config):
+    train_config = config['training']
+    base_lr = train_config['learning_rate']
+    weight_decay = train_config['weight_decay']
+    
+    layer_decay = train_config.get('layer_decay', 0.8) 
+    
+    # 获取所有的 blocks 数量用于计算深度
+    # 假设 model 是 HieraClassifier，其 encoder blocks 在 self.blocks 中
+    num_layers = len(model.blocks) + 1 # +1 处理 patch_embed
+    
+    parameter_groups = []
+    
+    # 1. 专门处理 Head (分类头通常使用最大的 base_lr)
+    head_lr = train_config.get('head_lr', base_lr)
+    parameter_groups.append({
+        "params": [p for n, p in model.named_parameters() if "head" in n],
+        "lr": head_lr,
+        "weight_decay": weight_decay
+    })
+
+    # 2. 处理 Encoder Blocks (按层衰减)
+    for i, block in enumerate(model.blocks):
+        # 深度越深（靠近 head），学习率越高
+        # 最后一层 i = num_layers-2，缩放接近 1.0
+        # 第一层 i = 0，缩放为 layer_decay^(num_layers)
+        scale = layer_decay ** (num_layers - i - 1)
+        
+        parameter_groups.append({
+            "params": block.parameters(),
+            "lr": base_lr * scale,
+            "weight_decay": weight_decay
+        })
+
+    # 3. 处理 Patch Embed 和其他初始层 (最低的学习率)
+    earliest_params = []
+    for n, p in model.named_parameters():
+        if "patch_embed" in n or "encoder_norm" in n:
+            earliest_params.append(p)
+    
+    if earliest_params:
+        parameter_groups.append({
+            "params": earliest_params,
+            "lr": base_lr * (layer_decay ** num_layers),
+            "weight_decay": weight_decay
+        })
+
+    if train_config['optimizer'].lower() == 'adamw':
+        optimizer = torch.optim.AdamW(
+            parameter_groups,
+            betas=tuple(train_config['betas']),
+            weight_decay=train_config['weight_decay']
+        )
+    elif train_config['optimizer'].lower() == 'sgd':
+        optimizer = torch.optim.SGD(
+            parameter_groups,
+            momentum=train_config.get('momentum', 0.9),
+            weight_decay=train_config['weight_decay']
+        )
+    else:
+        raise ValueError(f"Unsupported optimizer: {train_config['optimizer']}")
+
+    return optimizer
+
+
+def create_lr_scheduler(optimizer, config, steps_per_epoch):
+    """Create learning rate scheduler"""
+    train_config = config['training']
+    total_steps = train_config['epochs'] * steps_per_epoch
+    warmup_steps = train_config['warmup_epochs'] * steps_per_epoch
+
+    if train_config['lr_scheduler'].lower() == 'cosine':
+        def lr_lambda(current_step):
+            if current_step < warmup_steps:
+                # Linear warmup
+                return float(current_step) / float(max(1, warmup_steps))
+            else:
+                # Cosine annealing
+                progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
+                return max(train_config['min_lr'] / train_config['learning_rate'],
+                          0.5 * (1.0 + np.cos(np.pi * progress)))
+
+        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+    else:
+        raise ValueError(f"Unsupported scheduler: {train_config['lr_scheduler']}")
+
+    return scheduler

utils/utils.py

@@ -0,0 +1,237 @@
+import torch
+import datetime
+import time
+import torch.distributed as dist
+import yaml
+import os
+
+class MetricLogger:
+    """Metric logger for training"""
+    def __init__(self, delimiter="\t"):
+        self.meters = {}
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            if k not in self.meters:
+                self.meters[k] = SmoothedValue()
+            self.meters[k].update(v)
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(f"{name}: {meter}")
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ''
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt='{avg:.4f}')
+        data_time = SmoothedValue(fmt='{avg:.4f}')
+        space_fmt = ':' + str(len(str(len(iterable)))) + 'd'
+        log_msg = [
+            header,
+            '[{0' + space_fmt + '}/{1}]',
+            'eta: {eta}',
+            '{meters}',
+            'time: {time}',
+            'data: {data}'
+        ]
+        log_msg = self.delimiter.join(log_msg)
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available() and dist.get_rank() == 0:
+                    print(log_msg.format(
+                        i, len(iterable), eta=eta_string,
+                        meters=str(self),
+                        time=str(iter_time), data=str(data_time)))
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print(f'{header} Total time: {total_time_str} ({total_time / len(iterable):.4f} s / it)')
+
+
+class SmoothedValue:
+    """Track a series of values and provide access to smoothed values"""
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = []
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+        self.window_size = window_size
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        if len(self.deque) > self.window_size:
+            self.deque.pop(0)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """Synchronize across all processes"""
+        if not dist.is_available() or not dist.is_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device='cuda')
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = sorted(self.deque)
+        n = len(d)
+        if n == 0:
+            return 0
+        if n % 2 == 0:
+            return (d[n // 2 - 1] + d[n // 2]) / 2
+        return d[n // 2]
+
+    @property
+    def avg(self):
+        if len(self.deque) == 0:
+            return 0
+        return sum(self.deque) / len(self.deque)
+
+    @property
+    def global_avg(self):
+        if self.count == 0:
+            return 0
+        return self.total / self.count
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            max=max(self.deque) if len(self.deque) > 0 else 0,
+            value=self.deque[-1] if len(self.deque) > 0 else 0
+        )
+    
+
+
+def load_config(config_path):
+    """Load configuration from YAML file"""
+    with open(config_path, 'r') as f:
+        config = yaml.safe_load(f)
+    return config
+
+
+def log_to_file(log_file, message):
+    """Write message to log file"""
+    if log_file is not None:
+        with open(log_file, 'a') as f:
+            timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
+            f.write(f"[{timestamp}] {message}\n")
+            f.flush()
+
+
+def count_parameters(model, verbose=True):
+    """Count model parameters"""
+    def count_params(module):
+        return sum(p.numel() for p in module.parameters() if p.requires_grad)
+
+    def format_number(num):
+        if num >= 1e9:
+            return f"{num/1e9:.2f}B"
+        elif num >= 1e6:
+            return f"{num/1e6:.2f}M"
+        elif num >= 1e3:
+            return f"{num/1e3:.2f}K"
+        else:
+            return str(num)
+
+    # If DDP model, get original model
+    if hasattr(model, 'module'):
+        model = model.module
+
+    total_params = count_params(model)
+
+    if verbose:
+        print("\n" + "="*80)
+        print("Model Parameter Statistics")
+        print("="*80)
+
+        # Count encoder parameters
+        encoder_params = 0
+        for name in ['patch_embed', 'blocks', 'encoder_norm']:
+            if hasattr(model, name):
+                module = getattr(model, name)
+                params = count_params(module)
+                encoder_params += params
+                print(f"{name:.<35} {params:>15,} ({format_number(params):>8})")
+
+        # Count head parameters
+        if hasattr(model, 'head'):
+            head_params = count_params(model.head)
+            print(f"{'Classification/Regression Head':.<35} {head_params:>15,} ({format_number(head_params):>8})")
+
+        print("\n" + "="*80)
+        print(f"{'Encoder Parameters':.<35} {encoder_params:>15,} ({format_number(encoder_params):>8})")
+        print(f"{'TOTAL TRAINABLE PARAMETERS':.<35} {total_params:>15,} ({format_number(total_params):>8})")
+        print("="*80 + "\n")
+
+    return total_params
+
+
+
+def save_checkpoint(state, is_best, checkpoint_dir, filename='checkpoint.pth'):
+    """Save checkpoint"""
+    checkpoint_path = os.path.join(checkpoint_dir, filename)
+    torch.save(state, checkpoint_path)
+    if is_best:
+        best_path = os.path.join(checkpoint_dir, 'checkpoint_best.pth')
+        torch.save(state, best_path)
+
+
+def load_checkpoint(checkpoint_path, model, optimizer, scheduler, scaler=None):
+    """Load checkpoint"""
+    if not os.path.isfile(checkpoint_path):
+        print(f"No checkpoint found at '{checkpoint_path}'")
+        return 0, 0.0, 0.0
+
+    print(f"Loading checkpoint '{checkpoint_path}'")
+    checkpoint = torch.load(checkpoint_path, map_location='cpu')
+
+    start_epoch = checkpoint['epoch']
+    best_metric = checkpoint.get('best_metric', 0.0)
+    best_loss = checkpoint.get('best_loss', float('inf'))
+
+    model.load_state_dict(checkpoint['model_state_dict'])
+    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+    scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+
+    if scaler is not None and 'scaler_state_dict' in checkpoint:
+        scaler.load_state_dict(checkpoint['scaler_state_dict'])
+
+    print(f"Loaded checkpoint from epoch {start_epoch}")
+    return start_epoch, best_metric, best_loss
+
+
+
+class LabelScaler:
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def transform(self, labels):
+        """标准化: (y - mean) / std"""
+        return (labels - self.mean) / self.std