High-Level API for Autograd | PyTorch Developer Day 2020

Key Takeaways

[Music] hi everyone my name is alban i'm a research engineer here at facebook and i'm working on the front-end team for pytorch in new york more particularly on the autograd subsystem and so today i'm going to talk to you about the brand new high-level api that we added for the autograd in particular i'm going to give you a quick overview of like the motivations for this new api what it actually looks like and how to use it and finally some future works for what's going to come for it so to begin with the motivation for it is that the current autograd api we have is not for everyone in particular the dot backward function is really geared toward torch.n users as it populates the dot grad fields of all the tensors and it's really associated with states and the optimizers that we use with h m similarly autograph.grad is gears toward neural network user in the sense that it is built to do backpropagation and all the namings and all the api design has been done thinking about neural networks and so we think that now nowadays that we have more general users that use autograd to do more optimization than just neural network we needed a good generic api that these people can use to actually do everything they need moreover for higher level functions such as computing a jacobian that i'm going to use here as like an example for many things it's very nice to have a reference implementation within pi torch the reason for that is up to fairly recently if you wanted to compute the jacobian with spy torch you had to copy-paste some code from a guest that was made a few years ago and so people have basically a stale version of that gist potentially in their code and we can't really make any improvement to the way jacobian are computed for everyone to basically get automatically this new upgraded version and so this new api is going to help us bring more performance improvements to the users via this generic api and following this idea we also added recently a whole benchmark system around this api to make sure that existing models and models that are used very often by our users have good performance and we can also measure improvements we're making by different change but also we can measure regressions and make sure they are not significant and we can catch them and fix them before they hit main releases so what does this new api live so it lives under the torch.autograd.fun submodule you can find it in the autograd documentation under like the functional higher level api name and it has a slightly different spirit comparing to the existing apis mainly that it takes function as inputs and not the result of the forward pass the main reason for that is are twofold one is because it's closer to the mathematical formulation and people are getting used to differentiate functions directly not do the forward pass and then ask for the gradient second it also allows us to get more freedom respect to what happens during the forward pass and in particular for some of the optimization we have planned we will need to do some special things during the forward pass that we don't want the user to have to worry about and so this new api is going to allow us to do that very efficiently um so what does this new api contain the first big part is first order gradients so as i mentioned before we have a function that directly computed jacobian for you given a function and some input points we also have a vector jacobian product which corresponds to backward mode automatic differentiation and in the neural network world it corresponds to the back propagation algorithm so this is actually very close to the existing autograd.grad function and finally we also provide a jacobian vector product so the other way around and this one corresponds more to forward mode automatic differentiation and can be used to compute directional derivatives and we also have the second part of this api corresponds to second order gradients and very similar functions so a hessian function that computes all the second order derivatives a vector hessian product which allows us to efficiently compute with backward mode automatic differentiation a product between the hessian and a given vector and the similarly the other way around a hessian vector product to that corresponds more to four mode automatic differentiation so now an example on how to use this api so again as an example for the jacobian function now you don't have to copy paste code from that old guest you can just import it from torch and as you can see in the first example you have some input you just call jacobian on it and you get the value of the jacobian and that's easy that's that's everything you need to do what is very nice as well is that you can compose this new api with the existing autocad api and here for example if you input requires gradients you can ask the jacobian computation to create graph and so to be able to backgrade through it and as you can see you can then compute norm of the jacobian you just got and backward through that and then you can compose that with the rest of your training to get all the quantities that you need um there are many more examples and a few examples here are like gradient penalties based on jacobian computations uh jacobian vector product computations that corresponds to the forward mode automatic differentiation this is especially interesting because they actually compute directional derivatives in higher dimension these derivatives are very useful for quite a few optimization algorithms and similarly the second order methods that compute hessian or hessian vector products allow you to do newton step methods or approximate neutron stamp methods much more efficiently and much in a much easier way than with the current autorad api so to finish with some future work uh the first part is actually things that we are already currently working on so the first one is for node 18. so the idea would be for both jacobian vector and hessian vector products to replace that with an actual forward mode automatic differentiation so this is working progress hopefully we'll be able to release that very soon and get very good performance improvements based on that you so one interesting thing about this new api is that you will get the benefits of this new forward mode id for free if you already use that api when it is released similarly you can get batch gradients so this is a collaboration with the vmap features that you heard about and the idea here would be to speed up competition for jacobian and hessian computations and finally torch.nbn composibility is another topic we're working on the reason for that is currently and then modules hold a lot of states and they are not functional so they don't work very out of the box with the new api that we design so the idea here is going to be to work on torch and to try and provide a functional version of the indian module so that we can work with them more efficiently and for longer time work we are looking for ideas and please do share your feedback on how this work how it helps you do what you want to do do open issues on github or on the pytorch forum if you have questions or concerns about what is happening here and please help us make this api exactly what you're looking for so thank you very much for listening and have a good day