Posted inPython

Tăng tốc Python gấp 100 lần với Numba: Giải cứu vòng lặp khỏi tốc độ “rùa bò”

Python tutorial - IT technology blog
Python tutorial - IT technology blog

Nỗi ám ảnh mang tên “vòng lặp for” trong Python

Mình thường dùng Python để viết tool automation, từ script deploy đến hệ thống cảnh báo monitoring. Python rất tuyệt vì cú pháp trong sáng, dễ đọc. Tuy nhiên, điểm yếu lớn nhất của nó lộ ra ngay khi bạn cần tính toán số học nặng hoặc duyệt hàng triệu bản ghi: Tốc độ.

Nếu bạn từng viết vòng lặp lồng nhau để xử lý ma trận, bạn sẽ thấy CPU nhảy lên 100% nhưng chờ mãi không xong kết quả. Trước đây, cách duy nhất là viết lại đoạn code đó bằng C++ hoặc Cython. Nhưng thú thật, chẳng mấy ai muốn vật lộn với quản lý bộ nhớ thủ công hay con trỏ phức tạp. Đó là lý do mình chọn Numba cho hệ thống production suốt 6 tháng qua.

Numba là gì và tại sao nó lại nhanh đến thế?

Nói đơn giản, Numba là một trình biên dịch JIT (Just-In-Time). Thay vì để Python interpreter “vừa đọc vừa chạy” từng dòng code chậm chạp, Numba sẽ dịch thẳng code của bạn sang mã máy (machine code) ngay khi thực thi. Nó tận dụng hạ tầng LLVM để tối ưu hóa hiệu năng cực sâu.

Sức mạnh của Numba phát huy tốt nhất với các tác vụ xử lý số liệu và mảng NumPy. Khi bạn thêm decorator @jit, Numba sẽ phân tích kiểu dữ liệu và tối ưu hàm đó tương tự cách trình biên dịch C hay Fortran vận hành.

Bắt đầu thực hành: Cài đặt và ví dụ đầu tiên

Cài đặt Numba rất nhanh qua pip. Bạn nên dùng môi trường ảo (venv) để giữ hệ thống sạch sẽ.

pip install numba numpy

Hãy thử bài toán kinh điển: Tính tổng bình phương của 10 triệu số. Chúng ta sẽ so sánh Python thuần và Numba để thấy sự khác biệt.

import time
import numpy as np
from numba import njit

# Hàm Python thuần
def sum_sq_python(n):
    result = 0
    for i in range(n):
        result += i**2
    return result

# Hàm dùng Numba - thêm đúng 1 dòng này
@njit
def sum_sq_numba(n):
    result = 0
    for i in range(n):
        result += i**2
    return result

# Test với 10 triệu phần tử
n = 10_000_000

start = time.time()
sum_sq_python(n)
print(f"Python thuần: {time.time() - start:.4f}s")

# Chạy lần 1: Numba cần thời gian để biên dịch
start = time.time()
sum_sq_numba(n)
print(f"Numba lần 1 (compile): {time.time() - start:.4f}s")

# Chạy lần 2: Đã có sẵn mã máy trong cache
start = time.time()
sum_sq_numba(n)
print(f"Numba lần 2 (chạy ngay): {time.time() - start:.4f}s")

Kết quả thực tế trên máy mình rất sốc. Python thuần mất khoảng 0.72 giây. Numba ở lần chạy thứ hai chỉ tốn 0.00001 giây. Tốc độ đã tăng hơn 70.000 lần!

Chế độ nopython=True: Bí mật của hiệu suất

Trong ví dụ trên, mình dùng @njit, thực chất là viết tắt của @jit(nopython=True). Đây là chế độ tốt nhất. Ở chế độ này, Numba biên dịch toàn bộ hàm mà không cần sự can thiệp của Python interpreter.

Nếu hàm chứa kiểu dữ liệu lạ, Numba sẽ báo lỗi ngay. Điều này tốt hơn việc âm thầm chạy chậm ở chế độ “object mode”. Mẹo của mình là luôn dùng @njit. Nếu lỗi, hãy sửa code cho chuẩn thay vì chấp nhận tốc độ rùa bò.

NumPy và Numba: Cặp bài trùng hoàn hảo

Numba sinh ra để dành cho NumPy. Đôi khi, Numba còn tối ưu vòng lặp trên mảng tốt hơn cả các hàm vectorization có sẵn của NumPy trong các công thức tùy chỉnh phức tạp.

@njit
def process_array(arr):
    rows, cols = arr.shape
    out = np.empty_like(arr)
    for i in range(rows):
        for j in range(cols):
            # Công thức phức tạp khó dùng vectorization
            out[i, j] = np.sin(arr[i, j]) + np.cos(arr[i, j])
    return out

data = np.random.randn(5000, 5000)

Trước đây, chúng ta thường được khuyên “tránh dùng vòng lặp for trong Python”. Với Numba, triết lý này đảo ngược hoàn toàn. Bạn cứ viết vòng lặp cho dễ hiểu, Numba sẽ lo phần tốc độ.

Tận dụng đa nhân CPU với Parallel=True

Khả năng tự động song song hóa (auto-parallelization) là tính năng mình cực kỳ thích. Chỉ cần thêm parallel=True và đổi range thành prange, Numba sẽ tự chia nhỏ công việc lên toàn bộ các nhân CPU bạn có.

from numba import njit, prange

@njit(parallel=True)
def parallel_sum(arr):
    s = 0
    for i in prange(arr.shape[0]):
        s += arr[i]
    return s

Lưu ý: Chỉ dùng parallel=True cho các mảng dữ liệu đủ lớn. Với tác vụ nhỏ, chi phí quản lý luồng (overhead) sẽ làm code chạy chậm hơn bình thường.

Những lưu ý thực tế sau 6 tháng sử dụng

Dù mạnh mẽ, Numba không phải là công cụ vạn năng. Dưới đây là 4 bài học mình rút ra:

  • Hỗ trợ thư viện hạn chế: Numba không hiểu các object phức tạp từ Pandas hay Scikit-learn. Bạn cần chuyển dữ liệu về mảng NumPy trước khi xử lý.
  • Độ trễ lần đầu: Lần đầu gọi hàm sẽ chậm vì Numba phải compile. Nếu script của bạn chỉ chạy một lần rồi tắt (như tool CLI), tổng thời gian có thể không cải thiện nhiều.
  • Kỷ luật về kiểu dữ liệu: Vì biên dịch xuống mã máy, Numba rất khắt khe. Một biến không thể vừa là số nguyên, vừa là chuỗi được.
  • Khó debug: Việc đặt breakpoint trong hàm đã compile rất khó khăn. Hãy debug logic bằng Python thuần trước khi gắn @njit.

Lời kết

Numba giúp mình giải quyết bài toán hiệu năng mà không cần học ngôn ngữ mới hay thay đổi cấu trúc code. Nếu script xử lý số liệu của bạn đang chậm, hãy thử gắn @njit vào những hàm tốn thời gian nhất. Tốc độ của C nằm ngay trong file .py của bạn, ngại gì không thử?

Share:
Tags: