Spaces:
Sleeping
Sleeping
| import time | |
| import torch | |
| import pandas as pd | |
| from fvcore.nn import FlopCountAnalysis, ActivationCountAnalysis | |
| def profile_model(model, input_tensor, display=False): | |
| flops = calculate_flops(model, input_tensor[0:1, ...]) # (1, n_mels, n_frames) | |
| acts = calculate_activations(model, input_tensor[0:1, ...]) # (1, n_mels, n_frames) | |
| params = calculate_params(model) | |
| speed = calculate_speed(model, input_tensor[0:1, ...]) # (1, n_mels, n_frames) | |
| memory = calculate_memory(model, input_tensor) # (B, n_mels, n_frames) | |
| profile_data = { | |
| "Metric": [ | |
| "FLOPs (G)", | |
| "Activations (M)", | |
| "Params (M)", | |
| "Memory (GB)", | |
| "Speed (A/S)", | |
| ], | |
| "Value": [flops, acts, params, memory, speed], | |
| } | |
| profile_df = pd.DataFrame(profile_data).set_index("Metric").T | |
| if display: | |
| print(profile_df.to_markdown(index=False, tablefmt="grid")) | |
| return profile_df | |
| def calculate_speed(model, input_tensor, num_runs=100, warmup_runs=5): | |
| model.eval() | |
| if torch.cuda.is_available(): | |
| # Warm-up iterations | |
| with torch.no_grad(): | |
| for _ in range(warmup_runs): | |
| _ = model(input_tensor) | |
| # Create CUDA events for timing | |
| start = torch.cuda.Event(enable_timing=True) | |
| end = torch.cuda.Event(enable_timing=True) | |
| # Actual timing | |
| start.record() | |
| with torch.no_grad(): | |
| for _ in range(num_runs): | |
| _ = model(input_tensor) | |
| end.record() | |
| # Synchronize to wait for the events to be recorded | |
| torch.cuda.synchronize() | |
| # Calculate elapsed time | |
| elapsed_time = start.elapsed_time(end) # in milliseconds | |
| latency = elapsed_time / num_runs / 1000.0 # convert to seconds | |
| else: | |
| # Warm-up iterations | |
| with torch.no_grad(): | |
| for _ in range(warmup_runs): | |
| _ = model(input_tensor) | |
| # Actual timing | |
| start = time.time() | |
| with torch.no_grad(): | |
| for _ in range(num_runs): | |
| _ = model(input_tensor) | |
| end = time.time() | |
| # Calculate elapsed time | |
| latency = (end - start) / num_runs | |
| return 1.0 / latency | |
| def calculate_flops(model, input_tensor): | |
| """Calculate FLOPs in GigaFLOPs. | |
| Models often reports MACs as FLOPs e.g. ConvNeXt, timm library | |
| Reference: | |
| 1. https://github.com/huggingface/pytorch-image-models/blob/main/benchmark.py#L206 | |
| 2. https://github.com/facebookresearch/fvcore/issues/69 | |
| """ | |
| flops = FlopCountAnalysis(model, input_tensor).total() | |
| return flops / 1e9 # in GigaFLOPs | |
| def calculate_activations(model, input_tensor): | |
| acts = ActivationCountAnalysis(model, input_tensor).total() | |
| return acts / 1e6 # in Millions | |
| def calculate_params(model): | |
| return sum(p.numel() for p in model.parameters()) / 1e6 # in Millions | |
| def calculate_memory(model, input_tensor): | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| torch.cuda.reset_peak_memory_stats(device=None) | |
| start_memory = torch.cuda.max_memory_allocated(device=None) | |
| model.train() | |
| _ = model(input_tensor) | |
| end_memory = torch.cuda.max_memory_allocated(device=None) | |
| torch.cuda.empty_cache() | |
| torch.cuda.reset_peak_memory_stats(device=None) | |
| memory = (end_memory - start_memory) / (1024**3) # in GB | |
| else: | |
| memory = 0 | |
| return memory | |