implement Model-Contrastive Federated Learning (Koukyosyumei#164)

* implement Model-Contrastive Federated Learning

* update supported algorithms

* rm unused var

README.md

@@ -174,12 +174,12 @@ You can also find more examples in our tutorials and documentation.
 
 | | | |
 | ------------- | ---------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| Collaborative | Horizontal FL | [FedAVG](https://arxiv.org/abs/1602.05629), [FedProx](https://arxiv.org/abs/1812.06127), [FedKD](https://arxiv.org/abs/2108.13323), [FedGEMS](https://arxiv.org/abs/2110.11027), [FedMD](https://arxiv.org/abs/1910.03581), [DSFL](https://arxiv.org/abs/2008.06180) |
+| Collaborative | Horizontal FL | [FedAVG](https://arxiv.org/abs/1602.05629), [FedProx](https://arxiv.org/abs/1812.06127), [FedKD](https://arxiv.org/abs/2108.13323), [FedGEMS](https://arxiv.org/abs/2110.11027), [FedMD](https://arxiv.org/abs/1910.03581), [DSFL](https://arxiv.org/abs/2008.06180), [MOON](https://arxiv.org/abs/2103.16257) |
 | Collaborative | Vertical FL | [SplitNN](https://arxiv.org/abs/1812.00564), [SecureBoost](https://arxiv.org/abs/1901.08755) |
 | Attack | Model Inversion | [MI-FACE](https://dl.acm.org/doi/pdf/10.1145/2810103.2813677), [DLG](https://papers.nips.cc/paper/2019/hash/60a6c4002cc7b29142def8871531281a-Abstract.html), [iDLG](https://arxiv.org/abs/2001.02610), [GS](https://proceedings.neurips.cc/paper/2020/hash/c4ede56bbd98819ae6112b20ac6bf145-Abstract.html), [CPL](https://arxiv.org/abs/2004.10397), [GradInversion](https://openaccess.thecvf.com/content/CVPR2021/papers/Yin_See_Through_Gradients_Image_Batch_Recovery_via_GradInversion_CVPR_2021_paper.pdf), [GAN Attack](https://arxiv.org/abs/1702.07464) |
 | Attack | Label Leakage | [Norm Attack](https://arxiv.org/abs/2102.08504) |
 | Attack | Poisoning | [History Attack](https://arxiv.org/abs/2203.08669), [Label Flip](https://arxiv.org/abs/2203.08669), [MAPF](https://arxiv.org/abs/2203.08669), [SVM Poisoning](https://arxiv.org/abs/1206.6389) |
-| Attack | Backdoor | [DBA](https://openreview.net/forum?id=rkgyS0VFvr) |
+| Attack | Backdoor | [DBA](https://openreview.net/forum?id=rkgyS0VFvr), [Model Replacement](https://proceedings.mlr.press/v108/bagdasaryan20a.html) |
 | Attack | Free-Rider | [Delta-Weight](https://arxiv.org/pdf/1911.12560.pdf) |
 | Attack | Evasion | [Gradient-Descent Attack](https://arxiv.org/abs/1708.06131), [FGSM](https://arxiv.org/abs/1412.6572), [DIVA](https://arxiv.org/abs/2204.10933) |
 | Attack | Membership Inference | [Shaddow Attack](https://arxiv.org/abs/1610.05820) |

src/aijack/collaborative/__init__.py

@@ -21,5 +21,6 @@
  MPIFedMDClientManager,
  MPIFedMDServerManager,
 )
+from .moon import MOONClient # noqa :F401
 from .optimizer import AdamFLOptimizer, SGDFLOptimizer # noqa: F401
 from .splitnn import SplitNNAPI, SplitNNClient # noqa: F401

src/aijack/collaborative/moon/__init__.py

@@ -0,0 +1,3 @@
+from .client import MOONClient # noqa: F401
+
+__all__ = ["MOONClient"]

src/aijack/collaborative/moon/client.py

@@ -0,0 +1,91 @@
+import copy
+
+import torch
+import torch.nn.functional as F
+
+from ..fedavg import FedAVGClient
+
+
+class MOONClient(FedAVGClient):
+ """Client of MOON for single process simulation
+ (Li, Qinbin, Bingsheng He, and Dawn Song. "Model-contrastive
+ federated learning." Proceedings of the IEEE/CVF conference
+ on computer vision and pattern recognition. 2021.)
+
+ Args:
+ model (torch.nn.Module): local model
+ mu (float): weight of model-contrastive loss
+ tau (float): tempreature within model-contrastive loss
+ """
+
+ def __init__(
+ self,
+ model,
+ mu=0.1,
+ tau=1.0,
+ **kwargs,
+ ):
+ super(MOONClient, self).__init__(model, **kwargs)
+ self.mu = mu
+ self.tau = tau
+ self.global_model = copy.deepcopy(model)
+ self.prev_model = copy.deepcopy(model)
+
+ def local_train(
+ self,
+ local_epoch,
+ criterion,
+ trainloader,
+ optimizer,
+ communication_id=0,
+ ):
+ if communication_id != 0:
+ for param, glob_param in zip(
+ self.global_model.parameters(), self.model.parameters()
+ ):
+ if param is not None:
+ param = glob_param
+ for param, prev_param in zip(
+ self.prev_model.parameters(), self.prev_parameters
+ ):
+ if param is not None:
+ param = prev_param
+
+ for i in range(local_epoch):
+ running_loss = 0.0
+ running_data_num = 0
+ for _, data in enumerate(trainloader, 0):
+ inputs, labels = data
+ inputs = inputs.to(self.device)
+ labels = labels.to(self.device)
+
+ optimizer.zero_grad()
+ self.zero_grad()
+
+ outputs = self(inputs)
+ loss = criterion(outputs, labels)
+
+ if communication_id != 0:
+ glob_outputs = self.global_model(inputs)
+ prev_outputs = self.prev_model(inputs)
+
+ exp_sim_cg = torch.exp(
+ F.cosine_similarity(outputs, glob_outputs) / self.tau
+ )
+ exp_sim_cp = torch.exp(
+ F.cosine_similarity(outputs, prev_outputs) / self.tau
+ )
+ loss_con = -1 * torch.log(exp_sim_cg / (exp_sim_cg + exp_sim_cp))
+ loss = loss + self.mu * loss_con
+
+ loss.backward()
+
+ optimizer.step()
+
+ running_loss += loss.item()
+ running_data_num += inputs.shape[0]
+
+ print(
+ f"communication {communication_id}, epoch {i}: client-{self.user_id+1}",
+ running_loss / running_data_num,
+ )

test/collaborative/moon/test_moon.py

@@ -0,0 +1,71 @@
+def test_fedkd():
+ import torch
+ import torch.nn as nn
+ import torch.optim as optim
+ from torch.utils.data import DataLoader, TensorDataset
+
+ from aijack.collaborative import FedAVGAPI, FedAVGServer, MOONClient
+
+ torch.manual_seed(0)
+
+ lr = 0.01
+ client_num = 2
+
+ class Net(nn.Module):
+ def __init__(self):
+ super(Net, self).__init__()
+ self.conv = nn.Sequential(
+ nn.Conv2d(1, 32, 5),
+ nn.Sigmoid(),
+ nn.MaxPool2d(3, 3, 1),
+ nn.Conv2d(32, 64, 5),
+ nn.Sigmoid(),
+ nn.MaxPool2d(3, 3, 1),
+ )
+
+ self.lin = nn.Sequential(nn.Linear(256, 10))
+
+ def forward(self, x):
+ x = self.conv(x)
+ self.hidden_states = x.reshape((-1, 256))
+ x = self.lin(self.hidden_states)
+ return x
+
+ def get_hidden_states(self):
+ return [self.hidden_states]
+
+ x = torch.load("test/demodata/demo_mnist_x.pt")
+ x.requires_grad = True
+ y = torch.load("test/demodata/demo_mnist_y.pt")
+
+ local_dataloaders = [DataLoader(TensorDataset(x, y)) for _ in range(client_num)]
+
+ clients = [
+ MOONClient(
+ Net(),
+ user_id=i,
+ lr=lr,
+ )
+ for i in range(client_num)
+ ]
+ local_optimizers = [optim.SGD(client.parameters(), lr=lr) for client in clients]
+
+ global_model = Net()
+ server = FedAVGServer(clients, global_model, lr=lr)
+
+ criterion = nn.CrossEntropyLoss()
+
+ api = FedAVGAPI(
+ server,
+ clients,
+ criterion,
+ local_optimizers,
+ local_dataloaders,
+ num_communication=2,
+ local_epoch=1,
+ use_gradients=True,
+ custom_action=lambda x: x,
+ device="cpu",
+ )
+
+ api.run()