Xla: Should the num_workers in the colab examples be higher than 4?

[Santosh-Gupta]

❓ Questions and Help

In the colab examples, I noticed the num_workers is set to 4. I am wondering if it should be set to at least 8, which is the number of cores a TPU has. So a dataloader for each core. The num_workers is a parameter for the dataloader

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=num_workers)

I am wondering if this is based on colab CPUs having 4 cores in the CPU and GPU instances; however the TPU colab instance's CPU has 40 cores instead of 4.

[dlibenzi]

Tuning the number of workers depends on the amount of work the input pipeline is doing, and the available CPU cores.
Some CPU cores are also needed to convert tensors to device format, and some for running model's Python code, so we can imagine the maximum number of workers to be about NUM_CPU_CORES - NUM_TPU_CORES.
There is also no reason to set num_workers to a number higher than necessary.
So given the 40 CPU cores, you can try to bump it up progressively and see up to which point model's throughput increases.