compute_flops.py

import torch
from ptflops import get_model_complexity_info
from timm.models import create_model
from mamba_ssm.modules.mamba_simple import Mamba
import models_mamba_ecg 

# 1. Define the robust hook to compute exact BiMamba-v2 MACs
def mamba_flops_counter_hook(module, input, output):
    # Extract input tensor shape: [Batch, Sequence_Length, Hidden_Dim]
    x = input[0]
    B, L, D = x.shape
    
    macs = 0
    
    # 1. Forward Linear Projections
    if hasattr(module, 'in_proj'):
        macs += L * module.in_proj.in_features * module.in_proj.out_features
        
    # 2. Convolutions (Forward & Optional Backward)
    if hasattr(module, 'conv1d'):
        macs += L * module.conv1d.in_channels * module.conv1d.kernel_size[0]
    if hasattr(module, 'conv1d_b'):
        macs += L * module.conv1d_b.in_channels * module.conv1d_b.kernel_size[0]
        
    # 3. State Space Projections (Forward & Optional Backward)
    if hasattr(module, 'x_proj'):
        macs += L * module.x_proj.in_features * module.x_proj.out_features
    if hasattr(module, 'x_proj_b'):
        macs += L * module.x_proj_b.in_features * module.x_proj_b.out_features
        
    if hasattr(module, 'dt_proj'):
        macs += L * module.dt_proj.in_features * module.dt_proj.out_features
    if hasattr(module, 'dt_proj_b'):
        macs += L * module.dt_proj_b.in_features * module.dt_proj_b.out_features
        
    # 4. Dynamic Selective Scan Recurrences (Hardware-level operations)
    d_inner = module.dt_proj.out_features if hasattr(module, 'dt_proj') else (D * 2)
    d_state = getattr(module, 'd_state', 16) 
    
    # Base Forward Scan Pass (Always runs)
    macs += L * (2 * d_inner * d_state) 
    
    # PARAMETER-COUNT SIGNATURE CHECK (100% immune to library variable names)
    num_params = sum(p.numel() for p in module.parameters())
    
    # If parameters exceed the unidirectional baseline (~963k), it is running a bidirectional scan
    is_bidirectional_ssm = (num_params > 970000)
    
    if is_bidirectional_ssm:
        macs += L * (2 * d_inner * d_state) # Add the backward scan execution pass
        
    # 5. Final Output Projection
    if hasattr(module, 'out_proj'):
        macs += L * module.out_proj.in_features * module.out_proj.out_features
        
    # Register the complete accumulated MACs back to the ptflops profiler
    module.__flops__ += int(macs * B)


def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    model_name = 'ecg_vim_small_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2'
    input_shape = (12, 8192) 

    print(f"Creating model: {model_name} with depth=16...")
    
    model = create_model(
        model_name,
        pretrained=False,
        num_classes=26,             
        drop_rate=0.0,              
        drop_path_rate=0.0,         
        drop_block_rate=None,
        block='VisionMamba',        
        depth=16,                   
        fused_add_norm=False,       
        if_divide_out=False,        
        use_middle_cls_token=False, 
        img_size=8192
    )

    model.to(device)
    model.eval() 

    print("Computing MACs and Parameters. This might take a moment...")

    # 2. Register the hook via custom_modules_hooks
    with torch.no_grad():
        macs, params = get_model_complexity_info(
            model, 
            input_shape, 
            as_strings=True, 
            print_per_layer_stat=True, 
            verbose=True,
            custom_modules_hooks={Mamba: mamba_flops_counter_hook} # <--- CRITICAL INCLUSION
        )

    print("\n" + "="*50)
    print(f"Architecture: {model_name}")
    print(f"Input Shape:  {input_shape}")
    print(f"Computational Complexity (MACs): {macs}")
    print(f"Total Parameters:                {params}")
    print("="*50)

if __name__ == '__main__':
    main()


