Scale By The Bay 2019: Oscar Boykin, Scoring ONNX ML Models with Scala

Scale By The Bay 2019: Oscar Boykin, Scoring ONNX ML Models with Scala

Recording: Scale By The Bay 2019: Oscar Boykin, Scoring ONNX ML Models with Scala

[Music] I'm really excited to be here thanks to Lexi and all the organizers and so cool to see so much awesome work being done I'm gonna talk today about a format you might have not have heard of called onyx think it stands for open neural network exchange format and how we're using it a little bit of Skype and after that kind of like what the future holds for this yeah it's a murky and uncertain world so we're gonna grapple with that existential unknowns you know so let's see if I can move forward at all okay so here's what we're gonna talk about what what problem were we trying to solve why do we use this at all it's like so many buzzwords and tech now or it's like a buzzword business like how did this check our buzzword boxes so we gotta like talk about that I'm going to explain what onyx is hopefully they'll learn a little bit if you take nothing else away from the talk maybe you all I'd be able to speak intelligently about onyx then I'm going to talk about our take on that which is agate this is joint work I should say with a good friend of mine and colleague Eric awesome as well as team members on the machine learning infrastructure team as described and Eric has a quite a penchant for naming and if you don't know agate is a type of rock very similar to onyx spelled with the Y which is also a type of rock but I don't really know much about rock so then I'll give you a short demo of this part of the library a command-line tool and finally I'll talk about what we think is coming next for our team for the industry for you personally in your own life okay so the problem statement was like you know path dependency rules everything so if you're working in email these days as happenstance has it Python has won the training Wars so people train machine learning models with Python but we have a team that does a lot of Scala services and we have a really you know what nice I don't know you can talk to a pretty nice machine learning scoring server that we really like it has a lot of features its sharded it's highly reliable it has low latency and so the problem we are facing though is we want to adopt new models trained in different Python packages but we want to keep using the infrastructure that we already know how to operate that works well we want to retain our batch scoring we can do batch scoring with spark and scalding and so those are things we liked and so the models are being trained with sky scikit-learn or pi torch and some with tensorflow and so this is kind of the problem we were facing so the approach that we took before deep learning was like you know cool was we had decision trees we're in a forest --xg boost and it's not that hard to implement those algorithms to score them training them is more complicated but to score them especially decision trees is relatively easy so our approach was rather than put Python in production and deal with that we weren't really experts in operating the issues with memory and everything else and the packaging is a disaster from our point of view etc we wanted to operate a GVM service Escala service we just implemented serializers for the scikit-learn models and then scored those models in our DVM service and it was fine it was good everyone was happy with that so as we wanted to deploy more deep learning models we want we wanted to investigate at least like that same approach scoring and PI training and python scoring on the GBM so that was kind of the motivation so we looked around and here's onyx so someone asked me earlier what is onyx a company onyx is not a company it's a group a consortium of folks that are working together a lot of names someone heard of before I guess I feel like the audio just did a weird thing but maybe that was my imagination so this looked pretty good to us you know we heard of Microsoft we know those guys no idea this might be actually okay so we took a look to see if it was any good and so also you know it's like protobuf you know turtle is cool right everybody likes that so we took a look at it like what is it all about and so onyx really is just relatively straightforward I mean it's like super straightforward set of protobuf files that describe along with some other documentation what machine learning models were now if you're Scala programmer and you like you know correctness and good stuff you're probably imagining like what this thing is that and it's not that it's like take it down three or four levels it's like a bunch of strings and a bunch of proto but like I was really disappointed I was sad of apologies to any onyx people watching later I thought you'd have these nice messages that are basically att's and they fully specify all the parameters it's not that it's a bunch of strings but it's fine so you know we look through it we're like okay we can we can live our life this way so here is an example you know the previous slide you were showing this is what how it represents a tensor so for those of you who worked with if you don't know if you don't like fancy names tensor dudas multi-dimensional arrays and it could be one of several data types they have they have a bunch of different data types pretty obvious you know it's a bunch of numbers is deep-learning big deal here are the notes each operator in a graph so it onyx is trying to be like an intermediate representation if you've worked with compilers for like the kinds of things you would do in a machine learning models so those nodes those instructions individual instructions are nodes in this graph so that's a node proto now I'll here's where I was talking about all the sadness look at all the strings they're like you would think there'd be like a big list of all the possible operators but no it's like string and you better look up in a table on a piece of paper or someone wrote down is like what the operator means so that was a little bit sad but okay and also the inputs like just a string actually the inputs can refer to like a couple of different ways inputs can come into the system but you better just learn how to interpret those strings to mean the right thing so a little bit of bummers but you know they didn't totally derail us it was it was kind of okay and so finally there's kind of three layers to this you've got your tensors you've got your operations they're called nodes and then they're just the graph in a model kind of has some metadata points to a graph and the graph you'll issed of repeated nodes they're supposed to be in topologically sorted order but you know even if they weren't you could have sorted them and then it gives you some documentation strings and some initializers for the any kind of static data or like kind of constant data if you've ever worked with like a compiler that like maybe initializes some constants so that's onyx so you know I think you basically know about as much as everybody about onyx at this point so if you go to the Onyx operator page you can just find a list of some hundred or so operators and you know the string that you saw on that little table of what the operator name is you come and find it on this page and here's gem a pretty important one generalized matrix multiplication and then there's just some like text that you read is like a human and look back at those shrinks and interpret what they mean and you go and implement them it's not actually that hard the nice thing about these operators is that they're all pretty straightforward mathematical operations that are working on numbers so we don't have to do any kind of like you know too many like nothing's very complicated and they're very well specified so you go through and you look to see which ones you want to implement so this is kind of an attempt it's not like the most sophisticated as a bowl you over is like like like amazing but it actually works you can you can you can save onyx models from PI torch from tensorflow from other things and I think cloud providers are kind of interested in this idea that you have portable models that you could send around and we were interested to see what we could do with it so that's onyx so we decided to you know implement a library we called agate we weren't interested in training the models training is ML models is relatively complex it's you really care about performance we weren't really interested in that at all we were only interested in scoring and we wanted something that was really easy to deploy we wanted just like a jar we could ship it off to a spark cluster and just go load up a file and score the models so our hope was that we could score them fast enough that it wouldn't really matter that it was pure JVM and not like hitting the GPU you definitely want to use GPU for training but you probably don't need it in a service because you can probably answer in like a few milliseconds the the queries that in this course you need to make so the next part of this talk we're going to talk about a little bit about the design of this library and then we're gonna have a little bit of a demo I'll show you some some I don't know show you some real onyx models so the first thing that we did here's by the way if you're watching along and you're like maybe mostly interested in you know machine learning I kind of like like the pedagogy of like like functional programming I guess I mean the functional thing so this is just like totally orthogonal to any machine learning concern so you can get it out if you like here but the first question we want to say is how would you design a generic tensor class so the naive thing once you get in on board with all the functors and the type constructors is like Oh a tensor should be like this parent trick the parametric class that's like generalizing these multi-dimensional arrays so if you want to get really fancy you might even have like the size of the tensor at the type level that would be pretty cool that would prevent you from like combining tensors that have the wrong shape and you can definitely go down that route we'll talk a little bit about that more in a minute but the next realized is we do care about performance a lot and so when you start to go open that box it usually runs counter quite a bit to performance and you put yourself behind the eight-ball and performance and one of the things that we really knew that we would need if this was it all gonna be successful is that we would need to have unboxed data backing these tensors so if you're gonna have doubles floats whatever we had to have that there's no way we get away with that and there are some tricks to get that to work in scala with total parametric types but what we eventually wound up introducing was the notion that tensors could only be paralyzed over override what we call a data type so and not everything is a data type so what is a data type so data type is like I'm marker for the kind of like semantic notion of what the machine learning model is working over and it connects you to physical type on the JVM so for instance we have you int eight right well what back say youant eight well it's the GBM type of a byte and once we've got this data type we also have a class tag it's a reference to a class and that allows us to create an unboxed array of that type so that's kind of one of the core design choices that we made we you know use some nice Scala tricks here you know path-dependent this is called a path dependent type if you haven't worked on it so the the the type of the underlying array depends on the value that you have and we have a bunch of these values so why would you do that kind of thing so the first reason is that the backing type on the JVM is different than the semantics type the JVM only has like say short right it doesn't have unsigned short we also like float 16s in machine learning right 32 bits is too accurate for us so we need a 16-bit float that's backed by a short as well but there are actually two types of fluid sixteens google has one that's actually a little better and there's I Triple E specification all these are backed by the same underlying type so if we just parameterised on the kind of you know float or whatever that's not general enough to cover what onyx gives you onyx has many many more types that will be represented by the same physical type so that was a really nice thing so if you haven't used path dependent types think about how you can use them and the the kind of the take-home message of a path dependent type is if you have a type that depends on another type it's often going to be well modeled by a path dependent type so here ye the element type depends on the input type here it's a good case for a path dependent ID so now I'm just moving the slot around rather than go in the next page okay so the next thing that we we we did was that we we noticed that there's a concept coming up over and over again that like you have multi-dimensional arrays a new way you talk about those there are several things associated with them one thing is the the total size of the array like you know from you might have a two by three by one array right that's referring to the total size you might be talking about the coordinates within the array so you might be talking about point 0 0 0 or 0 1 0 and lastly you have the notion of how we step through the array because a cop very common optimization that you're going to do when you have a multi-dimensional array is really back it by a single dimensional array one dimensional linear array and you have some way of stepping through these but all these things are very related to the shape so we were able to come up with a very nice generic representation of shape that we could then parameterize on these three concepts coordinates like pointing within them axes describing just the kind of the bounds of the shape like what are the upper bounds and then what we call the the Dems the dimensions like how how are the dimensions within the shape represented internally so what does that mean so for instance this is really the fun stuff if you want to represent the matrix on the right but we want to sort and the underlying representation in a linear array on the left how could we do that well you can talk about each dimension can have an offset into the array and astride how big of a step do you take each time and by what's really cool is a lot of the operations on these arrays to transpose other operations broadcasting if you've worked at Python I actually don't change the array they just change which these sizes are with these shapes are so you don't touch the array at all you just change these steps that's basically an order one operation and you might have millions and millions of values but you've changed the shape one of the coolest ones that we that we noticed probably other people who do this all the time I'm sure it's a well-known optimization if you have a single value and you want to expand it out to fill a matrix you can actually do that with a stride of zero that's an interesting I yeah rather than thinking one step every time I could take zero steps every time and that's a way to fill an array and it just works within the abstraction so the shape abstraction we got a lot of mileage out of we were able to represent coal column major row major etc in these things you support to implement onyx or something like this so the last abstraction I want to talk about a little bit is we have a storage abstraction so why don't want this storage abstraction there's basically you know we knew that we probably we would often want for testing to deal with linear arrays just Java arrays but there's also if you work with the JVM a lot you've got byte buffers float buffers if you get a network you know some object from Jim from Neddie it might be a byte buffer you might want to directly load your your matrix from that and not have to copy it out so we wanted to have a generic representation of the backing but we also knew that we wanted to have memory map files now not a lot of Java programmers are like maybe often work with that but our use case was in machine learning embeddings are a very powerful tool and an example for us might be striped merchants people who use stripe to accept credit cards we might want to embed them in some space well there's a ton of merchants but at any given moment when we're scoring a model we only need like one column from that array so it might be a gigabyte of embeddings but to score this element we just need to fetch one column out we don't want to load a gigabyte it into memory just to keep that we can do that with a memory map file the JVM can help us with that so with this storage abstraction allowed us to separate these two things and by writing most of our operations on the storage array we knew each implementation of the storage array they were concrete here's a float buffer here's an array of floats we didn't have to do boxing again we wanted to always avoid boxing and but leverage some generic nature of our code did we get some share so so the storage abstraction really helped us with that so well once we had that the next thing that we needed to do is go through and implement operators there's something like 160 some operators and you don't want to just start at the top and go down and implement instead what we chose to do is we looked at the most common operators that we saw on the models people were training and we just implemented those in order so we started going from there so we needed them out as needed pole basis it's not kind of a we didn't guess and we have sufficient coverage now most of the neural nets for transactional data so stripe deals mostly with like you know categorical data or trains actually I don't know how you want to call this but it's not like we're not scoring image it you know we're not taking pictures and identifying cats or not with them or anything like that for those you would need a few operators we just didn't bother to get to yet namely convolution okay so the next thing you might ask is well it's functional programming can't we you know we really just want to be functional programming all day right but we have to like figure out some way to get paid and so it's like I guess I'll do this transaction thing for you if I can write some functional programming code so that's kind of the deal and so but it's like since we're functional programmers who want everything to be static right but like this sounds like a lot of dynamic stuff and that was like that was really tough for Eric and I we really wanted to write a compiler and a type checker and everything else cuz that's really what we wanted to be doing with our lives but instead we do this and so we really had to answer this question why are we doing this dynamic thing is this the right way to go what do I mean by dynamic we look at this model and then we dynamically walk through it and score it is that the right thing could it be static you could build a compiler for onyx models where you give me the compiler you give me the onyx model I will statically emit some code and spit it out on disk we could do that that would not be useful for back to our use case because we have a server that's already deployed and we want to take from users they want to throw it runtime all the time new models try this model let's take that model out through another model so we were already dealing with dynamic model loading we could like dynamically do code gin and go really off the rails with it but we were already in a dynamic setting the next thing was generally the monix models have already been run through a compiler that's the whole point they're an intermediate representation again the onyx model is kind of like well suppose tensorflow is like totally broken well if it's broken it's like you're probably hosed anyways I don't really know why you need to double-check that at the scoring time so we ultimately decided to just dynamically check we walked through the graph as we're going we knew this cheap that we have in our hand type that we have in our hand if we hit an operator that is inconsistent we fail at that point we've never seen that happen you know just like we don't emit bad graphs so here's an example of the way that we do that you know we you know here's batch normalization it's an interesting operation for deep learning we have a bunch of registers that are gonna hold pointers to tensors and we go through and we get the first one out we this gettin unifies unifies a term from type checkers but we we're checking dynamically that these types are the right type they all have to match and then we just call them and that turns out to work fine the next thing is that often dynamic type checking is slow but in in deep learning in tensor math the problem that you're really facing is multiplying large matrices of numbers together and checking a type is a very fast compared to multiplying a million you know floats with another million fluids so the performance we couldn't convince ourselves that perform that we would get a performance left now the last question you might ask if you were like well should I build this thing and no one asked this of us because they're like hurricane Oscar of course they could do this and make sure it's right but it kind of occurred to us if you're another programmer was doing this and someone didn't think that they could do it perfectly it's actually a beautiful space because you already have some training code the training code knows how to train it once you've run the training code to build your model you could just take the exact same training data and ship it through this other library and assert that the values are the same so no one has to write tests the training itself generates trusted test data for every model so what we do is when we train a model it goes through and then after we've got our model we take all the training data you gave us and we verify that it exactly matches Epsilon you know the -6 or whatever with the jvm implementation and so we have we sleep well at night we know that it's right it's exactly right now you don't have to use any code that we ever wrote or any other idea but I think that's an interesting idea to take home with you that that notion there that machine learning offers this like self test ability that I still don't think that we have gotten all the mileage out up from a systems perspective yet I think there can be lots of nice systems work around that property that can still be done so now we're going to move over to the demo portion of the talk so you know I don't know it's not maybe again I'm like a systems guy so like I'm not like I don't have the cat picture and classifying oh that's a cat not a dog that would probably be a cool demo but instead I'm just going to show you a command-line interface so we have a nice command-line interface for agate so it can help you understand onyx models a little bit better so now I'm going to shift over to that if you want to play with this if you wanted to implement your own onyx you could go and check out this what's called the onyx models you and that's what we're gonna do right now so let's take a look probably make that for folks okay so we have this little tool agate it gives a help command and I have a bunch of models that I've downloaded here and these are a bunch of onyx models and some of them are quite large as you can see you know well relatively large and so agate I just ran it and you kind of give you some help let me take a look I can let me just show I don't know this amnesty is kind of interesting so here you can see a little bit about how how onyx works it's telling you you this this node had this model has six distinct operators 12 total operators so some are used more than once in this particular you know like I said we only implemented operators that we need and on an as-needed basis two of them are unsupported so two out of these six distinct operators here we're unsupported and three totes so and one of those is repeated so 25% of the total operators aren't supported and then we have a list of them we have a drill you so this is uh if you know deep learning that's a convolution as I mentioned not supported matrix multiplication is interesting it should be very easy to support but we did gem which is a more that's another thing it's like they keep adding things and Jim as more of a general version of matrix multiplication but that old instructions still around so you may encounter models with it but anyways and so you can see here's some initializers you've got a shape of a tensor you know 4 by 4 by 10 with like some floats and and then here are the inputs etc so one cool thing that we can do with this let's see does anybody know by the way is it is it suspect did that command run so fast yeah anybody know how its running that fast yes core all VM it's it's because it's written in Scala and it's like everything static it works really well it's crawl VM native image so it's like super fast so yeah there's a help mode let me take the let me dump this output mode graphviz and let's do a zoo em nest and this and now I can get dot to turn it into a picture for us that might be cool let's take a look so yeah and then this kind of makes this you know relatively pretty graph of the nodes that we have in this thing so this is the direction these are inputs they do convolution they add some things together with a constant they do r lu which is like a sigmoid function like kind of you know ish they do max pool they do convolution again they add it to some other constant r lu again do some more mac school they reshape it multiply some issues add it together boom that's the output of their thing so that's how onyx works so you can count on there so that's it so now you can kind of ask what's the future of all this i would say that the course that we took wasn't the most obvious you could kind of like you know i like back to Heather's talk today software engineering is like is becoming like your job is to download software from the internet and install it so you didn't do that what's going on you did not just download it software from the internet so we decided to do this we gave it a try it's been really great for batch scoring it's really easy to deploy that wares deploying you know tensorflow on our hadoop cluster or whatever you know that's wasn't something that we wanted to do we're not really sure how much you know like is this definitely the right path is onyx gonna win I think it's still early days for intermediate representations of machine learning you know we'll see you know in a year too like we're still keeping our options open we'll probably you know we're still considering maybe we should just go all-in on downloading software from the internet and like just like let's just download tensorflow and like stop programming just be ops people that could be like a thing so we'll see what's going on but that's about it I want to think of ready for coming to my talk I would especially think my collaborators air across time and other members of my team who collaborated on this I want to call it that I do have members on my team who were here as well in the audience thank you guys for coming and that's it that's all I got wait I you have a question here who are hanging we're going to go ahead hiya Scott thanks for the talk this is great that one now but is this open-source can I use it oh oh my gosh what a question well it's not open source but actually why don't I just come over here and then make the repository public so you can give it a try and I got to put my password in okay anyways but I'll put my password in and you'll find it and it's github.com / stripe / yeah I probably am out of time I don't know if I have any more questions for question good time for one or two more okay I'm happy to answer any other questions okay we got one right thing hey-ya really cool project so all the operators implemented in Pierce Carla operations on arrays they're apparently they're implemented in pure scalar operations but we have this abstraction between like a little bit of an extraction between the underlying thing so they're not directly on the arrays they yep so kind of yes ish yeah I guess the the follow-up or real question is you mentioned things like convolutions are not yet implemented and those are typically pretty heavy and you'd normally drop down to like a native implementation so is there a plan to be able to do that yes yes that's a good question so we're definitely have no interest in making it fast enough for training so in our works we can score relatively large models in few milliseconds some of those so like matrix multiplication is also heavy enough that you want to like not go through an abstraction layer so for that we wind up we do the operation directly on arrays for convolution we'll probably do something similar internally we looked at like JP there are some libraries you can take a look at other like deep learning 4j is talked about doing onyx support but they didn't really do it yet there there's like India Ray 4j you can do we may integrate more with other libraries to help us with that very open to the but they just weren't it didn't need to get done immediately so it hasn't been done open the poll request now though thanks to Jorge's question yeah it is his image actually I totally should have put that on the thing there's like there's me and there's Eric awesome and that's Tom Switzer who worked on it these are like team members we were once the people who interviewed the most people in a period of time at stripe so we they they put that silly thing up and I kept it and I like to be self-deprecating so that's my github picture now no it is so you can see a little of the Hat yeah yeah all right thank you very much [Applause] [Music]