The computation graph is breaking in the outer loop of meta-learning, Meta gradients are None (FOMAML)

username: DeepLearner17
cooked: <p>Gradient computation in meta-learning: the computation graph is breaking in the outer loop. First Order MAML</p>
<p>I am relying on this pytorch lightning tutorial  <a href="https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/tutorial16/Meta_Learning.html" rel="noopener nofollow ugc">https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/tutorial16/Meta_Learning.html</a></p>
<p><code>In outer_loop()</code></p>
<p><code>p_global</code>  is <code>None</code> (all <code>param.grad</code> in<code> self.model.parameters()</code> are<code> None</code>, while those of <code>local_model_parameters()</code> are not). As a consequence <code> p_global.grad += p_local.grad</code> can’t be achieved</p>
<pre><code class="lang-auto"> for p_global, p_local in zip(self.model.parameters(), local_model.parameters()):
                    p_global.grad += p_local.grad  # First-order approx. -&gt; add gradients of finetuned and base model
</code></pre>
<p>It  could be related to</p>
<p><code>local_model = deepcopy(self.model)</code><br>
In <code>adapt_few_shot()</code> but <code>deepcopy()</code> is needed, or may be <code>create_graph=True</code> is needed ?</p>
<p>I tried the following but l am not sure that it’s true</p>
<pre><code class="lang-auto">                for p_global, p_local in zip(self.model.parameters(), local_model.parameters()):

                      if p_global.grad is None:
                          p_global.grad = torch.zeros_like(p_local.grad)
                    
                      p_global.grad += p_local.grad
</code></pre>
<p>Doing this leads to :</p>
<p><code> p_global.grad = p_local.grad</code> instead of <code>p_global.grad += p_local.grad</code> since after each gradient update,             <code>opt.step(), opt.zero_grad()</code> every <code>p_global.grad</code> in <code>self.model.parameters()</code> becomes <code>None</code></p>
<p>Code:</p>
<pre><code class="lang-auto">    def adapt_few_shot(self, support_imgs, support_targets):
        # Determine prototype initialization
        support_feats = self.model(support_imgs)
        prototypes, classes = ProtoNet.calculate_prototypes(support_feats, support_targets)
        support_labels = (classes[None,:] == support_targets[:,None]).long().argmax(dim=-1)
        # Create inner-loop model and optimizer
        local_model = deepcopy(self.model)
        local_model.train()
        local_optim = optim.SGD(local_model.parameters(), lr=self.hparams.lr_inner)
        local_optim.zero_grad()
        # Create output layer weights with prototype-based initialization
        init_weight = 2 * prototypes
        init_bias = -torch.norm(prototypes, dim=1)**2
        output_weight = init_weight.detach().requires_grad_()
        output_bias = init_bias.detach().requires_grad_()

        # Optimize inner loop model on support set
        for _ in range(self.hparams.num_inner_steps):
            # Determine loss on the support set
            loss, _, _ = self.run_model(local_model, output_weight, output_bias, support_imgs, support_labels)
            # Calculate gradients and perform inner loop update
            loss.backward()
            local_optim.step()
            # Update output layer via SGD
            # (https://discuss.pytorch.org/t/the-difference-between-torch-tensor-data-and-torch-tensor/25995/4):
            with torch.no_grad():
                output_weight.copy_(output_weight - self.hparams.lr_output * output_weight.grad)
                output_bias.copy_(output_bias - self.hparams.lr_output * output_bias.grad)

            # Reset gradients
            local_optim.zero_grad()
            output_weight.grad.fill_(0)
            output_bias.grad.fill_(0)

        # Re-attach computation graph of prototypes
        output_weight = (output_weight - init_weight).detach() + init_weight
        output_bias = (output_bias - init_bias).detach() + init_bias

        return local_model, output_weight, output_bias, classes

    def outer_loop(self, batch, mode="train"):
        accuracies = []
        losses = []
        self.model.zero_grad()

        # Determine gradients for batch of tasks
        for task_batch in batch:
            imgs, targets = task_batch
            support_imgs, query_imgs, support_targets, query_targets = split_batch(imgs, targets)
            # Perform inner loop adaptation
            local_model, output_weight, output_bias, classes = self.adapt_few_shot(support_imgs, support_targets)
            # Determine loss of query set
            query_labels = (classes[None,:] == query_targets[:,None]).long().argmax(dim=-1)
            loss, preds, acc = self.run_model(local_model, output_weight, output_bias, query_imgs, query_labels)
            # Calculate gradients for query set loss
            if mode == "train":
                loss.backward()

                for p_global, p_local in zip(self.model.parameters(), local_model.parameters()):
                    p_global.grad += p_local.grad  # First-order approx. -&gt; add gradients of finetuned and base model

            accuracies.append(acc.mean().detach())
            losses.append(loss.detach())

        # Perform update of base model
        if mode == "train":
            opt = self.optimizers()
            opt.step()
            opt.zero_grad()

        self.log(f"{mode}_loss", sum(losses) / len(losses))
        self.log(f"{mode}_acc", sum(accuracies) / len(accuracies))

    def training_step(self, batch, batch_idx):
        self.outer_loop(batch, mode="train")
        return None  # Returning None means we skip the default training optimizer steps by PyTorch Lightning
