Multiple forward passes, tuple mining

FrAnY: <p>Hello,<br>
i want to perform multiple forward passes of full batches with a GPU.<br>
Therefore each batch uses the complete VRAM.<br>
I want to store them temporarily and perform tuple mining with the concatenad batches.<br>
This is my approach but i run out of VRAM:</p>
<pre><code class="lang-auto">stored_embeddings = []
stored_labels = []
for i, (inputs, labels) in enumerate(loader):
                # Perform the forward pass
                outputs = model(inputs)
                # Mining on CPU if required
                if cpu_mining:
                    outputs_cpu = outputs.cpu()
                    del outputs
                    torch.cuda.empty_cache()
                    if (i + 1) % cpu_mining == 0:
                        # concat embeddings and pass to miner then compute loss based on them
                        outputs = torch.cat([torch.cat(stored_embeddings), outputs_cpu])
                        labels = torch.cat([torch.cat(stored_labels), labels])
                    else:
                        # fetch more data before mining
                        stored_embeddings.append(outputs_cpu)
                        stored_labels.append(labels)
                        continue
</code></pre>
<p>What is the correct way of doing this?<br>
Any help is appreciated.</p>

ptrblck: <p>You are attaching the computation graph with the intermediate tensors on the GPU via <code>outputs_cpu</code>.<br>
Moving <code>outputs</code> to the CPU and deleting <code>outputs</code> will not free the computation graph, since <code>outputs_cpu</code> is still attached to the computation graph as seen here:</p>
<pre><code class="lang-python">model = models.resnet50().cuda()
inputs = torch.randn(8, 3, 224, 224).cuda()

print(torch.cuda.memory_allocated()/1024**2)
# 102.32177734375

outputs = model(inputs)
print(torch.cuda.memory_allocated()/1024**2)
# 759.55859375

outputs_cpu = outputs.cpu()
del outputs # this does not delete the intermediates, as outputs_cpu still holds to them
print(torch.cuda.memory_allocated()/1024**2)
# 759.52783203125 !!!

# here you are attaching the outputs_cpu tensor with the entire computation graph
# with intermediates on the GPU
outputs = torch.cat([torch.cat(stored_embeddings), outputs_cpu])
</code></pre>
<p><code>.detach()</code> the <code>outputs</code> before moving them to the CPU and it should properly free the graph in <code>del outputs</code>:</p>
<pre><code class="lang-python">outputs_cpu = outputs.detach().cpu()
</code></pre>

FrAnY: <p>Thanks, you are a gem in this forum. Just to be sure this won’t break later backpropagation? Because the docs say: The result will never require gradient. But i guess this does not mean that it removes the old ones?</p>

ptrblck: <p>Detaching the <code>output</code> will break the computation graph and you will not be able to call <code>.backward()</code> on <code>output_cpu</code> or any loss calculated from it.<br>
If you need to call <code>backward</code> and calculate the gradients on <code>output_cpu</code> or any created tensor from it, you would need to keep the computation graph alive but would also store it in each iteration thus increasing the GPU memory.</p>

FrAnY: <p>Then there is no option to keep the gradients and free the GPU?</p>

ptrblck: <p>I’m not sure we are taking about the same issue.<br>
Are you referring to already calculated gradients which are stored in the <code>.grad</code> attributes of the parameters by “keep the gradients”? If so, these <code>.grad</code> attributes will not be changed in any way.<br>
Detaching the output will cut if from the computation graph, which stores the intermediate forward activations (not gradients), which are needed to calculate the gradients in the <code>backward</code> call.</p>

FrAnY: <p>Sorry for beeing imprecise, i mean the forward activations to calculate the gradients in the <code>backward</code><br>
call.</p>

ptrblck: <p>Ah OK. Yes, you will not be able to delete the forward activations and still be able to calculate the gradients (they are needed in the backpropagation).</p>