functorch

เครื่องราง

ทำไม Functorch? - คู่มือติดตั้ง การเปลี่ยนแปลง เอกสาร แผนการในอนาคต

ห้องสมุดนี้อยู่ระหว่างการพัฒนาอย่างหนัก - หากคุณมีคำแนะนำเกี่ยวกับ API หรือกรณีการใช้งานที่คุณต้องการได้รับการคุ้มครองโปรดเปิดปัญหา GitHub หรือเข้าถึง เราชอบที่จะได้ยินเกี่ยวกับวิธีการใช้ห้องสมุด

functorch เป็นฟังก์ชั่นที่เหมือนกันกับ JAX ที่แปลงเป็น Pytorch

มันมีจุดมุ่งหมายเพื่อให้ vmap และ grad การศึกษาสามารถแปลงได้ซึ่งทำงานกับโมดูล Pytorch และ Pytorch Autograd ที่มีประสิทธิภาพโหมดกระตือรือร้นที่ดี

นอกจากนี้ยังมีฟังก์ชั่นการทดลองเพื่อติดตามการเปลี่ยนแปลงเหล่านี้โดยใช้ FX เพื่อจับผลลัพธ์ของการแปลงเหล่านี้ล่วงหน้า สิ่งนี้จะช่วยให้เราสามารถรวบรวมผลลัพธ์ของ VMAP หรือผู้สำเร็จการศึกษาเพื่อปรับปรุงประสิทธิภาพ

ทำไมฟังก์ชั่นที่สามารถปรับแต่งได้?

มีหลายกรณีการใช้งานที่ยากที่จะทำใน Pytorch วันนี้:

• การคำนวณต่อตัวอย่างตัวอย่าง (หรือปริมาณอื่น ๆ ต่อตัวอย่าง)
• กำลังใช้ชุดรุ่นบนเครื่องเดียว
• การผสมผสานงานร่วมกันอย่างมีประสิทธิภาพในวงในของ MAML
• คำนวณ Jacobians และ Hessians ได้อย่างมีประสิทธิภาพ
• การคำนวณอย่างมีประสิทธิภาพจาโคเบียและเฮสเซียแบทช์

การเขียน vmap , grad , vjp และ jvp Transforms ช่วยให้เราสามารถแสดงผลข้างต้นได้โดยไม่ต้องออกแบบระบบย่อยแยกต่างหากสำหรับแต่ละระบบ ความคิดของฟังก์ชั่นการแต่งกายแบบคอมโพสิตนี้มาจากเฟรมเวิร์ก JAX

ติดตั้ง

มีสองวิธีในการติดตั้ง Functorch:

1. functorch จากแหล่งที่มา
2. Functorch Beta (เข้ากันได้กับการเปิดตัว Pytorch ล่าสุด)

เราขอแนะนำให้ลองใช้เบต้า Functorch ก่อน

การติดตั้ง functorch จากแหล่งที่มา

คลิกเพื่อขยาย

ใช้ colab

ทำตามคำแนะนำในสมุดบันทึก colab นี้

ในระดับท้องถิ่น

ณ วันที่ 9/21/2022 functorch ได้ติดตั้งพร้อมกับไบนารี Pytorch ทุกคืน โปรดติดตั้งตัวอย่าง (ทุกคืน) Pytorch Binary; ดู https://pytorch.org/ สำหรับคำแนะนำ

เมื่อคุณทำเสร็จแล้วให้ใช้การตรวจสุขภาพอย่างรวดเร็วใน Python:

 import torch
from functorch import vmap
x = torch . randn ( 3 )
y = vmap ( torch . sin )( x )
assert torch . allclose ( y , x . sin ())

การตั้งค่าการพัฒนา Functorch

ณ วันที่ 9/21/2022 functorch ได้ติดตั้งไว้ข้าง Pytorch และอยู่ในต้นไม้ต้นฉบับ Pytorch โปรดติดตั้ง pytorch จากแหล่งที่มาจากนั้นคุณจะสามารถ import functorch ได้

