Gradient for nn.Parameter

Ahmed_Khaled: <p>I am trying to implement this algorithm:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/3/d396678226d917348c2b8fa0ba72691da45d960e.png" data-download-href="https://discuss.pytorch.org/uploads/default/d396678226d917348c2b8fa0ba72691da45d960e" title="l2c algo"><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/3/d396678226d917348c2b8fa0ba72691da45d960e.png" alt="l2c algo" data-base62-sha1="ubMZ8ovVNT126lckJVMWqhW2sI6" width="342" height="500" data-dominant-color="EEE5E4"><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">l2c algo</span><span class="informations">404×589 40 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></p>
<p>So I created a model like this:</p>
<pre><code class="lang-auto">class CNNCifar(nn.Module):
    def __init__(self):
        super(CNNCifar, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
        self.alpha = nn.Parameter(torch.randn(100, 100), requires_grad=True)
        self.w = torch.randn((100, 100), requires_grad=True)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return F.log_softmax(x, dim=1)

</code></pre>
<p>with a training loop that goes like this:</p>
<pre><code class="lang-auto">k = len(neighbour_sets)
    device = torch.device("cuda" if not torch.cuda.is_available() else "cpu")
    model = CNNCifar().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.1)
    l2c_optimizer = optim.Adam([model.alpha], lr=beta, weight_decay=0.01)

    test_accuracies = [[] for _ in range(k)]

    theta = [model.state_dict().copy() for _ in range(k)]
    theta_half = [model.state_dict().copy() for _ in range(k)]

    # w = torch.randn(k, k, requires_grad=True)
    delta_theta = [model.state_dict().copy() for _ in range(k)]

    with tqdm_output(tqdm(range(T))) as trange:
        for t in trange:
            for i in range(k):
                # Local SGD step
                log.info(f'Started training a Local SGD at node {i + 1}')

                model.load_state_dict(theta[i])
                for m in range(S):
                    for _, data in enumerate(train_loaders[i]):
                        inputs, labels = data
                        inputs, labels = inputs.to(device), labels.to(device)
                        optimizer.zero_grad()
                        outputs = model(inputs)
                        loss = criterion(outputs, labels)
                        loss.backward()
                        optimizer.step()

                log.info(f'Finished training a Local SGD at node {i + 1}')



                # Change capturing
                log.info(f'Computing change capturing at node {i + 1}')
                for name, param in model.named_parameters():
                    delta_theta[i][name] = theta[i][name] - theta_half[i][name]

                log.info(f'Computing mixing weights at node {i + 1}')
                # Mixing weights calculation
                model.w = model.w.clone()
                model.w[i] = compute_mixing_weights(model.alpha[i], neighbour_sets[i])

                # Aggregation
                log.info(f'Aggergating at node {i + 1}')
                theta_next = {}
                for name, param in model.named_parameters():
                    theta_next[name] = theta[i][name].clone()

                for j in neighbour_sets[i]:
                    for name, param in model.named_parameters():
                        theta_next[name] -= model.w[i][j].item() * delta_theta[i][name][j].clone()


                # Update L2C
                log.info(f'Updating L2C at node {i + 1}')
                model.load_state_dict(theta_next)
                model.train()
                # a training loop to find alpha that minimizes the validation loss
                for _, data in enumerate(val_loaders[i]):
                    inputs, labels = data
                    inputs, labels = inputs.to(device), labels.to(device)
                    
                    l2c_optimizer.zero_grad()
                    model.alpha.requires_grad_(True)
                    
                    log.info(f'Forward pass check')
                    outputs = model(inputs)
                    loss = criterion(outputs, labels)
                    model.alpha.retain_grad()
                    loss.backward()
                    print(f'gradient of alpha is {model.alpha.grad}')
                    import pdb; pdb.set_trace()
                    l2c_optimizer.step()

                    # Update α[i]
                    # import pdb; pdb.set_trace()
                    # alpha_grad = model.alpha.grad  # Access the computed gradients
                    # model.alpha.data[i] -= beta * alpha_grad[i]
                

                # Remove edges for sparse topology
                if t == T_0:
                    for _ in range(K_0):
                        j = min(neighbour_sets[i], key=lambda x: w[i][x])
                        neighbour_sets[i].delete(j)

                theta[i] = model.state_dict().copy()
                theta_half[i] = model.state_dict().copy()

                # Compute test accuracy for each local model
                test_accuracies = compute_test_acc(model, test_loaders[i], device, test_accuracies, i)
            
        log.info(f'Test accuracies atiteration at Comm_round {t} =  {sum(test_accuracies) / k}')
    
    return theta, test_accuracies
