Swift for TensorFlow: The Next-Generation Machine Learning Framework (TF Dev Summit '19)

Key Takeaways

[Music] hi everyone I'm Chris and this is Brendon and we're super excited to tell you about a new approach to machine learning so here in the tensorflow team it is our jobs to push the state of the art of the machine learning forward and we've learned a lot over the last few years with deep learning and we've incorporated most of that all into tensorflow - and we're really excited about it but here we're looking a little bit further beyond that's flow - and what do I mean by further well eager mode makes it really easy to train a dynamic model but deploying it still requires you take that and then write a bunch of C++ code to help drive it and that's that could be better similarly some researchers are interested in taking machine learning models and integrating them into larger applications that also often requires writing C++ code we always want more flexible and expressive auto differentiation mechanisms and one things we're excited about is being able to define reusable types that then can be put into new new places and used with automatic differentiation and we always love improving your developer workflow we want to make you more productive by taking errors in your code and bringing them to your source and also by just improving your iteration time now what we're really trying to do here is lift tents load entirely new heights and to do that we need to be able to innovate at all levels of the stack this includes the compiler and the language and that's what Swift for tensorflow is all about we think that applying new solutions to old problems can help push machine learning even further than before but let's jump into some code so first what is Swift Swift is a modern and cross-platform programming language that's designed to be easy to learn and use so if two uses types and types are great because they can help you catch errors earlier and also they encourage good API design now Swift uses type inference so it's really easy to use and very elegant but it's also open source and it has an open language evolution process which allows us to change the language and make it better for machine learning which is really great let's jump into a more relevant example this is how you define a simple model in Swiffer tensorflow as you can see we're laying out our layers here and then we can find a four word function which composes them together in a linear sequence now you've probably noticed that this looks a lot like Karis and that's no accident of course we want you to be able to take what you know about Karis and bring it forward into this world as well now once we have a simple model let's train it how do we do that all we have to do is instantiate our model pick an optimizer and some random input data and then pick a training loop and here we'll write it by hand one of the reasons we like writing it by hand is that it gives you the maximum flexibility to play with different kinds of constructs and you can do whatever you want which is really great but some of the major advantages of Swift for tensorflow are the work flow and so instead of telling you about it what do you think Brian should be sure let's do it all right the team has thought long and hard about what's the easiest way for people to get started using Swift for tensorflow and what could be easier than just opening up a browser tab this is Google collab hosted Jupiter notebooks and it comes with Swift for tensorflow built right in let's see it in action here is the layer model the model that Chris just showed you a couple slides ago and we're gonna run it using some random training data right here in the browser so we're going to stand she ate the model we're gonna use the stochastic gradient descent SGD optimizer and here we go we have now just trained a model using swift for tensorflow in our browser on some training data right here now we can see the training loss is decreasing over time so that's great but if you're ever like me and whenever I try and use machine learning in any application I start with a simple model and I've got to iterate I've got to tweak the model to make it fit better to the the task at hand so since we're trying to show you the workflow let's actually edit this model let's make it more accurate so here we are now let's think a little for a moment what changes do we want to make to our model well this is deep learning after all so the answer is always to go deeper right but you may have been following the recent literature in state-of-the-art and that not just sequential layers but skip connections or residual connections are a really good idea to make sure your model continues to train effectively so let's go through and actually add make add an extra layer to our model let's add some skip connections and we're gonna do it all right now in under 90 seconds are you ready all right here we go so the first thing that we want to do is we need to define our additional layer so we're gonna fill in this dense layer whoops and one thing you can see is that we're using tab autocomplete to help fill in code as we're trying to develop and modify our model now we're gonna fix up the shapes right here really quick so that the residual connections will all work if I can type properly that would go better all right great we've now defined our model with the additional layers all we need to do is modify the forward pass so that we add those skipped connections so here we go the first thing we need to do is we need to store in a temporary variable the output of the dense are the output of the flatten layer then we're gonna feed the output of the flatten layer to our first dense layer so dense dot applied to temp in context now for the coup de Grasse here is our residual connection so dense to applied 210 plus temp to in context run that and yes that works we have now just defined a new model that has residual connections and is one additional layer deeper let's see how it does so we're gonna reinstate our model and rerun the training loop and if you recall from the loss that we saw before this one is now substantially lower this is great this is an example of what it's like to use Swiffer tensorflow to develop and iterate as you apply models to applications and challenges but swift for tensor book thank you but Swift for tensor flow was designed for researchers and researchers often need to do more than just change models and change the way the architecture fits together researchers often need to define entirely new abstractions or layers and so let's actually see that live right now let's define a new custom layer so let's say we had the brilliant idea that we wanted to modify the standard dense layer that takes a weights and biases and we wanted to add an additional bias set of parameters okay so we're gonna define this double by double bias dense layer right here so I'm going to type this really quickly standby 15 seconds here we go whoo all right that was great so let's actually walk through the code so that you can see what's going on so the first thing that we have is we define our parameters so these are W like our weights for our neurons and B 1 bias 1 and B 2 our second bias we define an initializer that takes an input size and an output size just like dents does and we use that to initialize our parameters the forward pass is very simple to write so here's just applied two and we just take the matrix multiplication of input by our weights and we add in our bias terms that's it we've now just defined a custom layer right in collab in just a few lines of code all right let's see how it goes here's model 2 and so we're going to use our double bias dense layer and we're going Stan she ate we're gonna train it using again our custom handwritten training loop here's an example of another way that we think so for tensorflow makes your life easier because swift for tensorflow can statically analyze your code it can be really helpful to you I don't know about you but I've put typos in my code I don't know if you saw me typing earlier and so if for tensorflow here is helping you out right it's saying look you mistyped softmax cross-entropy this should be labels okay all right so we run it we train it and our loss isn't as good this was not the right idea but this is an example of how easy it is for researchers to experiment with new ideas really easily in Swift for tensorflow but let's go deeper swifter tensorflow is again designed for researchers and researchers need to be able to customize everything right that's the whole point of research and so let's show an example of how to customize something other than just a model or a layer so you may have heard that large GPU clusters or TPU super pods are like delivering massive breakthroughs in research and and the state of the advancing the state of the art in certain applications and domains and you may have also heard that as you scale up to effectively utilize these massive hardware pools you need to increase your batch size and so let's say you're a researcher and you want to try and figure out what are the best ways to train deep neural networks at larger batch sizes well if you're a researcher you probably can't buy a whole GPU cluster or rent a whole TPU super pod all the time for your experiments but you often have a GPU under your desk so let's see how we can simulate running on a super large data parallel GP or TPU cluster on a single machine we're gonna do it all in a few lines of code right here so here's our custom training loop here's our well here's the standard part right this is our one to ten training epochs and what we're gonna do is instead of just applying our model forward once we have an additional inner loop right so we're gonna run our forward pass we're gonna run our model let's run our model four times and we're gonna take the gradients for each step and we're gonna aggregate them in this grads variable okay this simulates running on four independent accelerators four GPUs or 40 P use in a data parallel fashion on a batch that's actually four times as large as what we actually run we're going to then use our optimizer to update our model along these aggregated gradients again simulating a data parallel synchronous training process that's it that's all there is to it we're really excited by this sort of flexibility and capabilities that Swiffer tensorflow brings to researchers back over to you Chris thanks so we think that the focus on catching errors early and also productivity enhancements like code completion can help you in a lot of ways and it's not just about like automating typing of code but it can also be about discovery of api's so another thing that's really cool about swift as a language is it has really good interoperability with C code and so in Swift you can literally just import a C header file and call symbols directly from see without wrappers without boilerplate or anything involved it just works so we've taken this approach in the tensorflow team we've taken this approach and brought to the world of Python and one of the cool things about this is it that allows you to combine the power of Swift retentive flow with all the advantages of the Python ecosystem how about we take a look thanks Chris the Python data science ecosystem is incredibly powerful and vibrant and we wanted to make sure that as you start using Swift for tensorflow you didn't miss all your favorite libraries and utilities that you're used to and so we've built a seamless Python interoperability capability to swiffer tensorflow and let's see how it works in the context of my favorite Python data science library numpy so the first thing you need to do is import tensorflow and import python and one should do that that defines this Python object that allows you to import arbitrary Python libraries so here we import PI plot from the matplotlib library and numpy and we assign it to NP ok after that we can just use NP just as if we were in Python so here we call linspace we're gonna call sine and cosine and we're gonna pass those values to PI plot when we run this cell it just works exactly as you'd expect thank you now this sort of kind of looks like the Python code you're used to writing but this is actually pure Swift it just works seamlessly but this is maybe a bit of a toy example so let's see this a little bit more in context open AI has done a lot of work in the area of reinforcement learning and in order to help that along they developed a library a Python library called open AI jim-jim contains a collection of environments that are very useful when you're trying to train a reinforcement learning agent across a variety of different challenges let's use opening eye gym to train a reinforcement learning agent in Swift for tensorflow right in our browsers so the first thing we need to do is we need to import gym we're going to define a few hyper parameters here and now we define our neural network in this case we're gonna pick a simple two layer dense Network and and it's just a sequential model okay after that we have some helper code to filter out bad or short episodes and whatnot but here's the real meat of it we're going to use open tonight we're gonna use Jim to instantiate the cart pole v-0 environment so that's our em we're gonna then instantiate our network right here in our optimizer and here's our training loop there we go so we're going to get a bunch of episodes we're gonna run our model I get the gradients and we're gonna apply those to our optimizer and we're gonna record the mean rewards as we train okay it's all very simple straightforward Swift and here you can see by straining a Swifter tensorflow model in an open-air gym environment using the Python bridge totally seamless and of course afterwards you can keep track of the the parameters of the rewards in this case we're going to plot the mean rewards as the model train using Python numpy totally seamless you can get started using swift for tensorflow using all the libraries you know and love and take advantage of the take advantage of what swift for tensorflow brings to the table back over to you Chris thanks Brian so one of the things I love about this is it's not just about being able to leverage big important libraries like numpy we're working on the ability to integrate Swift for tensorflow and Python for tensorflow code together which we think will provide a nice transition path to make you able to incremental emove code from one world to the other now I think it's fair to say that calculus is an integral part of the machine learning and we think that differentiable programming is so important we've built it right into the language this has a number of huge advantages including enabling more flexible and custom work with with differentials with derivatives and we think this is really cool so I'd like to take a look so we've been using Swift for tensor flows differential programming capabilities throughout all of our demos so far but let's really break it down and see what's going on at a fundamental level so here we define my function that takes two doubles and returns a double based on some products and sums and quotients if we want Swift for tensor flow to automatically compute the derivative for us I would just annotate it at differential so if our tensorflow will then derive the derivative for this function right when we run the cell to use this auto-generated derivative use gradient so gradient takes two things it takes a closure to evaluate and the point that you want to evaluate your closure at so here we go we this is what it is to take the derivative of a function at a particular point so we can change this around this one's my favorite tasty number and that works nicely now one thing to note we've just been taking the partial derivatives of my function with respect to a but of course you can take the partial derivatives you get the full gradient of my function like so often with neural networks however you want to get not just the gradients for your network as you're trying to train it and optimize your loss function you often want what the network predicted right this is really useful to compute accuracy or other debugging sort of information and for that you can use value with gradient okay and that returns a tuple containing both the value and the gradient shockingly enough now one thing to note in Swift tuples can actually have named parameters they aren't just ordered and so you can actually see that it prints out really nicely and you can access values we think this is again another nice little thing that helps making makes writing and debugging code and more importantly reading it and understanding it later a little bit easier but the one thing that I want to call out is that throughout this we've been using just normal types these aren't tensor of something it's just plain old double this is because automatic differentiation is built right into the language and so for tensorflow it makes it really easy to express your thoughts very clearly but even though it's built into the language we've actually worked very hard to make sure that automatic differentiation is totally flexible so that you can customize it to whatever needs you have and instead of telling you about that let's show you so let's say you want to define an algebra in 2d space well you're certainly going to need a point datatype so here we define a point struct with X&Y we just mark it differentiable we can define helper functions on it like dot or other helper functions and Swift for tensorflow when you try and use your code it will often automatically infer when you need gradients to be computed for you automatically computed for you by the compiler but often it's a good idea to document your intentions and so you can annotate your helper functions as at differentiable the other reason why we recommend doing this is because this helps catch errors so here so if for tensorflow is actually telling you that hey you can only differentiate functions that return values that conform to differentiable but it doesn't conform to differentiable right what this is telling you is that my helper function returns an int + int is all about taking infinitesimally small small steps as you optimize and take gradients right and integers just aren't very discreet and so so for tensorflow is helping to catch errors you know right when you write the code very easily and tell you what's going on so the solution of course is just to not mark that is that differentiable the cell runs just fine but let's say we also wanted to go beyond just defining the dot product let's say we also wanted to define the magnitude helper function that is defined is the magnitude of the vector defined by the origin to the point in question so to do that we can use the distance formula if you're going to do Euclidean distance and we can define an extension on point that does this but we're going to pretend for a moment that swift doesn't include a square root function because I want a good excuse for you to see the interoperability with C okay so we're actually going to use a C's square root function that operates on doubles okay so based on the definition of Euclidean distance we can define the magnitude and it totally just no it doesn't quite work okay let's see what's going on so we wanted magnitude to be differentiable and it's saying that you can't differentiate the square root function because this is an external function that has been marked as differentiable okay what's that saying well the square root it's a see function it was compiled by the C compiler and as of today the C compiler can't automatically compute derivatives for you so surf tencel is saying like hey this isn't going to work this is excellent because it gives me a great excuse to show you how to write custom gradients all right so here we define a wrapper function my square root that just calls down in the forward pass to the C square root function and the backwards pass we take our double and we return a tuple of two values rather the first element in the tuple is the normal value in the forward pass and the second is a pull back closure and this is what's where you define the backwards pass capturing whatever values you need from the forward pass okay so we're going to run that we're going to go back up to our definition of magnitude and change it from square root to my square root rerun the cell and it works we've now defined point and two methods on it dot and magnitude and we can now combine these in arbitrary other silly differentiable functions so here I've defined the silly function and we've marked it as differentiable and we're gonna take two points we're also gonna take a double right you can mix and match differentiable datatypes totally fluidly we return doubles we're going to take magnitudes and do dot products it's a silly function after all and we can then use it compute the gradient of this function at arbitrary data points just like you'd expect you can get the value of the function get full partial get full gradients in addition to partial derivatives with respect to individual values that's been a quick run-through of how to use customization custom gradients custom datatypes with the language integrated automatic differentiation built into swift for tensor flow but let's go one one step further let's put all this together and show how you can write your own debuggers as an example of how this power is all in your hands so often when you're developing models you want to or debugging models you often want to be able to see the gradients at different points within your your model and so here we can just define in regular Swift code a gradient debugger now it's gonna take as input a double and it's going to return it just like normal for the forward pass right it's an identity function on the backwards pass we're going to get the gradient we're gonna print the gradient and then we're gonna return it so we're just passing it through just printing it out now that we've defined this gradient debugger ourselves we can use it in our silly function to see what's going on as we take derivatives so a gradient bugger there we go we can rerun that and when we take the gradients we can now see that for that point and the silly function of a dot B the gradient is three point eight zero that's how there's been a brief tour through how automatic differentiation works in Swift for tensor flow and how its customizable so that you can harness the power in whatever abstractions or systems you need to build back over to you Chris thanks Brendon so the funny thing about all this is that the algorithms that we're building on were defined back in the 1970s and so it really took language integration to be able to bring these things forward into the world of machine learning there's a tremendous amount of depth here and I'm really excited to see what you all can do with it and we think that this is going to enable new kinds of research which we're very excited about there's also a ton of depth and if you're interested in learning more we have a bunch of detailed design documents available online well let's talk about performance a little bit now Swift is fast this comes from a number of different things one of which is that the language itself has really good low-level performance there's also no Gil to get in the way of concurrency it's much more tense full also has some advanced compiler techniques to automatically identify graphs for you and extract them so you don't have to think about them the consequence of all this together is that we think Swift has the world's most advanced eager mode now you may not care that much about performance you may wonder like why do we care about this stuff well we're seeing various trends in the industry where people are defining neural nets and then want to integrate them into other larger applications and typically what this requires this is this requires you to export graphs and then write a bunch of C++ core code to load and orchestrate them in various ways so let's take an look at an example of this alphago zero is really impressive work that combines three major classes of techniques of course you have deep learning on the one hand but it also drives it through Monte Carlo tree search to actually find and evaluate these spaces and then it runs them at scale on industry-leading TPU accelerators and so it's the combination of all three of these signs that make alphago zahra possible now this is this is possible today and if you're an advanced team like deep mind you can totally do this but it's much more difficult than it should be and we think that breaking down barriers like this can lead to new breakthroughs in science and we think that this is what can drive progress for so instead of talking about it again let's take a look Mineko is an open-source go player inspired by deep minds alphago zero project it's available on github and you can certainly check out the code and I encourage you to they're also going to be here and they have some other presentations tomorrow but the mini go project when they started out they were writing everything in normal tensorflow and it was working great until they started trying to run at scale on large clusters of TP use there they ran into performance problems and had to rewrite things like Monte Carlo seeped research in does he plus plus in order to effectively utilize modern accelerators here we've reimplemented Monte Carlo tree search and the rest of the mini go self play in pure Swift and we're gonna let you see it running right here in collab so here we define a helper function where we take in the game configuration and a couple participants these are our white and black players and we're going to run basically play the game until we have a winner or loser and so let's actually run this here we define a game configuration we're gonna play between a Monte Carlo tree search powered by neural networks versus just a random player just to see how easy it is to flip back and forth or mix and match between deep learning and and other arbitrary machine learning algorithms right here in Swift so here you go you can see them playing white black playing different moves back and forth and it just goes we think that Swift for tensorflow is going to unlock whole new classes of algorithms and research because of how easy it is to do everything in one language with no barriers no having to rewrite things into C++ back over to you Chris thanks Brenda the cool thing about this of course is that you can actually do something like this in a workbook which is which is pretty phenomenal and we've seen many different families of new techniques that can be combined together and fused in different ways and bringing this to more people we think will lead to new kinds of interesting research now our work on usability and design is not just about high-end researchers so we love them but Swift is also widely used to teach new programmers how to code an education is very close to our hearts and so I'm very excited to announce the collaboration that we're embarking on with none other than Jeremy Howard but instead of talking about this I'd rather have Jeremy speak about it now at FASTA I we're always looking to push the boundaries of what's possible with deep learning especially pushing to make recent advances more accessible we've been involved with setting imagenet speed records at our cost of just $25 and building the world's best document classifier hundreds of thousands have become deep learning practitioners through our courses and are producing state-of-the-art results with our library we think that with Swift for tensorflow we can go even further so we're announcing today that our next course will include a big Swift component Coe taught by someone that knows Swift pretty well Chris I think he means you yeah we'll see how this goes so I'm super excited to be able to help teach the next generation of learners but I'm also really excited that Jeremy will be bringing his expertise in API design and helping us shape the high level API isn't Swift for tensorflow so we've talked about many things but the most important part is that Swift for dense flow is really tensorflow at its core and we think this is super important because we've worked really hard to make sure that it integrates with all the things going on the big tensorflow family and and we're very excited about that now you may be wondering where you could get this so sweat for tents flow is open source you can find out in github now and you can join our community it also works great in collab as you've seen today we have tutorials we have examples and all the demos you saw today are available now in collab straight we've also released our 0.2 release which includes all the basic infrastructure and underlying technology to power these demos and examples and we're actively working on high-level api's right now so this is not ready for production yet as you could guess but we're very excited about shaping this future building this out exploring this new programming model this is a great opportunity for for advanced researchers to get involved and help shape the future of this platform so we'd love it for you to try it out and let us know what you think thank you [Applause] [Music]