Menulis operasi khusus, kernel, dan gradien di TensorFlow.js

Ringkasan

Panduan ini menguraikan mekanisme untuk menentukan operasi kustom (operasi), kernel, dan gradien di TensorFlow.js. Hal ini bertujuan untuk memberikan gambaran umum tentang konsep utama dan petunjuk kode yang mendemonstrasikan konsep dalam tindakan.

Untuk siapa panduan ini?

Ini adalah panduan tingkat lanjut yang membahas beberapa bagian internal TensorFlow.js, dan mungkin berguna khususnya untuk kelompok orang berikut:

  • Pengguna TensorFlow.js tingkat lanjut yang tertarik untuk menyesuaikan perilaku berbagai operasi matematika (misalnya peneliti mengganti implementasi gradien yang ada atau pengguna yang perlu menambal fungsi yang hilang di perpustakaan)
  • Pengguna membuat perpustakaan yang memperluas TensorFlow.js (misalnya, perpustakaan aljabar linier umum yang dibangun di atas primitif TensorFlow.js atau backend TensorFlow.js baru).
  • Pengguna yang tertarik untuk menyumbangkan operasi baru ke tensorflow.js dan ingin mendapatkan gambaran umum tentang cara kerja mekanisme ini.

Ini bukan panduan untuk penggunaan umum TensorFlow.js karena panduan ini masuk ke dalam mekanisme implementasi internal. Anda tidak perlu memahami mekanisme ini untuk menggunakan TensorFlow.js

Anda harus terbiasa (atau ingin mencoba) membaca kode sumber TensorFlow.js agar dapat memanfaatkan panduan ini secara maksimal.

Terminologi

Untuk panduan ini, ada beberapa istilah penting yang berguna untuk dijelaskan terlebih dahulu.

Operasi (Ops) — Operasi matematika pada satu atau lebih tensor yang menghasilkan satu atau lebih tensor sebagai keluaran. Operasi adalah kode 'tingkat tinggi' dan dapat menggunakan operasi lain untuk menentukan logikanya.

Kernel — Implementasi spesifik dari operasi yang terkait dengan kemampuan perangkat keras/platform tertentu. Kernel bersifat 'tingkat rendah' ​​dan spesifik untuk backend. Beberapa operasi memiliki pemetaan satu-ke-satu dari operasi ke kernel sementara operasi lainnya menggunakan banyak kernel.

Gradien / GradFunc — Definisi 'mode mundur' dari sebuah op/kernel yang menghitung turunan dari fungsi tersebut sehubungan dengan beberapa masukan. Gradien adalah kode 'tingkat tinggi' (tidak spesifik untuk backend) dan dapat memanggil operasi atau kernel lain.

Kernel Registry - Peta dari tuple (nama kernel, nama backend) ke implementasi kernel.

Gradient Registry — Peta dari nama kernel ke implementasi gradien .

Organisasi kode

Operasi dan Gradien didefinisikan dalam tfjs-core .

Kernel bersifat spesifik untuk backend dan ditentukan di folder backend masing-masing (misalnya tfjs-backend-cpu ).

Operasi khusus, kernel, dan gradien tidak perlu didefinisikan di dalam paket ini. Namun akan sering menggunakan simbol serupa dalam implementasinya.

Menerapkan Operasi Kustom

Salah satu cara untuk memikirkan operasi khusus adalah sebagai fungsi JavaScript yang mengembalikan beberapa keluaran tensor, sering kali dengan tensor sebagai masukan.

  • Beberapa operasi dapat didefinisikan sepenuhnya berdasarkan operasi yang ada, dan sebaiknya mengimpor dan memanggil fungsi-fungsi ini secara langsung. Ini sebuah contoh .
  • Implementasi operasi juga dapat dikirim ke kernel backend tertentu. Hal ini dilakukan melalui Engine.runKernel dan akan dijelaskan lebih lanjut di bagian “implementing custom kernels”. Ini sebuah contoh .

Menerapkan Kernel Khusus

Implementasi kernel spesifik backend memungkinkan implementasi logika yang dioptimalkan untuk operasi tertentu. Kernel dipanggil oleh panggilan operasi tf.engine().runKernel() . Implementasi kernel ditentukan oleh empat hal

  • Nama kernel.
  • Backend tempat kernel diimplementasikan.
  • Input: Argumen tensor ke fungsi kernel.
  • Atribut: Argumen non-tensor pada fungsi kernel.

Berikut adalah contoh implementasi kernel . Konvensi yang digunakan untuk mengimplementasikan bersifat khusus untuk backend dan paling mudah dipahami dengan melihat implementasi dan dokumentasi masing-masing backend.

Umumnya kernel beroperasi pada level yang lebih rendah dari tensor dan sebagai gantinya langsung membaca dan menulis ke memori yang pada akhirnya akan dibungkus menjadi tensor oleh tfjs-core.

Setelah kernel diimplementasikan, kernel dapat didaftarkan ke TensorFlow.js dengan menggunakan fungsi registerKernel dari tfjs-core. Anda dapat mendaftarkan kernel untuk setiap backend yang Anda inginkan agar kernel tersebut berfungsi. Setelah terdaftar, kernel dapat dipanggil dengan tf.engine().runKernel(...) dan TensorFlow.js akan memastikan untuk dikirim ke implementasi di backend aktif saat ini.

Menerapkan Gradien Kustom

Gradien umumnya ditentukan untuk kernel tertentu (diidentifikasi dengan nama kernel yang sama dengan yang digunakan dalam panggilan ke tf.engine().runKernel(...) ). Hal ini memungkinkan tfjs-core menggunakan registri untuk mencari definisi gradien untuk kernel apa pun saat runtime.

Menerapkan gradien khusus berguna untuk:

  • Menambahkan definisi gradien yang mungkin tidak ada di perpustakaan
  • Mengganti definisi gradien yang ada untuk menyesuaikan komputasi gradien untuk kernel tertentu.

Anda dapat melihat contoh implementasi gradien di sini .

Setelah Anda mengimplementasikan gradien untuk panggilan tertentu, gradien tersebut dapat didaftarkan ke TensorFlow.js dengan menggunakan fungsi registerGradient dari tfjs-core.

Pendekatan lain untuk menerapkan gradien khusus yang melewati registri gradien (dan dengan demikian memungkinkan komputasi gradien untuk fungsi arbitrer dengan cara arbitrer adalah menggunakan tf.customGrad .

Berikut adalah contoh operasi dalam perpustakaan menggunakan customGrad