Saving the autograd state of a model for later backprop

neonbjb: <p>I’m looking to be able to save all the state data generated during the forward pass so that I can perform backprop at any time. The reason I want to do this is to be able to train large, composite models on a home computer with one GPU.</p>
<p>Here’s the general scheme I have in mind:</p>
<ol>
<li>Load and compute forward() on one model, collect the outputs.</li>
<li>Load another model, compute forward() with the outputs from the first.</li>
<li>Compute backward on second model and get the gradient for it’s inputs (updating its parameters along the way)</li>
<li>Re-load the first model and compute the backward pass with the gradients obtained from (3).</li>
</ol>
<p>Ideally, I’d keep the parameters for both models in GPU memory and just swap the intermediate and optimizer states for the models in and out of memory. I realize this might be quite slow, but I’d really like to experiment with it.</p>
<p>I’ve done a bit of tinkering trying to get this to work already. It seems like the intermediate states that I am trying to get at are likely being stored in the tensors returned by each operation. E.g.:</p>
<pre><code class="lang-auto">x, y = torch.randn(&lt;&gt;), torch.randn(&lt;&gt;)
t1 = x + y
t2 = t1 * x
....
output = tfinal.softmax()
</code></pre>
<p>In this example, the ‘t’ variables seem to contain the state within their ‘data’ property. I found this by building two identical graphs with different inputs, performing a forward on each, then copying that data property for each intermediate tensor from one graph to the other. The backprop then computes the same grads for each graph.</p>
<p>The thing is, this isn’t a very robust approach. First of all, it is pretty much impossible to get these intermediate tensors for some operations. It’s also a real pain to record them. Finally, you cannot do this with the nn package because it’ll throw an error when you try to do backprop on the doctored model.</p>
<p>I think the solution to this might involve tinkering with the C++ code. It seems like the grad_fn objects built up by autograd must contain references to this state data. If I could access this state data from these grad_fn objects, I could just crawl the graph backward and copy the state data to and from them. Can anyone give me any tips on where in the cpp code I should look to find these states? Maybe some tips on alternative approaches?</p>

ptrblck: <aside class="quote no-group" data-username="neonbjb" data-post="1" data-topic="71506">
<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/neonbjb/48/21111_2.png" class="avatar"> neonbjb:</div>
<blockquote>
<p>Maybe some tips on alternative approaches?</p>
</blockquote>
</aside>
<p>I would try to either:</p>
<ul>
<li>reduce the batch size and try to train both models on the device</li>
<li>if you have two GPUs, use each GPU for a single model</li>
<li>try to use <code>torch.utils.checkpoint</code> to trade compute for memory</li>
</ul>

neonbjb: <p>Hi ptrblock,<br>
I appreciate the tips. I think these would work for what I’m currently looking at (well, minus the two GPUs… <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 would still like to look in to this deferred backprop approach, though. If for no other reason than that I’m curious what kind of performance penalties it would entail. If anyone has any tips on where the backprop graph’s states are stored, I’d really appreciate it.</p>

neonbjb: <p>For anyone else who finds this question, I managed to get something like what I wanted by modding the code on my own. You’ll need to build your own libpytorch if you want to use it (for now):<br>
<aside class="onebox githubcommit">
  <header class="source">
      <a href="https://github.com/neonbjb/pytorch/commit/8ddcba20fd5744d877bcf03101cc86fd1bfff74e" target="_blank" rel="nofollow noopener">github.com/neonbjb/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Commit">
    <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M10.86 7c-.45-1.72-2-3-3.86-3-1.86 0-3.41 1.28-3.86 3H0v2h3.14c.45 1.72 2 3 3.86 3 1.86 0 3.41-1.28 3.86-3H14V7h-3.14zM7 10.2c-1.22 0-2.2-.98-2.2-2.2 0-1.22.98-2.2 2.2-2.2 1.22 0 2.2.98 2.2 2.2 0 1.22-.98 2.2-2.2 2.2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/neonbjb/pytorch/commit/8ddcba20fd5744d877bcf03101cc86fd1bfff74e" target="_blank" rel="nofollow noopener">Add ability to access saved variables from Python</a>
    </h4>

    <div class="github-info">
      <div class="date">
        committed <span class="discourse-local-date" data-format="ll" data-date="2020-03-06" data-time="03:52:45" data-timezone="UTC">03:52AM - 06 Mar 20 UTC</span>
      </div>

      <div class="user">
        <a href="https://github.com/neonbjb" target="_blank" rel="nofollow noopener">
          <img alt="neonbjb" src="https://avatars0.githubusercontent.com/u/833082?v=4" class="onebox-avatar-inline" width="20" height="20">
          neonbjb
        </a>
        
      </div>

      <div class="lines" title="changed 4 files with 40 additions and 1 deletions">
        <a href="https://github.com/neonbjb/pytorch/commit/8ddcba20fd5744d877bcf03101cc86fd1bfff74e" target="_blank" rel="nofollow noopener">
          <span class="added">+40</span>
          <span class="removed">-1</span>
        </a>
      </div>
    </div>

  </div>
</div>


  <div class="github-row">
    <pre class="github-content" style="white-space: normal;">Capability to restore them efficiently still needed. You can
do this in C++ by messing with the data_ field directly,
but a nice...</pre>
  </div>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>I’ve filed a feature request in case the mods are interested in bringing something like this into mainline:<br>
<aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/34273" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/34273" target="_blank" rel="nofollow noopener">Allow accessing or save/restore of SavedVariables</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-03-05" data-time="00:08:09" data-timezone="UTC">12:08AM - 05 Mar 20 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/neonbjb" target="_blank" rel="nofollow noopener">
          <img alt="neonbjb" src="https://avatars0.githubusercontent.com/u/833082?v=4" class="onebox-avatar-inline" width="20" height="20">
          neonbjb
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">🚀 Feature
Allow users to save and restore the state of the intermediate variables used to compute the backward pass of a...</p>
</div>

<div class="labels">
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">needs research</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">triaged</span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>