Penzai - ชุดเครื่องมือ JAX สำหรับสร้าง แก้ไข และแสดงภาพโครงข่ายประสาท

 (github.com/google-deepmind)
7 คะแนน โดย xguru 2024-04-23 | 1 ความคิดเห็น | แชร์ทาง WhatsApp
  • ไลบรารี JAX จาก Google DeepMind ที่เปิดตัวมาเพื่อเขียนโมเดลในรูปแบบโครงสร้างข้อมูล Pytree เชิงฟังก์ชันที่อ่านเข้าใจได้ง่าย
  • มุ่งเน้นให้สามารถทำงานกับโมเดลได้อย่างง่ายดายแม้หลังจากฝึกเสร็จแล้ว
    • เหมาะกับงานวิจัยที่เกี่ยวข้องกับการทำรีเวิร์สเอนจิเนียริงหรือการถอดส่วนประกอบของโมเดล รวมถึงการตรวจสอบและสำรวจ internal activations, model surgery, การดีบักสถาปัตยกรรม เป็นต้น
    • มีเครื่องมือสำหรับการแสดงภาพ แก้ไข และวิเคราะห์โมเดล จึงสามารถใช้เพื่อสร้างและฝึกโมเดลได้เช่นกัน
  • ประกอบด้วยชุดเครื่องมือแบบโมดูลาร์ที่ออกแบบมาคู่กัน แต่แต่ละโมดูลก็สามารถใช้งานแยกได้อย่างอิสระ

เครื่องมือที่มีมาให้

penzai.nn (pz.nn)

  • ไลบรารีโครงข่ายประสาทแบบ declarative combinator-based ที่เป็นทางเลือกแทนไลบรารีโครงข่ายประสาทอื่น ๆ เช่น Flax, Haiku, Keras และ Equinox
  • เปิดเผยโครงสร้างทั้งหมดของ forward pass ของโมเดลให้เห็นใน model pytree ทำให้สามารถดูทุกการทำงานของโมเดลผ่าน pretty printing และแทรก runtime logic ใหม่ด้วย jax.tree_util ได้
  • ภายในแล้ว โมเดลเป็นเพียง pytree ที่เรียกใช้งานได้เท่านั้น

penzai.treescope (pz.ts)

  • Python pretty-printer แบบโต้ตอบประสิทธิภาพสูงที่ทำงานแบบ drop-in replacement สำหรับเรนเดอเรอร์ IPython/Colab ปกติ
  • ออกแบบมาเพื่อช่วยให้เข้าใจโมเดล Penzai และ JAX pytree อื่น ๆ ที่ซ้อนลึก และรองรับการแสดงภาพ NDArray มิติใดก็ได้โดยอัตโนมัติ

penzai.core.selectors (pz.select)

  • เครื่องมืออเนกประสงค์สำหรับ pytree ที่ขยายไวยากรณ์ .at[...].set(...) ของ JAX ให้เป็นการท่องโครงสร้าง pytree ตามชนิดข้อมูลแบบทั่วไป
  • ช่วยให้เขียนการรีไรต์ที่ซับซ้อนหรือแพตช์แบบ on-the-fly สำหรับโมเดล Penzai และโครงสร้างข้อมูลอื่น ๆ ได้ง่าย

penzai.core.named_axes (pz.nx)

  • ระบบ named axes แบบน้ำหนักเบาที่ช่วย vectorize ฟังก์ชัน JAX ปกติให้ทำงานกับ named axes และสลับไปมาระหว่างสไตล์การเขียนโปรแกรมแบบใช้ชื่อกับแบบใช้ตำแหน่งได้อย่างราบรื่น โดยไม่ต้องเรียนรู้ array API ใหม่

penzai.data_effects (pz.de)

  • ระบบแบบ opt-in สำหรับ side inputs, random numbers และ state variables ที่สร้างบนพื้นฐานของการท่อง pytree และยังคงให้การควบคุมได้โดยไม่รบกวนการเขียนหรือใช้งานโมเดล

บทความที่เกี่ยวข้อง

1 ความคิดเห็น

 
brainer 2024-04-23

มีอะไรคล้ายๆ กันสำหรับ Pytorch ไหมนะ?