Unable to manually assign Adam's parameter gradients

hwaseem04: <p>Hi,</p>
<p>I am trying to implement a minimalistic federated SGD algorithm for MNIST classification.<br>
For which I am trying to accumulate the gradients and then make optimizer.step.</p>
<p>Here is my gradient accumulation code</p>
<pre><code class="lang-auto">def server_aggregate_gradients(n_clients, device): # For FedSGD
    grads_info = []
    for c in range(n_clients):
        grads_info.append(torch.load(f'./model_cache/client_{c}.pkl', map_location=device))

    total_grads = {}
    n_total_samples = 0
    # acc_gradient = summation[ (number_samples_per_client/total_samples) *  gradient_per_client]
    for info in grads_info:
        n_samples = info['n_samples']
        for name, grad_value in info['named_grads'].items():
            if name not in total_grads:
                total_grads[name] = grad_value * n_samples
            else:
                total_grads[name] += (grad_value * n_samples) 
        n_total_samples += n_samples

    gradients = {}
    for name, grad_value in total_grads.items():
        gradients[name] = grad_value/ n_total_samples
    return gradients
</code></pre>
<p>After accumulating gradients(weighted average of gradients) from all clients, I am setting those gradients to the main model as below and the perform gradient descent</p>
<pre><code class="lang-auto">def server_step(model, optimizer, gradients=None, weights=None):
    model.train()
    optimizer.zero_grad()
    for name, parameter in model.named_parameters():
        print(gradients[name].shape) # torch.Size([32, 1, 5, 5])
        print(parameter.shape) # torch.Size([32, 1, 5, 5])
        print(gradients[name].dtype) # torch.float32
        print(parameter.dtype) # torch.float32

        parameter.grad = gradients[name] 
    optimizer.step()
    torch.save(model.state_dict(), './model_cache/global_model_state.pkl')
</code></pre>
<p>But I get error at line <code>parameter.grad = gradients[name] </code></p>
<p>Here is the error I get;</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/Users/muhammadwaseem/Documents/GitHub/CV-projects/federated-learning/server.py", line 86, in &lt;module&gt;
    fedSGD(n_clients, n_server_epochs, batch_size, device=device)
  File "/Users/muhammadwaseem/Documents/GitHub/CV-projects/federated-learning/server.py", line 73, in fedSGD
    server_step(model, optimizer, gradients=gradients)
  File "/Users/muhammadwaseem/Documents/GitHub/CV-projects/federated-learning/server.py", line 41, in server_step
    parameter.grad = gradients[name]
RuntimeError: assigned grad has data of a different type
</code></pre>
<p><strong>Why</strong>: Shape and dtype of both   <code>parameter.grad</code> and <code>gradients[name]</code> is same, then why do I get this error?</p>
<p>Edit: Also <code>parameter.grad</code> seems to be having the value <code>None</code>, does it have anything to do with the error</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> , any thoughts on this issue?</p>

hwaseem04: <p>I figured out that <code>parameter.grad</code> and <code>gradients[name]</code> were on different device. Once I put them in same device (either CPU or cuda) then it worked.</p>