Enabling multiple worker support for DataLoaders in inference scripts

[msoham123]

PyTorch's DataLoader can use python multiprocessing to "multithread" batch loading from storage. This can linearly speed up the data loading time and improve overall inference performance.

To add multiple worker support, we need to pass the num_workers arg to the DataLoader.

cifar10_dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=4)

Python multiprocessing creates additional python processes to accelerate data loading. On Linux, this uses the fork method. However, recently on MacOS, starting with Python 3.8, this is done using the spawn method (https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods).

However, when using spawn, this can cause the process to repeat by recursing on itself. To avoid this, we must call any code using multiple workers from a main function (see https://pytorch.org/docs/stable/notes/windows.html#multiprocessing-error-without-if-clause-protection).

e.g

from inference_cifar10_resnet import run_inference_cifar10_resnet

if __name__ == '__main__':
    run_inference_cifar10_resnet()

This is important when running inference over large datasets (e.g ImageNet), as we can reduce the time to begin inference significantly.