พยายามทำการทดสอบบางอย่างเพื่อให้แน่ใจว่าทั้งหมดเป็นโอเค:

pytest test/test_vmap.py -v
pytest test/test_eager_transforms.py -v

Aotautograd มีข้อกำหนดเพิ่มเติมเพิ่มเติม คุณสามารถติดตั้งผ่าน:

pip install networkx

ในการเรียกใช้การทดสอบ Functorch โปรดติดตั้งการทดสอบการทดสอบของเรา ( expecttest , pyyaml )

การติดตั้ง functorch beta (เข้ากันได้กับการเปิดตัว Pytorch ล่าสุด)

คลิกเพื่อขยาย

ใช้ colab

ทำตามคำแนะนำที่นี่

ปิ๊ก

วิชาบังคับก่อน: ติดตั้ง pytorch

pip install functorch

ในที่สุดเรียกใช้การตรวจสุขภาพอย่างรวดเร็วใน Python:

 import torch
from functorch import vmap
x = torch . randn ( 3 )
y = vmap ( torch . sin )( x )
assert torch . allclose ( y , x . sin ())

การแปลงคืออะไร?

ตอนนี้เราสนับสนุนการแปลงต่อไปนี้:

• grad , vjp , jvp ,
• jacrev , jacfwd , hessian
• vmap

นอกจากนี้เรามีสาธารณูปโภคบางอย่างสำหรับการทำงานกับโมดูล Pytorch

• make_functional(model)
• make_functional_with_buffers(model)

VMAP

หมายเหตุ: vmap กำหนดข้อ จำกัด เกี่ยวกับรหัสที่สามารถใช้งานได้ สำหรับรายละเอียดเพิ่มเติมโปรดอ่านเอกสาร

vmap(func)(*inputs) เป็นการแปลงที่เพิ่มมิติให้กับการทำงานของเทนเซอร์ทั้งหมดใน func vmap(func) ส่งคืนฟังก์ชั่นใหม่ที่แมป func ในบางมิติ (ค่าเริ่มต้น: 0) ของแต่ละเทนเซอร์ใน inputs

vmap มีประโยชน์สำหรับการซ่อนมิติแบตช์: หนึ่งสามารถเขียนฟังก์ชั่น func ที่ทำงานบนตัวอย่างแล้วยกขึ้นไปยังฟังก์ชั่นที่สามารถนำตัวอย่างของตัวอย่างด้วย vmap(func) นำไปสู่ประสบการณ์การสร้างแบบจำลองที่ง่ายขึ้น:

 from functorch import vmap
batch_size , feature_size = 3 , 5
weights = torch . randn ( feature_size , requires_grad = True )

def model ( feature_vec ):
    # Very simple linear model with activation
    assert feature_vec . dim () == 1
    return feature_vec . dot ( weights ). relu ()

examples = torch . randn ( batch_size , feature_size )
result = vmap ( model )( examples )

ผู้สำเร็จการศึกษา

grad(func)(*inputs) ถือว่า func ส่งคืนเทนเซอร์องค์ประกอบเดียว มันคำนวณการไล่ระดับสีของเอาต์พุตของ func WRT ไปยัง inputs[0]

 from functorch import grad
x = torch . randn ([])
cos_x = grad ( lambda x : torch . sin ( x ))( x )
assert torch . allclose ( cos_x , x . cos ())

# Second-order gradients
neg_sin_x = grad ( grad ( lambda x : torch . sin ( x )))( x )
assert torch . allclose ( neg_sin_x , - x . sin ())

เมื่อประกอบกับ vmap grad สามารถใช้ในการคำนวณต่อตัวอย่างตัวอย่าง:

 from functorch import vmap
batch_size , feature_size = 3 , 5

