speed up activations retrieval #41
Description
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. brain-score/vision#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).