ECG-Mamba-V2 / compute_throughput.py

Upload 80 files

32051f1 verified 3 months ago

3.09 kB

import torch
import time
from timm.models import create_model

# Import your model registry to ensure custom modules are recognized
import models_mamba_ecg 

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

    # Parameters mirroring your exact bash script configuration
    model_name = 'ecg_vim_small_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2'
    batch_size = 20
    channels = 12
    depth = 16
    embed_dim = 384
    seq_len = 8192
    drop_path_rate=0
    uniform_drop_path_rate=0.1

    print(f"Creating model: {model_name} (depth={depth}, embed_dim={embed_dim}, seq_len={seq_len})...")
    model = create_model(
        model_name,
        pretrained=False,
        num_classes=27,
        drop_rate=0.0,
        drop_path_rate=drop_path_rate,
        drop_block_rate=None,
        block='VisionMamba',
        depth=depth,
        fused_add_norm=False,
        if_divide_out=False,
        use_middle_cls_token=False,
        img_size=seq_len,
        embed_dim = embed_dim,
        uniform_drop_path_rate=uniform_drop_path_rate,
        bimamba_type="v2"
    )

    
    model.to(device)
    model.eval() # Ensure dropout/norms are locked in inference mode

    # Create dummy data directly on the GPU to isolate pure model throughput
    print(f"Allocating dummy input tensor (Batch: {batch_size}, Channels: {channels}, Length: {seq_len})...")
    dummy_input = torch.randn(batch_size, channels, seq_len, device=device)

    warmup_iters = 20
    benchmark_iters = 100

    # 1. WARM-UP PHASE
    print(f"Executing {warmup_iters} warm-up iterations (discarding timing data)...")
    with torch.no_grad():
        for _ in range(warmup_iters):
            _ = model(dummy_input)

    # 2. BENCHMARK PHASE
    print(f"Executing {benchmark_iters} official benchmark iterations...")
    
    # CRITICAL: Force Python to wait until all warm-up kernels finish before starting the clock
    if device.type == 'cuda':
        torch.cuda.synchronize()
        
    start_time = time.time()

    with torch.no_grad():
        for _ in range(benchmark_iters):
            _ = model(dummy_input)

    # CRITICAL: Force Python to wait until the final forward pass finishes before stopping the clock
    if device.type == 'cuda':
        torch.cuda.synchronize()

    end_time = time.time()
    total_time = end_time - start_time

    # 3. METRICS CALCULATION
    total_files_processed = benchmark_iters * batch_size
    throughput = total_files_processed / total_time
    ms_per_batch = (total_time / benchmark_iters) * 1000

    print("\n" + "="*50)
    print(" THROUGHPUT BENCHMARK RESULTS")
    print("="*50)
    print(f"Batch Size:            {batch_size}")
    print(f"Total Time Elapsed:    {total_time:.4f} seconds")
    print(f"Total Files Processed: {total_files_processed} files")
    print(f"Model Throughput:      {throughput:.2f} files/s")
    print(f"Latency per Batch:     {ms_per_batch:.2f} ms")
    print("="*50)

if __name__ == '__main__':
    main()

	import torch
	import time
	from timm.models import create_model

	# Import your model registry to ensure custom modules are recognized
	import models_mamba_ecg

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

	# Parameters mirroring your exact bash script configuration
	model_name = 'ecg_vim_small_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2'
	batch_size = 20
	channels = 12
	depth = 16
	embed_dim = 384
	seq_len = 8192
	drop_path_rate=0
	uniform_drop_path_rate=0.1

	print(f"Creating model: {model_name} (depth={depth}, embed_dim={embed_dim}, seq_len={seq_len})...")
	model = create_model(
	model_name,
	pretrained=False,
	num_classes=27,
	drop_rate=0.0,
	drop_path_rate=drop_path_rate,
	drop_block_rate=None,
	block='VisionMamba',
	depth=depth,
	fused_add_norm=False,
	if_divide_out=False,
	use_middle_cls_token=False,
	img_size=seq_len,
	embed_dim = embed_dim,
	uniform_drop_path_rate=uniform_drop_path_rate,
	bimamba_type="v2"
	)


	model.to(device)
	model.eval() # Ensure dropout/norms are locked in inference mode

	# Create dummy data directly on the GPU to isolate pure model throughput
	print(f"Allocating dummy input tensor (Batch: {batch_size}, Channels: {channels}, Length: {seq_len})...")
	dummy_input = torch.randn(batch_size, channels, seq_len, device=device)

	warmup_iters = 20
	benchmark_iters = 100

	# 1. WARM-UP PHASE
	print(f"Executing {warmup_iters} warm-up iterations (discarding timing data)...")
	with torch.no_grad():
	for _ in range(warmup_iters):
	_ = model(dummy_input)

	# 2. BENCHMARK PHASE
	print(f"Executing {benchmark_iters} official benchmark iterations...")

	# CRITICAL: Force Python to wait until all warm-up kernels finish before starting the clock
	if device.type == 'cuda':
	torch.cuda.synchronize()

	start_time = time.time()

	with torch.no_grad():
	for _ in range(benchmark_iters):
	_ = model(dummy_input)

	# CRITICAL: Force Python to wait until the final forward pass finishes before stopping the clock
	if device.type == 'cuda':
	torch.cuda.synchronize()

	end_time = time.time()
	total_time = end_time - start_time

	# 3. METRICS CALCULATION
	total_files_processed = benchmark_iters * batch_size
	throughput = total_files_processed / total_time
	ms_per_batch = (total_time / benchmark_iters) * 1000

	print("\n" + "="*50)
	print(" THROUGHPUT BENCHMARK RESULTS")
	print("="*50)
	print(f"Batch Size: {batch_size}")
	print(f"Total Time Elapsed: {total_time:.4f} seconds")
	print(f"Total Files Processed: {total_files_processed} files")
	print(f"Model Throughput: {throughput:.2f} files/s")
	print(f"Latency per Batch: {ms_per_batch:.2f} ms")
	print("="*50)

	if __name__ == '__main__':
	main()