def model ( weights , feature_vec ):
    # Very simple linear model with activation
    assert feature_vec . dim () == 1
    return feature_vec . dot ( weights ). relu ()

def compute_loss ( weights , example , target ):
    y = model ( weights , example )
    return (( y - target ) ** 2 ). mean ()  # MSELoss

weights = torch . randn ( feature_size , requires_grad = True )
examples = torch . randn ( batch_size , feature_size )
targets = torch . randn ( batch_size )
inputs = ( weights , examples , targets )
grad_weight_per_example = vmap ( grad ( compute_loss ), in_dims = ( None , 0 , 0 ))( * inputs )

VJP

การแปลง vjp ใช้ func กับ inputs และส่งคืนฟังก์ชั่นใหม่ที่คำนวณ VJPs ให้ cotangents บางตัวเทนเซอร์

 from functorch import vjp
outputs , vjp_fn = vjp ( func , inputs ); vjps = vjp_fn ( * cotangents )

JVP

jvp แปลงการคำนวณ Jacobian-vector-Products และเป็นที่รู้จักกันในชื่อ "โฆษณาไปข้างหน้า" มันไม่ได้เป็นฟังก์ชั่นที่มีลำดับสูงกว่าการแปลงอื่น ๆ ส่วนใหญ่ แต่จะส่งคืนผลลัพธ์ของ func(inputs) รวมถึง jvp S

 from functorch import jvp
x = torch . randn ( 5 )
y = torch . randn ( 5 )
f = lambda x , y : ( x * y )
_ , output = jvp ( f , ( x , y ), ( torch . ones ( 5 ), torch . ones ( 5 )))
assert torch . allclose ( output , x + y )

Jacrev, Jacfwd และ Hessian

การแปลง jacrev ส่งคืนฟังก์ชั่นใหม่ที่ใช้ใน x และส่งคืน Jacobian of torch.sin เมื่อเทียบกับ x โดยใช้โฆษณาย้อนกลับโหมด

 from functorch import jacrev
x = torch . randn ( 5 )
jacobian = jacrev ( torch . sin )( x )
expected = torch . diag ( torch . cos ( x ))
assert torch . allclose ( jacobian , expected )

ใช้ jacrev เพื่อคำนวณ Jacobian สิ่งนี้สามารถแต่งด้วย VMAP เพื่อผลิตจาโคเบียแบทช์:

 x = torch . randn ( 64 , 5 )
jacobian = vmap ( jacrev ( torch . sin ))( x )
assert jacobian . shape == ( 64 , 5 , 5 )

jacfwd เป็นการแทนที่ jacrev ที่คำนวณ Jacobians โดยใช้ AD Mode Forward:

 from functorch import jacfwd
x = torch . randn ( 5 )
jacobian = jacfwd ( torch . sin )( x )
expected = torch . diag ( torch . cos ( x ))
assert torch . allclose ( jacobian , expected )

การแต่ง jacrev ด้วยตัวเองหรือ jacfwd สามารถผลิต Hessians:

 def f ( x ):
  return x . sin (). sum ()

x = torch . randn ( 5 )
hessian0 = jacrev ( jacrev ( f ))( x )
hessian1 = jacfwd ( jacrev ( f ))( x )

hessian เป็นฟังก์ชั่นความสะดวกสบายที่รวม jacfwd และ jacrev :

 from functorch import hessian

def f ( x ):
  return x . sin (). sum ()

x = torch . randn ( 5 )
hess = hessian ( f )( x )

ติดตามการเปลี่ยนแปลง

นอกจากนี้เรายังสามารถติดตามผ่านการแปลงเหล่านี้เพื่อจับภาพผลลัพธ์เป็นรหัสใหม่โดยใช้ make_fx นอกจากนี้ยังมีการรวมการทดลองกับคอมไพเลอร์ NNC (ใช้งานได้กับ CPU เท่านั้นในตอนนี้!)

 from functorch import make_fx , grad
def f ( x ):
    return torch . sin ( x ). sum ()
