Scaling Tensorflow data processing with tf.data (TF Dev Summit '20)

Key Takeaways

[Music] hi all I'm Rohan and I'm here to talk to you about how you can scale up an input data processing the TF data so let's start with a high-level view of your ML training job typically your ML training step will have two phases to it the first is data pre-processing where you're going to look at the input files and do all kinds of transformations on them to make them ready for the next phase which is model computation while you're doing data pre-processing which happens on the CPU you might be doing some kind of things such as for images you're cropping them for videos you may be sampling them and whatnot so if your training speed is slow you could have a bottleneck in either one of these two places and I hope that the talk on profiling would give you an indication on how to figure out which one of the two phases you are getting slow at and I'm here to talk to you about the first kind of pre-processing bottle like the bottleneck which is data pre-processing so let's try to look into what this bottleneck really is so the in the last few years we've done a fantastic job making accelerators which do the ML operations really fast and so the amount of time it takes us to do a matrix operation on all the linear algebra operations is is a lot smaller but the hosts and the CPUs that feed the data to these accelerators have not been able to keep up with them and so there ends up being a bottleneck we thought that we could mitigate this by making the models more complex but what happens is that the accelerators have constraints on how much RAM they have and more importantly where you deploy these models tends to be something like a mobile device or something like that which tends to restrict the amount of complexity you can introduce into your model so that hasn't really panned out the second approach people take is that they try to attain larger batch sizes but larger batch require a large amount of pre-processing to assemble the batch so then that puts further pressure on them so that's why this is becoming an increasingly larger problem within alphabet and even externally and I'm going to talk to you about how you can solve it using TF data D F data is tensor flows data pre-processing framework it's fast it's flexible and it's easy to use and you can learn more about it at our guide to sort of for background for the rest of the talk I think I'm going to go through like a typical T of data pipeline and that will help us in the later later stages so suppose you have some data in some TF record files which are your training data so you can now start off with a TF record data set with that data and then after that you start doing your pre-processing this is typically the bulk of the logic so if it's images you're doing cropping maybe flipping or all sorts of things there after that you will shuffle the data so that you don't train to the order in which you see the the examples and the the input and that helps you with your training accuracy and after that we would batch it so that the model the accelerator can now take use of vectorized computations finally you want to do some software pipelining so that you ensure that while the model is of working on one batch of data you can then the pre-processing side can produce the next batch so that everything works very efficiently finally you can then feed this TF data data set to a karas model so that you can now start doing your training so given that sort of basic pipeline and suppose you have a bottleneck the first thing I'd recommend you to do is to go through our single host performance guide and try to utilize every trick and transformation that is available in TF data to be able to extract the maximum possible performance so that you using all the cores and and and whatever there's excellent information at the guide that we have here and even Giri did a great talk at the ML Tokyo Summit which you can take a look at to learn more about this so that's the first thing I'd recommend you do but suppose you've done that and you've tried all the different recommendations that we have here but you're still bottlenecked on that on that data pre-processing part and don't worry you're not alone this is very common we've increasingly seen this with a lot of internal customers and so now I'm going to I'm very pleased to present a couple of solutions that we've been working on on the team to help you solve that problem so the first idea is that why why don't we just do use the computation so suppose you're playing around with different model architectures your input pre-processing sort of part kind of remains the same and if it's expensive and time consuming why don't we just do it once save it and then every subsequent time we just read from it and do that quickly so we noticed a bunch of internal customers teams within alphabet who were trying to do this on their own outside of TF data and we decided to bring it into TF data and make it incredibly fast flexible and easy to use and so this is what we call snapshot the idea is what I explained to you it's you materialize the output of your data pre-processing once and then you can use it many many times this is incredibly useful for playing around with different model architectures and if we have settle down on an architecture doing hyper parameter tuning and so you you you can get that speed up using snapshot next I'm gonna go through the pipeline that we talked about before and see how you can add snapshot to it to make it faster so that's the original pipeline that we had and so notice that there's this pre-processing step which is expensive so now with snapshot you just add a snapshot transformation right after that with a with the directory path and with this everything that is before the snapshot will now be written to disk the first time it's run and then every subsequent time we would just read from it and we would go through the rest of the steps as usual one thing I'd like to point out is that we place the snapshot at a particular location before the shuffle because if it's after the shuffle everything gets frozen so all the randomization that you get out of shuffle you lose because every subsequent time you're just going to be reading the same exact order again and again so that's why we introduced it at that at that stage in the pipeline so snapshot we've been we've developed it internally there are internal users and teams that are using it and and driving benefit out of it and now we're bringing it to the open source world we published an RFC which has more information about it and some of the technical details and this should be available in tensorflow 2.3 but I believe it would be available in the nightly shortly so remember I talked about two ideas so the second idea is that now not all computation is reusable so because suppose you had some randomized crops in there and if you wrote that to disk and read them back you again use that lose that randomization and so snapshot is probably not applicable in that scenario so the second idea is to be able to distribute the computation so the initial setup is that you have one host CPU which is driving a bunch of this accelerator but now you can offload this computation from this horse to maybe a cluster and now you can utilize the ability and the computational power that you have for all these different workers to be able to feed the host so that you're not bottleneck on the input pre-processing anymore and things move fast this is TF data service it's it's a TF data feature that allows you to scale your workload horizontally so if you're seeing a slowness in your input pre-processing you can start adding workers and it'll just scale up it's got a master worker architecture where the master drives the work for the different workers and it gives you fault tolerance so if one of the workers fails you're still good and you still can make progress so let's see how you can use the TF data service for your for the example that we have so here instead of having sort of an expensive pre-processing let's say you have some random SP processing so now this is not snapshot able because if you snapshot then the you lose the randomization so we will provide you a binary which allows you to run the data service on the cluster set of manager that you like whether skew burn at ease or cloud or something like that and then you can once you have that up and running you can just add or distribute transformation to your TF data pipeline and provide the master address anything before the distribute transformation would now get run on the cluster that you have set up and everything after will run on the host and so this allows you to sort of scale up again note that because we are not doing any kind of freezing of the data we can now put this transformation as late as possible in there so notice that I've put it after shut the shuffle transformation the service like snapshot has been developed with internal users they've been using it and it's been like a game-changer in terms of TPU utilization and now again we're bringing it to to you and so we published an RFC which was well-received and this should be available in in 2.3 for you to play around with so to summarize what did I talk about today so as with various trends in hardware and software we've ended up in a scenario where a lot of input machine learning jobs are getting bottlenecked on input pre-processing and I've told about two solutions that TF data team has been working on to help you solve this bottleneck first a snapshot which allows you to reuse your pre-processing so that you don't have to do it multiple times and the second is the TF data service which allows you to distribute this computation to a cluster so that you you get the scale up that you need I hope you play around with these and give us feedback and thank you for your time [Music]