Posted on
• Artificial Intelligence

Artificial Intelligence Python Performance

Author
  • User
    linuxbash
    Posts by this author
    Posts by this author

Squeezing Performance out of AI Python on Linux: A Practical, Bash-First Guide

If your AI models are slow, you’re probably leaving free speed on the table. The good news: you can often get 2–10× faster training and inference on the same Linux box—no algorithm changes—by fixing your Python and math stack, profiling the right hotspots, and tuning threads. This post is a hands-on, Bash-friendly guide to do exactly that.

What you’ll get:

  • Why AI performance in Python is usually I/O and math-backend bound—not “just Python”

  • 3–5 actionable steps you can run today

  • Copy-pasteable commands for apt, dnf, and zypper

  • Minimal examples you can benchmark

Why this matters (and why Python is not the bottleneck you think)

  • AI workloads spend the majority of time inside native code (BLAS/LAPACK, FFTs, convolution kernels). If those aren’t optimized or are misconfigured, you waste cores and cache.

  • Python overhead (the GIL, interpreter loops) bites you mostly when you write loop-heavy code in pure Python. Use vectorization or JIT to escape it.

  • Threading defaults are not universal. NumPy, SciPy, PyTorch, OpenBLAS/MKL/BLIS, and OpenMP all have their own knobs. Mis-set threads can slow you down dramatically.

  • Linux gives you full control over compilers, math libraries, and process affinity—use it.

Below are five core actions with commands and small code snippets you can try immediately.

0) Prerequisites: System packages and a clean Python environment

Install a base toolchain and fast math libraries so Python wheels can link against optimized BLAS/LAPACK.

  • Debian/Ubuntu (apt):
sudo apt update
sudo apt install -y \
  python3 python3-venv python3-pip \
  build-essential gcc g++ gfortran make pkg-config \
  libopenblas-dev liblapack-dev
  • Fedora/RHEL/CentOS Stream (dnf):
sudo dnf groupinstall -y "Development Tools"
sudo dnf install -y \
  python3 python3-virtualenv python3-pip \
  gcc gcc-c++ gcc-gfortran make pkgconf-pkg-config \
  openblas-devel lapack-devel
  • openSUSE Leap/Tumbleweed (zypper):
sudo zypper refresh
sudo zypper install -y \
  python3 python3-virtualenv python3-pip \
  gcc gcc-c++ gcc-fortran make pkg-config \
  openblas-devel lapack-devel

Create and activate a fresh virtual environment, then install common AI/scientific packages:

python3 -m venv ~/.venvs/aiperf
source ~/.venvs/aiperf/bin/activate

pip install -U pip wheel
pip install -U numpy scipy scikit-learn numba numexpr line-profiler py-spy

Optional (CPU-only PyTorch wheels):

pip install --index-url https://download.pytorch.org/whl/cpu torch torchvision torchaudio

Optional (NVIDIA GPU with CUDA 12.1 wheels):

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install cupy-cuda12x

1) Make sure NumPy/SciPy are using a fast BLAS (and tune its threads)

Check which BLAS/LAPACK backend NumPy is actually using:

python - <<'PY'
import numpy as np, os
np.show_config()
print("OPENBLAS_NUM_THREADS =", os.getenv("OPENBLAS_NUM_THREADS"))
print("MKL_NUM_THREADS      =", os.getenv("MKL_NUM_THREADS"))
print("OMP_NUM_THREADS      =", os.getenv("OMP_NUM_THREADS"))
PY

If you don’t see OpenBLAS, MKL, or BLIS, you’re likely on a slow reference BLAS. Installing the dev packages above and then reinstalling NumPy usually fixes it:

pip install --force-reinstall --no-binary=:all: numpy scipy

Control threads to avoid over-subscription (common cause of slowdowns):

  • For OpenBLAS:
export OPENBLAS_NUM_THREADS=$(nproc)   # try full cores first
# or pin to fewer:
export OPENBLAS_NUM_THREADS=4
  • For MKL (if you use Intel MKL via custom repos):
export MKL_NUM_THREADS=8
export MKL_DYNAMIC=FALSE
  • OpenMP-using libs:
export OMP_NUM_THREADS=8

Quick matrix-multiply benchmark to see scaling and best thread count:

python - <<'PY'
import os, time, numpy as np
np.random.seed(0)
for t in [1,2,4,8,16]:
    os.environ["OPENBLAS_NUM_THREADS"]=str(t)
    A = np.random.randn(3000,3000).astype(np.float64)
    B = np.random.randn(3000,3000).astype(np.float64)
    s = time.perf_counter()
    C = A @ B
    dt = time.perf_counter() - s
    print(f"threads={t:>2}  time={dt:.2f}s  GFLOP/s≈{(2*3000**3/1e9)/dt:.1f}")
PY

Pick the thread count with the lowest time on your machine and export it in your shell profile.

2) Profile first: find Python hotspots before optimizing

You can’t optimize what you haven’t measured.

  • Quick, built-in profiler:
python -m cProfile -o prof.pstats your_script.py
python - <<'PY'
import pstats
p = pstats.Stats('prof.pstats').strip_dirs().sort_stats('tottime')
p.print_stats(30)
PY
  • Line-level profiling:
pip install line-profiler
# Add @profile decorator to functions of interest, then:
kernprof -l your_script.py
python -m line_profiler your_script.py.lprof
  • Low-overhead sampling (great for production):
pip install py-spy
py-spy top -- python your_script.py
# or record a flamegraph
py-spy record -o profile.svg -- python your_script.py

Use this to decide: vectorize, JIT, or move hot loops to native code.

