Is it possible to preserve the structure of computational graph without tensor data?

sh0416: <p>Hello,</p>
<p>I am struggling with implementing reversible residual network.</p>
<p>When I implement this, the computational graph is consisted of some nodes, and this nodes are constructed using Tensor.</p>
<p>The strength of reversible network is that it can construct earlier activation with current activation. However, the computational graph use Tensor as nodes, which makes them always store the earlier activation.</p>
<p>I want to use gpu memory as efficient as I could. So, I want to retain the structure of graph while the tensor data is freed. Can I do this? If I can, how to do this?</p>
<p>One possible solution is to delete tensor.data. However, I worried that it makes unexpected result.</p>
<p>Thanks in advance.</p>

sh0416: <p>I found out that deleting tensor.data is impossible.</p>
<p>My test code:</p>
<pre><code class="lang-auto">a = torch.zeros([32, 64], requires_grad=True).cuda()
b = torch.zeros([64, 128], requires_grad=True).cuda()
c = torch.zeros([128, 256], requires_grad=True).cuda()
print("Calculate result four times: ", torch.cuda.memory_allocated())
d = torch.matmul(a, b)
e = torch.matmul(d, c)
print("Calculate result four times: ", torch.cuda.memory_allocated())
del d.data
print("Calculate result four times: ", torch.cuda.memory_allocated())
e.mean().backward()
print("Calculate result four times: ", torch.cuda.memory_allocated())
print(a.grad)
print(b.grad)
</code></pre>

albanD: <p>Hi,</p>
<p><code>.data</code> is not a thing anymore. It has been removed from the doc and will be deleted.</p>
<p>The graph is actually composed of Node (the function that needs to run in the backward) and not Tensor. Only the Tensors required to compute the backward are saved by the Nodes that needs them.</p>

sh0416: <p>Can I ask how to deal with this situation??</p>
<ul>
<li>I construct a computational graph with tensor.</li>
<li>I <strong>don’t want to hold data, i.e. activation</strong>, in the intermediate nodes.</li>
<li>When I call <code>backward</code>, I want to calculate the gradient of intermediate nodes with my custom <code>backward</code>.</li>
</ul>
<p>Clearly, the first and third requirements can be accomplished by using normal tensor operation and the custom <code>autograd.Function</code>, respectively.<br>
The main concern is how to remove <strong>only data</strong>, not the computation node itself.</p>
<p>Is there any suggestion about implementing this?? I think I might to hack <code>ctx</code> in the <code>autograd.Function</code>, but there is lack of documentation… <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=9" title=":frowning:" class="emoji" alt=":frowning:"></p>
<p>Some code template I want is to this.</p>
<pre><code class="lang-auto">class CustomFunction(autograd.Function):
  def forward(ctx, x):
    output = f(input)
    # TODO: remove input data or deallocate gpu memory for input data. but how?
    return output

  def backward(ctx, grad_output):
    return torch.ones_like(grad_output)
</code></pre>

albanD: <p>Hi,</p>
<p>If you don’t want any Tensor to be saved, just don’t save anything in the <code>ctx</code> and no Tensor will be saved for that Function.</p>
<p>Every Tensor is deleted as soon as it is not referenced by anything. So the input data will be deleted as soon as you don’t use it in your forward function.</p>

sh0416: <p>Thanks for the answer. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>I tested it by my custom code.</p>
<pre><code class="lang-python">class CustomMM(autograd.Function):
  def forward(ctx, x, y):
    out = torch.mm(x, y)
    return out

  def backward(ctx, grad_out):
    return magic_grad_for_x(grad_out), magic_grad_for_y(grad_out)

print(torch.cuda.memory_allocated())
x = torch.rand(10000, 10, requires_grad=True).cuda()
y = torch.rand(10, 10000).cuda()
z = torch.rand(10000, 100).cuda()
print(torch.cuda.memory_allocated())

output = CustomMM.apply(CustomMM.apply(x, y), z)
print(torch.cuda.memory_allocated())

# output.backward()  # It doesn't work because there is no magic function :)
</code></pre>