Viết các op, hạt nhân và độ dốc tùy chỉnh trong TensorFlow.js

Tổng quan

Hướng dẫn này phác thảo các cơ chế xác định các hoạt động tùy chỉnh (ops), hạt nhân và độ dốc trong TensorFlow.js. Nó nhằm mục đích cung cấp một cái nhìn tổng quan về các khái niệm chính và các con trỏ tới mã thể hiện các khái niệm đang hoạt động.

hướng dẫn này dành cho ai?

Đây là hướng dẫn khá nâng cao đề cập đến một số nội dung bên trong của TensorFlow.js, nó có thể đặc biệt hữu ích cho các nhóm người sau:

  • Người dùng nâng cao của TensorFlow.js quan tâm đến việc tùy chỉnh hành vi của các phép toán khác nhau (ví dụ: các nhà nghiên cứu ghi đè việc triển khai gradient hiện có hoặc người dùng cần vá chức năng bị thiếu trong thư viện)
  • Người dùng xây dựng các thư viện mở rộng TensorFlow.js (ví dụ: thư viện đại số tuyến tính tổng quát được xây dựng dựa trên các nguyên hàm TensorFlow.js hoặc chương trình phụ trợ TensorFlow.js mới).
  • Người dùng quan tâm đến việc đóng góp các hoạt động mới cho tensorflow.js muốn có cái nhìn tổng quan về cách hoạt động của các cơ chế này.

Đây không phải là hướng dẫn sử dụng chung TensorFlow.js vì nó đi sâu vào các cơ chế triển khai nội bộ. Bạn không cần phải hiểu các cơ chế này để sử dụng TensorFlow.js

Bạn cần phải cảm thấy thoải mái (hoặc sẵn sàng thử) đọc mã nguồn TensorFlow.js để tận dụng tối đa hướng dẫn này.

Thuật ngữ

Đối với hướng dẫn này, một số thuật ngữ chính rất hữu ích để mô tả trước.

Hoạt động (Ops) - Một phép toán trên một hoặc nhiều tensor tạo ra một hoặc nhiều tensor làm đầu ra. Ops là mã 'cấp cao' và có thể sử dụng các ops khác để xác định logic của chúng.

Kernel - Việc triển khai cụ thể một op gắn liền với các khả năng phần cứng/nền tảng cụ thể. Hạt nhân là 'cấp thấp' và phụ trợ cụ thể. Một số op có ánh xạ một-một từ op đến kernel trong khi các op khác sử dụng nhiều kernel.

gradient / GradFunc - Định nghĩa 'chế độ lùi' của op/kernel tính toán đạo hàm của hàm đó liên quan đến một số đầu vào. Gradien là mã 'cấp cao' (không dành riêng cho chương trình phụ trợ) và có thể gọi các op hoặc kernel khác.

Sổ đăng ký hạt nhân - Bản đồ từ một bộ (tên hạt nhân, tên chương trình phụ trợ) tới quá trình triển khai hạt nhân.

Sổ đăng ký gradient - Bản đồ từ tên kernel đến triển khai gradient .

Tổ chức mã

Các hoạt độngđộ dốc được xác định trong tfjs-core .

Hạt nhân là dành riêng cho phần phụ trợ và được xác định trong các thư mục phụ trợ tương ứng của chúng (ví dụ: tfjs-backend-cpu ).

Các hoạt động, hạt nhân và độ dốc tùy chỉnh không cần phải được xác định bên trong các gói này. Nhưng thường sẽ sử dụng các ký hiệu tương tự khi thực hiện.

Triển khai hoạt động tùy chỉnh

Một cách để nghĩ về một op tùy chỉnh giống như một hàm JavaScript trả về một số đầu ra tensor, thường có các tensor làm đầu vào.

  • Một số hoạt động có thể được xác định hoàn toàn theo các hoạt động hiện có và chỉ cần nhập và gọi trực tiếp các hàm này. Đây là một ví dụ .
  • Việc triển khai một op cũng có thể được gửi tới các hạt nhân cụ thể phụ trợ. Việc này được thực hiện thông qua Engine.runKernel và sẽ được mô tả sâu hơn trong phần “triển khai hạt nhân tùy chỉnh”. Đây là một ví dụ .

Triển khai hạt nhân tùy chỉnh

Việc triển khai kernel cụ thể ở phần phụ trợ cho phép triển khai logic được tối ưu hóa cho một hoạt động nhất định. Hạt nhân được gọi bởi các op gọi tf.engine().runKernel() . Việc triển khai kernel được xác định bởi bốn điều

  • Một tên hạt nhân.
  • Phần phụ trợ mà kernel được triển khai.
  • Đầu vào: Đối số tensor cho hàm kernel.
  • Thuộc tính: Các đối số không phải tensor cho hàm kernel.

Đây là một ví dụ về triển khai kernel . Các quy ước được sử dụng để triển khai là dành riêng cho phần phụ trợ và được hiểu rõ nhất khi xem xét cách triển khai và tài liệu của từng phần phụ trợ cụ thể.

Nói chung, các hạt nhân hoạt động ở mức thấp hơn các tenxơ và thay vào đó trực tiếp đọc và ghi vào bộ nhớ mà cuối cùng sẽ được tfjs-core gói thành các tenxơ.

Sau khi kernel được triển khai, nó có thể được đăng ký với TensorFlow.js bằng cách sử dụng hàm registerKernel từ tfjs-core. Bạn có thể đăng ký kernel cho mọi chương trình phụ trợ mà bạn muốn kernel đó hoạt động. Sau khi đăng ký, kernel có thể được gọi bằng tf.engine().runKernel(...) và TensorFlow.js sẽ đảm bảo gửi đến quá trình triển khai trong phụ trợ hoạt động hiện tại.

Triển khai các gradient tùy chỉnh

Các gradient thường được xác định cho một hạt nhân nhất định (được xác định bằng cùng tên hạt nhân được sử dụng trong lệnh gọi tới tf.engine().runKernel(...) ). Điều này cho phép tfjs-core sử dụng sổ đăng ký để tra cứu định nghĩa độ dốc cho bất kỳ hạt nhân nào trong thời gian chạy.

Việc triển khai các gradient tùy chỉnh rất hữu ích cho:

  • Thêm định nghĩa độ dốc có thể không có trong thư viện
  • Ghi đè định nghĩa độ dốc hiện có để tùy chỉnh tính toán độ dốc cho hạt nhân nhất định.

Bạn có thể xem các ví dụ về triển khai gradient tại đây .

Khi bạn đã triển khai một gradient cho một lệnh gọi cụ thể, nó có thể được đăng ký với TensorFlow.js bằng cách sử dụng hàm registerGradient từ tfjs-core.

Cách tiếp cận khác để triển khai các gradient tùy chỉnh bỏ qua sổ đăng ký gradient (và do đó cho phép tính toán các gradient cho các hàm tùy ý theo những cách tùy ý là sử dụng tf.customGrad .

Đây là một ví dụ về một op trong thư viện sử dụng customGrad