Understanding numpyro

numpyro の基本的な使い方から内部構造の把握．
plate, sample, HMC, MCMC の挙動を中心に．

Transcript

1. numpyro Λཧղ͢Δ _zak3 2022/10
2. Agenda • numpyro ͱ͸ • Πϯετʔϧ • ࢖͍ํͷཧղ • ಺෦ߏ଄ͷཧղ 2
3. nympyro ͱ͸ • ܰྔͳ֬཰ϓϩάϥϛϯάϥΠϒϥϦʢPPLʣɽ • jax ϕʔεɽJITίϯύΠϧɼࣗಈඍ෼ɼGPU ߴ଎ԽʹରԠɽ 3 Github: https://github.com/pyro-ppl/numpyro
4. Πϯετʔϧ • pip install --upgrade pip • pip install numpyro 4
5. ࢖͍ํͷཧղ
6. 8 Schools • ίʔνϯά͕ SAT*1 ͷ੒੷ʹ༩͑ΔӨڹʹ͍ͭͯͷσʔλɽ • y : ॲஔޮՌɼsigma : ඪ४ޡࠩɽ 6 *1: [Wikipedia] SAT͸ɺຊདྷScholastic Aptitude TestʢେֶదੑࢼݧʣΛҙຯ͢ΔུޠͰ͕͋ͬͨɺ1990೥ʹScholastic Assessment TestʢେֶೳྗධՁࢼݧʣʹมΘΓɺݱࡏ͸ུͱͯ͠Ͱ͸ͳ͘SATͦͷ΋ͷ໊͕લʹ࢖ΘΕ͍ͯΔɻ
7. ϞσϦϯά • Ϟσϧ͸ؔ਺Ͱදݱ͢Δʢ͜͜Ͱ͸ eight_schoolsʣɽ • άϥϑΟΧϧϞσϧͱରൺ͠ͳ͕ΒͩͱΘ͔Γ΍͍͢ɽ 7
8. αϯϓϦϯά • mcmc.run ͰαϯϓϦϯάΛ࣮ߦɽ • ͜ͷ࣌ʹϞσϧͷҾ਺΋ࢦఆ͢Δʢ͕͜͜গ͠෼͔ΓͮΒ͍ʣɽ 8
9. αϯϓϧͷ֬ೝ (1 / 3) • arviz Λར༻͢Δͱศརɽ • mcmc ͔Β InferenceData Λऔಘɽ 9
10. αϯϓϧͷ֬ೝ (2 / 3) • summary 10
11. αϯϓϧͷ֬ೝ (3 / 3) • trace plot 11
12. ಺෦ߏ଄ͷཧղ plate, sample, NUTS, MCMC ͷϓϩάϥϜతͳڍಈʹ͍ͭͯɽ
13. plate • খจࣈ͕ͩ Messenger Λܧঝͨ͠Ϋϥεɽ • with ۟Ͱར༻͞ΕΔ context managerɽ • __enter__, __exit__ Λ࣋ͭɽ • plate ίϯςΩετ಺Ͱͷ sample ϝιου͸ ϓϨʔταΠζʹϒϩʔυΩϟετ͞ΕΔɽ 13 ར༻ͷ؍఺Ͱ͸͜ͷڍಈ͚ͩཧղ͓͚ͯ͠͹ ྑ͍ͱࢥΘΕΔɽ
14. Messenger • ؔ਺ fn ΛҾ਺ʹΠϯελϯεԽɽ • with ۟ʹೖͬͨͱ͖ɼ __enter__͕࣮ߦ ͞Ε _PYRO_STACK ʹ appendɽ • with ͔۟Βग़Δͱ͖ɼ __exit__ ͕࣮ߦ ͞Εɼ_PYRO_STACK ͔Β popɽ • process_message / postprocess_message ͸ܧঝઌͰ࣮૷ɽ 14
15. primitives.sample • plate ίϯςΩετͰ͸ _PYRO_STACK ͕ଘ ࡏ͢ΔͨΊӈͷ࿮಺ͷίʔυ͕࣮ߦ͞ΕΔɽ • apply_stackʹͯ _PYRO_STACK ্Ͱ msg[“fn”] ͕ධՁ͞ΕΔɽ 15 ελοΫ=จ຺ɽ ॲཧΛจ຺ʹΑͬͯ੾Γସ͑Δػߏɽ
16. apply_stack() • msg ͸ dict ܕͷ஋ɼhandler ͸ Messenger ͷαϒΫϥεͷ Πϯελϯεɽ • _PYRO_STACK ্ʹ͸ɼ plate, trace, substitute, seed ͳͲͷෳ਺ͷ handlers ͕ొ࿥͞Ε͍ͯΔɽ • see: https://num.pyro.ai/en/latest/handlers.html • ҎԼͷॱͰ࣮ߦ͞Ε msg[“value”]ͳͲ༷ʑͳ৘ใ͕෇Ճ ͞ΕΔɽ • handler#process_message • default_process_message • hander#postprocess_message 16 see also: Poutine http://pyro.ai/examples/effect_handlers.html
17. plate#process_message plate#postprocess_message 17 plate ʹࢦఆ͞Εͨ size ͰϒϩʔυΩϟετ͞Ε͍ͯΔՕॴɽ .expand ʹΑΓ Distribution Ϋϥε͔Β ExpandedDistribution Ϋϥε ʹม׵͞Ε͍ͯΔɽ
18. default_process_message • hander#process_message ʹΑΓม׵͞Εͨ msg[“fn”] ͕࣮ߦ͞ΕΔɽ • msg[“args”], msg[“kwargs”] ͕Ҿ਺ɽ • msg[“value”] ʹฦΓ஋͕֨ೲ͞ΕΔɽ 18
19. NUTS • __init__ ͷΈ͕ఆٛ͞Ε͍ͯΔͳͷͰɼ ࣮ଶ͸ HMCɽ 19
20. HMC • MCMCKernel ͷαϒΫϥεʹͳ͍ͬͯΔɽ • ந৅ϝιου init, sample Λ࣮૷ɽ 20
21. HMC#init • self._sample_fn ͕ηοτ͞Εɼhmc_init_fn ͷ࣮ߦ݁Ռ init_state ͕ return ͞ΕΔɽ 21
22. HMC#sample • self._sample_fn Λ࣮ߦ 22
23. MCMC#run • _single_chain_mcmc ͷҾ਺Λ෦෼ద༻ͨ͠ partial_map_fn Λ࡞੒ɽ • num_chains, chain_method ͷछผʹԠͯ͡ partial_map_fn Λ࣮ߦͯ݁͠Ռstates/ last_state Λ֫ಘ͢Δɽ • states : dict ܕ • states[‘z’] ʹύϥϝʔλͷαϯϓϧ͕ ֨ೲ͞Ε͍ͯΔɽ • last_state : HMCState ܕ 23
24. HMCState • namedtuple Ͱఆٛ͞Ε͍ͯΔɽ 24