5 Steps to Speed Up Your Data-Analysis on a Single Core

Speaker 1 [00:03]

Thank you very much. Hi and welcome. My name is Jonathan. If you have any questions afterwards, feel free to catch me outside in the Discord, via email, on Twitter, whatever. I'm working for Scalable Minds. We are a small software company in Potsdam, close to Berlin, and I just want to explain some parts of my domain, why I want to speak about the speedups, and why this actually was quite important for us. So we've built Webgnosis. Webgnosis is a 3D viewer and annotation tool for 3D image data. So here you can see some electron microscopy data of brain tissue and what we are doing besides that is also the analysis of those segments. So we're tracing which part is a cell in the brain tissue, and this is what's colored here in the picture. If you want to see some of those cool data sets that biological researchers acquired, you can check out WebNozzles.org. We also built a Python library to interact with it that's WebNozzles.py, and you can just do pip install WebNozzles to try it yourself. This is all also open source. When we are doing data analysis on those data sets, we are speaking about petabyte scale image data sets. The resolution of those data sets is in the nanometer scheme, so there is quite a lot of data we are just handling. Then we have some machine learning systems that we train on this data and run on it, and we have some post-processing steps on this. We do segmentation based on the predictions we get from our machine learning systems, and we do agglomeration. so we can trace longer cells, especially because the neurons have very thin processes that's quite hard to track. And what we get out basically is those neuron reconstructions where we can see a cell, and not just one cell, but possibly all of the cells in the volume, and where they connect, where the interactions are, the synapses of the neurons, so that the biological researchers can do the biological research in the end and not the boring data analysis. And, those experiments are running for weeks or even months in high-performance computing clusters. So, that's just to give a bit of a context where actually the problem we encounter are coming from. So, we're talking about speedups with data analysis, so why to speed up when it's slow? And especially it's important that you speed up then, and also it's important not to over-optimize everything beforehand, which I know many of us are prone to, and I think it's all fair if it's totally working to hack something, and if the data grows, then it's getting too slow, then you can tune it even more. And that's actually what we did. So here's an example of one of the problems we encountered. So we had a pipeline running and where we combine statistics of those segments of chunk data sets. So we have the brain tissue, we chunk it into smaller parts and we gather statistics for those cells in those chunks and afterwards we want to combine them across the different chunks. And after the data grew even more, it just couldn't run at all anymore because it crashed with out of memory after two days running and while the machine was having a terabyte of memory. So that's quite annoying. We thought well that cannot possibly be good, let's do something. And then we did some optimizations, brought it down to 18 hours and then even after some more optimizations down to seven hours. And that's quite important because when we do the experiments now they run first and And we can do more iterations per week when we try different combinations of parameters that was really important. So we had a speedup of something 500 to 700% in total for this single analysis step. And I've prepared a toy example where we will see all the effects we did there, just in a really small example. And we will have a speedup of roughly 200% there, but it's quite comparable if the data just grows and it's also a 2D example and not 3D, then this will be similar. So the example you can see here on the top right, there we have chunked segments. So all of the circles are different segments and in this case the statistics are really simple. We just look at the center of all of the different segments and the size or the volume of the segments. We have this per chunk and we want to combine it across all the chunks so that we have a single center per segment and the volume of the segment. The volume is just denoted by the size of the circle here. So that's about why we would speed up. Why should we speed up on a single core? Because that's somehow the title of the presentation, right? So there are some parts where you want to parallelize soon, but think first maybe about if you want to parallelize, if it pays off, if you can do much more optimization on a single core first. So because every optimization you do on a single core usually also pays off afterwards if you have parallelized it. If you have a speed up on a single core, you probably will have the same speed up also after parallelization. Parallelization needs resources, so you need to have the cores, the memory, cluster nodes if you are in high performance clusters or cloud nodes. It usually costs some money, time, time to run, time to build it. And some code might just not be easy to parallelize or it's purely impossible. So exactly in this example, it would be quite hard to do more parallelization. Since we did gather the statistics per chunk of our data and we just want to combine it all afterwards, that's not easily parallelizable. I brought you five steps. The first one is profiling, and I think that's the most important one where we'll also spend some more time. Then we'll talk about efficient I.O., about data formats, about vectorization, memory and precision trade-offs and cheating with number, just as a very last point. So let's have a look at the actual code. So I brought you the notebooks. I'm not doing any live coding. I would be too nervous for that, to be honest. So this is just the example I showed you before. So we do some imports and then we just have some code and we know, well, it's too slow. want to optimize it and we have some load data code some combined statistics code and there's just a call that calls those both and maybe we just have a brief look what this roughly does so load data seems to take a couple of JSON files and has some lists and then just takes the data of the JSON files in a pencil. Okay. And combine statistics. Okay. There's something happens. Let's see later what it exactly does. We don't need to know for sure right now. Um, and then center is somewhere multiplied with account. Okay. That might make sense because smaller segments need to have a different weight when combining the centers, and there we return the results for the different things, okay. But that's just somebody wrote it, we want to optimize it, just skim through it, no deep understanding yet, totally fine. So we run it, and I have this handy plot function, seems like it does what it should be doing. So now we want to profile it. are the bottlenecks and I think that's the most important part you can optimize many things about this code and you could dive in and think oh that seems to be too slow but let's check actually first what are the slow parts so what we could do is run time it that's built in in the Jupiter or you can also use directly with Python and I have prepared this large data set as well so this one was just a small sample okay and this seems to take like 15 seconds more or less okay interesting now we know at the time but let's get into the details so there are multiple profilers in Python and I'm going to show you two here exactly one is yuppie yet another Python profiler I think something similar and And okay, we can either measure CPU time or wall clock time, and we start YAPI, and we just call the function we're interested in, YAPI stop, and print the results. Then we can see here how often a specific method or function was called. And the thing is that YAPI is instrumenting all of our code to exactly get those information and get the statistics for it, and this just takes time, which is problematic. So if you have lots of calling functions, like here, some things are called 10,000 times, that actually might be a problem because the profiler is just slowing down everything and and also just giving different results since those are instrumented much more often than others. Another profiler I quite like is PySpy. PySpy has the nice property that it's sampling based and it's just attaching to the Python process you're running. You can also just, if you use a Python script, do PySpy and run your Python script afterwards, or you just attach to a running process by getting the process ID here and just attaching this. This is just a shell script, basically, what we're doing here. We are opening a sub-process, starting PySpy to watch exactly on the thing that is currently running, and gather the statistics there. This is much more efficient. We still have a slowdown to roughly 20 seconds, something like this, but since it's just sampling from the outside and also it's really about efficiency here, I think PySpy is written in Rust and just attaching to the process from the outside, so we don't have that much deviation from the actual runs. I've configured some more things here, so we can also check native code if we have native extensions, C extensions, Rust extensions, whatever. We can also watch sub-processors. That's just showing how everything is grouped later in the graph. And we have an output that is a flame graph in this case. So PySpy can do chop-like outputs, if you know what that is. Also different things and flame graphs. And let's have a look at it here. So, this is a flame graph for this program. Can you read it some? Or is it too small? Okay. Perfect. So, here we have lots of calling functions just from the outside because it's in a Jupyter notebook, but we're actually interested in load and combine stats. I can click on it and just resize on this. And we can see combined stats is taking just a little portion of the time, and load data here seems to be quite a lot. So OK, load data definitely seems to be a bottleneck. And we can also have a look more here, JSON decode, OK, that takes quite long. We're actually talking about numbers, so maybe there's something more efficient than just decoding strings from a JSON. So with this information, we can go on. And I've written a small decorator here, measure and compare, which does exactly the same. We just get the flame graph afterwards and also it just gives us the time of the step. So profiling, I would say, is the most important of those steps because it helps you where you want to go deeper and actually spend time on optimization. So just a small overview, we've looked at PySpy, it is sampling based, we have flame graphs, it can also do speed scopes if you know what that is, there's YAPI that does it linewise, Cprofile is built in with Python, it's directly there, PyInstrument is quite popular, Pylontere I think rather young, I didn't try it myself but also really popular, And for memory profiling, I usually use memory profiler, or GAPI 3, also quite popular. For speed, I think except for PySpy, all of those are directly instrumenting your code, so they all suffer from the same downsides as YAPI. So, we've already seen that data loading seems to be a bottleneck, okay, because we're loading lots of JSON files. So a typical thing is we write everything as a JSON, or as a CSV, or maybe as a YAML if you're into that, and that oftentimes might not scale. I mean, we implement it as writing it for JSON when the datasets were just orders of magnitude smaller, and what usually is much more efficient is directly writing into binary data format. So there are multiple, HDF5 is quite known, also compatible with other languages, there are many implementations, NPy is from NumPy directly, Parquet is quite known in Pickle for Python things themselves, you can also use an SQLite database, why not? Or SAR for multidimensional arrays, also possible, or many others. So what we did in our case was using HDF5 and let's have a look. So I also stored all the input chunks directly as HDF5 files here just for the preparation of this talk. And now we're just loading the HDF5 files. Still the rest looks pretty similar. We are pending the different segment numbers, counts, and centers, and everything else is just the same. We use the same combined stats function as before, so the result seems to be fine. We have now just much smaller times, so this is already 55% faster than before. And have a look at the flame graph again. So now load data is still the majority but it looks more sane to be honest. Also I mean it's still in different files. If we would be able to just squeeze everything in one single file that would probably still be much more efficient. I didn't put this in here, but also just iterating across 10,000 files in this case isn't a good idea. So if you're able to just squeeze everything in one, that's usually more efficient as well. So the next step is vectorization. So now we go from the data loading to actually the combined stats parts. So let's have a look at the original implementation once again. So what we do here, this is centers per segment and counts per segment. So this might be a group by across the segments, just guessing from the names. And then we go across all segments and sum up things. Okay, so after loading this from HD5, we already get NumPy arrays. Why not do the math with NumPy as well? Let's go down here again. So we just do this. Here the group by code is still the same. And after we did the group by, we put everything in a single array and use NumPy methods directly. So here the sum on this and here calculations for the center. And this obviously is much more efficient already. So the number here is not that much more impressive, but still it gives us a bit. And that's probably what all of you know already. But the important part is we can also do the group by with NumPy. So what we are doing here is we concatenate all the input arrays directly. And then to do a group by a NumPy, we sort all our arrays by the segment numbers, get the first index of the different segments and split all the arrays by those indices. And this directly gives us the different groups for each segment. So if you want to have a look more deeply at the code, I think you can just check out the GitHub later. But this now actually gives us 100% more, which is crazy. So NumPy, in comparison to, for example, just Python, is really compact in the representation because all of your data is just continuous, more or less, in memory, and you have optimized functions directly for those things. Then maybe for the groupby, you think, well, I could have done that with pandas as well. Yeah, definitely. But also, for example, especially groupby with pandas isn't that fast. It is faster than the Python loops we saw before, but still the NumPy version is much faster. There are also better possibilities. That's just the last row. There's so many alternatives, and some of them are quite efficient. But I'd also recommend if you want to really get down into the optimization part, just sticking to NumPy is also a good case because what can you do if the group by a slow in some library, probably not that much. So if you just use NumPy for the different things you have, you can also tune it much better than with anything that's just pre-made. point is memory and precision trade-offs. When we're talking about speed, we also have to talk about memory and precision. Usually, we can have some trade-offs in this triangle. For example, we can change the data type, and we're just doing that in a second. This will take less memory directly, but also the precision will be less, but it might just be faster because we're iterating about smaller portions of data. Then there are other things like iterative methods. You can consider divide and conquer things, reduce. This usually also uses smaller memory, but you might lack numerical precision in the end. A typical example also are look-out tables or just compression. When we check this here, this is basically the same code as before, but we do as type uint32 for the segment IDs first because they just all fit into 32 bits. Before it was 64, and the same for the counts and the centers, and the rest is totally the same as before. We even get 20% more speedup, which is great, and just by reducing the precision, which we don't need. We don't need the full precision at all. The last point I'll show is if that is not enough, you can use chitting. So far, we've used NumPy, but if you use NumPy, you can also consider Numba, which is chitting your code and optimizing the different calls, because for NumPy code, you still have to consider that all the arrays can be n-dimensional, can have different types, they're different checks, and you can just put number n chit here before. The n is for the no Python mode, so it forces the chitting that is just all native code afterwards. Everything is still the same, and we just run it again. We will run into an error here, okay, failed in no Python mode, let's see, can we find something here? No implementation of function get item for the signature, okay, interesting, the slice. If we check this again, show it here. This is just this line, and this just doesn't work with numbers. There are some parts that won't work that well, and this is just one of those. You will get different error messages, but you can go through all of those, and this will be fine. We fix this by using reshape instead, and then it works. The first time we're running it, it will be slower than before because the chitting itself takes also the compilation time, but the second time it's a bit faster than before. In this example, because just averaging all those centers and the counts, there's not much that number can speed up, but I've also prepared here a different example. We are looking also at the covariance metrics of those and there after just also using number chit here, now we get 5% more than before. So that's really simple. Just add number and chit and you have to fix some errors probably when you use number, but afterwards it might just be faster. So besides what I've shown you, there are obviously more things. So you can upgrade your dependencies in Python itself. That might just be faster afterwards. Just check the change logs if that has something to do with your parts. You can use faster Python runtimes if that's less data centric and more Python loops, etc. You can optimize critical code paths with something that's closer to the metal like Python, PyBind 11, CFFI for Rust, there's PyO3. I just watched a talk yesterday about ONNX, that's why I also put that in here. And then also you can think about parallelization, but optimizing your code first on a single core just pays off afterwards, so that's what I'd recommend. The most important thing for you to take away is just do profiling, look where is the bottleneck, and then you have more steps how to optimize it. Thanks a lot, and you can find the code in this repository. If you're interested in working with me at Scalable Minds, we're looking for engineers. to help me up outside.

