Spaces:
Sleeping
Sleeping
File size: 1,732 Bytes
6c514a7 d304e58 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | import torch
import time
from cnn1 import UNet
from thop import profile
import numpy as np
def benchmark(device="cpu"):
print(f"--- Benchmarking on {device.upper()} ---")
# 1. Setup Model
model = UNet(text_dim=512).to(device)
model.eval()
# 2. Setup Dummy Inputs (Standard 512x512 image + CLIP embedding)
dummy_image = torch.randn(1, 3, 512, 512).to(device)
dummy_text = torch.randn(1, 512).to(device) # Normalized text embedding
# --- METRIC 1: FLOPs & Parameters ---
if device == "cpu": # Only need to calc this once
print("Calculating FLOPs and Params...")
macs, params = profile(model, inputs=(dummy_image, dummy_text), verbose=False)
print(f"Parameters: {params / 1e6:.2f} M")
print(f"GFLOPs: {macs / 1e9:.2f} G")
# --- METRIC 2: Inference Time ---
print("Measuring Inference Speed...")
# Warmup (get cache ready)
for _ in range(5):
with torch.no_grad():
_ = model(dummy_image, dummy_text)
# Measure
latencies = []
with torch.no_grad():
for _ in range(50): # Run 50 times
start = time.time()
_ = model(dummy_image, dummy_text)
if device == "cuda":
torch.cuda.synchronize() # Wait for GPU to finish
end = time.time()
latencies.append(end - start)
avg_time = np.mean(latencies)
print(f"Avg Inference Time: {avg_time:.4f} seconds")
print(f"FPS: {1/avg_time:.2f}")
print("-" * 30)
if __name__ == "__main__":
# Run on CPU
benchmark("cpu")
# Run on GPU (If available, for comparison)
if torch.cuda.is_available():
benchmark("cuda") |