</code></pre>
<p>I have ckecked the following discussions:</p>
<ul>
<li>
<p><a href="https://discuss.pytorch.org/t/gradient-computation-in-meta-learning-algorithms/6039/4@rinuboney">https://discuss.pytorch.org/t/gradient-computation-in-meta-learning-algorithms/6039/4@rinuboney</a>  <a class="mention" href="/u/asura">@asura</a></p>
</li>
<li>
<p><a href="https://discuss.pytorch.org/t/calculating-meta-gradients-with-approximate-inner-loop-gradient/141025" class="inline-onebox">Calculating meta-gradients with approximate inner-loop gradient</a> <a class="mention" href="/u/alband">@albanD</a> <a class="mention" href="/u/tinfoilhat0">@TinfoilHat0</a></p>
</li>
<li>
<p><a href="https://discuss.pytorch.org/t/meta-learning-gradient-computation-error-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/181454" class="inline-onebox">#meta learning# #gradient computation error# RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation</a> <a class="mention" href="/u/ptrblck">@ptrblck</a> <a class="mention" href="/u/jsktt01">@jsktt01</a></p>
</li>
<li>
<p><a href="https://discuss.pytorch.org/t/how-to-free-the-graph-after-create-graph-true/58476" class="inline-onebox">How to free the graph after create_graph=True</a> <a class="mention" href="/u/alband">@albanD</a></p>
</li>
</ul>

username: Harbar-Inbound
cooked: <p>It seems like you are implementing first-order MAML (foMAML).</p>
<p>To achieve this, keep in mind that PyTorch does not initialize the gradient attributes of a module/parameter during instantiation, so these will be set to None. As a result, when you attempt to add the inner-loop grads to the meta-model, you attempt to add a tensor to <code>None</code> (the un-initialized meta-models gradient attributes).</p>
<p>You can resolve this, by manually setting the meta-models parameter’s attributes to a zero tensor with the same shape as the learnable weights, e.g. using <code>torch.zeros_like</code>. Or create an if statement to set the outer-loop grad to <code>p_local.grad</code> if <code>p_global.grad</code> is <code>None</code>.</p>

username: DeepLearner17
cooked: <p>Thank you <a class="mention" href="/u/harbar-inbound">@Harbar-Inbound</a> for your response. Indeed it is FOMAML.</p>
<p>Yes, l have tried the following but l am not sure that it’s true</p>
<pre><code class="lang-auto">                for p_global, p_local in zip(self.model.parameters(), local_model.parameters()):

                      if p_global.grad is None:
                          p_global.grad = torch.zeros_like(p_local.grad)
                    
                      p_global.grad += p_local.grad
</code></pre>
<p>Doing this leads to :</p>
<p><code> p_global.grad = p_local.grad</code> instead of <code>p_global.grad += p_local.grad</code> since after each gradient update,             <code>opt.step(), opt.zero_grad()</code> every <code>p_global.grad</code> in <code>self.model.parameters()</code> becomes <code>None</code></p>

username: Harbar-Inbound
cooked: <p>You are describing expected behavior, as <code>opt.zero_grad()</code> will set the <code>.grad</code> attribute of the optimized model to <code>None</code>.</p>
<p>See also <a href="https://discuss.pytorch.org/t/what-step-backward-and-zero-grad-do/33301">the difference between <code>step</code>, <code>backward</code>, and <code>zero_grad</code></a>.</p>
<p>If you want to change the behavior (i.e., set grads to zero), see the <code>set_to_none</code> parameter of <code>torch.optim.Optimizer</code> <a href="https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html#torch.optim.Optimizer.zero_grad" rel="noopener nofollow ugc">here</a>.</p>
<p>If you want to check that gradients are properly applied, you can implement the following:</p>
<ol>
<li>retain a copy of the meta-model (e.g., deep-copy),</li>
<li>step with an SGD optimizer (without any momentum, etc.),</li>
<li>compare the difference between the updated meta-model and copied model, and the gradient of the local-model.</li>
</ol>
<p>If the result of step 3 is <code>torch.allclose</code>, then gradients were applied properly*.</p>
<p>*These results may differ slightly if you use an optimizer without momentum, etc., and a learning rate of <code>1.0</code>.</p>

username: DeepLearner17
cooked: <p>To be sure l understand your trick <a class="mention" href="/u/harbar-inbound">@Harbar-Inbound</a> , you mean testing the following:</p>
<pre><code class="lang-auto">opt.step()
torch.allclose(p_global.grad,p_local.grad) 
torch.allclose(p_global,p_local)
opt.zero_grad()
</code></pre>
<p>?</p>