How to save computation graph of a gradient?

craigyang: <p>Hi, how should I save the computation graph of a gradient vector computed from <code>torch.autograd.grad(loss, model.parameters(), create_graph=True)</code>?</p>
<p>The background is that I want to compute the Hessian-vector products of <code>k</code> vectors: <code>H V</code>, in which <code>H</code> is the Hessian of a neural network with <code>n</code> parameters, and <code>V</code> is a constant matrix with <code>n</code> rows and <code>k</code> columns.  To do that, I compute the gradient of the inner product between gradient of the network forward function <code>g</code> and <code>V</code>, with respect to the network parameters. An example that works for a tiny network is</p>
<pre><code class="lang-auto">import torch

# define the tiny "network"
class quadratic_fun(torch.nn.Module):
    def __init__(self):
        super(quadratic_fun, self).__init__()
        self.x = torch.nn.Parameter(torch.ones(5, requires_grad=True))
        self.y = torch.nn.Parameter(torch.ones(5, requires_grad=True))

    def forward(self):
        loss = torch.norm(self.x) ** 2 + torch.norm(self.y) ** 2
        return loss

# compute the flattened gradient with create_graph=True
model = quadratic_fun()
loss_quad = model.forward()
grad_ft = torch.autograd.grad(loss_quad, model.parameters(), create_graph=True)
flat_grad = torch.cat([g.contiguous().view(-1) for g in grad_ft])

# generate the constant matrix V, and compute the matrix-gradient product
torch.manual_seed(0)
V = torch.randn((10, 3))
h = torch.matmul(flat_grad, V)

# compute the matrix-Jacobian product by iterating over the columns of the constant matrix
for i in range(3):
    hvp = torch.autograd.grad(h[i], model.parameters(), retain_graph=True)
    hvp_flat = torch.cat([g.contiguous().view(-1) for g in hvp])
    print(hvp_flat)
</code></pre>
<p>which gives</p>
<pre><code class="lang-auto">tensor([-2.2517, -0.8678, -0.6320, -2.5267,  0.2397, -0.2232, -0.9854,  0.2248,
        -0.2046,  0.1050])
tensor([-2.3047,  1.6974, -4.2304,  0.7000,  2.4753, -1.2272,  0.4968, -1.6821,
         1.5849,  1.0457])
tensor([-0.5012,  1.3840,  0.6445,  0.6163, -0.2869,  0.0632,  0.8794, -4.6321,
        -0.5793,  4.6044])
</code></pre>
<p>However, this is not feasible on CUDA when <code>H</code> is the Hessian of a large neural network: with <code>retain_graph=True</code> in the third from last line, the CUDA memory will quickly be filled up. While if I don’t retain the graph, the graph will be freed after one iteration of the for loop. In that case, I would need to compute the gradient again, which is time-consuming. Thus I wonder if I can save the not only the gradient value, but also its associated computation graph (both generated from <code>grad_ft = torch.autograd.grad(loss_quad, model.parameters(), create_graph=True)</code>) to a file or buffer, and reload it in a later iteration of the for loop.</p>
<p>Some other posts I looked into but didn’t find an answer:</p>
<ul>
<li>
<a href="https://discuss.pytorch.org/t/how-to-save-computation-graph-in-pytorch/97787">This post</a> suggests using JIT, but it is not clear to me how to use the API for the graph of a gradient vector.</li>
<li>A reply in <a href="https://discuss.pytorch.org/t/get-the-jacobian-using-autograd/47567">this post</a> suggests to compute the matrix-Jacobian product with <a href="https://pytorch.org/docs/stable/generated/torch.autograd.functional.jacobian.html#torch.autograd.functional.jacobian" rel="noopener nofollow ugc"><code>torch.autograd.functional.jacobian</code></a>, but it looks like the API only works when the function to compute Jacobian is explicitly defined.)</li>
</ul>
<p>Thanks!</p>

