Problem with Autograd

username: Andrea_Rosasco
cooked: <p>Hello,<br>
the problem I’m facing is quite complex and I’m not sure if I can reduce it to a minimal working example without cutting out some detail that might be the source of my problem, so I’ll try to be as clear as possible.</p>
<p>The problem I’m working on is the re-implementation of the paper “Dataset Distillation”. The overall objective is to modify a small subset of the training set so that training a network on this subset will lead to a good accuracy (potentially comparable to the one obtained by training the model on the whole dataset).</p>
<p>The algorithm consists of doing one or more SGD steps on the model using the small subsets then evaluating the model with a batch sampled from the training set and propagating the gradient back to the subset modifying the input samples. To implement this I had to use the library higher as I had to deal with second-order derivatives.</p>
<pre><code class="lang-auto">def distill(model, buff_data, img_lr, criterion, train_loader):
    model.train()

    buff_imgs, buff_trgs = buff_data

    model_lr = 0.1

    buff_opt = torch.optim.SGD([buff_imgs], lr=img_lr)
    model_opt = torch.optim.SGD(model.parameters(), lr=model_lr)

    for i in range(config['outer_steps']):
        for step, (ds_imgs, ds_trgs) in enumerate(train_loader):
            ds_imgs = ds_imgs.cuda()
            ds_trgs = ds_trgs.cuda()
 
            with higher.innerloop_ctx(model, model_opt) as (fmodel, diffopt):
								acc_loss = None
                for j in range(config['inner_steps']):
                    # First step modifies the model
                    buff_out = fmodel(buff_imgs)
                    buff_loss = criterion(buff_out, buff_trgs)
                    diffopt.step(buff_loss)

		            # Second step modifies the images
		            ds_out = fmodel(ds_imgs)
		            ds_loss = criterion(ds_out, ds_trgs)
		            acc_loss = acc_loss + ds_loss if acc_loss is not None else ds_loss

                acc_loss.backward()
                buff_opt.step()
                buff_opt.zero_grad()
</code></pre>
<p>It worked fine but I then decided to implement the learning of the learning rates. It works like this: if I’m doing 10 inner SGD steps I’ll have a list of ten tensors containing a single element.<br>
Each learning rate is specific to an SGD step and it’s learned by backpropagating on the training loss specific to that step.</p>
<pre><code class="lang-auto">def distill(model, buff_data, meta_lr, criterion, train_loader):
    model.train()

    buff_imgs, buff_trgs = buff_data

    model_lr = 0.1

    buff_opt = torch.optim.SGD([buff_imgs], lr=meta_lr)
    # DIFF
    model_opt = torch.optim.SGD(model.parameters(), lr=1)
    lr_list = []
    lr_opts = []
    for _ in range(config['inner_steps']):
        lr = torch.tensor([config['model_lr']], requires_grad=True, device='cuda')
        lr_list.append(lr)
        lr_opts.append(torch.optim.SGD([lr], meta_lr,))
    # END DIFF

    for i in range(config['outer_steps']):
        for step, (ds_imgs, ds_trgs) in enumerate(train_loader):
            ds_imgs = ds_imgs.cuda()
            ds_trgs = ds_trgs.cuda()
 
            with higher.innerloop_ctx(model, model_opt) as (fmodel, diffopt):
								acc_loss = None
                for j in range(config['inner_steps']):
                    # First step modifies the model
                    buff_out = fmodel(buff_imgs)
                    buff_loss = criterion(buff_out, buff_trgs)
                    diffopt.step(buff_loss)

		            # Second step modifies the images
		            ds_out = fmodel(ds_imgs)
		            ds_loss = criterion(ds_out, ds_trgs)
		            acc_loss = acc_loss + ds_loss if acc_loss is not None else ds_loss

                    # DIFF
		            lr_opts[j].zero_grad()
                    ds_loss.backward(retain_graph=True)
                    lr_opts[j].step()
                    # END DIFF

                acc_loss.backward()
                buff_opt.step()
                buff_opt.zero_grad()
</code></pre>
<p>If I don’t use <code>retain_graph=True</code> when calling backward to update the learning rate I get an error on the second iteration. If I set it to True it starts working but it’s extremely slow.</p>
<p>Since the update to a specific learning rate is computed using local information I can’t understand why I have to set <code>retain_graph=True</code>. I’m surely doing something wrong but I don’t understand what. Could someone give me some help?</p>

username: albanD
cooked: <aside class="quote no-group" data-username="Andrea_Rosasco" data-post="1" data-topic="103151">
<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/andrea_rosasco/48/14939_2.png" class="avatar"> Andrea_Rosasco:</div>
<blockquote>
<p>If I don’t use <code>retain_graph=True</code> when calling backward to update the learning rate I get an error on the second iteration.</p>
</blockquote>
</aside>
<p>Is that on the <code>ds_loss</code> backward or <code>acc_loss</code> backward? Because the <code>acc_loss</code> will backprop in the same graph as the <code>ds_loss</code> so I think the retain_graph is required here.<br>
For the slowdown, I guess it depends how many times you unroll your loop as you differentiate through all these iterations.</p>

