Issues when trying to implement FedAVG using torch.multiprocessing.Pool

username: Zixun_Xiong
cooked: <p>After finishing a single process version of FedAVG(a federated learning), I tried to apply torch.multiprocessing.Pool() and use apply_async to share and update the model. However, I got a way lower accuracy （from 0.8 to 0.1）in each client’s local model.</p>
<p>my original code is like this:<br>
def update_server(self, T):<br>
‘’’<br>
FedAVG<br>
‘’’<br>
client_acc = []<br>
# 2: for t=0, …, T-1 do<br>
for round in range(T):<br>
print(f"Round {round+1} started…")<br>
client_models = []<br>
client_losses = []<br>
client_accs = []<br>
delta_is_t = []</p>
<pre><code>        # --multi-pro pool
        pool = mp.Pool(self._num_clients)
        results = [] 
        # --
        
        # 3: TODO 2 : select a subset of client
        # 4: x{i, 0}^t = x_t
        x_t = self._global_model.state_dict()
        # 5: for each client i \in S in parallel do
        for i in range(self._num_clients):
            #TODO 1 : do this in parallel
            
            #  x_{i, 0}^t = x_t  : x_t is model from last epoch 
            x_t_temp = copy.deepcopy(x_t)
            
            # cur_user._model.load_state_dict(x_t)
            # 6-8： x{i, K}^t = CLIENTOPT
            # --
            # result = pool.apply_async(self._clients[i].client_update, args=(round+1,i, x_t_temp))
            # x_i_K_t, client_loss, client_acc = \
            # pool.apply_async(self._clients[i].client_update, args=(round+1,i, x_t_temp)).get()
            # --
            x_i_K_t, client_loss, client_acc =  \
               self._clients[i].client_update(epoch=round+1, id=i, global_model=x_t_temp)
            # 9: \delta_{client i, epoch t} = x{i, K}^t - x_t
            # delta_i_t =  x_i_K_t - x_t
            # delta_i_t = {}
            # for key in x_i_K_t.keys():
            #     delta_i_t[key] = x_i_K_t[key] - x_t[key]
            # --
            # results.append(result)
            # --
            client_models.append(x_i_K_t)
            client_losses.append(client_loss)
            client_accs.append(client_acc)
            # delta_is_t.append(delta_i_t)
            # -- print(f"Client {i+1} loss: {client_loss:.4f}, accuracy {client_acc: .4f} ")
        # --
        
        # for result in results:
        #     client_state_dict, client_loss, client_acc = result.get()
        #     client_models.append(client_state_dict)
        #     client_losses.append(client_loss)
        #     client_accs.append(client_acc)
        # --

        # x_{t+1} = \frac{1}{S} \sum_{i \in S} \delta_i^t
        global_state_dict = self.server_update(client_models)
        self._global_model.load_state_dict(global_state_dict)
        print(f"Round {round+1} finished, global loss:  \
            {sum(client_losses)/len(client_losses):.4f},  \
                global accuracy: {sum(client_accs)/len(client_accs): .4f}")
        #--
        pool.close()
        pool.join()
        #--
        
    return sum(client_accs)/len(client_accs) 
</code></pre>
<p>my updated version only changed this:<br>
ool.apply_async(self._clients[i].client_update, args=(round+1,i, x_t_temp)).get()<br>
(I used .get() method for simplicity)</p>
<p>My client update functions stays same between two version:<br>
class ClientBase():<br>
‘’‘base class for FL learning’‘’<br>
def <strong>init</strong>(self, dataloader, model, optimizer, device, E, B):<br>
self._dataloader = dataloader<br>
self._model = model<br>
self._optimizer  = optimizer<br>
self._device = device<br>
self._E = E<br>
self._B = B</p>
<pre><code>def client_update(self, epoch, id, global_model):
    '''ClientUpdate in FedAVG;'''
    # print(f'client {id+1} is started to run.')
    self._model.train()
    self._model.load_state_dict(global_model)
    
    self._model.to(device=self._device)
    criterion = nn.CrossEntropyLoss()
    running_loss = 0
    num_euqal = 0
    acc = None
    
    for _ in range(self._E):
        for inputs, labels in self._dataloader:
            inputs, labels = inputs.to(self._device), labels.to(self._device)
            self._optimizer.zero_grad()
            outputs = self._model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            self._optimizer.step()
            running_loss += loss.item() * inputs.size(0)
    self._model.eval()
    acc_num = 0
    total_num = 0
    for inputs, labels in self._dataloader:
        inputs, labels = inputs.to(self._device), labels.to(self._device)
        test_output = self._model(inputs)
        pred_y = torch.max(test_output, 1)[1].data.squeeze()
        num_equal = (pred_y == labels).sum().item()
        acc_num += num_equal
        total_num += labels.size()[0] 
    print(f"Client {id+1} Ended-loss:  \
        {running_loss / len(self._dataloader.dataset)*self._E:.4f}, \
            accuracy {acc_num/total_num: .4f} ")
    return copy.deepcopy(self._model.cpu()).state_dict(), \
        running_loss / len(self._dataloader.dataset)*self._E, acc_num/total_num
</code></pre>
<p>My whole project can be found here:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/kitaharasetusna/FedAvg">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/kitaharasetusna/FedAvg" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/345;"><img src="https://opengraph.githubassets.com/518030d274b3c7ec799e910d14ef3d6989e1e9272eda17ba102beededd170642/kitaharasetusna/FedAvg" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/kitaharasetusna/FedAvg" target="_blank" rel="noopener nofollow ugc">GitHub - kitaharasetusna/FedAvg: Simple implementation of fedavg for learning...</a></h3>

  <p>Simple implementation of fedavg for learning and researching purpose - GitHub - kitaharasetusna/FedAvg: Simple implementation of fedavg for learning and researching purpose</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>I’m new to pytorch and this seems to be a silly question, but trying to solve it and got stuck for a month make me write this post. I’ll be grateful if anyone can help me solve it. Thanks for reading.</p>