Flexible ML Experiment Tracking System for Python Coders with DVC and Streamlit
There exist so many tools to do data science today that it is sometimes difficult to navigate. Many of them are AI platforms that “do everything by clicking on a UI” and do not leverage pre-existing tools e.g., GIT for versioning, or good old python IDE instead of Jupyter Notebooks. On the other hand, ML engineering is not classical software engineering:
- in addition to the code, the data should also be versioned;
- in its essence, ML engineering is an exploratory work: one can not know if the model is going to work before testing it;
- there is no clear way to guarantee the quality of the trained model: the data-scientist has to play with it to make it “talk”.
In this talk, we will build a fully customizable and complete system in Python to track Machine Learning experiments. For the purpose of this talk, we will train a neural network (Tensorflow) to classify images between cat and dog, though, the main focus is on the tooling and not the ML algorithm. We will use:
- DVC) (Data Version Control) to 1) version the data alongside the code with GIT 2) build training pipelines to orchestrate the python scripts 3) version experiments.
- Streamlit) to build data exploration apps to play with the trained models.
Both DVC and Streamlit are open-source libraries with python APIs. In the second part of the talk, we will focus on various ways of combining DVC and Streamlit. For instance, we will see how to build a Streamlit app that allows selecting any trained model tracked with DVC (provided its GIT commit), loading it, and testing it on given input images.
During the talk, I will provide actionable code samples and live demos.
This session took place in track PyData & Scientific Libraries Stack and was classified suitable for some domain / some python by the speaker.
Transcript (auto)
Auto-generated from the recording utilizing Open-Source AI. Speaker labels (Speaker 1, Speaker 2) reflect diarization, not identity. Timestamps refer to the recording.
Speaker 1 [00:04]
Hi, everyone. Thanks for being here. So Let me let me start so my name is Antoine. I'm French. I live in Paris and I work at Sikar it's a consulting firm we do AI for clients and I work there for like for four years ago and I am head of data science there So, at CKR, the official job title for tech is data scientist, but actually in my mind I think we are more like machine learning engineers, because what we do for our clients, it's in between data science and classical software engineering. We do we conceive models, we train, we do machine learning, and we also put it in production, build monitoring systems, and so on. So since it's software engineering, do not forget to rely on best practices and best tools. best tools. On the other hand, machine learning engineering, it's a lot more than just regular software engineering. Here you can see in the figure extracted from the hidden technical debts in machine learning systems. It's a famous paper in the RIPS. The code, it's a small black box and you have many systems around to have an AI project functioning. So the reasons, the main reasons why software, why machine learning is different from software engineering is that you need to track the data, not only the code, so you need to track the data and to version the data, and machine learning is an exploratory work, it's not linear like regular software engineering, you have to try and test and try again and test again, and you need solid tools for an investigation, like once the model is trained you want to touch it to feel it and see how it works. So let me talk about what is an ML experiment tracking system. In every machine learning project you have this loop where at the beginning you collect data and you define a model, then you train it, then you evaluate it, and at the end you have a feedback loop because you you learn something and you modify something. You modify either the code or the data. So when you run experiments, you want a tracking system to be able to log the experiment metadata, like which code did you use, which parameters did you use. You want to be able to compare experiments. When you run an AI project in the long term, you end up with lots of experiments, so you want to organise them in a nice way so that it's easy to search them. Reproducibility is very important, because when something is working, you want to reproduce in order to be able to iterate over it. Finally, we work as teams, so collaboration is very important. There exists many platforms or tools that do experiment tracking, and I'm not going to talk about them today. We use a lot of MLflow, and we also have a Polyaxon running on Kubernetes on premise, but I will not talk about this today. Instead, I will show you another approach which is building an entire experiment tracking system with smaller bricks, which are DVC and Streamlit, because it will allow a lot more flexibility and customisation. So let's start doing machine learning. For the purpose of this talk, I choose a simple example. So I will train a model to classify images of cats and dogs. It's like a one-on-one machine learning for computer vision. And so the goal is to have a model that should predict, given an image, whether it's a cat or a dog. A few words about the data set. It's available in the TensorFlow data set. It's 3,000 images. And here you have some examples of cats or dogs from the data set. So here you can see the training pipeline. So to build my model, I will execute this pipeline, so first I need to download the data to train the model, then I need to split because in TensorFlow, you only have train and test splits, you don't have validation splits, so I split the train into train and val splits. Once I did that, I can train the a model, and I get a trained model, and at the end, I evaluate this model against the test set. So I will go very quickly on the scripts. So to download the data set, I just run command line commands, download the data set, and put it in the right place. Let me me show you quickly the training script. So it's just adapted from the TensorFlow tutorial. So nothing fancy in here. What I do basically is that I import parameters, then I load the train data set and the validation data set. Here I define a model. It's a neural network based on a backbone which is a parameter. It can be mobile net or whatever. And I define some callbacks and here it goes. I train the model in here. And you see here the backbone, it's frozen. Here I unfroze the backbone and I train it again. So, okay. Now that I have those four scripts, I want to execute this training pipeline. To do this, I need to go in four steps. First, I need to set up my experiment. So define the parameter I want to use, eventually modify the code or the data. Then I need to run the pipeline. So So far I have no tool, so I need to run the command manually in the right order. So this is I'm not very happy with that. Then I need to save the results. The scripts I just showed you, they produce data, it downloads the data set, at the end When the model is trained, I have the neural network weights and the metric files. So I have to save it somewhere, either on my laptop, an S3 bucket, or somewhere. And finally, I need to have a record of the experiment. Which parameter did I use, which code, and so on. So usually when you have no tools, you put it somewhere in some spreadsheet. So to improve this, you could use DVC. How many people here know DVC? Okay. So DVC, it's a very popular tool to version data. So the most basic feature is that it will replace large files that you don't want to track with Git by small metadata file, and that you can track with Git. And it also has a it also supports a remote storage, so when you when you decide to track a large file with DVC, it's stored in the local DVC cache folder, and you can run DVC push pull to sync with remote data storage. If you are familiar with Git API, it's very similar. Like in Git, you have Git add, Git commit, in DVC you have DVC add, DVC commit, DVC check out. It's basically the same API. So this is nice, but so far, it's not more than older tools like Git large file system that do pretty much the same. What is very interesting with DVC is that it allows to build reproducible pipelines. So it's the next step. So here is a simple For example, I have a pipeline with two stages, let's say stage one, it consumes file A and B and produce file D, and I have another stage that will produce file E. So I define, I simply define the pipeline with this YAML file here. All I need to provide for each stage is the command to execute. It can be Python, but it can be anything, actually. And I have to declare the dependencies and the output file. And whenever I want to run the pipeline, I simply run dvc-repro-zcml-file, and it will resolve the DAGs to know in which order it should execute the stage. It will run the pipeline, so execute the stage. It will track all the data, inputs and outputs, and there is also a mechanism for caching stages, like if, say, stage 2, the dependencies didn't change, it will automatically restore the output, which speeds up a little bit the pipeline execution. So let's go back to my training pipeline. Now I have two more files. So the first one is the pipeline definition. So I have four stages. One for each green box here. So for instance, the train script, I just ran the Python train.py. I have the dependency, so the script itself, the train data set, and at the end, it produced the model base weights, the model itself, and tons of all the files. Okay. So now that I have the DVC pipeline, let's go back again to running an experiment. So to run the pipeline, now it's much simpler. I simply have to run DVC repro. Save the result. It's also done because DVC will automatically save all the data produced by your pipeline, and you can save it to the remote storage by running DVC push. For the last part, put the experiment in the right place, it's better than before because you can access the experiment in the commit history, but still it's not very practical. You don't want to search experiment by running Git log. So it's time to come again to differences between software engineering and data science. Data science, classical software engineering, it's mostly linear workflow. Even if you have branches in your projects, when as individual you code, you do one commit after you work linearly. Data science, it's not linear because usually what you do is you modify some code, then you run not only one experiment but usually many experiments, because you want to do, say, a grid search on some parameter, so from one commit, you will have several experiments, and sometimes from an experiment, you may want to refine it by running it again, so it's not linear. In DVC, they have this feature, DVC experiments, that allows you to track this nonlinear workflow. Basically, it relies on Git custom reference. A Git reference, if you are familiar with Git, you manipulate them every day. the Git reference you know are branches on tags. Basically a reference is just a commit name. Just a pointer to a commit. And DVC experiments, it will create a new kind of Git reference called experiment. And it will like this track every experiment you run. in a non-linear way. So let's talk about the API. Instead of running DVC, you just run DVCX run. What you will do is that it will, for you, edit the parameter.tml file, so you You can provide the parameters from the command line. It will run the pipeline like a regular DVC repo, and at the end, it will automatically create a commit and track it with a Git customer. So how do I see the experiment I run? So let me run this. It's not this. Here I will get the list of experiments available in my Git repository. So here you see I have a commit hash, which is the parent of this nine experiments. If If I want details about those experiments, I can sorry. To get the detail of this experiment, because here I don't have any information, I use DVCX show, and I enter the git commit parents. So here you see I got the nine experiments and I have a lot of columns. So let's say I want to only see a few ones, like the accuracy and some training parameters, I can filter the column like that. And in the end, I end up with an experiment table which is pretty fine. So in my experiment, I see the accuracy is close to, it's pretty good, it's 96%, because the problem is pretty easy, actually. And what changes between the experiment is the train seed and the depth of the fine-tuning. So if we get back to running an experiment, what we can see is that for the three first steps we are pretty much the same, but for the organising your experiments, we are a little bit better, thanks to the experiment table you can display running DVC exp show. So let's take a step back here and see what we got regarding the ML experiment tracking system. Thanks to DVC, we are very strong with logging experiments data and reproducibility. Because you can check out any commit from any experiments and run DVC repro and it will run just as before. Regarding organising and comparing experiments, I can do it, but from the command line, which is fine, but there is room for improvement, like how do we build a UI to visualise experiments. So that's where I introduced Streamlit. So Streamlit, it's a Python library to build data apps, so you don't have to know how to develop web applications, only code in Streamlit, and actually the slides that I'm just showing you are written in Streamlit, and it's about to become the most popular library for building a dashboard in Python, so on. I will go for a quick demo of the Streamlit API. So you just install it with pip, and you have a lot of function in the Streamlit API that allows you to display whatever you want. So let's say I want to display a latex formula. I can do it like this, and what is really important, because so far, display data, you can do it in regular Jupyter notebooks. What you can build easily with Streamlit is interaction. So let's say I want to have a select box like this. So here you see I declare a select box with two options, cats or dogs, And I store the result in this variable that I just print in here, so it renders this selector, and whenever I change the value here, it gets updated in here. And you have more advanced inputs, like, for instance, I could use my camera to get a picture. So, you see here, I use camera input, and I get the image here, and if there is an image, I just show it in here. So, it's a numpy array. I can do whatever I want with it. So, let's get back to the cats and dogs problem. So my neural network, it's a bare output, it's a probability of being a cat or a dog, and if I look at the raw prediction on the test set. It's here. So you see here I read the prediction using Pandas, and I display the data frame with streamlit in this line. This will produce this table. So you see the prediction, it's the probability of being a dog. I have the true label and the predicted label. The predicted label, by default, it's computed with 50 per cent threshold, but let's say I want to modify this threshold, for instance, the first line, it's a cat, but the prediction is close to 50, so let's say I move the slider below 43, and then, as you can see, the predicted with threshold, it changes from cats to dogs. Cats, dogs. And I get the accuracy updated in real time. And all this, the right part of my slide, is produced by this code, which is very, very simple to do. You could display image. Let's say I want to see the images for which the model is not sure, so close to 50-50. So with a slider with two entries, I can get a min and max threshold, and here I simply filter my data frame to get only the images whose confidence is below those two thresholds, and I use st.image to plot the images. For instance, for those two thresholds, I have seven unsure predictions, and this one is pretty interesting, because there is a cat and a dog in the same image, so it looks like the model is working fine, actually. And finally, what I could do also is to directly run the model in my app. So to do that, I simply need to load the model using TensorFlow. And I get an image from a file uploader and get it to the model and print the output prediction with trimlit. Let's try it. This cute cat, for instance. Okay. So, it's working fine because it's pretty sure it's a cat. And what is interesting here is that you can try the model with whatever you want. For example, let's try edge cases. Let's try a lion. And you see that the model is pretty confident the lion is a dog. So, it's very easy to build a UI to experiment the model, to try it, to you can do a lot of things with this. So, What's next? The Streamlit app I just showed you, it shows you the model that is in my current work space, and I cannot what I would like to do is to see the model from other past experiments. So actually it's like I have two dimensions, Streamlit it allows me to build very expressive data visualisation, try the model, run inference, and so on. Streamlit is very good at tracking experiment data, and now I want to combine them to build a UI that allows me to explore experiments. So how do we do that? The first thing to do this is to retrieve experiments in Python. So to do this, I will use the DVC Python API, which is very similar to Git Python API. So first I initialise a repo from my current directory, then I can get all experiment commits running this command. It's like a DVC list that I just showed you before, but in Python. So it gets me a JSON, this is the commit parent and this is the experiment names. Once I have this, I can get the experiment detail by running this, experiments.show, it's similar to DVCX show come online, but it's in Python and you see here I have a lot of information coming out. So, for instance, here it's an experiment, I have all the parameters, which file was used, which parameters, batch size, image size, and so on. I have the hashes of all dependencies. I have the hashes of all outputs, and I have the metric here, and I have the experiment name. So having this, it's very easy now to build a UI for a table of experiments. So this is the whole code. So first I retrieve an experiment metadata, just like I showed you before, then I flatten it this way so that I can wrap it into a data frame, and in the end I use streamlit to do the rendering. What I get with that is that I have this table here, it's similar to what we got with DVC Exo, except that it's in Streamlit, and I can do plots, like here I plot a bar chart of all the accuracies of all experiments, and actually you can do whatever you want at this stage. For instance, here, I want to plot the accuracy versus the level of fine tuning, so I use Plotly, because with Streamlit, you can plot graphs with Plotly, with Vega, whatever framework you are familiar with, it's very easy to plug it into Streamlit. And here it goes. So you have the accuracy against the level of fine-tuning, and as you can see, there is no correlation. It seems like it's a random seed that defines the level of accuracy of the model. And what can we do after that? Now that I I have the table of experiments, I have all the details of the experiment logs, but what I would like to do is to actually load the files from this experiment. Let's say I want to try the model from a past experiment, I need to retrieve the model files from that commit. To do that, you can use the DVC Python API, which provides an open function that is similar to the regular Python open function, except that it's like it's open to a new dimension because you can pass, you see here, the commit hash, and it will read the file from any commit. In particular, it could be an experiment commit that is not tracked by a Git branch. So here I just open the prediction.cv file and read it with Pandas. And once I know that I have this function, I need to feed it with commit ashes. So I just get the experiment data just as before and I build options for a select box that looks like this. So here I have the experiment commit hash and the experiment name. And when I have this, it becomes pretty easy to build something like that. So let's say I want to build an app to compare two models. So I simply put two selectors for selecting experiments. And so the two selectors are here and here. Once I have selected a commit hash, I simply load the predictions with the function using dvc.open function, and I merge them in a single data frame, and I can do whatever I want. So for instance, I can get the predictions on which both models disagree, and display those images in the Streamlit app. So let's try it. so it looks like these two experiments agree. Okay. So here it goes. I will hide this column. So these two experiments, you see that the model disagree on two images. And each of them is right in one case. And here you have the image where the model disagrees. So let's take a step back here and see what we got. Now we have all we need to build an experiment tracking system. So to compare experiments, sorting them, I can build a Streamlit app to organise my experiments. The demonstration I just showed you is pretty basic, but in Streamlit you can go... You can customise as much as you want. And you can compare experiments, You can have the experiment table, and since it's code, you can commit the Streamlit application in your project repository, so it's a common interface for the whole team. It goes through code review and so on. So let's talk a bit about the advantage of the drawbacks of this approach. The main advantage is that you can do whatever you want with the UI. Here I just show you example of images, but you can have an AI project of where you deal with video, audio, text, all of them. You can do pretty much what you want in Streamlit, and you have a lot of primitive function in Streamlit, and you also have some communities that develop a lot of custom components. You can even develop your own component in JavaScript if you want. And also another advantage, it's very easy to set up. You don't need, for instance, if you want to deploy a MLflow, you need to have a server, you need to run the infrastructure, you need to deploy it here. In five minutes, you are set up, you can initialise DVC and have your first Streamlit dashboard. On the other hand, DVC and Streamlit are quite recent technologies. So we cannot predict the future. They are very promising technologies, but still, it's a very recent technology. The way of developing this UI for tracking experiments, it may be hard to maintain, because it's the team who develop it, so you have to code review it, you have to deal with technical debt, and you also have to handle retro compatibility, because during the life cycle of a project, the pipeline, they may change a lot, and when you go deeper in customisation, maybe what was working before is not working now, so you have to fix continuously the web UI. You don't have this problem if you use a more classical tracking experiment system. What I didn't show you today is how to deploy a Streamlit app, so you can just dockerise what I've shown you. Something which is not trivial is how to handle with DVC in Docker. It's very easy to dockerise a Streamlit application. What is not easy here is that the UI, it requires a DVC repository to be able to search experiments. So usually when you deploy code, you only deploy the code, you do not deploy the Git repo with it. So you have to declarize this as well. And I didn't talk about the infrastructure, which is outside the scope of ML experiment tracking systems. And what you could do is have an integration with CICDs. Actually iterative.ai who is developing DVC, they propose a CML that allows you to say trigger the training of a model when you push some code on GitHub. So that's it for me.
Speaker 2 [37:24]
Yeah, thank you very much. That was a very, very interesting talk. I guess everyone was really impressed with the layout of your presentation. So we have many questions. We probably won't be able to go through all of them, but let us collect the most kind of all at once. So the first question is, what's your experience with using DVC for storing trained models? Maybe compared to MLflow, do they work well together?
Speaker 1 [37:51]
I didn't try DVC and MLflow at the same time, but I know I have colleagues that did that. And I know it's feasible, but it's not that easy. Because MLflow and DVC, there is an intersection between what the two tools do. They both track model parameters and so on. So it's not easy to define who is doing what. But I know it's feasible.
Speaker 2 [38:27]
All right. I think the second one you comment a little bit, but maybe you can expand on that. So can you include execution of DVC pipelines into a common CI-CD setup in order to run pipelines automatically with every new push?
Speaker 1 [38:41]
Okay, you can use a CML, and it's a wrapping, it's GitHub actions that allows you to trigger whatever pipeline you want when you push on some branches on GitHub. It's very easy to set up, like, in five minutes you can say whenever I merge on the master branch, it triggers training on EC2 instance on AWS.
Speaker 2 [39:10]
The next question is, how well does DVC handle GDPR compliance? Can you delete experiments or raw data from its history?
Speaker 1 [39:21]
I have no idea how to answer this one.
Speaker 2 [39:26]
All right, sorry No worries, so a One a a question, but you already may be answered during the presentation is what's the name of the new neat presentation tool?
Speaker 1 [39:40]
I'm sorry, what's the name of?
Speaker 2 [39:41]
what's the name of the presentation tool that you like how did you write the presentation oh yeah
Speaker 1 [39:45]
Oh, yeah, I use a component called the book, and it allows me to have this left panel here with subsections. So actually the main chapter selection, it's a regular Streamlit select box, but what you got in here, it's from a Streamlit book. So you just run a pip install streamlit book and you have a very nice and simple interface. And each slide, it's actually a separate Python file. It's a bit overkill to do a slide this way. But... Thank you.
Speaker 2 [40:30]
I would say it's not.
Speaker 1 [40:31]
But it's really nice because actually I did that because it was much easier to show the result of Streamlit code directly in the slides without sweeping from slides to web browser.
Speaker 2 [40:48]
Or maybe a related question is if you're planning to share the code of the generation.
Speaker 1 [40:52]
Yeah, sure.
Speaker 2 [40:53]
Yeah, sure.
Speaker 1 [40:53]
Okay.
Speaker 2 [40:54]
All right. The next question is, what's the advantage compared to solutions like Neptune AI?
Speaker 1 [41:01]
Actually, it's maybe the one I know the... I never tried Neptune AI. My understanding is that it's a service where you have a Python API to log parameters and artefacts like data, it's completely separated from the Git history. The main advantage of DVC is that it allows you to really link experiments with Git commits. And it's not logging. It's a real link. Like if I want to retrieve my experiment, I just check out any commit and run DVC pool and I get all the files on the code and so on. Whenever you use Neptune AI or MLflow, it's a logging system, so it's outside the code. So of course you can log the commit hash and so on, but let's say you do Git rebase and the commit hash is changed, you lose the commit reference. It's not possible with DVC.
Speaker 2 [42:18]
All right, another question is how does DVC experiments tracking work if you have multiple data scientists working at the same repo?
Speaker 1 [42:27]
That's a good question. It's actually just like a code review. When you have conflicts, you have to resolve them, and usually it's not that hard. It's simpler than resolving code conflicts, because when you have code conflicts, you have to look in the file and see, okay, it modifies this line, it modifies this line, how do I merge this? And when you run experiments, most conflict resolution just overwrites the files that track the hashes of the output file. In practice, it's fairly easy.
Speaker 2 [43:13]
So there's some other questions about the infrastructure. So how do you integrate DVC into infrastructure? Is there a Kubernetes scheduler or similar?
Speaker 1 [43:22]
And DVC do not provide, to my knowledge, a tool to do that, except CML. So you have to do this outside. And the most basic way of doing it is just SSH to a server and run DVC. What you could do is automate this a little bit, but actually it's maybe at CKR, we use a lot of DVC and Streamlit, and maybe this is what we should improve. All right. So I have no proper solution to do that.
Speaker 2 [44:04]
All right, there are many questions which are tied. So let me finalize with a question about selecting the tool for the visualization. So are they advantage or disadvantage? Or let's say, what's your take on Streamlit versus Dash?
Speaker 1 [44:21]
versus Dash? I used Dash like two or three years ago. Dash is really nice, but I think you need to be more familiar with web development. Python, Streamlit, what is really a strength of Streamlit is that it really makes you feel like you are coding a Python script. Like you have no callbacks, you have no components. It's a really Pythonic way of developing a web app. It makes you believe you are coding a Python script. But you are not. You are building a web app, actually. And when you develop Plotly Dash, you have to be familiar with some concept of web development, like callbacks, components, the interaction of components. It's a bit trickier.
Speaker 2 [45:24]
All right, thanks. So I think that's all for the questions, but I guess you can reach Antoine in the conference. So let's thanks Antoine again for this great presentation.