DataLoader num_workers drastically increases run time

dorkytorchy123: <p>Pytorch 2.0.1+cu117, Win10, powerful workstation with nvidia gpu and 16 cpus, full code below</p>
<p>I have a dataset of 240,000 28x28 images, but once I figure out this issue I plan on moving to a much larger dataset. I was simply loading these images and running them through the model, but wanted to augment my data by slightly changing the images at random.</p>
<p>Making this augmentation change (in <code>__getitem__</code>) about quadrupled my training time due to the <code>perturbImage</code> call. It was at this point I looked into using the <code>DataLoader</code> class to load and alter the data in the background and avoid the drastic slowdown I just introduced. However, passing in any value for <code>num_workers</code> besides 0 results in extreme slowdowns and if I set num_workers around 12 and let it run for a few minutes, it just freezes my computer until I manage to open task manager and kill the Python tasks. I’ve tried playing around with some settings like pin_memory and persistent_workers with no noticeable positive changes. What is going on here?</p>
<p>Time to load data and run 1 epoch:<br>
num_workers 0, no perturbImage: 20 seconds<br>
num_workers 0, with perturbImage: 91 seconds<br>
num_workers 1, with perturbImage: 4 minutes<br>
num_workers 3, with perturbImage: 14 minutes<br>
num_workers 12, with perturbImage: greater than 30 minutes, killed process</p>
<p>Code, with no perturbImage call and 0 workers:</p>
<pre><code class="lang-auto">from multiprocessing import freeze_support
import random
import torch.nn as nn
import torch.nn.functional as F
import torch
import torchvision
from torch.utils.data import Dataset
import torch.optim as optim
import time
import os

def perturbImage(img, jit, blurrer, pt, transpose = False):

    numOptions = 4
    pickedOption = random.randint(1, numOptions)

    if pickedOption == 1:
        img = jit(img)
    elif pickedOption == 2:
        img = blurrer(img)
    elif pickedOption == 3:
        img = pt(img)
    elif pickedOption == 4:
        rots = [90, 180, 270]
        pickedRot = rots[random.randint(0, 2)]
        img = torchvision.transforms.functional.rotate(img, pickedRot)

    if transpose:
        img = torch.transpose(img, 0, 2)
        img = torch.transpose(img, 0, 1)

    return img

class CustomMnistRotateDataset(Dataset):
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.jit = torchvision.transforms.ColorJitter(brightness=.5, hue=0.1)
        self.blurrer = torchvision.transforms.GaussianBlur(kernel_size=(5, 9), sigma=(0.1, .8))
        self.pt = torchvision.transforms.RandomPerspective(distortion_scale=0.4, p=1.0)
        
        print(f'loading CustomMnistRotateDataset from {data_dir}')
        self.data = torch.load(data_dir)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img = self.data[idx][0]
        label = self.data[idx][1]
        #img = perturbImage(img, self.jit, self.blurrer, self.pt)

        return img, label

def loadData(type = None, trainPath = None, testPath = None):
    if type == 'custfashmnist':
        raw_training_data = CustomMnistRotateDataset(trainPath)
        raw_test_data = CustomMnistRotateDataset(testPath)
    return (raw_training_data, raw_test_data)

class BaseNet(nn.Module):
    def __init__(self, imgSize, outputSize):
        super().__init__()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(imgSize, 5120),
            nn.ReLU(),
            nn.Linear(5120, 5120),
            nn.ReLU(),
            nn.Linear(5120, 5120),
            nn.ReLU(),
            nn.Linear(5120, outputSize),
        )

    def forward(self, x):
        x = self.linear_relu_stack(x)
        x = F.log_softmax(x, dim=1)
        return x

def train(net, epochs, trainset, losses, start_time, lossRecordEvery = 10, imgSize = (28*28), lr = 0.001, device = "cuda"):
    optimizer = optim.Adam(net.parameters(), lr = lr) 
    i = 0
    for epoch in range(epochs):
        for batchData in trainset:
            batchImg, batchLabel = batchData
            net.zero_grad()
            forwardOutput = net(batchImg.view(-1, imgSize).to(device))
            loss = F.nll_loss(forwardOutput, batchLabel.to(device))
            if (i == lossRecordEvery):
                losses.append(loss.item())
                i = 0
            i += 1
            loss.backward()
            optimizer.step()
        print(loss)
        print(f"--- {(time.time() - start_time)} epoch {epoch} time to run ---" )

def go():
    start_time = time.time()

    device = (
        "cuda"
        if torch.cuda.is_available()
        else "mps"
        if torch.backends.mps.is_available()
        else "cpu"
    )
    print(f"Using {device} device")

    dataset = 'custfashmnist'

    if dataset == 'custfashmnist':
        imgSize = 28*28 
        imgDimsenion = 1 
        outputSize = 10 
        raw_training_data, raw_test_data = loadData('custfashmnist', 'data/expanded-train', 'data/expanded-test')

    net = BaseNet(imgDimsenion* imgSize, outputSize).to(device)

    epochs = 1
    losses = []

    # num_workers=1 takes 4 minutes for 1 epoch(!!!), num_workers=3 takes 14 minutes for 1 epoch(!!!) num_workers=0 with perturbImage takes 91 seconds for 1 epoch, without perturbImage num_workers=0 takes 20 seconds
    trainset = torch.utils.data.DataLoader(raw_training_data, batch_size=500, shuffle=True, num_workers=0, persistent_workers=False, pin_memory=False)
    print('start train')
    train(net, epochs, trainset, losses, start_time, imgSize=(imgSize*imgDimsenion), lossRecordEvery=100, lr = 0.0001)

if __name__ == '__main__':
    #freeze_support()
    print(torch.__version__) # 2.0.1+cu117
    go()

</code></pre>