การเขียน ops เคอร์เนล และการไล่ระดับสีแบบกำหนดเองใน TensorFlow.js

ภาพรวม

คู่มือนี้จะสรุปกลไกในการกำหนดการดำเนินการแบบกำหนดเอง (ops) เคอร์เนล และการไล่ระดับสีใน TensorFlow.js มีจุดมุ่งหมายเพื่อให้ภาพรวมของแนวคิดหลักและตัวชี้ไปยังโค้ดที่แสดงให้เห็นถึงแนวคิดในทางปฏิบัติ

คู่มือนี้เหมาะสำหรับใคร?

นี่เป็นคำแนะนำขั้นสูงที่ค่อนข้างครอบคลุมเนื้อหาภายในบางส่วนของ TensorFlow.js ซึ่งอาจเป็นประโยชน์อย่างยิ่งสำหรับกลุ่มคนต่อไปนี้:

  • ผู้ใช้ขั้นสูงของ TensorFlow.js สนใจในการปรับแต่งพฤติกรรมของการดำเนินการทางคณิตศาสตร์ต่างๆ (เช่น นักวิจัยที่เอาชนะการใช้งานการไล่ระดับสีที่มีอยู่ หรือผู้ใช้ที่ต้องการแก้ไขฟังก์ชันการทำงานที่ขาดหายไปในไลบรารี)
  • ผู้ใช้สร้างไลบรารีที่ขยาย TensorFlow.js (เช่น ไลบรารีพีชคณิตเชิงเส้นทั่วไปที่สร้างขึ้นบนพื้นฐาน TensorFlow.js หรือแบ็กเอนด์ TensorFlow.js ใหม่)
  • ผู้ใช้ที่สนใจสนับสนุนการดำเนินการใหม่ๆ ให้กับ tensorflow.js ที่ต้องการดูภาพรวมทั่วไปเกี่ยวกับวิธีการทำงานของกลไกเหล่านี้

นี่ ไม่ใช่ คำแนะนำสำหรับการใช้งาน TensorFlow.js โดยทั่วไป เนื่องจากเกี่ยวข้องกับกลไกการใช้งานภายใน คุณไม่จำเป็นต้องเข้าใจกลไกเหล่านี้เพื่อใช้ TensorFlow.js

คุณต้องสบายใจ (หรือเต็มใจที่จะลอง) อ่านซอร์สโค้ดของ TensorFlow.js เพื่อใช้ประโยชน์สูงสุดจากคู่มือนี้

คำศัพท์เฉพาะทาง

สำหรับคู่มือนี้ มีคำศัพท์สำคัญสองสามคำที่เป็นประโยชน์ในการอธิบายอย่างตรงไปตรงมา

การดำเนินงาน (Ops) — การดำเนินการทางคณิตศาสตร์กับเทนเซอร์ตั้งแต่หนึ่งตัวขึ้นไปที่สร้างเทนเซอร์ตั้งแต่หนึ่งตัวขึ้นไปเป็นเอาท์พุต Ops คือโค้ด 'ระดับสูง' และสามารถใช้ ops อื่นเพื่อกำหนดตรรกะได้

เคอร์เนล — การใช้งานเฉพาะของ op ที่เชื่อมโยงกับความสามารถของฮาร์ดแวร์/แพลตฟอร์มเฉพาะ เคอร์เนลเป็น 'ระดับต่ำ' และเฉพาะแบ็กเอนด์ ops บางตัวมีการแมปแบบหนึ่งต่อหนึ่งจาก op ไปยังเคอร์เนล ในขณะที่ ops อื่น ๆ ใช้หลายเคอร์เนล

Gradient / GradFunc - คำจำกัดความ 'โหมดย้อนกลับ' ของ op/kernel ที่คำนวณอนุพันธ์ของฟังก์ชันนั้นโดยคำนึงถึงอินพุตบางตัว การไล่ระดับสีเป็นโค้ด 'ระดับสูง' (ไม่เฉพาะเจาะจงแบ็กเอนด์) และสามารถเรียก ops หรือเคอร์เนลอื่นได้

Kernel Registry - แผนที่จากทูเพิล (ชื่อเคอร์เนล ชื่อแบ็กเอนด์) ไปจนถึงการใช้งานเคอร์เนล

Gradient Registry - แผนที่จาก ชื่อเคอร์เนลไปจนถึงการใช้งานการไล่ระดับสี

องค์กรรหัส

การดำเนินการ และ การไล่ระดับสี ถูกกำหนดไว้ใน tfjs-core

เคอร์เนลเป็นแบ็กเอนด์เฉพาะและถูกกำหนดไว้ในโฟลเดอร์แบ็กเอนด์ที่เกี่ยวข้อง (เช่น tfjs-backend-cpu )

ไม่จำเป็นต้องกำหนด ops เคอร์เนล และการไล่ระดับสีแบบกำหนดเองภายในแพ็คเกจเหล่านี้ แต่มักจะใช้สัญลักษณ์ที่คล้ายกันในการใช้งาน

การนำ Custom Ops ไปใช้