</code></pre>
<p>The problem is: in step 18 of the algorithm, the gradient of the loss is computed with respect to alpha. but when I access the <strong>model.alpha.grad</strong> it’s None, an no gradient is available.</p>
<p>What am I doing wrong here?</p>

KFrank: <p>Hi Ahmed!</p>
<aside class="quote no-group quote-modified" data-username="Ahmed_Khaled" data-post="1" data-topic="185089" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/e36b37/48.png" class="avatar"> Ahmed_Khaled:</div>
<blockquote>
<pre><code class="lang-auto">                model.w[i] = compute_mixing_weights(model.alpha[i], neighbour_sets[i])
...
                for j in neighbour_sets[i]:
                    for name, param in model.named_parameters():
                        theta_next[name] -= model.w[i][j].item() * delta_theta[i][name][j].clone()

                # Update L2C
                log.info(f'Updating L2C at node {i + 1}')
                model.load_state_dict(theta_next)
</code></pre>
<p>when I access the <strong>model.alpha.grad</strong> it’s None</p>
</blockquote>
</aside>
<p>First, although <code>model.w</code> does depend (differentiably) on <code>model.alpha</code><br>
(depending, of course, on what <code>compute_mixing_weights()</code> actually<br>
does), calling <code>.item()</code> “breaks the computation graph” because <code>.item()</code><br>
returns a python scalar that is no longer tracked by autograd.  As a result,<br>
<code>theta_next[name]</code> does not, in the sense of autograd, depend on <code>alpha</code>.</p>
<p>Second, the more serious problem is that <code>load_state_dict()</code> also<br>
breaks the computation graph.  The short explanation is that pytorch is<br>
very fastidious about not letting you modify <code>Parameter</code>s (in a way that<br>
is tracked by autograd).</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.0.0'
&gt;&gt;&gt; lin1 = torch.nn.Linear (2, 3)
&gt;&gt;&gt; lin2 = torch.nn.Linear (2, 3)
&gt;&gt;&gt; loss = lin2 (torch.randn (2)).sum()
&gt;&gt;&gt; lin1.load_state_dict (lin2.state_dict())
&lt;All keys matched successfully&gt;
&gt;&gt;&gt; loss.backward()
&gt;&gt;&gt; lin2.weight.grad
tensor([[0.4859, 0.7030],
        [0.4859, 0.7030],
        [0.4859, 0.7030]])
