PyTorch Developer Podcast
PyTorch Developer Podcast
Edward Yang, Team PyTorch
torch.nn
14 minutes Posted Jun 7, 2021 at 1:00 pm.
0:00
14:18
Download MP3
Show notes

What goes into the implementation of torch.nn? Why do NN modules exist in the first place? What's the function of Parameter? How do modules actually track all the parameters in question? What is all of the goop in the top level NN module class? What are some new developments in torch.nn modules? What are some open problems with our modules?

Further reading:

Liner notes:

  • python for hackability (c++ is reimplemented)
  • parameters
    • parameter collection (for optimization)
    • buffers: not considered optimizable
  • modules
    • functorial operation (_apply)
    • jit script: staged computation (init is not scripted)
    • __call__ to forward (extra instrumentation)
    • serialization / state_dict
  • new stuff: device kwarg (joel schlosser)
  • new stuff: lazy modules (emcastillo)
  • open problems: parameter initialization