3 เทคนิคใช้ Numba เพิ่มความเร็ว Python Runtime ให้แรงกว่าเดิม

pip install numba

import time
import numpy as np
from numba import njit
 
def total_plain(x):
    total = 0.0
    for i in range(x.shape[0]):
        total += np.sqrt(x[i]) * np.sin(x[i])
    return total
 
@njit
def total_jit(x):
    total = 0.0
    for i in range(x.shape[0]):
        total += np.sqrt(x[i]) * np.sin(x[i])
    return total
 
def total_numpy(x):
    return np.sum(np.sqrt(x) * np.sin(x))
 
def benchmark(func, x, repeats=3):
    """Run func repeats times and return (best_time, mean_time, result)."""
    times = []
    result = None
    for _ in range(repeats):
        start = time.perf_counter()
        result = func(x)
        times.append(time.perf_counter() - start)
    return min(times), sum(times) / len(times), result
 
x = np.random.default_rng(0).random(10_000_000)
 
# Measure JIT compilation separately (first call compiles)
start = time.perf_counter()
total_jit(x)
compile_time = time.perf_counter() - start
print(f"Numba first call (includes compilation): {compile_time:.4f} s\n")
 
# Plain Python loop is slow on 10M elements, so run it only once
results = {
    "Plain Python loop": benchmark(total_plain, x, repeats=1),
    "Numba @njit":       benchmark(total_jit, x, repeats=5),
    "NumPy vectorized":  benchmark(total_numpy, x, repeats=5),
}
 
baseline = results["Plain Python loop"][0]
 
print(f"{'Method':<20} {'Best (s)':>10} {'Mean (s)':>10} {'Speedup':>10}  Result")
print("-" * 72)
for name, (best, mean, value) in results.items():
    print(f"{name:<20} {best:>10.4f} {mean:>10.4f} {baseline / best:>9.1f}x  {value:.6f}")
 
# Sanity check that all methods agree
values = [r[2] for r in results.values()]
print("\nResults match:", np.allclose(values, values[0]))

Numba first call (includes compilation): 0.2799 s
 
Method                 Best (s)   Mean (s)    Speedup  Result
------------------------------------------------------------------------
Plain Python loop        3.0789     3.0789       1.0x  3641603.675817
Numba @njit              0.0384     0.0385      80.2x  3641603.675817
NumPy vectorized         0.0552     0.0606      55.7x  3641603.675816
 
Results match: True

from numba import njit, prange
 
@njit(parallel=True)
def total_parallel(x):
    total = 0.0
    for i in prange(x.shape[0]):
        total += np.sqrt(x[i]) * np.sin(x[i])
    return total

Numba first call (includes compilation): 0.2431 s
 
Method                 Best (s)   Mean (s)    Speedup  Result
------------------------------------------------------------------------
Plain Python loop        3.0229     3.0229       1.0x  3641603.675817
Numba @njit              0.0374     0.0379      80.7x  3641603.675817
NumPy vectorized         0.0547     0.0608      55.3x  3641603.675816
Parallel @njit           0.0087     0.1222     347.3x  3641603.675816
 
Results match: True

@njit(parallel=True, cache=True)
def total_cached(x):
    total = 0.0
    for i in prange(x.shape[0]):
        total += np.sqrt(x[i]) * np.sin(x[i])
    return total

Numba first call (includes compilation): 0.1248 s
 
Method                 Best (s)   Mean (s)    Speedup  Result
------------------------------------------------------------------------
Plain Python loop        3.0025     3.0025       1.0x  3641603.675817
Numba @njit              0.0384     0.0384      78.2x  3641603.675817
NumPy vectorized         0.0550     0.0563      54.5x  3641603.675816
Parallel @njit           0.0086     0.0383     349.1x  3641603.675816
Cached @njit             0.0086     0.0094     349.9x  3641603.675816
 
Results match: True