Speaker 2 [26:12]

All right, so that was Jonathan Striebel with his talk. So we still have a few questions from Slido, and if we have time, we'll get to the audience. So the first question is, would you advise storing data in a single file if we eventually want to spread our computation over multiple cores?

Speaker 1 [26:31]

That's a good question, so it depends also on your file system, right? I mean if Reading from the same file is not a bottleneck for you Then this might be a good idea because if you change your parallelization strategy later, and you don't know yet From which files you will need to read and you have larger chunks for example in a later step Then you will need to read from more files again But if reading from the same file across different nodes, for example, is a bottleneck then you should definitely consider doing this in different files I mean, that's why there is chunking and sharding in different libraries exactly for that reason

Speaker 2 [27:08]

All right, so the next question is, have you thought about doing all of this in a database? It seems like the data is well-structured.

Speaker 1 [27:16]

Yes. Actually, so we didn't do it in production. I considered this when doing this talk and preparing it. So I think in this special case, it's quite possible, but also just loading it in a database first is quite an overhead. And having a single database where everybody is writing to first is just also not that efficient. I mean, if you load it into memory once, you can also just do the operations and then it's done. With a database you first need to get the data in and if you have a data analysis pipeline You do it usually just once and we don't Do the same statistics acclimation multiple times for the same data. So then it doesn't pay off I'd say

Speaker 2 [27:59]

All right, so that was the side of questions. Does the audience have any questions they want to ask? We have time for one or two. All right, then I'll do another side of question. So another question is, you mentioned performance evaluation at the beginning. So one of the audience asks, how do these performance evaluations compare to inline magic, like line profiler.

Speaker 1 [28:31]

So, I think the YAPI I've shown is quite something like line profiling. So, that's really similar. I'm not sure where the big difference would be, to be honest. So, PySpy that I've shown is sampling-based. So, you have a sampling rate. And, for example, in this case, it sampled just 100 times per second. and this helps you to not get too much overhead from the profiling itself since it just in regular intervals checks where the code is and if it's running for long enough you'll just get good data and those line based ones just instrument all your code and every single call in there will be noted somewhere and this just takes more time

Speaker 2 [29:21]

All right, thank you very much. So that was a great talk. And I'd like to thank you again for the session.