File size: 2,512 Bytes
a10ba7f | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | import torch
import time
import os
import psutil
import sys
from pathlib import Path
# Add project root to path
sys.path.append(str(Path(__file__).parent.parent))
from src.models.student import LIPEV2Student
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
def estimate_flops(model, input_patches, input_landmarks):
"""
Very rough estimate of FLOPs for the LIPEV2Student.
"""
flops = 0
# 1. MiniConvEmbedder (3 Conv layers + GAP)
# Conv1: 16 filters, 3x3 kernel, input 1x8x8 -> output 16x6x6
# FLOPs approx: 2 * Cin * K * K * Hout * Wout * Cout
flops += 4 * (2 * 1 * 3 * 3 * 6 * 6 * 16)
# Conv2: 16 -> 32, 6x6 -> 4x4
flops += 4 * (2 * 16 * 3 * 3 * 4 * 4 * 32)
# Conv3: 32 -> 64, 4x4 -> 2x2
flops += 4 * (2 * 32 * 3 * 3 * 2 * 2 * 64)
# 2. Geo MLP
# Linear 1: 956 -> 256
flops += 2 * 956 * 256
# Linear 2: 256 -> 256
flops += 2 * 256 * 256
# 3. Heads
# Pitch: 256 -> 64 -> 90
flops += 2 * 256 * 64 + 2 * 64 * 90
# Yaw: 256 -> 64 -> 90
flops += 2 * 256 * 64 + 2 * 64 * 90
return flops
def profile_model():
model = LIPEV2Student()
model.eval()
params = count_parameters(model)
print(f"Total Trainable Parameters: {params:,}")
dummy_patches = torch.randn(1, 4, 8, 8)
dummy_landmarks = torch.randn(1, 956)
# FLOPs Estimation
flops = estimate_flops(model, dummy_patches, dummy_landmarks)
print(f"Estimated FLOPs per frame (State A): {flops / 1e6:.2f} MFLOPs ({flops / 1e9:.4f} GFLOPs)")
# RAM Measurement
process = psutil.Process(os.getpid())
mem_before = process.memory_info().rss / 1024 / 1024 # MB
# Run some inferences
with torch.no_grad():
for _ in range(100):
_ = model(dummy_patches, dummy_landmarks, state='A')
mem_after = process.memory_info().rss / 1024 / 1024 # MB
print(f"Memory Usage (RSS): {mem_after:.2f} MB (Baseline: {mem_before:.2f} MB)")
# Latency
start_time = time.time()
num_iters = 500
with torch.no_grad():
for _ in range(num_iters):
_ = model(dummy_patches, dummy_landmarks, state='A')
end_time = time.time()
avg_latency = (end_time - start_time) / num_iters * 1000
print(f"Average Inference Latency (State A): {avg_latency:.4f} ms")
print(f"Estimated Max FPS (State A): {1000 / avg_latency:.2f}")
if __name__ == '__main__':
profile_model()
|