โครงสร้างภายในของ PyTorch (2019)

• คำอธิบายเกี่ยวกับโครงสร้างภายในของ PyTorch เป็นคู่มือสำหรับผู้ที่ต้องการมีส่วนร่วมกับโค้ดเบส C++ ของ PyTorch
• เป้าหมายของบทความนี้คือช่วยให้เข้าใจโครงสร้างของไลบรารีเทนเซอร์และเทคนิคการหาอนุพันธ์อัตโนมัติ (autograd) ของ PyTorch รวมถึงช่วยให้ค้นหาเส้นทางในโค้ดเบสได้

โครงสร้างพื้นฐานของเทนเซอร์ใน PyTorch

• ใน PyTorch เทนเซอร์คือโครงสร้างข้อมูลพื้นฐานที่สุด
• เทนเซอร์เป็นโครงสร้างข้อมูลแบบ n มิติ ซึ่งสามารถเก็บค่าสเกลาร์ เช่น จำนวนทศนิยมแบบลอยตัว (float), จำนวนเต็ม (int) เป็นต้น
• เทนเซอร์มีเมทาดาทาดังต่อไปนี้:
  • ขนาด (size): ข้อมูลเกี่ยวกับมิติของเทนเซอร์
  • dtype: ชนิดข้อมูลที่จัดเก็บ (เช่น float32, int64 เป็นต้น)
  • device: ตำแหน่งที่เก็บข้อมูล (CPU, CUDA เป็นต้น)
  • stride: ข้อมูลออฟเซ็ตของข้อมูลในหน่วยความจำจริง

• บทบาทของ Stride

  • stride ใช้สำหรับแปลงดัชนีเชิงตรรกะให้เป็นตำแหน่งในหน่วยความจำจริง
  • stride กำหนดออฟเซ็ตสำหรับแต่ละมิติ และคูณค่าดัชนีกับค่า stride เพื่อระบุตำแหน่งในหน่วยความจำจริง
  • ด้วย stride จึงสามารถมองข้อมูลชุดเดียวกันในรูปแบบที่ต่างออกไปเป็น view ได้ โดยไม่ต้องสร้างเทนเซอร์ใหม่

แนวคิดเรื่องเทนเซอร์และสตอเรจ (Storage)

• ใน PyTorch เทนเซอร์ไม่ได้เก็บข้อมูลจริงไว้โดยตรง → ข้อมูลถูกจัดการโดยสตอเรจ (Storage)
• Tensor = ขนาด + dtype + device + stride + offset
• เทนเซอร์หลายตัวสามารถแชร์สตอเรจเดียวกันได้ → รองรับแนวคิดของวิว (View)
• การแยกสตอเรจออกจากเทนเซอร์ช่วยให้ใช้หน่วยความจำได้อย่างมีประสิทธิภาพ

กระบวนการดิสแพตช์ (Dispatch) ของการคำนวณบนเทนเซอร์

• ใน PyTorch การคำนวณจะผ่านดิสแพตช์สองขั้นตอน:
  1. ดิสแพตช์ตามชนิดของอุปกรณ์และเลย์เอาต์
     • จะรันโค้ดติดตั้งใช้งานต่างกันตามว่าเป็นเทนเซอร์ CPU หรือ CUDA
  2. ดิสแพตช์ตาม dtype
     • จะเรียกเคอร์เนลต่างกันตามชนิดข้อมูล เช่น float กับ int

โมเดลการขยายเทนเซอร์ของ PyTorch

• องค์ประกอบหลัก 3 อย่างสำหรับการขยายเทนเซอร์:

  • Device: กำหนดวิธีจัดสรรหน่วยความจำบน CPU, GPU, TPU เป็นต้น
  • Layout: กำหนดวิธีจัดเก็บเทนเซอร์ในหน่วยความจำ (เช่น การจัดเก็บแบบต่อเนื่อง, การจัดเก็บแบบ sparse เป็นต้น)
  • dtype: กำหนดชนิดข้อมูลที่จะเก็บในแต่ละองค์ประกอบของเทนเซอร์

• ทางเลือกในการขยาย:

  • สามารถขยายเทนเซอร์ได้ด้วยการแก้โค้ดของ PyTorch โดยตรง
  • สามารถเขียน wrapper class ที่ครอบเทนเซอร์เดิมได้
  • หากต้องใช้ wrapper ระหว่างการหาอนุพันธ์อัตโนมัติ ต้องขยายโดยตรง