Varal7: <aside class="quote no-group" data-username="craigyang" data-post="1" data-topic="128286">
<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/craigyang/48/29894_2.png" class="avatar"> craigyang:</div>
<blockquote>
<p>Thus I wonder if I can save the not only the gradient value, but also its associated computation graph (both generated from <code>grad_ft = torch.autograd.grad(loss_quad, model.parameters(), create_graph=True)</code>) to a file or buffer, and reload it in a later iteration of the for loop.</p>
</blockquote>
</aside>
<p>The short answer is: I think so! Using saved tensors default hooks.</p>
<p>The docs are still under review (<a href="https://github.com/pytorch/pytorch/pull/62362" rel="noopener nofollow ugc">#62362</a> and <a href="https://github.com/pytorch/pytorch/pull/62361" rel="noopener nofollow ugc">#62361</a>) but the functionality is already merged to master as of today!</p>
<p>In particular, the first PR describes exactly your use case where you want to save a computation graph to the disk and retrieve it later when needed.<br>
I think in your case you would want to do something like:</p>
<pre data-code-wrap="python"><code class="lang-python"># compute the flattened gradient with create_graph=True and store the graph on disk
torch.autograd.graph.set_saved_tensors_default_hooks(pack_hook, unpack_hook)  
model = quadratic_fun()
loss_quad = model.forward()
grad_ft = torch.autograd.grad(loss_quad, model.parameters(), create_graph=True)
flat_grad = torch.cat([g.contiguous().view(-1) for g in grad_ft])

# generate the constant matrix V, and compute the matrix-gradient product
torch.manual_seed(0)
V = torch.randn((10, 3))
h = torch.matmul(flat_grad, V)
torch.autograd.graph.reset_saved_tensors_default_hooks()

# compute the matrix-Jacobian product by iterating over the columns of the constant matrix
for i in range(3):
    hvp = torch.autograd.grad(h[i], model.parameters(), retain_graph=True)
    hvp_flat = torch.cat([g.contiguous().view(-1) for g in hvp])
    print(hvp_flat)

</code></pre>
<p>where <code>pack_hook</code> and <code>unpack_hook</code> are defined <a href="https://github.com/pytorch/pytorch/pull/62362/files#diff-db7358efad18a9d12f25e26575ad27527fcb7438cfd7b888db6600c9cc6ea3fcR674-R681" rel="noopener nofollow ugc">here</a>.</p>
<p>You can control exactly which part of the graph should be saved to disk by adapting the position of the calls to <code>set_saved_tensors_default_hooks</code> and <code>reset_saved_tensors_default_hooks</code>.</p>
<p>Alternatively, you use the context manager <code>torch.autograd.graph.save_on_cpu</code>, cf <a href="https://github.com/pytorch/pytorch/pull/62410/files" rel="noopener nofollow ugc">#62410</a>.</p>

craigyang: <p>Thanks Victor for the pointers to the new functionality! Is there a concrete example of what <code>pack_hook</code> and <code>unpack_hook</code> should look like? I tried your example <a href="https://github.com/pytorch/pytorch/pull/62362/files#r682315984" rel="noopener nofollow ugc">here</a> but have no idea why you have <code>inc()</code> and <code>lambda x: x</code> as pack and unpack hooks, and do not really understand why <code>f("cpu")</code> and <code>f("cuda")</code> have seemingly arbitrarily large values between the set and reset hook functions.</p>
<p>Also, in your example</p>
<pre><code class="lang-auto">    class SelfClosingTempFile():
        def __init__(self):
            self.fp = tempfile.TemporaryFile()

        def __del__(self):
            self.fp.close()

    def pack_hook(tensor):
        sctf = SelfClosingTempFile()
        torch.save(tensor, sctf.fp)
        return sctf

    def unpack_hook(sctf):
        sctf.fp.seek(0)
        return torch.load(sctf.fp)
</code></pre>
<p>what should <code>tempfile.TemporaryFile()</code> be?</p>

Varal7: <p>Hi Craig,</p>
<p>Please ignore my example that you linked, that’s a POC of what an incorrect usage of the hooks would be!<br>
You can use the docs at <a href="https://docs-preview.pytorch.org/62362/notes/autograd.html#hooks-for-saved-tensors" class="inline-onebox" rel="noopener nofollow ugc">Autograd mechanics — PyTorch 1.10.0 documentation</a> and at <a href="https://docs-preview.pytorch.org/62361/autograd.html#saved-tensors-default-hooks" class="inline-onebox" rel="noopener nofollow ugc">Automatic differentiation package - torch.autograd — PyTorch 1.10.0 documentation</a>.<br>
To use the example you pasted, you need to <code>import tempfile</code> (<a href="https://docs.python.org/3/library/tempfile.html#tempfile.TemporaryFile" class="inline-onebox" rel="noopener nofollow ugc">tempfile — Generate temporary files and directories — Python 3.9.6 documentation</a>)</p>

craigyang: <p>Thanks Victor for the explanation. I imported <code>tempfile</code> as you suggested and have the following code that works:</p>
<pre><code class="lang-auto">import torch
import tempfile

class quadratic_fun(torch.nn.Module):
    def __init__(self):
        super(quadratic_fun, self).__init__()
        self.x = torch.nn.Parameter(torch.ones(5, requires_grad=True))
        self.y = torch.nn.Parameter(torch.ones(5, requires_grad=True))

    def forward(self):
        loss = torch.norm(self.x) ** 2 + torch.norm(self.y) ** 2
        return loss

class SelfClosingTempFile():
    def __init__(self):
        self.fp = tempfile.TemporaryFile()

    def __del__(self):
        self.fp.close()

def pack_hook(tensor):
    sctf = SelfClosingTempFile()
    torch.save(tensor, sctf.fp)
    return sctf

def unpack_hook(sctf):
    sctf.fp.seek(0)
    return torch.load(sctf.fp)

torch.autograd.graph.set_saved_tensors_default_hooks(pack_hook, unpack_hook)
model = quadratic_fun()
loss_quad = model.forward()
grad_ft = torch.autograd.grad(loss_quad, model.parameters(), create_graph=True)
flat_grad = torch.cat([g.contiguous().view(-1) for g in grad_ft])

torch.manual_seed(0)
V = torch.randn((10, 3))
h = torch.matmul(flat_grad, V)

torch.autograd.graph.reset_saved_tensors_default_hooks()

for i in range(3):
    hvp = torch.autograd.grad(h[i], model.parameters(), retain_graph=True)
    hvp_flat = torch.cat([g.contiguous().view(-1) for g in hvp])
    print(hvp_flat)
</code></pre>
<p>However, I want to set <code>retain_graph</code> to be <code>False</code> in the third from last line: I cannot retain the graph because of memory limits (and that’s exactly why I want to save the graph when I created it). If I remove <code>retain_graph=True</code> from the above code, I would still get</p>
<pre><code class="lang-auto">RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.
</code></pre>
<p>which is the same as not using the hook functions. Is it because I put <code>reset_saved_tensors_default_hooks</code> to the wrong position? How should I do what I want with all the hook functions here?</p>

Varal7: <p>What if you call <code>torch.autograd.graph.reset_saved_tensors_default_hooks()</code> after the <code>for</code> loop instead (but keep the <code>retain_graph</code> option on). Does that still exceed your memory requirements?</p>

craigyang: <p>Hi Victor, sorry for getting back to you late (it took me some time to install the latest PyTorch onto a machine with CUDA).</p>
<p>In my code that does the actual training (not the example code above), I tried to put <code>torch.autograd.graph.reset_saved_tensors_default_hooks()</code> to after the <code>for</code> loop, but got</p>
<pre><code class="lang-auto">...
AttributeError: 'SelfClosingTempFile' object has no attribute 'fp'
...
RuntimeError: OSError: [Errno 24] Too many open files: '/tmp/tmpw5thrpry'
</code></pre>
<p>I have no control over how many files the process can open at the same time, though, and I did not find how to control the number of created files in the <code>SelfClosingTempFile()</code> class (like in <code>tempfile.TemporaryFile()</code>). Do you know if there is a workaround? Thanks!</p>

Varal7: <p>Hi Craig,</p>
<p>Thanks for taking the time to test this new functionality! You can try to increase the number of files that can be opened. For example, here: <a href="https://stackoverflow.com/questions/16526783/python-subprocess-too-many-open-files" class="inline-onebox" rel="noopener nofollow ugc">Python Subprocess: Too Many Open Files - Stack Overflow</a>.</p>
<p>Edit: actually, I’ll provide you with another version of the hooks that should handle this issue.</p>

craigyang: <p>Thanks for the pointer! After setting <code>ulimit -Sn 500000</code>, now there seems to be a problem with writing and running the tmp files when I run my code with <code>torch.autograd.graph.reset_saved_tensors_default_hooks()</code> after the <code>for</code> loop:</p>
<pre><code class="lang-auto">terminate called after throwing an instance of 'c10::Error'
 what():  [enforce fail at inline_container.cc:300] . unexpected pos 320 vs 252
frame #0: c10::ThrowEnforceNotMet(char const*, int, char const*, std::string const&amp;, void const*) + 0x47 (0x7f627ac4e4b7 in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libc10.so)
frame #1: &lt;unknown function&gt; + 0x25844b0 (0x7f62c2c164b0 in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so)
frame #2: &lt;unknown function&gt; + 0x257fa8c (0x7f62c2c11a8c in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so)
frame #3: caffe2::serialize::PyTorchStreamWriter::writeRecord(std::string const&amp;, void const*, unsigned long, bool) + 0xb5 (0x7f62c2c196f5 in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so)
frame #4: caffe2::serialize::PyTorchStreamWriter::writeEndOfFile() + 0x173 (0x7f62c2c199e3 in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so)
frame #5: caffe2::serialize::PyTorchStreamWriter::~PyTorchStreamWriter() + 0x125 (0x7f62c2c19c55 in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so)
frame #6: &lt;unknown function&gt; + 0xb1cec3 (0x7f62d5797ec3 in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_python.so)
frame #7: &lt;unknown function&gt; + 0x558988 (0x7f62d51d3988 in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_python.so)
frame #8: &lt;unknown function&gt; + 0x559c8e (0x7f62d51d4c8e in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_python.so)
&lt;omitting python frames&gt;
frame #12: &lt;unknown function&gt; + 0x5548f5 (0x7f62d51cf8f5 in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_python.so)
frame #13: &lt;unknown function&gt; + 0xaa175 (0x7f62d65c0175 in /lib/x86_64-linux-gnu/libstdc++.so.6)
frame #14: &lt;unknown function&gt; + 0xfbf034 (0x7f62c1651034 in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so)
frame #15: &lt;unknown function&gt; + 0x37083df (0x7f62c3d9a3df in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so)
frame #16: std::_Sp_counted_base&lt;(__gnu_cxx::_Lock_policy)2&gt;::_M_release() + 0x46 (0x7f62d51cab66 in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_python.so)
frame #17: torch::autograd::Engine::thread_main(std::shared_ptr&lt;torch::autograd::GraphTask&gt; const&amp;) + 0x132 (0x7f62c3d98672 in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so)
frame #18: torch::autograd::Engine::thread_init(int, std::shared_ptr&lt;torch::autograd::ReadyQueue&gt; const&amp;, bool) + 0x89 (0x7f62c3d90589 in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so)
frame #19: torch::autograd::python::PythonEngine::thread_init(int, std::shared_ptr&lt;torch::autograd::ReadyQueue&gt; const&amp;, bool) + 0x53 (0x7f62d5729163 in /home/craig/anaconda3/lib/python3.7/site-packages/torch/lib/libtorch_python.so)
frame #20: &lt;unknown function&gt; + 0xd6de4 (0x7f62d65ecde4 in /lib/x86_64-linux-gnu/libstdc++.so.6)
frame #21: &lt;unknown function&gt; + 0x9609 (0x7f62e658b609 in /lib/x86_64-linux-gnu/libpthread.so.0)
frame #22: clone + 0x43 (0x7f62e64b2293 in /lib/x86_64-linux-gnu/libc.so.6)
</code></pre>
<p>Do you have any idea what the cause might be?</p>

Varal7: <p>This seems to be a serialize error. Do you get the same error with these hooks:</p>
<pre><code class="lang-auto">import torch
import os
import uuid

tmp_dir = "temp"

class quadratic_fun(torch.nn.Module):
    def __init__(self):
        super(quadratic_fun, self).__init__()
        self.x = torch.nn.Parameter(torch.ones(5, requires_grad=True))
        self.y = torch.nn.Parameter(torch.ones(5, requires_grad=True))

    def forward(self):
        loss = torch.norm(self.x) ** 2 + torch.norm(self.y) ** 2
        return loss

class SelfDeletingTempFile():
    def __init__(self):
        self.name = os.path.join(tmp_dir, str(uuid.uuid4()))

    def __del__(self):
        os.remove(self.name)

def pack_hook(tensor):
    temp_file = SelfDeletingTempFile()
    torch.save(tensor, temp_file.name)
    return temp_file

def unpack_hook(temp_file):
    return torch.load(temp_file.name)


torch.autograd.graph.set_saved_tensors_default_hooks(pack_hook, unpack_hook)

model = quadratic_fun()
loss_quad = model.forward()
grad_ft = torch.autograd.grad(loss_quad, model.parameters(), create_graph=True)
flat_grad = torch.cat([g.contiguous().view(-1) for g in grad_ft])

# generate the constant matrix V, and compute the matrix-gradient product
torch.manual_seed(0)
V = torch.randn((10, 3))
h = torch.matmul(flat_grad, V)

# compute the matrix-Jacobian product by iterating over the columns of the constant matrix
for i in range(3):
    hvp = torch.autograd.grad(h[i], model.parameters(), retain_graph=True)
    hvp_flat = torch.cat([g.contiguous().view(-1) for g in hvp])
    print(hvp_flat)

torch.autograd.graph.reset_saved_tensors_default_hooks()
</code></pre>

Varal7: <p>Here are two other thoughts:</p>
<p>Maybe you can keep the graph of <code>h</code> on GPU but only move to to disk the part that computes the matrix-Jacobian product.</p>
<pre><code class="lang-auto">model = quadratic_fun()
loss_quad = model.forward()
grad_ft = torch.autograd.grad(loss_quad, model.parameters(), create_graph=True)
flat_grad = torch.cat([g.contiguous().view(-1) for g in grad_ft])

# generate the constant matrix V, and compute the matrix-gradient product
torch.manual_seed(0)
V = torch.randn((10, 3))
h = torch.matmul(flat_grad, V)

torch.autograd.graph.set_saved_tensors_default_hooks(pack_hook, unpack_hook)
# compute the matrix-Jacobian product by iterating over the columns of the constant matrix
for i in range(3):
    hvp = torch.autograd.grad(h[i], model.parameters(), retain_graph=True)
    hvp_flat = torch.cat([g.contiguous().view(-1) for g in hvp])
    print(hvp_flat)

torch.autograd.graph.reset_saved_tensors_default_hooks()
</code></pre>
<p>or even, instead of moving to disk, moving to CPU:</p>
<pre><code class="lang-auto">model = quadratic_fun()
loss_quad = model.forward()
grad_ft = torch.autograd.grad(loss_quad, model.parameters(), create_graph=True)
flat_grad = torch.cat([g.contiguous().view(-1) for g in grad_ft])

# generate the constant matrix V, and compute the matrix-gradient product
torch.manual_seed(0)
V = torch.randn((10, 3))
h = torch.matmul(flat_grad, V)

with torch.autograd.graph.save_on_cpu(pin_memory=True):
    for i in range(3):
        hvp = torch.autograd.grad(h[i], model.parameters(), retain_graph=True)
        hvp_flat = torch.cat([g.contiguous().view(-1) for g in hvp])
        print(hvp_flat)
</code></pre>

craigyang: <p>EDIT:</p>
<p>Thanks Victor for the additional thoughts. These two code snippets do work!</p>
<p>There is one issue though: I tried the functions you provided with a recent nightly build (1.10.0.dev20210805+cu102). Probably because of some incompatibility issues between this version and my cudatoolkit, the <code>torch.autograd.grad()</code> is much slower than before (in one instance, 40 seconds vs 5 seconds). Is there a PyTorch+cudatoolkit version combination that you recommend me to try? Thanks!</p>

Varal7: <p>Sorry, I’m not very familiar with version of python and cudatoolkit but there should not be a slowdown with the more recent versions!<br>
Is <code>torch.autograd.grad()</code> becoming much slower with the latest version of pytorch even when you don’t use the hooks?<br>
It is expected that using the hooks will incur a performance penalty.</p>

craigyang: <p>Gotcha, thanks! Yes, <code>torch.autograd.grad()</code> has become much slower even when I don’t use the hooks. I am reinstalling everything and just wondered if you have any idea on which version combination would work.</p>
<p>I’ll tell you how much penalty the hooks will incur once I get the proper PyTorch+CUDA toolkit versions installed and up running.</p>

Yufei-Gu-451: <p>I am attempting a similar operation after 3 years with the new torch.autograd.graph.saved_tensors_hooks(<em>pack_hook</em> , <em>unpack_hook</em> ) API.</p>
<p>I am iteratively computing large amounts of hessian vector products and sum over all batches in my dataset (dataloader). Thus to optimize computation time and CUDA/CPU storage, I am attempting to save the first-order gradients and computation graphs of all data batches outside the iteration and retrieve them when computing the hessian vector products.</p>
<pre><code class="lang-auto">grad_ls = []
for batch in dataloader
            # Set Tensors Hook to save Gradients and Computation Graph on Disk
            with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
                grad_ft = torch.autograd.grad(loss, self.params, allow_unused=True, create_graph=True, retain_graph=True)

            # Pack also the Gradient on Disk/Tempfile; Save the hook to a list
            grad_ls.append((pack_hook(grad_ft), len(batch)))
</code></pre>
<p>However, when computing the hvps using torch.autograd.grad(), the function returns None for all model parameters, in contrast to computing gradients and hvps at the same time. The code is provided below where v_list are the target vector in</p>
<pre><code class="lang-auto">for (grad_ft, batch_size) in grad_ls:
      grad_ft = unpack_hook(grad_ft)

      Hv = torch.autograd.grad(outputs=grad_ft, inputs=self.params,
                                     grad_outputs=v_list, create_graph=False, allow_unused=True)
</code></pre>
<p>I have checked the grad_ft tensor tuple saved on disk is the same in each retrieval. The second torch.autograd.grad() is executed, making the computation graph likely to be existing. I am curious on the cause of this behaviour.</p>