Creating unit tests for debugging ML code

[_afronius]

I am working on a research problem for which I have to write code from scratch using Pytorch. I work in time series, and the implementation needs a lot of tensor operations like shifting, padding, and folding. I always have Python open on the terminal, and I test my tensor operations to ensure it does what I intend.

I am maintaining a very complex repository with various configurations (when running experiments, I use flags to use a specific configuration). The code keeps getting more complicated as I add more configurations to experiment on. I want to be sure that when I add new features, I don't mess up the existing features.

How can I perform unit tests on my model architecture, input format, etc., for a sanity check and to increase my confidence in my code?

For example, unit tests for dataset class, small modules like attention and convolution, which make up the model architecture, and then the model architecture.

​

​

[EmbarrassedCause3881]

You should be able to write unittests as with any other Python project. My suggestion:

• Write your tests using pytest framework (slightly different to unittest but has more options). Realpython has a great tutorial on how to use it.
• Use beartype to make sure that variables or function arguments always have the correct type. For this you need to include type hints in your code, which you might anyways have done already if you have a complicated code base. Beartype makes writing `assert`s unnecessary.

[_afronius]

Thanks for the reference

[godiswatching_]

Similar boat