หลักการทำงานของการหาอนุพันธ์อัตโนมัติ (Autograd)

• PyTorch ทำการหาอนุพันธ์อัตโนมัติโดยอาศัย reverse-mode differentiation
• ระหว่างการคำนวณไปข้างหน้า (forward) จะสร้างกราฟ → ตอนย้อนกลับ (backward) จะไล่กราฟเพื่อคำนวณอนุพันธ์
• Autograd จัดการข้อมูลเพิ่มเติมดังต่อไปนี้:
  • AutogradMeta: เมทาดาทาที่เชื่อมกับเทนเซอร์และใช้ในการย้อนกลับ
  • บันทึกผลลัพธ์ของการคำนวณ และคำนวณอนุพันธ์ในระหว่างการย้อนกลับ

โครงสร้างโค้ดของ PyTorch และตำแหน่งไฟล์

• ไดเรกทอรีหลักในโค้ดเบสของ PyTorch:
  • torch/ → โมดูล Python (โค้ด Python)
  • torch/csrc/ → โค้ด binding ระหว่าง Python กับ C++, เอนจิน autograd, คอมไพเลอร์ JIT เป็นต้น
  • aten/ → นิยามการคำนวณของเทนเซอร์ (รวมการคำนวณหลักส่วนใหญ่)
  • c10/ → นิยามโครงสร้างหลัก เช่น เทนเซอร์และสตอเรจ

ขั้นตอนการทำงานของการคำนวณใน PyTorch

• ตัวอย่าง: กระบวนการเมื่อเรียก torch.add() :
  1. แปลงอาร์กิวเมนต์จาก Python ไปเป็นโค้ด C++
  2. ดำเนินการดิสแพตช์ใน VariableType
  3. ดำเนินการดิสแพตช์ตาม Device/เลย์เอาต์
  4. รันเคอร์เนลสุดท้าย

ขั้นตอนและเครื่องมือในการเขียนเคอร์เนล

• ใน PyTorch การเขียนเคอร์เนลมีขั้นตอนดังนี้:
  1. เขียนเมทาดาทาของการคำนวณ: กำหนดฟังก์ชันซิกเนเจอร์ อุปกรณ์และชนิดข้อมูลที่รองรับ
  2. ตรวจสอบอินพุต: ตรวจสอบอินพุต เช่น มิติและชนิด
  3. จัดสรรเอาต์พุตเทนเซอร์
  4. ดิสแพตช์ตาม dtype: รันเคอร์เนลตามชนิดข้อมูล
  5. การประมวลผลแบบขนาน: บน CPU ใช้ OpenMP, บน CUDA ใช้การขนานในตัว
  6. เข้าถึงข้อมูลและคำนวณ: ใช้ TensorAccessor, TensorIterator เป็นต้น

แมโครดิสแพตช์หลัก

• AT_DISPATCH_ALL_TYPES → ใช้ดิสแพตช์ตาม dtype
• รองรับแมโครสำหรับชนิดข้อมูลหลากหลายประเภท → ช่วยเพิ่มประสิทธิภาพได้

เคล็ดลับด้านการเพิ่มประสิทธิภาพและประสิทธิผลในการทำงาน

• แก้ไขไฟล์เฮดเดอร์ให้น้อยที่สุด → เพราะการแก้ไขจะทำให้ต้องรีบิลด์โค้ดทั้งหมด
• ตั้งค่าสภาพแวดล้อมพัฒนาในเครื่อง → ช่วยลดเวลาที่ต้องเสียไปกับการใช้ CI
• ใช้ ccache → ช่วยประหยัดเวลาในการคอมไพล์ซ้ำ
• ใช้เซิร์ฟเวอร์ที่มีประสิทธิภาพสูง → ช่วยลดเวลาในการคอมไพล์ C++ และบิลด์ CUDA

คู่มือการมีส่วนร่วมกับ PyTorch

• งานที่เหมาะสำหรับเริ่มต้นมีส่วนร่วม:
  • issue ที่มีป้าย triaged → เป็น issue ที่นักพัฒนา PyTorch ตรวจสอบแล้ว
  • การปรับปรุงเอกสารและช่วยทำซ้ำบั๊ก
  • การแสดงความคิดเห็นต่อ RFC (ข้อเสนอฟีเจอร์) ของ PyTorch
• PyTorch เติบโตขึ้นมาด้วยผู้ร่วมพัฒนาโอเพนซอร์ส และยินดีต้อนรับการมีส่วนร่วมจากชุมชน