3) Vectorize and JIT: escape the GIL and pure-Python loops

  • Replace Python loops with NumPy operations whenever possible:
# Slow: pure Python loop
def relu_py(x):
    out = []
    for v in x:
        out.append(v if v > 0 else 0)
    return out

# Fast: vectorized
import numpy as np
def relu_np(x):
    x = np.asarray(x)
    return np.maximum(x, 0, out=x.copy())
  • Use Numba for tight loops that resist vectorization:
from numba import njit, prange
import numpy as np

@njit(parallel=True, fastmath=True)
def pairwise_l2(X):
    n, d = X.shape
    D = np.empty((n, n), dtype=np.float32)
    for i in prange(n):
        for j in range(n):
            s = 0.0
            for k in range(d):
                diff = X[i,k] - X[j,k]
                s += diff*diff
            D[i,j] = s
    return D

X = np.random.rand(3000, 64).astype(np.float32)
D = pairwise_l2(X)  # JIT-compiled on first call
  • NumExpr can speed up certain array expressions by multi-threading and avoiding temporaries:
import numexpr as ne
import numpy as np

a = np.random.rand(10_000_000)
b = np.random.rand(10_000_000)
c = ne.evaluate("2*a + 3*b")

4) Control parallelism and memory to prevent slowdowns

  • Avoid thread explosions: set only one library in charge of threads. If PyTorch is your main engine on CPU:
python - <<'PY'
import torch
torch.set_num_threads(8)
torch.set_num_interop_threads(1)
print(torch.get_num_threads(), torch.get_num_interop_threads())
PY
  • Pin processes/threads (advanced, useful on dual-socket or noisy systems):
# Run with CPU affinity to first 8 cores
taskset -c 0-7 python your_script.py
  • NUMA-aware runs (multi-socket boxes):
numactl --cpunodebind=0 --membind=0 python your_script.py
  • Keep arrays contiguous and dtypes consistent (float32 vs float64) to reduce copies:
X = np.ascontiguousarray(X, dtype=np.float32)
  • Batch smartly to fit caches and GPU/CPU memory; tiny batches underutilize hardware, giant batches thrash memory.

5) Use the right accelerator at the right time (CPU vs GPU)

  • CPU first: modern OpenBLAS/MKL with tuned threads is surprisingly strong for classical ML, smaller models, and preprocessing.

  • GPU when you have large, parallel workloads (CNNs, transformers, big matmuls). Installation can be as simple as:

# NVIDIA CUDA-enabled PyTorch
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
# Or CPU-only wheels if no GPU:
pip install --index-url https://download.pytorch.org/whl/cpu torch torchvision torchaudio
  • Try CuPy for NumPy-like GPU acceleration:
pip install cupy-cuda12x
python - <<'PY'
import cupy as cp
x = cp.random.randn(10_000, 10_000, dtype=cp.float32)
y = x @ x.T
cp.cuda.Stream.null.synchronize()
print(y.shape)
PY

Measure end-to-end, including data transfer, to make sure the GPU gives net wins.

Real-world micro-benchmarks to copy-paste

1) Pure Python vs NumPy vectorization:

python - <<'PY'
import numpy as np, time

n = 20_000_000
x = np.random.randn(n).astype(np.float32)

def relu_py(v):
    out = []
    for z in v:
        out.append(z if z > 0 else 0)
    return out

s=time.perf_counter(); relu_py(x); print("python loop:", time.perf_counter()-s, "s")

s=time.perf_counter(); y=np.maximum(x,0); print("numpy vect:", time.perf_counter()-s, "s")
PY

2) Numba speedup:

python - <<'PY'
import numpy as np, time
from numba import njit

@njit(fastmath=True)
def accum(a):
    s=0.0
    for i in range(a.size):
        s += a[i]*a[i]
    return s

x = np.random.randn(50_000_000).astype(np.float32)
s=time.perf_counter(); r=accum(x); print("numba first call (compile):", time.perf_counter()-s, "s")
s=time.perf_counter(); r=accum(x); print("numba second call:", time.perf_counter()-s, "s")
PY

3) Thread tuning for matmul (as above) to find best OPENBLAS_NUM_THREADS.

Common pitfalls checklist

  • Installing NumPy/SciPy before system BLAS dev packages ⇒ slow reference BLAS. Fix by reinstalling with the dev libs present.

  • Too many threads across libraries ⇒ over-subscription. Set one of OPENBLAS_NUM_THREADS, MKL_NUM_THREADS, or OMP_NUM_THREADS; limit others.

  • Pure-Python loops in hot paths ⇒ vectorize or use Numba/Cython.

  • Mixed dtypes (float64 vs float32) ⇒ unnecessary conversions and memory pressure.

  • Measuring only kernel time, not end-to-end latency ⇒ misleading conclusions.

Conclusion and Next Steps (CTA)

You don’t need a new server to go faster—you need a tuned stack.

  • Step 1: Install the optimized math libraries with your package manager (apt/dnf/zypper) and confirm NumPy’s BLAS.

  • Step 2: Profile your code with cProfile, line-profiler, or py-spy to find hotspots.

  • Step 3: Vectorize or JIT the heavy loops; manage threads with OPENBLAS_NUM_THREADS/OMP_NUM_THREADS or torch.set_num_threads.

  • Step 4: Re-run the benchmarks and record gains. Iterate.

Want a quick win? Copy the matmul and Numba snippets, try 1, 2, 4, 8, 16 threads, and pick the best setting for your shell profile. Then profile your workload and knock out the biggest hotspot today.

Related posts