Penzai - Bộ công cụ JAX để xây dựng, chỉnh sửa và trực quan hóa mạng nơ-ron

• Thư viện JAX do Google DeepMind công bố, dùng để viết mô hình dưới dạng cấu trúc dữ liệu Pytree mang tính hàm và dễ đọc
• Tập trung vào việc giúp có thể làm việc với mô hình một cách dễ dàng ngay cả sau khi mô hình đã được huấn luyện
  • Phù hợp cho nghiên cứu liên quan đến reverse engineering hoặc loại bỏ các thành phần mô hình, cũng như kiểm tra và khảo sát kích hoạt bên trong, phẫu thuật mô hình, gỡ lỗi kiến trúc, v.v.
  • Cung cấp các công cụ để trực quan hóa, chỉnh sửa và phân tích mô hình, nên cũng có thể dùng đơn giản cho mục đích xây dựng và huấn luyện mô hình
• Gồm một bộ công cụ mô-đun được thiết kế cùng nhau nhưng mỗi mô-đun đều có thể dùng độc lập

Các công cụ đi kèm

penzai.nn (pz.nn)

• Thư viện mạng nơ-ron dạng declarative dựa trên combinator, là lựa chọn thay thế cho các thư viện mạng nơ-ron khác như Flax, Haiku, Keras, Equinox
• Phơi bày toàn bộ cấu trúc của forward pass của mô hình trong model pytree, nên có thể xem mọi thao tác của mô hình bằng pretty printing và chèn logic runtime mới bằng jax.tree_util
• Về bản chất bên trong, mô hình chỉ là một pytree có thể gọi được

penzai.treescope (pz.ts)

• Trình pretty-printer Python tương tác mạnh mẽ, hoạt động theo kiểu drop-in để thay thế renderer IPython/Colab thông thường
• Được thiết kế để giúp hiểu các mô hình Penzai và những JAX pytree lồng sâu khác, đồng thời mặc định hỗ trợ trực quan hóa NDArray với số chiều bất kỳ

penzai.core.selectors (pz.select)

• Công cụ đa năng cho pytree, tổng quát hóa cú pháp .at[...].set(...) của JAX thành cơ chế duyệt pytree dựa trên kiểu dữ liệu tùy ý
• Có thể dễ dàng thực hiện các thao tác viết lại phức tạp hoặc vá nóng on-the-fly cho mô hình Penzai và các cấu trúc dữ liệu khác

penzai.core.named_axes (pz.nx)

• Hệ thống trục được đặt tên gọn nhẹ, có thể vector hóa các hàm JAX thông thường theo các trục được đặt tên và cho phép chuyển đổi liền mạch giữa phong cách lập trình dùng tên và phong cách lập trình theo vị trí mà không cần học API mảng mới

penzai.data_effects (pz.de)

• Hệ thống opt-in cho các đối số phụ, số ngẫu nhiên và biến trạng thái, được xây dựng trên cơ sở duyệt pytree và vẫn giữ quyền kiểm soát mà không cản trở việc viết hoặc sử dụng mô hình

[brainer]

Không biết có cái nào tương tự dành cho PyTorch không nhỉ?