Torch dataloader is slow

yuanc: <p>I’m practicing using torch.utils.data.DataLoader to load batches of data into GPU, but find it’s significantly slower than just indexing the data, like 10 times slower. I wonder why or there is something I did wrong.<br>
Here is my code:</p>
<pre><code class="lang-auto">for train_index, test_index in kf.split(X, Y):
    # Split train-test
    X_train,X_test=X[train_index],X[test_index]
    y_train,y_test=y[train_index],y[test_index]
    train_dataset = torch.utils.data.TensorDataset(torch.from_numpy(X_train), torch.from_numpy(y_train))
    train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size, shuffle=True)
    
    model = NNTrain()
    model.cuda()
    optimizer = torch.optim.Adam(model.parameters(), lr = LEARNING_RATE)

    model.train()
    for i in range(epochs):
        for X_batch, y_batch in train_dataloader:
            X_batch = X_batch.float().cuda()
            y_batch = y_batch.float().cuda()
            optimizer.zero_grad()
            y_batch_pred = model(X_batch)
            loss = nn.MSELoss()(y_batch_pred, y_batch)
            loss.backward()
            optimizer.step()
</code></pre>
<p>and the faster one:</p>
<pre><code class="lang-auto">for train_index, test_index in kf.split(X, Y):
    # Split train-test
    X_train,X_test=X[train_index],X[test_index]
    y_train,y_test=y[train_index],y[test_index]

    model = NNTrain()
    model.cuda()
    optimizer = torch.optim.Adam(model.parameters(), lr = LEARNING_RATE)

    model.train()
    for i in range(epochs):
        for j in range(batch_num):
            X_batch = torch.from_numpy(X_train[batch_size*j:batch_size*(j+1),:]).float().cuda()
            y_batch = torch.from_numpy(y_train[batch_size*j:batch_size*(j+1),:]).float().cuda()
            y_batch_pred = model(X_batch)            
            loss = nn.MSELoss()(y_batch_pred, y_batch)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
</code></pre>

ptrblck: <p>This might be expected, since the <code>DataLoader</code> will add overhead to the already loaded tensors by e.g.:</p>
<ul>
<li>shuffling the indices using a sampler</li>
<li>indexing the <code>TensorDataset</code> with each index separately</li>
<li>creating the batch by calling into the <code>collate_fn</code>
</li>
</ul>
<p>This workflow is useful, if you want to lazily load the dataset and process each sample via data augmentation techniques.<br>
Since your dataset is already loaded, indexing would be the faster approach.</p>

yuanc: <p>Thanks for the reply! I found a custom dataloader later and it’s much faster so its all good now.</p>

yjguo: <p>hi <a class="mention" href="/u/yuanc">@yuanc</a> ,  is the later dataloader you found public? If so, could you share, thanks.</p>

Lizhen_Ji: <p>Hi, similar problem here, I have a very large dataset (1200000 images) and it cannot fit into memory and using indexing like in this example.<br>
I also tried using h5py file to create dataset.</p>
<p>The training goes well with shuffle=False and getting slower and slower with shuffle=True<br>
Can you give me some advice on how to solve the problem?<br>
Thanks a lot</p>
<p>Here is <a href="https://discuss.pytorch.org/t/same-gpu-same-pytorch-training-script-training-time-sometimes-fast-sometimes-slow/134776">my question</a></p>