client_with_pgfed.py 5.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  1. import copy
  2. import gc
  3. import logging
  4. import time
  5. from collections import Counter
  6. import numpy as np
  7. import torch
  8. import torch._utils
  9. import torch.nn as nn
  10. import torch.nn.functional as F
  11. import model
  12. import utils
  13. from communication import ONLINE, TARGET, BOTH, LOCAL, GLOBAL, DAPU, NONE, EMA, DYNAMIC_DAPU, DYNAMIC_EMA_ONLINE, SELECTIVE_EMA
  14. from easyfl.client.base import BaseClient
  15. from easyfl.distributed.distributed import CPU
  16. from client import FedSSLClient
  17. logger = logging.getLogger(__name__)
  18. L2 = "l2"
  19. def model_dot_product(w1, w2, requires_grad=True):
  20. """ Return the sum of squared difference between two models. """
  21. dot_product = 0.0
  22. for p1, p2 in zip(w1.parameters(), w2.parameters()):
  23. if requires_grad:
  24. dot_product += torch.sum(p1 * p2)
  25. else:
  26. dot_product += torch.sum(p1.data * p2.data)
  27. return dot_product
  28. class FedSSLWithPgFedClient(FedSSLClient):
  29. def __init__(self, cid, conf, train_data, test_data, device, sleep_time=0):
  30. super(FedSSLWithPgFedClient, self).__init__(cid, conf, train_data, test_data, device, sleep_time)
  31. self._local_model = None
  32. self.DAPU_predictor = LOCAL
  33. self.encoder_distance = 1
  34. self.encoder_distances = []
  35. self.previous_trained_round = -1
  36. self.weight_scaler = None
  37. self.latest_grad = copy.deepcopy(self.model)
  38. self.lambdaa = 1.0 # PGFed learning rate for a_i, Regularization weight for pFedMe
  39. self.prev_loss_minuses = {}
  40. self.prev_mean_grad = None
  41. self.prev_convex_comb_grad = None
  42. self.a_i = None
  43. def train(self, conf, device=CPU):
  44. start_time = time.time()
  45. loss_fn, optimizer = self.pretrain_setup(conf, device)
  46. if conf.model in [model.MoCo, model.MoCoV2]:
  47. self.model.reset_key_encoder()
  48. self.train_loss = []
  49. self.model.to(device)
  50. old_model = copy.deepcopy(nn.Sequential(*list(self.model.children())[:-1])).cpu()
  51. for i in range(conf.local_epoch):
  52. batch_loss = []
  53. for (batched_x1, batched_x2), _ in self.train_loader:
  54. x1, x2 = batched_x1.to(device), batched_x2.to(device)
  55. optimizer.zero_grad()
  56. if conf.model in [model.MoCo, model.MoCoV2]:
  57. loss = self.model(x1, x2, device)
  58. elif conf.model == model.SimCLR:
  59. images = torch.cat((x1, x2), dim=0)
  60. features = self.model(images)
  61. logits, labels = self.info_nce_loss(features)
  62. loss = loss_fn(logits, labels)
  63. else:
  64. loss = self.model(x1, x2)
  65. loss.backward()
  66. if self.prev_convex_comb_grad is not None:
  67. for p_m, p_prev_conv in zip(self.model.parameters(), self.prev_convex_comb_grad.parameters()):
  68. p_m.grad.data += p_prev_conv.data
  69. dot_prod = model_dot_product(self.model, self.prev_mean_grad, requires_grad=False)
  70. self.update_a_i(dot_prod)
  71. optimizer.step()
  72. batch_loss.append(loss.item())
  73. if conf.model in [model.BYOL, model.BYOLNoSG, model.BYOLNoPredictor] and conf.momentum_update:
  74. self.model.update_moving_average()
  75. current_epoch_loss = sum(batch_loss) / len(batch_loss)
  76. self.train_loss.append(float(current_epoch_loss))
  77. # get loss_minus and latest_grad
  78. self.loss_minus = 0.0
  79. test_num = 0
  80. self.optimizer.zero_grad()
  81. for i, (x, y) in enumerate(self.train_loader):
  82. if type(x) == type([]):
  83. x[0] = x[0].to(self.device)
  84. else:
  85. x = x.to(self.device)
  86. y = y.to(self.device)
  87. test_num += y.shape[0]
  88. output = self.model(x)
  89. loss = self.criterion(output, y)
  90. self.loss_minus += (loss * y.shape[0]).item()
  91. loss.backward()
  92. self.loss_minus /= test_num
  93. for p_l, p in zip(self.latest_grad.parameters(), self.model.parameters()):
  94. p_l.data = p.grad.data.clone() / len(self.train_loader)
  95. self.loss_minus -= model_dot_product(self.latest_grad, self.model, requires_grad=False)
  96. self.train_time = time.time() - start_time
  97. # store trained model locally
  98. self._local_model = copy.deepcopy(self.model).cpu()
  99. self.previous_trained_round = conf.round_id
  100. if conf.update_predictor in [DAPU, DYNAMIC_DAPU, SELECTIVE_EMA] or conf.update_encoder in [DYNAMIC_EMA_ONLINE, SELECTIVE_EMA]:
  101. new_model = copy.deepcopy(nn.Sequential(*list(self.model.children())[:-1])).cpu()
  102. self.encoder_distance = self._calculate_divergence(old_model, new_model)
  103. self.encoder_distances.append(self.encoder_distance.item())
  104. self.DAPU_predictor = self._DAPU_predictor_usage(self.encoder_distance)
  105. if self.conf.auto_scaler == 'y' and self.conf.random_selection:
  106. self._calculate_weight_scaler()
  107. if (conf.round_id + 1) % 100 == 0:
  108. logger.info(f"Client {self.cid}, encoder distances: {self.encoder_distances}")
  109. def update_a_i(self, dot_prod):
  110. for clt_j, mu_loss_minus in self.prev_loss_minuses.items():
  111. self.a_i[clt_j] -= self.lambdaa * (mu_loss_minus + dot_prod)
  112. self.a_i[clt_j] = max(self.a_i[clt_j], 0.0)
  113. def set_prev_mean_grad(self, mean_grad):
  114. if self.prev_mean_grad is None:
  115. self.prev_mean_grad = copy.deepcopy(mean_grad)
  116. else:
  117. self.set_model(self.prev_mean_grad, mean_grad)
  118. def set_prev_convex_comb_grad(self, convex_comb_grad, momentum=0.0):
  119. if self.prev_convex_comb_grad is None:
  120. self.prev_convex_comb_grad = copy.deepcopy(convex_comb_grad)
  121. else:
  122. self.set_model(self.prev_convex_comb_grad, convex_comb_grad, momentum=momentum)