Pratik Bhadane
106 lines
2.8 KiB
Python
106 lines
2.8 KiB
Python
"""
|
||
GPU vs CPU benchmark for ferro_ta.gpu (SMA, EMA, RSI).
|
||
|
||
Requires:
|
||
pip install "ferro-ta[gpu]" # or pip install torch
|
||
|
||
Run:
|
||
python benchmarks/bench_gpu.py
|
||
|
||
The script compares wall-clock time for 1M-element arrays and prints a
|
||
summary table. If PyTorch is not installed or no GPU is found, GPU columns are skipped.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import time
|
||
|
||
import numpy as np
|
||
|
||
# Try to import PyTorch
|
||
try:
|
||
import torch
|
||
|
||
TORCH_AVAILABLE = True
|
||
if torch.cuda.is_available():
|
||
DEVICE = "cuda"
|
||
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||
DEVICE = "mps"
|
||
else:
|
||
DEVICE = None
|
||
except ImportError:
|
||
torch = None # type: ignore[assignment]
|
||
TORCH_AVAILABLE = False
|
||
DEVICE = None
|
||
|
||
from ferro_ta.gpu import ema, rsi, sma
|
||
|
||
N = 1_000_000
|
||
REPEATS = 10
|
||
|
||
|
||
def _time_fn(fn, *args, **kwargs) -> float:
|
||
"""Return minimum wall time (seconds) over REPEATS calls."""
|
||
times = []
|
||
for _ in range(REPEATS):
|
||
t0 = time.perf_counter()
|
||
fn(*args, **kwargs)
|
||
if DEVICE == "cuda":
|
||
torch.cuda.synchronize()
|
||
elif DEVICE == "mps":
|
||
torch.mps.synchronize()
|
||
times.append(time.perf_counter() - t0)
|
||
return min(times)
|
||
|
||
|
||
def main() -> None:
|
||
rng = np.random.default_rng(42)
|
||
close_cpu = rng.uniform(100.0, 200.0, N)
|
||
|
||
print(f"Array size: {N:,} elements")
|
||
print(f"Repeats: {REPEATS}")
|
||
print(f"Device: {DEVICE if DEVICE else 'CPU'}")
|
||
print()
|
||
|
||
header = f"{'Indicator':<20} {'CPU (ms)':>10}"
|
||
if DEVICE:
|
||
header += f" {'GPU (ms)':>10} {'Speedup':>10}"
|
||
print(header)
|
||
print("-" * len(header))
|
||
|
||
for name, fn, kwargs in [
|
||
("sma(period=30)", sma, {"timeperiod": 30}),
|
||
("ema(period=30)", ema, {"timeperiod": 30}),
|
||
("rsi(period=14)", rsi, {"timeperiod": 14}),
|
||
]:
|
||
cpu_time = _time_fn(fn, close_cpu, **kwargs) * 1000 # ms
|
||
|
||
row = f"{name:<20} {cpu_time:>10.3f}"
|
||
if DEVICE:
|
||
dtype = torch.float32 if DEVICE == "mps" else torch.float64
|
||
close_gpu = torch.tensor(close_cpu, dtype=dtype, device=DEVICE)
|
||
# Warm-up
|
||
fn(close_gpu, **kwargs)
|
||
if DEVICE == "cuda":
|
||
torch.cuda.synchronize()
|
||
elif DEVICE == "mps":
|
||
torch.mps.synchronize()
|
||
gpu_time = _time_fn(fn, close_gpu, **kwargs) * 1000 # ms
|
||
speedup = cpu_time / gpu_time
|
||
row += f" {gpu_time:>10.3f} {speedup:>10.2f}×"
|
||
print(row)
|
||
|
||
if not TORCH_AVAILABLE:
|
||
print()
|
||
print("PyTorch not available — GPU columns skipped.")
|
||
print("Install with: pip install 'ferro_ta[gpu]'")
|
||
elif not DEVICE:
|
||
print()
|
||
print(
|
||
"PyTorch found, but no CUDA or MPS device detected — GPU columns skipped."
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|