import torch
import [Link] as nn
import [Link] as optim
from torchvision import datasets, transforms
import numpy as np
import [Link] as plt
# Configuration
NUM_DEVICES = 50
EPOCHS = 2
ROUNDS = 100
BATCH_SIZE = 64
LEARNING_RATE = 0.01
DEVICE_PER_ROUND = 10 # For round robin
# Energy, Bandwidth, SNR initialization
device_energy = [Link](30, 100, NUM_DEVICES)
device_bandwidth = [Link](1.0, 5.0, NUM_DEVICES)
device_snr = [Link](10, 40, NUM_DEVICES)
# Dataset loading (split among devices)
transform = [Link]([[Link]()])
full_dataset = [Link]('./data', train=True, download=True,
transform=transform)
datasets_split = [Link].random_split(full_dataset,
[int(len(full_dataset)/NUM_DEVICES)]*NUM_DEVICES)
# CNN Model with separable representation and classifier (for FedRep)
class CNN_FedRep([Link]):
def __init__(self):
super(CNN_FedRep, self).__init__()
[Link] = [Link](
nn.Conv2d(1, 10, 5),
[Link](),
nn.MaxPool2d(2),
[Link]()
)
[Link] = [Link](1440, 10)
def forward(self, x):
rep = [Link](x)
return [Link](rep)
# Initialize models and optimizers
models = [CNN_FedRep() for _ in range(NUM_DEVICES)]
optimizers_rep = [[Link]([Link](), lr=LEARNING_RATE)
for model in models]
optimizers_cls = [[Link]([Link](), lr=LEARNING_RATE) for
model in models]
loss_fn = [Link]()
# Training function for FedRep
def train_rep(model, data_loader, optimizer_rep, optimizer_cls=None,
train_classifier=False):
[Link]()
for epoch in range(EPOCHS):
for data, target in data_loader:
if train_classifier:
optimizer_cls.zero_grad()
optimizer_rep.zero_grad()
output = model(data)
loss = loss_fn(output, target)
[Link]()
optimizer_rep.step()
if train_classifier:
optimizer_cls.step()
return [Link].state_dict()
# Model averaging (FedRep: only classifier aggregated)
def average_classifiers(classifier_states):
new_state = {}
for key in classifier_states[0].keys():
new_state[key] = sum([state[key] for state in classifier_states]) /
len(classifier_states)
return new_state
# Evaluation
def evaluate_model(model, test_loader):
[Link]()
correct = 0
total = 0
with torch.no_grad():
for data, target in test_loader:
output = model(data)
_, predicted = [Link]([Link], 1)
total += [Link](0)
correct += (predicted == target).sum().item()
return correct / total
# Test loader
test_loader = [Link]([Link]('./data', train=False,
download=True, transform=transform), batch_size=1000)
# === Federated Training ===
accuracy_round_robin = []
accuracy_fedrep = []
for round_idx in range(ROUNDS):
print(f"\n--- Round {round_idx+1} ---")
# Round Robin Selection
start = (round_idx * DEVICE_PER_ROUND) % NUM_DEVICES
selected_devices = list(range(start, start + DEVICE_PER_ROUND))
print(f"Selected Devices (Round Robin): {selected_devices}")
# === FedRep Training ===
classifier_states = []
for i in selected_devices:
data_loader = [Link](datasets_split[i],
batch_size=BATCH_SIZE, shuffle=True)
state_dict = train_rep(models[i], data_loader, optimizers_rep[i],
optimizers_cls[i], train_classifier=True)
classifier_states.append(state_dict)
if classifier_states:
avg_classifier = average_classifiers(classifier_states)
for model in models:
[Link].load_state_dict(avg_classifier)
acc_fedrep = evaluate_model(models[0], test_loader)
accuracy_fedrep.append(acc_fedrep)
print(f"FedRep Accuracy: {acc_fedrep*100:.2f}%")
# === Device Info ===
print("\n--- Device Initialization Info ---")
for i in range(NUM_DEVICES):
print(f"Device {i}: Energy = {device_energy[i]:.2f}, Bandwidth =
{device_bandwidth[i]:.2f}, SNR = {device_snr[i]:.2f}")
# === Plotting ===
[Link](range(1, ROUNDS+1), [round(a*100, 2) for a in accuracy_fedrep],
label='FedRep', marker='o')
[Link]("Rounds")
[Link]("Accuracy (%)")
[Link]("FedRep Accuracy (Round Robin Scheduling)")
[Link](True)
[Link]()
plt.tight_layout()
[Link]()