&gt;&gt;&gt; lin1.weight.grad
&gt;&gt;&gt;
</code></pre>
<p>If I understand what you are trying to do, the best approach might be<br>
something like:</p>
<pre data-code-wrap="python"><code class="lang-python">class CNNCifar(nn.Module):
    def __init__(self):
        ...
        self.fc1_weight = torch.randn (120, 16 * 5 * 5)
        self.fc1_bias = torch.randn (120)
    ...
    def forward(self, x):
        ...
        x = F.relu (torch.nn.functional.linear (x, self.fc1_weight, self.fc1_bias)
        ...
</code></pre>
<p>Here, you replace <code>Linear</code> (and the other layers in <code>model</code>) whose<br>
<code>Parameter</code>s are difficult to modify in a way that autograd can track<br>
with ordinary tensors (that can be tracked by autograd) that you then<br>
pass to the functional version of <code>linear()</code>.</p>
<p>Your <code>compute_mixing_weights()</code> would then mix together the <code>fc1_weight</code><br>
and <code>fc1_bias</code> tensors from the various <code>neighbour_sets</code> (and similarly for<br>
the other modifiable tensors that replace <code>self.conv1</code>, etc.).</p>
<p>You would then do something like:</p>
<pre data-code-wrap="python"><code class="lang-python">                for j in neighbour_sets[i]:
                    theta_next.fc1_weight -= model.w[i][j] * delta_theta[i, j].fc1_weight.clone()
                    theta_next.fc1_bias -= model.w[i][j] * delta_theta[i, j].fc1_bias.clone()
                    ...
</code></pre>
<p>where <code>theta_next</code> and <code>delta_theta</code> are no longer state dictionaries, but<br>
just something that holds copies of the <code>model</code> tensors (i.e., <code>fc1_weight</code><br>
and <code>fc1_bias</code>).</p>
<p>Note that <code>fc1_weight</code> is not a <code>Parameter</code> and doesn’t have <code>requires_grad</code><br>
<em>explicitly</em> set to <code>True</code>.  It will “inherit” <code>requires_grad = True</code> when it is<br>
updated using <code>w</code>, which does have <code>requires_grad = True</code>.</p>
<p>(If I understand your use case correctly, you want to separately optimize<br>
both the “ordinary” <code>parameters()</code> of <code>model</code> and <code>model.alpha</code>.  So you<br>
may, indeed, need to endow <code>fc1_weight</code> with <code>requires_grad = True</code>.<br>
Note that it’s fine to set up an <code>Optimizer</code> to optimizer non-<code>Parameter</code><br>
tensors (such as <code>fc1_weight</code>).</p>
<p>Best.</p>
<p>K. Frank</p>

Ahmed_Khaled: <p>Thank you K. Frank very much for this detailed explanation.</p>
<p>So I adjusted everything, as much as I understood from your comment, as follows:</p>
<pre><code class="lang-auto">k = len(neighbour_sets)
    device = torch.device("cuda" if not torch.cuda.is_available() else "cpu")
    model = CNNCifar().to(device)
    criterion = nn.CrossEntropyLoss()
    # optimizer for all params except alpha and w
    params = [model.conv1_weight, model.conv1_bias, model.conv2_weight, model.conv2_bias, model.fc1_weight, model.fc1_bias, model.fc2_weight, model.fc2_bias, model.fc3_weight, model.fc3_bias]
    optimizer = optim.SGD(params, lr=0.01, momentum=0.9, weight_decay=0.01)
    l2c_optimizer = optim.Adam([model.alpha], lr=beta, weight_decay=0.01)

    test_accuracies = [[] for _ in range(k)]

    theta = [model for _ in range(k)]
    theta_half = [model for _ in range(k)]
    delta_theta = [model for _ in range(k)]
    theta_next = [model for _ in range(k)]

    w = torch.zeros(k, k, dtype=theta[0].alpha.dtype, device=theta[0].alpha.device)

    with tqdm(range(T)) as trange:
        for t in trange:
            for i in range(k):
                # Local SGD step
                log.info(f'Started training a Local SGD at node {i + 1}')

                # theta_half[i] = theta[i]
                # for m in range(S):
                #     for _, data in enumerate(train_loaders[i]):
                #         inputs, labels = data
                #         inputs, labels = inputs.to(device), labels.to(device)
                #         optimizer.zero_grad()
                #         outputs = theta_half[i](inputs)
                #         loss = criterion(outputs, labels)
                #         loss.backward()
                #         optimizer.step()

                log.info(f'Finished training a Local SGD at node {i + 1}')

                # Change capturing
                log.info(f'Computing change capturing at node {i + 1}')
                import pdb; pdb.set_trace()
                delta_theta[i].fc1_weight = theta[i].fc1_weight - theta_half[i].fc1_weight.clone()
                delta_theta[i].fc1_bias = theta[i].fc1_bias - theta_half[i].fc1_bias.clone()
                delta_theta[i].fc2_weight = theta[i].fc2_weight - theta_half[i].fc2_weight.clone()
                delta_theta[i].fc2_bias = theta[i].fc2_bias - theta_half[i].fc2_bias.clone()
                delta_theta[i].fc3_weight = theta[i].fc3_weight - theta_half[i].fc3_weight.clone()
                delta_theta[i].fc3_bias = theta[i].fc3_bias - theta_half[i].fc3_bias.clone()
                delta_theta[i].conv1_weight = theta[i].conv1_weight - theta_half[i].conv1_weight.clone()
                delta_theta[i].conv1_bias = theta[i].conv1_bias - theta_half[i].conv1_bias.clone()
                delta_theta[i].conv2_weight = theta[i].conv2_weight - theta_half[i].conv2_weight.clone()
                delta_theta[i].conv2_bias = theta[i].conv2_bias - theta_half[i].conv2_bias.clone()
                               
                log.info(f'Computing mixing weights at node {i + 1}')
                # Mixing weights calculation like this :\highlightcyan{$w_{i,j}=\frac{\exp(\alpha_{i,j})}{\sum_{\ell\in i\cup N(i)}\exp(\alpha_{i,\ell})}$}
                for j in neighbour_sets[i]:
                    w[i][j] = torch.exp(theta[i].alpha[i][j])
                w[i] /= w[i].sum()
                
                theta[i].w.data = w

                # Aggregation
                log.info(f'Aggergating at node {i + 1}')

                for j in neighbour_sets[i]:
                    theta_next[i].fc1_weight = theta[i].fc1_weight -  theta[i].w[i][j] * delta_theta[i, j].fc1_weight.clone()
                    theta_next[i].fc1_bias = theta[i].fc1_bias - theta[i].w[i][j] * delta_theta[i, j].fc1_bias.clone()
                    theta_next[i].fc2_weight = theta[i].fc2_weight - theta[i].w[i][j] * delta_theta[i, j].fc2_weight.clone()
                    theta_next[i].fc2_bias = theta[i].fc2_bias - theta[i].w[i][j] * delta_theta[i, j].fc2_bias.clone()
                    theta_next[i].fc3_weight = theta[i].fc3_weight - theta[i].w[i][j] * delta_theta[i, j].fc3_weight.clone()
                    theta_next[i].fc3_bias = theta[i].fc3_bias - theta[i].w[i][j] * delta_theta[i, j].fc3_bias.clone()
                    theta_next[i].conv1_weight = theta[i].conv1_weight - theta[i].w[i][j] * delta_theta[i, j].conv1_weight.clone()
                    theta_next[i].conv1_bias = theta[i].conv1_bias - theta[i].w[i][j] * delta_theta[i, j].conv1_bias.clone()
                    theta_next[i].conv2_weight = theta[i].conv2_weight - theta[i].w[i][j] * delta_theta[i, j].conv2_weight.clone()
                    theta_next[i].conv2_bias = theta[i].conv2_bias - theta[i].w[i][j] * delta_theta[i, j].conv2_bias.clone()

                # Update L2C
                log.info(f'Updating L2C at node {i + 1}')
                # a training loop to find alpha that minimizes the validation loss
                for _, data in enumerate(val_loaders[i]):
                    inputs, labels = data
                    inputs, labels = inputs.to(device), labels.to(device)
                    
                    l2c_optimizer.zero_grad()
                    model.alpha.requires_grad_(True)
                    outputs = theta_next[i](inputs)
                    loss = criterion(outputs, labels)
                    loss.backward()
                    print(f'gradient of alpha is {theta_next[i].alpha.grad}')
                    l2c_optimizer.step()

                # Remove edges for sparse topology
                if t == T_0:
                    for _ in range(K_0):
                        j = min(neighbour_sets[i], key=lambda x: w[i][x])
                        neighbour_sets[i].delete(j)

                # theta[i] = model.state_dict().copy()
                # theta_half[i] = model.state_dict().copy()

                # Compute test accuracy for each local model
                test_accuracies = compute_test_acc(model, test_loaders[i], device, test_accuracies, i)
            
        log.info(f'Test accuracies atiteration at Comm_round {t} =  {sum(test_accuracies) / k}')
    
    return theta, test_accuracies

</code></pre>
<p>But still got None  from this line</p>
<pre><code class="lang-auto">print(f'gradient of alpha is {theta_next[i].alpha.grad}')

</code></pre>
<p>My model (based on your suggestions) is as follows:</p>
<pre><code class="lang-auto">
from torch import nn
import torch.nn.functional as F
import torch
from utils import compute_mixing_weights

class CNNCifar(nn.Module):
    def __init__(self):
        super(CNNCifar, self).__init__()
        self.pool = nn.MaxPool2d(2, 2)

        self.conv1_weight = torch.randn(6, 3, 5, 5, requires_grad=True)
        self.conv1_bias = torch.randn(6, requires_grad=True)

        self.conv2_weight = torch.randn(16, 6, 5, 5, requires_grad=True)
        self.conv2_bias = torch.randn(16, requires_grad=True)

        self.fc1_weight = torch.randn(120, 16 * 5 * 5, requires_grad=True)
        self.fc1_bias = torch.randn(120, requires_grad=True)

        self.fc2_weight = torch.randn(84, 120, requires_grad=True)
        self.fc2_bias = torch.randn(84, requires_grad=True)

        self.fc3_weight = torch.randn(10, 84, requires_grad=True)
        self.fc3_bias = torch.randn(10, requires_grad=True)

        self.alpha = nn.Parameter(torch.randn(100, 100), requires_grad=True)
        self.w = nn.Parameter(torch.randn(100, 100), requires_grad=True)


    def forward(self, x):
        x = self.pool(F.relu(F.conv2d(x, self.conv1_weight, self.conv1_bias)))
        x = self.pool(F.relu(F.conv2d(x, self.conv2_weight, self.conv2_bias)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(torch.nn.functional.linear(x, self.fc1_weight, self.fc1_bias))
        x = F.relu(torch.nn.functional.linear(x, self.fc2_weight, self.fc2_bias))
        x = torch.nn.functional.linear(x, self.fc3_weight, self.fc3_bias)
        return F.log_softmax(x, dim=1)


</code></pre>

KFrank: <p>Hi Ahmed!</p>
<aside class="quote no-group" data-username="Ahmed_Khaled" data-post="3" data-topic="185089" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/e36b37/48.png" class="avatar"> Ahmed_Khaled:</div>
<blockquote>
<p>So I adjusted everything, as much as I understood from your comment, as follows:</p>
</blockquote>
</aside>
<p>You have quite a few things going on here.  I would suggest that you work<br>
with a simplified, “toy” model where you figure out how to optimize both<br>
<code>alpha</code> and some simple model parameter in the presence of “weight mixing.”</p>
<p>Get that working, and then try to extend it to your more-complicated use case.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">    theta = [model for _ in range(k)]
    theta_half = [model for _ in range(k)]
    delta_theta = [model for _ in range(k)]
    theta_next = [model for _ in range(k)]
</code></pre>
</blockquote>
</aside>
<p>One immediate problem is that at this point <code>theta</code>, <code>theta_half</code>, etc.,<br>
are all lists that contain references to the <em>same</em> instance of <code>model</code>.  So,<br>
if your modify, say, <code>theta[2]</code>, you have also modified <code>theta[0]</code> and<br>
<code>theta_half[3]</code>, because these are all actually references to the same<br>
object, namely <code>model</code>.  This surely isn’t what your want and could likely<br>
be the cause of your problem.</p>
<p>Best.</p>
<p>K. Frank</p>