username: Andrea_Rosasco
cooked: <p>It’s on the <code>ds_loss</code> backward (you can see it in the second snippet inside the inner and after the <code> # DIFF</code> comment).</p>
<blockquote>
<p>Because the <code>acc_loss</code> will backprop in the same graph as the <code>ds_loss</code> so I think the retain_graph is required here.</p>
</blockquote>
<p>I thought about it and I had also tried to change <code>ds_loss.backward(retain_graph=True)</code> to <code>criterion(fmodel(ds_imgs, ds_trgs)).backward()</code> to create a separate path on the graph which does not interfere with the following backpropagation on <code>acc_loss</code>. Unfortunately, that didn’t help much as the problem occur earlier when I call for the second time <code>ds_loss.backward()</code>. Since every iteration optimizes a different learning rate I don’t understand why should I retain the graph there.</p>
<blockquote>
<p>For the slowdown, I guess it depends how many times you unroll your loop as you differentiate through all these iterations.</p>
</blockquote>
<p>That’s fair but for the same number of inner iterations, the second snippet takes a lot more. I don’t think it should since I’m learning literally just one new parameter and it doesn’t have to backprop through several steps (just through one since I update it after I evaluate the model on the training set).</p>
<p>I think that whatever the bug is the two “symptoms” are correlated i.e. it’s taking so much cause it’s doing something which requires <code>retain_graph=True</code> while what I’m actually trying to do shouldn’t need it and should do a much more simple (and cost efficient) update</p>

username: albanD
cooked: <p>That sounds good.</p>
<p>Could you try to use <a href="https://github.com/szagoruyko/pytorchviz/"><code>torchviz</code></a> to print the graph and see where it overlaps (give it as input the loss from the first and second iteration and see where the two graphs start to overlap)? As it is a bit hard to say from just looking at the code indeed.</p>

username: Andrea_Rosasco
cooked: <p>Sure, if it’s hard for me that I’ve written it, it’s even harder for you since you are not familiar with it.<br>
I didn’t know about torchviz and I had some problem using it on windows but here is the first graph:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/f/7f412432e9bbbe819dd9fd0a3358451d458d4643.png" data-download-href="https://discuss.pytorch.org/uploads/default/7f412432e9bbbe819dd9fd0a3358451d458d4643" title="first"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/f/7f412432e9bbbe819dd9fd0a3358451d458d4643_2_458x500.png" alt="first" data-base62-sha1="i9KcCtiK1ed7Zuv8gJsQQN7Evbd" width="458" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/f/7f412432e9bbbe819dd9fd0a3358451d458d4643_2_458x500.png, https://discuss.pytorch.org/uploads/default/optimized/3X/7/f/7f412432e9bbbe819dd9fd0a3358451d458d4643_2_687x750.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/f/7f412432e9bbbe819dd9fd0a3358451d458d4643_2_916x1000.png 2x"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">first</span><span class="informations">5604×6113 999 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
I couldn’t update the second because it’s too big but it basically the second iteration with the first one stacked above it (since to backpropagate on acc loss all the previous iterations are needed).</p>
<p>Now, I don’t know if you can download it in full resolution but on the center-left of the image, you can see a one-element tensor i.e. the learning rate for that step.</p>
<p>In the graph at the second iteration, there are both the learning rate for the second step and, the one for the previous step. Basically, what I think is happening is that, at a given inner iteration, I’m not only updating the current learning rate but also all of the past ones.</p>
<p>Is there any way I can avoid this?</p>

username: albanD
cooked: <p>The image you shared if full resolution if you click on it. Thanks!</p>
<p>Could it be that the <code>j</code> index is a typo in your code? <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: Andrea_Rosasco
cooked: <p>No, sorry I just mistyped it in the snippet (I edited now); j is the index of the inner loop so it makes sense.</p>
<p>I think that, since <code>fmodel</code> is functional and the graph of all inner iterations is kept in order to backpropagate to <code>buff_imgs</code>, when I call backward on <code>ds_loss</code> to update the current learning rate, the gradient actually flows back trough all the graph.</p>
<p>What could I do to prevent this? I need the whole graph to update <code>buff_imgs</code> but when I update the learning rate I just need the gradient to flow to the current learning rate tensor.</p>
<p>I think one solution could be to calculate the output of <code>fmodel</code> two times: one to update <code>buff_imgs</code> after all iterations and the other with <code>no_grad()</code> to update the learning rate. Would that work?</p>

username: albanD
cooked: <aside class="quote no-group" data-username="Andrea_Rosasco" data-post="7" data-topic="103151">
<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/andrea_rosasco/48/14939_2.png" class="avatar"> Andrea_Rosasco:</div>
<blockquote>
<p>What could I do to prevent this?</p>
</blockquote>
</aside>
<p>If you use nightly build, you can pass <code>inputs=</code> to <code>.backward()</code> to specify what you want to compute gradients for. In this case <code>ds_loss.backward(inputs=(lr_list[j],), retrain_graph=True)</code>.</p>
<p>Otherwise, you can do <code>grad, = autograd.grad(ds_loss, lr_list[j], retain_graph=True); lr_list[j].grad = grad</code> but it is not as nice <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>

username: Andrea_Rosasco
cooked: <p>Thank you so much, that’s just what I needed!<br>
Yeah, the nightly build version is cooler but I’ll stick with the standard one for now.</p>
<p>Side question: do you plan to add graph visualization tools directly to PyTorch?</p>

username: albanD
cooked: <p>torchviz is more or less the official way to do it <img src="https://discuss.pytorch.org/images/emoji/apple/stuck_out_tongue.png?v=9" title=":stuck_out_tongue:" class="emoji" alt=":stuck_out_tongue:"><br>
The main drawback is that it is very hard to read from a user point of view in general.<br>
Note that the tensorboard implementation has a nice tool to visualize graphs at the torch.nn level as well <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>