Numba - Python on steroid
1. Numba là gì
Numba là trình biên dịch dành cho các hàm Python thực thi trên dữ liệu dạng số và mảng. Nó cho phép viết các chương trình thuần Python mà lại mang tốc độ ngang ngửa các ngôn ngữ biên dịch khác.
Numba thực hiện điều đó bằng cách tạo ra mã máy tối ưu từ code Python, sử dụng LLVM, một framework biên dịch open-source nổi tiếng. Chỉ với một vài thay đổi nhỏ đơn giản tới codebase, các đoạn code Python xử lý nặng dữ liệu kiểu mảng và số sẽ được tối ưu với tốc độ ngang ngửa các ngôn ngữ như C, C++ hay Fortran, mà không cần sử dụng ngôn ngữ khác hay trình thông dịch Python khác như PyPy.
Những tính năng chính của numba là:
- Sinh code ngay lập tức (runtime hoặc khi import)
- Sinh code native cho CPU và GPU
- Tích hợp với các thư viện tính toán khoa học của Python như Numpy
2. Cài đặt Numba
Chúng ta có thể cài đặt Numba thông qua Conda, một package manager và environment manager được viết bằng python và có thể sử dụng nhiều ngôn ngữ khác nhau. Nó có thể:
- Tạo và quản lý các môi trường
- Tìm kiếm, cài đặt packages vào một môi trường có sẵn giúp dễ dàng quản lý, kiểm soát các packages
$ conda install numba
Ngoài ra chúng ta có thể cài đặt numba bằng pip
$ pip install numba
Câu lệnh sẽ cài toàn bộ dependencies cần thiết, bao gồm cả LLVM (thông qua llvmlite)
Sau khi cài đặt, chúng ta có thể kiểm tra xem numba đã được cài đặt thành công hay chưa bằng cách import thẳng numba từ python
$ python
Python 3.12.2 (main, Apr 6 2024, 18:55:05) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import numba
>>> numba.__version__
'0.60.0'
Hoặc bằng cách xử dụng CLI của numba
$ numba -s
System info:
--------------------------------------------------------------------------------
__Time Stamp__
Report started (local time) : 2024-09-30 10:57:37.293762
Running time (s) : 0.862805
__Hardware Information__
Machine : x86_64
CPU Name : znver3
CPU Count : 12
Number of accessible CPUs : 12
List of accessible CPUs cores : 0-11
CFS Restrictions (CPUs worth of runtime) : None
3. Biên dịch Python code với Numba
Numba cung cấp một vài hàm utilities dùng cho việc sinh mã, tuy nhiên chức năng chính nằm ở decorator @numba.jit. Bằng cách sử dụng decorator này, chúng ta có thể đánh dấu một hàm Python cần tối ưu bằng trình biên dịch JIT của Numba. Cách chúng ta gọi decorator này sẽ thay đổi cách thức cũng như các options của compiler.
Lazy compilation
Cách gọi @numba.jit decorator được khuyến khích là cho phép Numba quyết định khi nào nên tối ưu và nên tối ưu như thế nào:
from numba import jit
@jit
def f(x, y):
# A somewhat trivial example
return x + y
Quá trình biên dịch sẽ được đẩy lùi về đến khi hàm trên được gọi lần đầu. Numba sẽ suy diễn kiểu dữ liệu tại thời điểm được gọi và sinh code dựa vào thông tin được suy diễn. Điều này cũng cho phép Numba biên dịch các trường hợp đặc biệt dựa trên kiểu dữ liệu đầu vào, ví dụ khi gọi hàm f() với 2 kiểu dữ liệu khác nhau, integer và complex, sẽ cho ra kết quả khác nhau:
>>> f(1, 2)
3
>>> f(1j, 2)
(2+1j)
Eager compilation
Bên cạnh đó chúng ta có thể định nghĩa signature của hàm cần tối ưu, Numba lúc này sẽ chỉ sinh code dành cho kiểu dữ liệu đã được định nghĩa mà thôi.
from numba import jit, int32
@jit(int32(int32, int32))
def f(x, y):
# A somewhat trivial example
return x + y
Ở đây int32(int32, int32) là signature của hàm f(x, y). Cách này cho phép chúng ta kiểm soát được kiểu dữ liệu đầu vào của hàm tối ưu theo ý thích (ví dụ chỉ thực thi trên kiểu float32).
Nếu ta bỏ đi kiểu dữ liệu trả về, (int32, int32) thay vì int32(int32, int32), Numba sẽ tự động suy diễn kiểu trả về phù hợp.
Gọi và inline các hàm khác
Các hàm đẫ được biên dịch bởi Numba có thể gọi các hàm đã được biên dịch khác, thậm chí code gọi hàm có thể được inline trong native code, tuỳ vào chiến thuật 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))
Decorator @numba.jit phải được thêm ở tất cả các hàm được gọi, nếu không Numba sẽ sinh ra code chậm hơn rất nhiều.
Tuỳ chọn biên dịch
Một vài tham số có thể truyền vào decorator @numba.jit nhằm thay đổi cách trình biên dịch hoạt động:
nopython
Numba có hai chế độ biên dịch chính: nopython và object. Mặc định Numba sẽ biên dịch ở chế độ nopython và sẽ sinh code nhanh hơn so với object.
@jit # same as @jit(nopython=True) or @njit since Numba 0.59
def f(x, y):
return x + y
nogil
Bất cứ khi nào Numba tối ưu Python code sang native code với kiểu dữ liệu và biến native (thay vì kiểu Python objects), việc sử dụng Python GIL là không còn cần thiết nữa. Numba sẽ không còn giữ GIL khi hàm đã biên dịch được gọi khi flag trên được bật.
@jit(nogil=True)
def f(x, y):
return x + y
Đoạn code được gắn flag trên sẽ chạy song song với các threads đang chạy code Python và Numba khác. Điều này sẽ không xảy ra nếu ta chạy ở object. Lưu ý, khi đặt nogil=true, chúng ta phải chú ý đến những vấn đề xảy ra khi chạy song song như race conditions, synchronization, ...)
cache
Để tránh biên dịch lại mỗi lần chúng ta chạy chương trình, Numba sẽ viết kết quả biên dịch ra một file cache riêng:
@jit(cache=True)
def f(x, y):
return x + y
Lưu ý, cache có các vấn đề sau:
- Cache không được lưu the từng hàm mà lưu tại hàm jit chính, tất cả các hàm được gọi bởi hàm jit đó đều sẽ được cache trong cùng một file
- Không phát hiện sự thay đổi ở các hàm import từ các module khác. Ví dụ khi ta gọi một hàm từ module khác, và hàm đó có thay đổi, cache sẽ không được cập nhật và sẽ ảnh hưởng đến kết quả cuối cùng
- Biến toàn cục được coi là hằng số. Cache sẽ nhớ giá trị của biến toàn cục tại compile time. Khi load cache, Numba sẽ không gán lại biến toàn cục với giá trị mới
parallel
Cho phép tự động song song hoá (và các tối ưu liên quan) dành cho các operations có khả năng song song. Chi tiết các toán tử hỗ trợ song song hoá Danh sách các toán tử hỗ trợ:
@jit(nopython=True, parallel=True)
def f(x, y):
return x +
4. Hiệu năng của Numba
Chúng ta sẽ cùng test thử hiệu năng của Numba so với Python thuần (sử dụng thư viện numpy). Thuật toán được chọn là phép tích chập hai chiều - convolution. Thuật toán này là nền tảng của các mạng deep dùng trong xử lý hình ảnh như ResNet, ConvNet, VGG, v.v. Cả hai đều sử dụng cùng một implementation, khác biệt duy nhất là một bên sử dụng jit compilation của Numba, bên còn lại là Cython. Source code phép tính (được lấy từ bài blog):
def convolve_numpy(img, kernel):
img_height = img.shape[0]
img_width = img.shape[1]
kernel_height = kernel.shape[0]
kernel_width = kernel.shape[1]
H = (kernel_height - 1) // 2
W = (kernel_width - 1) // 2
out = np.zeros((img_height, img_width))
for i in np.arange(H, img_height - H):
for j in np.arange(W, img_width - W):
sum = 0
for k in np.arange(-H, H + 1):
for l in np.arange(-W, W + 1):
a = img[i + k, j + l]
w = kernel[H + k, W + l]
sum += w * a
out[i, j] = sum
return out
Đối với Numba, ta chỉ việt thêm decorator như sau là xong:
@numba.jit(nopython=True, parallel=True)
def convolve_numba(img, kernel):
img_height = img.shape[0]
img_width = img.shape[1]
kernel_height = kernel.shape[0]
kernel_width = kernel.shape[1]
H = (kernel_height - 1) // 2
W = (kernel_width - 1) // 2
out = np.zeros((img_height, img_width))
for i in np.arange(H, img_height - H):
for j in np.arange(W, img_width - W):
sum = 0
for k in np.arange(-H, H + 1):
for l in np.arange(-W, W + 1):
a = img[i + k, j + l]
w = kernel[H + k, W + l]
sum += w * a
out[i, j] = sum
return out
Ta sử dụng stdlib timeit đi kèm với python để test hiệu năng, thực hiện chạy 100 lần:
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ả chạy:
numpy: 58.314815886998986
numba: 1.2203670519993466
Có thể thấy chỉ với 1 dòng code đơn giản, thuật toán của chúng ta đã giảm đến gần 60 lần tốc độ thực thi.
5. Kết luận
Numba cho phép ta speed-up code Python một cách nhanh chóng mà dễ dàng mà không cần phải thay đổi interpreter, không phải thay đổi quá nhiều code base.
Ngoài ra Numba còn có thể hỗ trợ sinh native code CUDA, cho phép ta tận dụng khả năng xử lý đa luồng mạnh mẽ của NVIDIA GPU để tăng tốc độ thực thi của thuật toán. Mọi người có thể tham khảo documentations chi tiết của Numba để áp dụng vào dự án của bản thân.
Tài liệu tham khảo
All rights reserved