วิธีหนึ่งที่จะนึกถึง op ที่กำหนดเองก็คือฟังก์ชัน JavaScript ที่ส่งคืนเอาต์พุตเทนเซอร์บางตัว ซึ่งมักจะมีเทนเซอร์เป็นอินพุต

  • ops บางอย่างสามารถกำหนดได้อย่างสมบูรณ์ในแง่ของ ops ที่มีอยู่ และควรนำเข้าและเรียกใช้ฟังก์ชันเหล่านี้โดยตรง นี่คือตัวอย่าง
  • การใช้งาน op ยังสามารถส่งไปยังแบ็กเอนด์เคอร์เนลเฉพาะได้ ซึ่งดำเนินการผ่าน Engine.runKernel และจะมีการอธิบายเพิ่มเติมในส่วน “การใช้เคอร์เนลแบบกำหนดเอง” นี่คือตัวอย่าง

การใช้เคอร์เนลแบบกำหนดเอง

การใช้งานเคอร์เนลเฉพาะของแบ็กเอนด์ช่วยให้สามารถปรับใช้ตรรกะสำหรับการดำเนินการที่กำหนดได้อย่างเหมาะสม เคอร์เนลถูกเรียกใช้โดย ops โทร tf.engine().runKernel() การใช้งานเคอร์เนลถูกกำหนดโดยสี่สิ่ง

  • ชื่อเคอร์เนล
  • แบ็กเอนด์เคอร์เนลถูกนำมาใช้
  • อินพุต: อาร์กิวเมนต์เทนเซอร์สำหรับฟังก์ชันเคอร์เนล
  • คุณลักษณะ: อาร์กิวเมนต์ที่ไม่ใช่เทนเซอร์สำหรับฟังก์ชันเคอร์เนล

นี่คือตัวอย่าง การใช้งานเคอร์เนล แบบแผนที่ใช้ในการนำไปใช้นั้นมีความเฉพาะเจาะจงของแบ็กเอนด์และเป็นที่เข้าใจได้ดีที่สุดจากการดูการใช้งานและเอกสารประกอบของแบ็กเอนด์แต่ละรายการ

โดยทั่วไปแล้ว เคอร์เนลจะทำงานในระดับที่ต่ำกว่าเทนเซอร์ และจะอ่านและเขียนโดยตรงไปยังหน่วยความจำแทน ซึ่งในที่สุดจะถูกห่อหุ้มเป็นเทนเซอร์ด้วย tfjs-core

เมื่อนำเคอร์เนลไปใช้แล้ว จะสามารถลงทะเบียนกับ TensorFlow.js ได้โดยใช้ ฟังก์ชัน registerKernel จาก tfjs-core คุณสามารถลงทะเบียนเคอร์เนลสำหรับทุกแบ็กเอนด์ที่คุณต้องการให้เคอร์เนลนั้นทำงาน เมื่อลงทะเบียนแล้ว เคอร์เนลจะสามารถเรียกใช้ด้วย tf.engine().runKernel(...) และ TensorFlow.js จะทำให้แน่ใจว่าได้ส่งไปยังการใช้งานใน แบ็กเอนด์ที่ใช้งานอยู่ในปัจจุบัน

การใช้การไล่ระดับสีแบบกำหนดเอง

โดยทั่วไปการไล่ระดับสีถูกกำหนดไว้สำหรับเคอร์เนลที่กำหนด (ระบุโดยชื่อเคอร์เนลเดียวกันกับที่ใช้ในการเรียก tf.engine().runKernel(...) ) สิ่งนี้ทำให้ tfjs-core สามารถใช้รีจิสทรีเพื่อค้นหาคำจำกัดความของการไล่ระดับสีสำหรับเคอร์เนลใด ๆ ในขณะรันไทม์

การใช้การไล่ระดับสีแบบกำหนดเองมีประโยชน์สำหรับ:

  • การเพิ่มคำจำกัดความของการไล่ระดับสีที่อาจไม่มีอยู่ในไลบรารี
  • การแทนที่คำจำกัดความการไล่ระดับสีที่มีอยู่เพื่อปรับแต่งการคำนวณการไล่ระดับสีสำหรับเคอร์เนลที่กำหนด

คุณสามารถดูตัวอย่าง การใช้งานการไล่ระดับสีได้ที่นี่

เมื่อคุณใช้การไล่ระดับสีสำหรับการโทรที่กำหนดแล้ว จะสามารถลงทะเบียนกับ TensorFlow.js ได้โดยใช้ ฟังก์ชัน registerGradient จาก tfjs-core

อีกวิธีหนึ่งในการนำการไล่ระดับสีแบบกำหนดเองไปใช้โดยผ่านการลงทะเบียนการไล่ระดับสี (และช่วยให้สามารถคำนวณการไล่ระดับสีสำหรับฟังก์ชันที่กำหนดเองในรูปแบบที่กำหนดเองได้คือการใช้ tf.customGrad

นี่คือ ตัวอย่างของ op ภายในไลบรารี ของการใช้ customGrad