brain-score / model-tools

Helper functions to extract model activations and translate from Machine Learning to Neuroscience
MIT License
8 stars 27 forks source link

speed up activations retrieval #41

Open mschrimpf opened 3 years ago

mschrimpf commented 3 years ago

when we started writing model-tools, we were primarily thinking of neuroscience stimulus sets with only a couple thousand images. Speed was therefore less of an issue because even with a suboptimal implementation, a few k images are quickly passed through the network.

We are now evaluating models on increasingly large ML benchmarks (e.g. https://github.com/brain-score/brain-score/pull/232) and due to the slow activations retrieval, the evaluation takes very long (days), sometimes timing out on the cluster. We therefore need to speed up the activations extraction.

Models are already using cuda when possible, I believe the main bottleneck is actually the loading of images which is currently single-threaded. We should profile the code to confirm this and (if true) use multiple workers to load the images to pass into the model (e.g. here and here for pytorch, ideally using existing tools such as DataLoader).