1. Numba là gì

Numba là một trình biên dịch dành cho các hàm Python có khả năng xử lý dữ liệu số và mảng một cách hiệu quả. Numba giúp các chương trình viết bằng Python đạt được tốc độ tương đương với các ngôn ngữ biên dịch khác như C, C++, hoặc Fortran mà không cần phải viết lại mã nguồn hay sử dụng trình thông dịch Python khác.

Để đạt được điều này, Numba chuyển đổi mã Python thành mã máy tối ưu thông qua LLVM, một framework biên dịch mã mở nổi tiếng. Chỉ cần vài thay đổi nhỏ trong mã nguồn, các đoạn code Python có khối lượng tính toán lớn liên quan đến mảng và số sẽ được tối ưu hóa với hiệu suất cao.

Những tính năng chính của Numba bao gồm:

• Sinh mã ngay tại thời điểm chạy (runtime) hoặc khi import.
• Tạo mã native cho CPU và GPU.
• Tích hợp dễ dàng với các thư viện tính toán khoa học trong Python như NumPy.

2. Cài đặt Numba

Để cài đặt Numba, bạn có thể sử dụng Conda - một công cụ quản lý gói và môi trường được viết bằng Python, hỗ trợ nhiều ngôn ngữ khác nhau. Công cụ này cho phép tạo và quản lý các môi trường cũng như cài đặt các gói cần thiết một cách dễ dàng.

$ conda install numba

Ngoài ra, bạn có thể cài đặt Numba bằng pip:

$ pip install numba

Lệnh trên sẽ cài đặt tất cả các dependencies cần thiết, bao gồm cả LLVM (thông qua llvmlite). Sau khi cài đặt thành công, bạn có thể kiểm tra phiên bản Numba bằng đoạn mã sau:

$ python
Python 3.12.2 (main, Apr  6 2024, 18:55:05) [GCC 11.4.0] on linux
>>> import numba
>>> numba.__version__
'0.60.0'

Hoặc bằng cách sử dụng CLI của Numba:

$ numba -s

3. Biên dịch Python code với Numba

Numba cung cấp một số công cụ tiện ích cho việc sinh mã, nhưng chức năng chính nằm ở decorator @numba.jit. Bằng cách sử dụng decorator này, bạn có thể đánh dấu một hàm Python mà bạn muốn tối ưu hóa bằng trình biên dịch JIT của Numba.

Biên dịch lười (Lazy Compilation)

Cách gọi decorator @numba.jit nên được sử dụng để cho phép Numba quyết định thời điểm và cách thức tối ưu:

from numba import jit

@jit
def f(x, y):
    return x + y

Quá trình biên dịch sẽ diễn ra khi hàm được gọi lần đầu. Numba sẽ tự động suy diễn kiểu dữ liệu và sinh mã trên cơ sở thông tin đó.

Biên dịch ngay (Eager Compilation)

Ngoài ra, bạn có thể định nghĩa rõ chữ ký của hàm cần tối ưu. Khi đó, Numba sẽ chỉ sinh mã cho các kiểu dữ liệu đã được xác định.

from numba import jit, int32

@jit(int32(int32, int32))
def f(x, y):
    return x + y

Gọi và inline hàm khác

Các hàm được biên dịch bởi Numba có thể gọi nhau, và mã của các hàm này có thể được inline vào mã native tùy thuộc vào chiến lược tối ưu của trình biên dịch.

@jit
def square(x):
    return x ** 2

@jit
def hypot(x, y):
    return math.sqrt(square(x) + square(y))

Tùy chọn biên dịch

Có một số tham số có thể truyền vào decorator @numba.jit để điều chỉnh cách hoạt động của trình biên dịch:

nopython

Mặc định, Numba biên dịch ở chế độ nopython và sẽ sinh mã nhanh hơn so với chế độ object:

@jit
def f(x, y):
    return x + y

nogil

Khi tối ưu hóa mã Python sang mã native, nếu không cần sử dụng GIL, bạn có thể thêm tham số nogil:

@jit(nogil=True)
def f(x, y):
    return x + y

cache

Để tránh việc biên dịch lại mỗi lần chương trình chạy, Numba có thể lưu trữ kết quả biên dịch vào cache:

@jit(cache=True)
def f(x, y):
    return x + y

parallel

Cho phép tự động song song hóa các thao tác có thể thực hiện song song.

@jit(nopython=True, parallel=True)
def f(x, y):
    return x + y

4. Hiệu năng của Numba

Để so sánh hiệu năng của Numba và Python thuần (sử dụng thư viện NumPy), chúng ta sẽ kiểm tra thuật toán tích chập hai chiều (convolution), là nền tảng của các mạng deep learning trong xử lý hình ảnh như ResNet, ConvNet, VGG, v.v.

So sánh hiệu năng

Dưới đây là cách implement cho cả NumPy và Numba:

# Phép tích chập với NumPy

def convolve_numpy(img, kernel):
    # ...
    return out

Đối với Numba, bạn chỉ cần thêm một decorator:

@numba.jit(nopython=True, parallel=True)
def convolve_numba(img, kernel):
    # ...
    return out

Sử dụng timeit để kiểm tra hiệu năng:

in_img = np.random.random((100, 100))
kernel = np.random.random((10, 10))
print(f"numpy: {timeit.timeit(lambda: convolve_numpy(in_img, kernel), number=100)*(10**-3)}")
print(f"numba: {timeit.timeit(lambda: convolve_numba(in_img, kernel), number=100)*(10**-3)}")

Kết quả cho thấy rằng mã Numba nhanh hơn tới gần 60 lần:

numpy: 58.314
numba: 1.220

5. Kết luận

Numba cho phép ta tăng tốc mã Python một cách nhanh chóng và dễ dàng mà không cần thay đổi interpreter hay mã nguồn quá nhiều. Thêm vào đó, Numba còn hỗ trợ sinh mã native CUDA, giúp tận dụng tối đa khả năng xử lý của GPU NVIDIA. Bạn có thể tham khảo tài liệu của Numba để áp dụng vào dự án của mình.
source: viblo