Google JAX hoặc Just After Execution là một khung do Google phát triển để tăng tốc các tác vụ học máy.
Bạn có thể coi nó như một thư viện dành cho Python để giúp bạn hoàn thành công việc nhanh hơn, tính toán khoa học, chuyển đổi hàm, học sâu, mạng lưới thần kinh, v.v.
Giới thiệu về Google JAX
Gói tính toán cơ bản nhất trong Python là gói NumPy, chứa tất cả các hàm như tập hợp, phép toán vectơ, đại số tuyến tính, thao tác với mảng và ma trận n chiều và nhiều hàm nâng cao khác.
Điều gì sẽ xảy ra nếu chúng ta có thể thực hiện các phép tính NumPy nhanh hơn nữa – đặc biệt là với các bộ dữ liệu khổng lồ?
Chúng tôi có thứ gì đó có thể hoạt động tốt như nhau trên các loại CPU khác nhau như GPU hoặc TPU mà không có bất kỳ thay đổi mã nào không?
Điều gì sẽ xảy ra nếu hệ thống có thể thực hiện chuyển đổi các chức năng có thể kết hợp một cách tự động và hiệu quả hơn?
Google JAX là một thư viện (hoặc khung như Wikipedia nói) thực hiện điều đó và có thể hơn thế nữa. Nó được xây dựng để tối ưu hóa hiệu suất và thực hiện các nhiệm vụ học máy (ML) và học sâu một cách hiệu quả. Google JAX cung cấp các tính năng chuyển đổi sau đây, giúp nó khác biệt với các thư viện ML khác và trợ giúp tính toán khoa học nâng cao cho mạng lưới thần kinh và học sâu:
- tự động phân biệt
- Vector hóa tự động
- song song tự động
- Biên dịch đúng lúc (JIT)
Các tính năng Google JAX độc đáo
Tất cả các phép biến đổi đều sử dụng XLA (Đại số tuyến tính tăng tốc) để có hiệu suất cao hơn và tối ưu hóa bộ nhớ. XLA là một công cụ biên dịch tối ưu hóa dành riêng cho miền, thực hiện đại số tuyến tính và tăng tốc các mô hình TensorFlow. Sử dụng XLA trong mã python không yêu cầu bất kỳ thay đổi mã quan trọng nào!
Chúng ta hãy xem xét chi tiết từng tính năng này.
Các tính năng của Google JAX
Google JAX đi kèm với các tính năng chuyển đổi quan trọng mà bạn có thể soạn thảo để cải thiện hiệu suất và thực hiện các tác vụ học sâu hiệu quả hơn. Ví dụ: tự động phân biệt để lấy độ dốc của hàm và tìm đạo hàm của bất kỳ thứ tự nào. Tương tự như vậy tự động song song và JIT để chạy song song nhiều tác vụ. Những biến đổi này rất quan trọng đối với các ứng dụng như chế tạo rô-bốt, chơi game và thậm chí là nghiên cứu.
Hàm biến đổi có thể kết hợp là một hàm thuần túy biến đổi một tập hợp dữ liệu thành một dạng khác. Chúng được gọi là có thể kết hợp được vì chúng độc lập (nghĩa là các hàm này không phụ thuộc vào phần còn lại của chương trình) và chúng không trạng thái (tức là cùng một đầu vào sẽ luôn tạo ra cùng một đầu ra).
Y(x) = T: (f(x))
Trong phương trình trên, f(x) là hàm nguyên thủy mà phép biến đổi được áp dụng. Y(x) là hàm kết quả sau khi áp dụng phép biến đổi.
Ví dụ: nếu bạn có một hàm gọi là “total_bill_amt” và bạn muốn kết quả là một phép biến đổi hàm, thì bạn có thể sử dụng bất kỳ phép biến đổi nào bạn muốn, chẳng hạn như độ dốc (grad):
grad_total_bill = grad(total_bill_amt)
Bằng cách biến đổi các hàm số với các hàm như grad(), chúng ta có thể dễ dàng nhận được các đạo hàm bậc cao hơn của chúng, chúng ta có thể sử dụng rộng rãi trong các thuật toán tối ưu hóa học sâu như giảm độ dốc, làm cho các thuật toán nhanh hơn và hiệu quả hơn. Tương tự, sử dụng jit(), chúng ta có thể biên dịch các chương trình Python ngay lập tức (lười biếng).
#1. tự động phân biệt
Python sử dụng autograd để tự động phân biệt giữa NumPy và mã Python gốc. JAX sử dụng phiên bản sửa đổi của autograd (tức là grad) và kết hợp XLA (Đại số tuyến tính tăng tốc) để tự động phân biệt và tìm các đạo hàm theo bất kỳ thứ tự nào cho GPU (đơn vị xử lý đồ họa) và TPU (đơn vị xử lý tensor).]
Lưu ý nhanh về TPU, GPU và CPU: CPU hoặc bộ xử lý trung tâm quản lý tất cả các hoạt động trên máy tính. GPU là một bộ xử lý bổ sung giúp tăng sức mạnh tính toán và cho phép bạn thực hiện các thao tác nâng cao. TPU là một đơn vị mạnh mẽ được phát triển đặc biệt cho khối lượng công việc nặng và phức tạp như trí tuệ nhân tạo và thuật toán học sâu.
Giống như hàm autograd, có thể phân biệt giữa các vòng lặp, đệ quy, nhánh, v.v., JAX sử dụng hàm grad() cho các gradient ở chế độ đảo ngược (lan truyền ngược). Ngoài ra, chúng ta có thể thay đổi chức năng theo bất kỳ thứ tự nào bằng cách sử dụng phân cấp:
mưa đá(grad(grad(sin θ)))) (1,0)
Tự động phân biệt bậc cao hơn
Như đã đề cập trước đó, grad rất hữu ích trong việc tìm đạo hàm riêng của các hàm. Chúng ta có thể sử dụng đạo hàm riêng để tính toán độ dốc của hàm chi phí trên các tham số mạng thần kinh trong học sâu để giảm thiểu tổn thất.
Tính đạo hàm riêng
Giả sử một hàm có nhiều biến x, y và z. Việc tìm đạo hàm của một biến bằng cách giữ nguyên các biến khác được gọi là đạo hàm riêng. Giả sử chúng ta có một chức năng
f(x,y,z) = x + 2y + z2
Một ví dụ cho thấy đạo hàm riêng
Đạo hàm riêng của x sẽ là ∂f/∂x, cho chúng ta biết hàm thay đổi như thế nào đối với một biến khi các biến khác không đổi. Nếu làm thủ công, chúng ta cần viết một chương trình vi phân, áp dụng nó cho từng biến, sau đó tính độ dốc giảm dần. Điều này sẽ trở thành một vấn đề phức tạp và tốn thời gian cho nhiều biến số.
Vi phân tự động chia một hàm thành một tập hợp các phép toán cơ bản như +, -, *, / hoặc sin, cos, tan, exp, v.v., sau đó áp dụng quy tắc dây chuyền để tính đạo hàm. Chúng ta có thể làm điều đó cả về phía trước và phía sau.
Nó không phải là nó! Tất cả những tính toán này diễn ra rất nhanh (tốt, hãy nghĩ về một triệu phép tính như trên và thời gian có thể mất!). XLA quan tâm đến tốc độ và hiệu quả.
#2. gia tốc đại số tuyến tính
Hãy lấy phương trình trước. Nếu không có XLA, quá trình tính toán sẽ yêu cầu ba (hoặc nhiều hơn) hạt nhân, với mỗi hạt nhân thực hiện một nhiệm vụ nhỏ hơn. Ví dụ,
Hạt nhân k1 -> x * 2y (nhân)
k2 -> x * 2y + z (bổ sung)
d3 -> Giảm
Nếu cùng một nhiệm vụ được thực hiện bởi XLA, thì một hạt nhân duy nhất sẽ xử lý tất cả các hoạt động trung gian bằng cách kết hợp chúng. Kết quả trung gian của các hoạt động nguyên thủy được truyền trực tuyến thay vì lưu trữ trong bộ nhớ, giúp tiết kiệm bộ nhớ và tăng tốc độ.
#3. Biên dịch đúng lúc
JAX nội bộ sử dụng trình biên dịch XLA để tăng tốc độ thực thi. XLA có thể tăng tốc độ CPU, GPU và TPU. Tất cả điều này được thực hiện bằng cách thực thi mã JIT. Để sử dụng cái này, chúng ta có thể sử dụng jit thông qua nhập:
from jax import jit def my_function(x): …………some lines of code my_function_jit = jit(my_function)
Một cách khác là trang trí jit trên định nghĩa chức năng:
@jit def my_function(x): …………some lines of code
Mã này nhanh hơn nhiều vì biến đổi sẽ trả lại phiên bản mã đã biên dịch cho người gọi thay vì sử dụng trình thông dịch python. Điều này đặc biệt hữu ích cho các đầu vào vectơ như mảng và ma trận.
Điều tương tự cũng áp dụng cho tất cả các chức năng python hiện có. Ví dụ: các hàm NumPy. Trong trường hợp này, chúng ta nên nhập jax.numpy dưới dạng jnp thay vì NumPy:
import jax import jax.numpy as jnp x = jnp.array([[1,2,3,4], [5,6,7,8]])
Khi bạn làm điều này, đối tượng mảng JAX bên dưới có tên là DeviceArray sẽ thay thế mảng NumPy tiêu chuẩn. DeviceArray lười biếng – các giá trị được lưu trữ trong bộ tăng tốc cho đến khi chúng cần thiết. Điều này cũng có nghĩa là chương trình JAX không đợi kết quả trả về cho người gọi (Python) và do đó xảy ra việc gửi đi không đồng bộ.
#4. Vector hóa tự động (vmap)
Trong một thế giới máy học điển hình, chúng tôi có bộ dữ liệu với một triệu điểm dữ liệu trở lên. Chúng tôi rất có thể sẽ thực hiện một số tính toán hoặc thao tác trên từng hoặc hầu hết các điểm dữ liệu này – đây là một công việc rất tốn thời gian và bộ nhớ! Ví dụ: nếu bạn muốn tìm bình phương của từng điểm dữ liệu trong tập dữ liệu, điều đầu tiên bạn nghĩ đến là tạo một vòng lặp và lấy từng bình phương một – argh!
Nếu chúng ta tạo các điểm này dưới dạng vectơ, chúng ta có thể thực hiện tất cả các ô vuông trong một lần bằng cách thực hiện các thao tác vectơ hoặc ma trận trên các điểm dữ liệu bằng NumPy yêu thích của mình. Và nếu chương trình của bạn có thể tự động làm điều đó – bạn còn đòi hỏi gì hơn nữa? Đây chính xác là những gì JAX làm! Nó có thể tự động véc tơ hóa tất cả các điểm dữ liệu, vì vậy bạn có thể dễ dàng thực hiện các thao tác tùy ý trên chúng – làm cho các thuật toán nhanh hơn và hiệu quả hơn nhiều.
JAX sử dụng chức năng vmap để vector hóa tự động. Hãy xem xét bảng sau:
x = jnp.array([1,2,3,4,5,6,7,8,9,10]) y = jnp.square(x)
Bằng cách chỉ làm như trên, phương pháp bậc hai sẽ được thực hiện cho từng điểm trong mảng. Nhưng nếu bạn làm như sau:
vmap(jnp.square(x))
Phương thức bình phương sẽ chỉ được thực thi một lần vì các điểm dữ liệu hiện được tự động véc tơ hóa bằng phương thức vmap trước khi thực thi hàm và vòng lặp được đẩy xuống mức hoạt động cơ bản – dẫn đến phép nhân ma trận thay vì phép nhân vô hướng để có hiệu suất tốt hơn.
#5. Lập trình SPMD (pmap)
SPMD – hoặc lập trình nhiều dữ liệu trong một chương trình – là điều cần thiết trong bối cảnh học sâu – thường các chức năng giống nhau được áp dụng cho các bộ dữ liệu khác nhau nằm trên nhiều GPU hoặc TPU. JAX có một tính năng gọi là máy bơm cho phép phát triển song song trên nhiều GPU hoặc bất kỳ máy gia tốc nào. Cũng giống như JIT, các chương trình sử dụng pmap sẽ được XLA biên dịch và thực thi đồng thời trên các hệ thống. Tính song song tự động này hoạt động cho cả tính toán tiến và lùi.
Cách pmap hoạt động
Chúng ta cũng có thể áp dụng nhiều phép biến đổi cùng một lúc theo bất kỳ thứ tự nào cho bất kỳ chức năng nào như:
pmap(vmap(jit(grad(f(x))))))
Nhiều sáng tác sáng tác
Giới hạn JAX của Google
Các nhà phát triển Google JAX đã nghĩ rất kỹ về việc tăng tốc các thuật toán học sâu bằng cách giới thiệu tất cả các biến đổi tuyệt vời này. Các gói và chức năng tính toán khoa học được mô phỏng theo NumPy, vì vậy bạn không phải lo lắng về đường cong học tập. Tuy nhiên, JAX có những hạn chế sau:
- Google JAX vẫn đang trong giai đoạn phát triển ban đầu và trong khi mục tiêu chính của nó là tối ưu hóa hiệu suất, nó không mang lại nhiều lợi ích cho việc tính toán CPU. NumPy dường như hoạt động tốt hơn và việc sử dụng JAX chỉ có thể tăng chi phí hoạt động.
- JAX vẫn đang trong giai đoạn nghiên cứu hoặc giai đoạn đầu và cần cải tiến nhiều hơn để đáp ứng các tiêu chuẩn cơ sở hạ tầng của các khung như TensorFlow, được thiết lập lâu hơn và có nhiều mô hình được xác định trước, dự án nguồn mở và tài liệu học tập.
- Hiện tại, JAX không hỗ trợ hệ điều hành Windows – một máy ảo là cần thiết cho hoạt động của nó.
- JAX chỉ hoạt động trên các chức năng thuần túy – những chức năng không có tác dụng phụ. Đối với các chức năng có tác dụng phụ, JAX có thể không phải là một giải pháp tốt.
Làm cách nào để cài đặt JAX trong môi trường Python của bạn?
Nếu bạn có thiết lập python trên hệ thống của mình và muốn chạy JAX trên máy cục bộ (CPU), hãy sử dụng các lệnh sau:
pip install --upgrade pip pip install --upgrade "jax[cpu]"
Nếu bạn muốn chạy Google JAX trên GPU hoặc TPU, hãy làm theo hướng dẫn được cung cấp trên trang GitHub JAX. Để thiết lập Python, hãy truy cập trang Tải xuống Python chính thức.
Đăng kí
Google JAX là công cụ tuyệt vời để viết các thuật toán nghiên cứu, robot và học sâu mạnh mẽ. Bất chấp những hạn chế của nó, nó được sử dụng rộng rãi với các framework khác như Haiku, Flax và nhiều framework khác. Bạn sẽ có thể đánh giá cao những gì JAX thực hiện khi chạy các chương trình và xem sự khác biệt về thời gian thực thi mã khi có và không có JAX. Bạn có thể bắt đầu bằng cách đọc tài liệu Google JAX chính thức khá rộng rãi.