x = torch . randn ( 100 )
grad_f = make_fx ( grad ( f ))( x )
print ( grad_f . code )

def forward ( self , x_1 ):
    sin = torch . ops . aten . sin ( x_1 )
    sum_1 = torch . ops . aten . sum ( sin , None );  sin = None
    cos = torch . ops . aten . cos ( x_1 );  x_1 = None
    _tensor_constant0 = self . _tensor_constant0
    mul = torch . ops . aten . mul ( _tensor_constant0 , cos );  _tensor_constant0 = cos = None
    return mul

การทำงานกับโมดูล NN: make_functional และเพื่อน

บางครั้งคุณอาจต้องการทำการแปลงด้วยความเคารพต่อพารามิเตอร์และ/หรือบัฟเฟอร์ของ nn.module สิ่งนี้สามารถเกิดขึ้นได้เช่นใน:

• Modelsembling ซึ่งน้ำหนักและบัฟเฟอร์ทั้งหมดของคุณมีมิติเพิ่มเติม
• การคำนวณต่อตัวอย่าง gradient ที่คุณต้องการคำนวณต่อตัวอย่างของการสูญเสียตามพารามิเตอร์โมเดล

การแก้ปัญหาของเราในตอนนี้คือ API ที่ได้รับ nn.module สร้างเวอร์ชันไร้สัญชาติที่สามารถเรียกได้ว่าเป็นฟังก์ชั่น

• make_functional(model) ส่งคืนรุ่นที่ใช้งานได้ของ model และ model.parameters()
• make_functional_with_buffers(model) ส่งคืนรุ่นที่ใช้งานได้ของ model และ model.parameters() และ model.buffers()

นี่คือตัวอย่างที่เราคำนวณต่อตัวอย่างตัวอย่างโดยใช้เลเยอร์ nn.linear:

 import torch
from functorch import make_functional , vmap , grad

model = torch . nn . Linear ( 3 , 3 )
data = torch . randn ( 64 , 3 )
targets = torch . randn ( 64 , 3 )

func_model , params = make_functional ( model )

def compute_loss ( params , data , targets ):
    preds = func_model ( params , data )
    return torch . mean (( preds - targets ) ** 2 )

per_sample_grads = vmap ( grad ( compute_loss ), ( None , 0 , 0 ))( params , data , targets )

หากคุณกำลังสร้างชุดโมเดลคุณอาจพบว่า combine_state_for_ensemble มีประโยชน์

เอกสาร

สำหรับเอกสารเพิ่มเติมดูเว็บไซต์เอกสารของเรา

การดีบัก

torch._C._functorch.dump_tensor : DUMPS DISPATCH KEYS บนสแต็ก torch._C._functorch._set_vmap_fallback_warning_enabled(False) หาก VMAP เตือนสแปมรบกวนคุณ

แผนการในอนาคต

ในท้ายที่สุดเราต้องการที่จะต้นน้ำสิ่งนี้เป็น Pytorch เมื่อเรารีดรายละเอียดการออกแบบ หากต้องการทราบรายละเอียดเราต้องการความช่วยเหลือของคุณ - โปรดส่งกรณีการใช้งานของคุณให้เราโดยเริ่มการสนทนาในตัวติดตามปัญหาหรือลองใช้โครงการของเรา

ใบอนุญาต

Functorch มีใบอนุญาตสไตล์ BSD ตามที่พบในไฟล์ใบอนุญาต

อ้างถึง functorch

หากคุณใช้ Functorch ในสิ่งพิมพ์ของคุณโปรดอ้างอิงโดยใช้รายการ BibTex ต่อไปนี้

 @Misc { functorch2021 ,
  author =       { Horace He, Richard Zou } ,
  title =        { functorch: JAX-like composable function transforms for PyTorch } ,
  howpublished = { url{https://github.com/pytorch/functorch} } ,
  year =         { 2021 }
}