How to Unit Test Deep Learning: Tests in TensorFlow, mocking and test coverage
How to Unit Test Deep Learning: Tests in TensorFlow, mocking and test coverage
Programming a deep learning model is not easy (I’m not going to lie) but testing one is even harder. That’s why most of the TensorFlow and PyTorch code out there does not include unit testing. But when your code is going to live in a production environment, making sure that it actually does what is intended should be a priority. After all, machine learning is not different from any other software.
Note that this post is the third part of the Deep Learning in Production course where we discover how to convert a notebook into production-ready code that can be served to millions of users.
In this article, we are going to focus on how to properly test machine learning code, analyze some best practices when writing unit tests and present a number of example cases where testing is kind of a necessity. We will start on why we need them in our code, then we will do a quick catch up on the basics of testing in python, and then we will go over a number of practical real-life scenarios.
Why we need unit testing
When developing a neural network, most of us don’t care about catching all possible exceptions, finding all corner cases, or debugging every single function. We just need to see our model fit. And then we just need to increase its accuracy until it reaches an acceptable point. That’s all good but what happens when the model will be deployed into a server and used in an actual public faced application? Most likely it will crash because some users may be sending wrong data or because of some silent bug that messes up our data preprocessing pipeline. We might even discover that our model was in fact corrupted all this time.
This is where unit tests come into play. To prevent all these things before they even occur. Unit tests are tremendously useful because they:
Find software bugs early
Debug our code
Ensure that the code does what is supposed to do
Simplify the refactoring process
Speed up the integration process
Act as documentation
Don’t tell me that you don’t want at least some of the above. Sure testing can take up a lot of your precious time but it’s 100% worth it. You will see why in a bit.
But what exactly is a unit test?
Basics of Unit testing
In simple terms, unit testing is just a function calling another function (or a unit) and checking if the values returned match the expected output. Let’s see an example using our UNet model to make it more clear (If you haven’t followed the series you can find the code in our GitHub repo .
In a few words, we took an official Tensorflow google colab that performs image segmentation and we try to convert it into a highly optimized production-ready code. Check the first two parts here and here ).
So we...
