Assign values in a tensor to 0 to stop building computation graph

Yujian_Liu: <p>Hi, I’m having CUDA out of memory issue now, and I’m thinking to reduce some useless variables in the computation graph to save memory.</p>
<p>Particularly, I have a tensor <code>sim_score</code> which is calculated in the forward pass by some previous encoders. e.g., <code>sim_score = tensor([[1, 2, 3, 4], [5, 6, 7, 8]])</code>. I also have a mask tensor, e.g., <code>mask = tensor([[0, 1, 0, 0], [1, 1, 1, 0]])</code>. I will do something like <code>output = torch.sum(sim_score * mask, dim=1)</code>. As you can see, 1, 3, 4, 8 in <code>sim_score</code> are useless in future computation. However, multiplying the mask tensor seems to keep all values in the computation graph. I wonder is there a way to don’t keep those useless values in the graph to save some gpu memory?</p>
<p>I’ve tried to directly assign the values in <code>sim_score</code> to 0 like <code>sim_score[~mask] = 0</code>, but I got the error: RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation.</p>
<p>Many thanks!</p>

tom: <p>Two comments here:<br>
The autograd graph if per-tensor not per-element, so you cannot mask within a tensor to free the computational graph.</p>
<p>For much the same reason assigning values will not cut the graph but instead insert a new node CopySlice:</p>
<pre data-code-wrap="python"><code class="lang-python">a = torch.randn(5, requires_grad=True)
b = torch.randn(5, requires_grad=True)
c = torch.zeros(4, 7, requires_grad=True)
d = c * 2
d[1, 1:-1] = a
d[2, 1:-1] = b
</code></pre>
<p>will give an autograd graph that looks like this:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/7/37f24fd009b8c275039d6597131a648d9471ac74.png" alt="image" data-base62-sha1="7YVpvAy2FyCvlRciI8tApjtvdze" width="663" height="347"></p>
<p>If you were to draw intermediate graphs, we would see that while the tensor d stays the same, the autograd node directly attached to it moves from MulBackward0 to the CopySlices with the assignments.</p>
<p>Best regards</p>
<p>Thomas</p>

Yujian_Liu: <p><a class="mention" href="/u/tom">@tom</a> Hi Thomas, thanks for your reply! I have a follow up question, what if I have:</p>
<pre><code class="lang-auto">a = torch.randn((4, 7), requires_grad=True)
c = torch.zeros(4, 7, requires_grad=True)
c[1, 1:3] = a[1, 1:3]
c[2, 3:4] = a[2, 3:4]
# instead of operating on a, we use c for future computation
</code></pre>
<p>And imagine <code>a</code> is calculated by some previous layers. I guess some values in <code>a</code> won’t receive gradient in backward pass. But can this save some gpu memory for me or not at all (compared to using <code>a * mask</code> for future computation)?</p>
<p>Best<br>
Yujian</p>

tom: <p>No, it cannot, because</p>
<ol>
<li>the tensors will still be there,</li>
<li>the building of the graph (what consumes the gpu memory) happens before the masking, so it would not help,</li>
<li>the autograd does not know which bits you don’t need because autograd only works with tensor-level information.</li>
</ol>
<p>So you could only save something if you had separate tensors and could avoid using some of them and let them go out of scope.</p>
<p>Best regards</p>
<p>Thomas</p>

Yujian_Liu: <p>Thanks! This is really helpful!</p>