Should Machine Learning Engineers Write Unit Tests?

If you are looking for a tl;dr, the answer is yes.

I somehow seem to run into engineers fairly frequently that think there is nothing to test in machine learning code or at least can’t imagine where one might need a test. Generally speaking, the argument is something like “It’s just a model. You just give it input and give it output so there is nothing to test”. Having worked in machine learning for several years now, I regularly run into situations where a test could save tons of time. Let me explain a few reasons why you might want to test your code.

Unit Tests Help Make Reusable Code

This should be obvious enough to most engineers, but if you can reuse your code, you can get more projects done more quickly. Unit tests help you think ahead for problems that may occur and make sure that your code works in all circumstances that you might expect it to. This will prevent you from scratching your head when things worked perfectly on your practice hotdog dataset, but not on your car dataset.

Unit Tests Catch Functionality-Breaking Changes

If you work with a group of people or if you pick up a project after a long time, there may be situations where you are working on code that you are unfamiliar with. In these situations it is always possible that a change you (or a teammate!) make could inadvertently break the functionality of your code, which in some cases may not be immediately clear to you. In some cases, you may not even notice that something isn’t working correctly until much later in development. If you test your functionality with unit tests, when changes are made to your code you can simply run the tests to make sure that the functionality still works how you expect.

Unit tests Save You Time

Have you ever written a training script with a dataloader and a model and run it on your data only to have an error happen that you then had to trace back and fix? Or maybe it was even worse, you launched an Azure cloud job with your script running inside it, only to come back 10 minutes later and see it failed. If you have then repeated this step multiple times as you keep going back and debugging, then hopefully you have already seen my point. Unit tests allow you to test the functionality of your ML components quickly without needing to put all the components together and load actual data.

When to Write Unit Tests

So now that we have established that it is beneficial to write unit tests, you might be wondering where exactly then can they be applied for ML code? This is not an exhaustive list, but here are some of my recommendations:

Anywhere there is an exception

When building a Pytorch dataset, you may want to add some exceptions for things like when an annotation file doesn’t exist or when there are no images in a folder. It is important that we know these exceptions work so that when someone uses the dataset, the exception properly tells them what they did wrong. Likewise, you may have an exception in a model that tells you when the configurations are not set properly. If you wait until you run the code to test whether it works or not, it may take a long time before you actually see the exception work (or figure out it isn’t working).

For inputs and outputs

Machine learning is all about inputs and outputs. If your data loader is outputting the wrong format, you’ll never get it to work with your model and if your visualization function doesn’t accept the shape of the output from your model, you won’t be visualizing anything. Take a minute to write unit tests that check the input and output of your various ML components.

When you have custom logic

If you ever write your own components, like a loss function, a special layer or maybe a post processing step, it is always a good idea to make some tests to check if it is working correctly. This doesn’t need to be exhaustive, because it can be next to impossible to write tests that cover every possibility. However, you can certainly cover a couple of cases to have some peace of mind that things are working the way you expect them to.

From one ML engineer to another, I highly encourage you to learn more about unit tests and look for other areas of your code that might benefit from it. There are tons of great resources out there on how to use, write, and debug using unit tests and this information is not exclusive to standard software engineering. Next time you write an dataset or a neural network, try some of the tips I listed above. If you benefited from it (or not), please let me know in the comments!

For more articles on Python, check out some of my other content.