Efficient GPU Data Movement?

jagraves21: <p>Normally I move all of my data to the GPU before training, but I have a dataset that is too big for the GPU’s memory (it easily fits into system RAM).  So rather than moving the entire dataset to the GPU, I changed by code to only move the mini batches to the GPU as needed.  I’ve done some testing with datasets that fit entirely on the GPU, and I find that moving the mini batches to the GPU as needed greatly increases training time (about 1.5x for my data).  I thought that if I moved larger chucks of data to the GPU at a time and then split the chucks into mini batches, it might help.  However, this seems to be twice as slow as just moving the mini batches as needed.  Any suggestions on a better way to move data to the GPU would be appreciated.</p>
<p>Original Approach:</p>
<pre><code>    X = torch.from_numpy(X).cuda()
    y = torch.from_numpy(y).cuda()
    tensor_dataset = torch.utils.data.TensorDataset(X, y)
    train_loader = torch.utils.data.DataLoader(tensor_dataset, 256, shuffle=True)
    for epoch in range(1, n_epochs+1):
        for batch, (data, target) in enumerate(train_loader, 1):
            model.train()
            optimizer.zero_grad()
            outputs = model.forward(data)
            loss = loss_function(outputs, target)
            loss.backward()
            optimizer.step()
</code></pre>
<p>Mini Batch Approach:</p>
<pre><code>    X = torch.from_numpy(X)
    y = torch.from_numpy(y)
    tensor_dataset = torch.utils.data.TensorDataset(X, y)
    train_loader = torch.utils.data.DataLoader(tensor_dataset, 256, shuffle=True)
    for epoch in range(1, n_epochs+1):
        for batch, (data, target) in enumerate(train_loader, 1):
            data = data.cuda()
            target = target.cuda()
            model.train()
            optimizer.zero_grad()
            outputs = model.forward(data)
            loss = loss_function(outputs, target)
            loss.backward()
</code></pre>
<p>Chunked Approach:</p>
<pre><code>    X = torch.from_numpy(X)
    y = torch.from_numpy(y)
    tensor_dataset = torch.utils.data.TensorDataset(X, y)
    chuck_loader = torch.utils.data.DataLoader(tensor_dataset, 100000, shuffle=True)
    for epoch in range(1, n_epochs+1):
        for chunk_data, chunk_target in chuck_loader:
            chunk_data = chunk_data.cuda()
            chunk_target = chunk_target.cuda()
            tmp_tensor_dataset = torch.utils.data.TensorDataset(chuck_data, chuck_target)
            train_loader = torch.utils.data.DataLoader(tmp_tensor_dataset, 256, shuffle=False)
            for batch, (data, target) in enumerate(train_loader, 1):
                model.train()
                optimizer.zero_grad()
                outputs = model.forward(data)
                loss = loss_function(outputs, target)
                loss.backward()
                optimizer.step()</code></pre>

ptrblck: <p>Could you try to use multiple workers in your <code>DataLoader</code> (<code>num_workers&gt;0</code>) as well as <code>pin_memory=True</code> and <code>data = data.to('cuda', non_blocking=True)</code>?</p>

jagraves21: <p>If the problem really is data movement, how would setting <code>num_workers</code> help since the <code>DataLoader</code> does not move the data to the GPU?</p>
<p>I’ve considered writing a custom <code>torch.utils.data.Dataset</code>, that moves the data to the GPU before returning it and then using a <code>DataLoader</code> with multiple workers, but this seems like it would cause a lot of memory fragmentation on the GPU.</p>

ptrblck: <aside class="quote no-group" data-username="jagraves21" data-post="3" data-topic="63707">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/a3d4f5/48.png" class="avatar"> jagraves21:</div>
<blockquote>
<p>If the problem really is data movement, how would setting <code>num_workers</code> help since the <code>DataLoader</code> does not move the data to the GPU?</p>
</blockquote>
</aside>
<p>Based on the description and provided code, I doubt the transfer is the bottleneck, but the data loading, but correct me, if I’m wrong.</p>
<aside class="quote no-group" data-username="jagraves21" data-post="3" data-topic="63707">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/a3d4f5/48.png" class="avatar"> jagraves21:</div>
<blockquote>
<p>I’ve considered writing a custom <code>torch.utils.data.Dataset</code> , that moves the data to the GPU before returning it and then using a <code>DataLoader</code> with multiple workers, but this seems like it would cause a lot of memory fragmentation on the GPU.</p>
</blockquote>
</aside>
<p>Why would this cause a lot of fragmentation?<br>
I wouldn’t recommend this approach in general, as multiprocessing and CUDA calls might be hard to handle.</p>

jagraves21: <aside class="quote no-group quote-modified" data-username="ptrblck" data-post="4" data-topic="63707">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<blockquote>
<p>If the problem really is data movement, how would setting <code>num_workers</code> help since the <code>DataLoader</code> does not move the data to the GPU?</p>
</blockquote>
<p>Based on the description and provided code, I doubt the transfer is the bottleneck, but the data loading, but correct me, if I’m wrong.</p>
</blockquote>
</aside>
<p>The data is loaded completely into RAM and stored in <code>X</code> and <code>y</code> initially. The problem is that the data does not fit entirely on the GPU.  The only difference between the first block of code and the second is when the data is moved to the GPU.  I’ve tested with subsets of my large dataset that will fit into memory and found that the second block of code is about 1.5 times slower than the first.</p>
<aside class="quote no-group quote-modified" data-username="ptrblck" data-post="4" data-topic="63707">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p><img src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/a3d4f5/40.png" alt="" width="20" height="20" role="presentation"> jagraves21:</p>
<blockquote>
<p>If the problem really is data movement, how would setting <code>num_workers</code> help since the <code>DataLoader</code> does not move the data to the GPU?</p>
<p>I’ve considered writing a custom <code>torch.utils.data.Dataset</code> , that moves the data to the GPU before returning it and then using a <code>DataLoader</code> with multiple workers, but this seems like it would cause a lot of memory fragmentation on the GPU.</p>
</blockquote>
<p>Why would this cause a lot of fragmentation?<br>
I wouldn’t recommend this approach in general, as multiprocessing and CUDA calls might be hard to handle.</p>
</blockquote>
</aside>
<p>I just assumed that it would not be the best use of GPU memory to have the <code>torch.utils.data.Dataset</code> move each sample to the GPU when it was requested.  Currently, 128 samples are moved to the GPU at a time.</p>

ptrblck: <p>Thanks the for information!<br>
If you are preloading all samples into the RAM, you might also try to simply index the tensors in your loop and push each chunk to the GPU. I assume you don’t apply any transformation etc. in your <code>Dataset</code>.<br>
If that’s the case, the “manual” indexing might be a bit faster, as the <code>DataLoader</code> might add some overhead to load each sample.<br>
Could you try that and check the performance?</p>
<p>Also, you could try to call <code>pin_memory()</code> on the data tensors and use <code>.to("cuda", non_blocking=True)</code>.</p>