"""
output:
(vim3) huawei@huawei-B460MDS3HV2:/media/huawei/93219077-4d31-405d-88e2-323585a4ba3e/effective-ECG-recognition/1DCNN-ECG-Mamba-For-versatile_result$ python compute_flops.py
Using device: cuda
Creating model: ecg_vim_small_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2 with depth=16...
This is the drop_path_rate: 0.0
This is whether the fused_add_norm: False
This is whether the if_divide_out: False
This is whether the use_middle_cls_token: False
This is whether the stochastic dropout: 0.0
Computing MACs and Parameters. This might take a moment...
Warning: module ECG_Patch_embedding is treated as a zero-op.
Warning: module Dropout is treated as a zero-op.
Warning: module Identity is treated as a zero-op.
Warning: module SiLU is treated as a zero-op.
Warning: module RMSNorm is treated as a zero-op.
Warning: module Block is treated as a zero-op.
Warning: module VisionMamba is treated as a zero-op.
VisionMamba(
  16.77 M, 97.670% Params, 17.49 GMac, 100.000% MACs, 
  (ECG_patch_embedding): ECG_Patch_embedding(
    74.11 k, 0.432% Params, 75.82 MMac, 0.433% MACs, 
    (multiple_cnn): Sequential(
      74.11 k, 0.432% Params, 75.82 MMac, 0.433% MACs, 
      (0): Conv1d(74.11 k, 0.432% Params, 75.82 MMac, 0.433% MACs, 12, 384, kernel_size=(16,), stride=(8,))
    )
  )
  (pos_drop): Dropout(0, 0.000% Params, 0.0 Mac, 0.000% MACs, p=0.0, inplace=False)
  (head): Linear(10.01 k, 0.058% Params, 10.01 KMac, 0.000% MACs, in_features=384, out_features=26, bias=True)
  (drop_path): Identity(0, 0.000% Params, 0.0 Mac, 0.000% MACs, )
  (layers): ModuleList(
    (0-15): 16 x Block(
      1.04 M, 6.074% Params, 1.09 GMac, 6.223% MACs, 
      (mixer): Mamba(
        1.04 M, 6.074% Params, 1.09 GMac, 6.223% MACs, 
        (in_proj): Linear(589.82 k, 3.435% Params, 0.0 Mac, 0.000% MACs, in_features=384, out_features=1536, bias=False)
        (conv1d): Conv1d(3.84 k, 0.022% Params, 0.0 Mac, 0.000% MACs, 768, 768, kernel_size=(4,), stride=(1,), padding=(3,), groups=768)
        (act): SiLU(0, 0.000% Params, 0.0 Mac, 0.000% MACs, )
        (x_proj): Linear(43.01 k, 0.250% Params, 0.0 Mac, 0.000% MACs, in_features=768, out_features=56, bias=False)
        (dt_proj): Linear(19.2 k, 0.112% Params, 0.0 Mac, 0.000% MACs, in_features=24, out_features=768, bias=True)
        (conv1d_b): Conv1d(3.84 k, 0.022% Params, 0.0 Mac, 0.000% MACs, 768, 768, kernel_size=(4,), stride=(1,), padding=(3,), groups=768)
        (x_proj_b): Linear(43.01 k, 0.250% Params, 0.0 Mac, 0.000% MACs, in_features=768, out_features=56, bias=False)
        (dt_proj_b): Linear(19.2 k, 0.112% Params, 0.0 Mac, 0.000% MACs, in_features=24, out_features=768, bias=True)
        (out_proj): Linear(294.91 k, 1.717% Params, 0.0 Mac, 0.000% MACs, in_features=768, out_features=384, bias=False)
      )
      (norm): RMSNorm(0, 0.000% Params, 0.0 Mac, 0.000% MACs, )
      (drop_path): Identity(0, 0.000% Params, 0.0 Mac, 0.000% MACs, )
    )
  )
  (norm_f): RMSNorm(0, 0.000% Params, 0.0 Mac, 0.000% MACs, )
)

==================================================
Architecture: ecg_vim_small_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2
Input Shape:  (12, 8192)
Computational Complexity (MACs): 17.49 GMac
Total Parameters:                17.17 M
==================================================

"""