Pin memory vs sending direct to GPU from dataset

DuaneNielsen: <p>I’m implementing an atari pong playing policy gradient agent.</p>
<p>The main loop at the moment is.</p>
<ol>
<li>Run the policy to gather some experience and save the experience (images, actions, rewards) in a dataset.</li>
<li>Run a single iteration of training over the gathered experience to update the policy.</li>
<li>Throw all the data away and start again.</li>
</ol>
<p>Up until now, I have just been using cpu training, but now I’d like to push the training to GPU.</p>
<p>When it comes to loading It seems I have a matrix of choices.</p>
<ol>
<li>Create tensors on CPU, then push them to GPU via pinned memory.</li>
<li>Create tensors directly on GPU.</li>
</ol>
<p>"1.  Create tensors on the <strong>get_item</strong>(index) of the DataSet<br>
2.  Create tensors in the collate_batch function of the DataLoader (or write my own DataLoader)</p>
<p>I thought I understood all this, but as it turns out, there are a few gaps in my understanding.  I have 3 questions.</p>
<p>Is there any downside to directly creating GPU tensors in the <strong>get_item</strong>(index) call on the DataSet.  Is this a bad idea?</p>
<p>What’s the best practice for creating tensors generally, should they created for each item in the dataset, or should the dataset just return what it returns, and the loader take care of creating tensors?</p>
<p>Where is the right place to decide which device a tensor goes to?  The DataSet?  The DataLoader, or the training loop itself?</p>
<p>Any insight appreciated!</p>

ptrblck: <p>The downside of creating GPU tensors in <code>__getitem__</code> or push CPU tensors onto the device is that your <code>DataLoader</code> won’t be able to use multiple workers anymore.<br>
If you try to set <code>num_workers &gt; 0</code>, you’ll get a CUDA error:</p>
<blockquote>
<p>RuntimeError: CUDA error: initialization error</p>
</blockquote>
<p>This also means that your host and device operations (most likely) won’t be able to overlap anymore, i.e. your CPUs cannot load and process the data while your GPU is busy training the model.</p>
<p>If you are “creating” the tensors, i.e. sampling them, this could still be a valid approach.<br>
However, if you are loading and processing some data (e.g. images), I would write the <code>Dataset</code> such that a single example is loaded, processed and returned as a CPUTensor.<br>
If you use <code>pin_memory=True</code> in you <code>DataLoader</code>, the transfer from host to device will be faster as described <a href="https://devblogs.nvidia.com/how-optimize-data-transfers-cuda-cc/">in this blogpost</a>.<br>
Inside the training loop you would push the tensors onto the GPU. If you set <code>non_blocking=True</code> as an argument in <code>tensor.to()</code>, PyTorch will try to perform the transfer asynchronously as decribed <a href="https://pytorch.org/docs/stable/tensors.html#torch.Tensor.to">here</a>.</p>
<p>The <code>DataLoader</code> might use a <code>sampler</code> or a custom <code>collate_fn</code>, but shouldn’t be responsible of creating the tensors.</p>

DuaneNielsen: <p>That’s incredibly helpful (as usual).  Thanks Peter!</p>

raunaks: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I had a question related to this. Actually I am loading image filenames in my dataset, and loading the images in <code>__get_item__</code>. I also use manually written data augmentation functions in <code>__get_item__</code>, so wouldn’t it be better to have these run on the GPU? Currently training is pretty slow.</p>
<p>Now I think I have 2 options - either I write the data augmentation functions for batched tensors outside the DataLoader, or I can create GPU tensors in <code>__get_item__</code>.</p>
<p>What would you recommend?</p>

ptrblck: <p>What kind of data augmentation methods are you currently using?<br>
If you are using <code>PIL</code>, you might want to install <code>PIL-SIMD</code>, which is a drop-in replacement and will use SIMD operations to speed up the transformations.<br>
Are you using multiple workers in your <code>DataLoader</code>? This should usually speeding up the data loading and processing.</p>
<p>If you are bottlenecked by the CPU and have enough GPU resources, you might also try to use <a href="https://github.com/NVIDIA/DALI" rel="nofollow noopener">NVIDIA/DALI</a>. <a class="mention" href="/u/januszl">@JanuszL</a> might give you some more information about it. Have a look at <a href="https://discuss.pytorch.org/t/how-to-prefetch-data-when-processing-with-gpu/548/31">his post here</a> for some additional information.</p>