class FedAvg5Trainer(BaseTrainer):
"""
Original Scheme
"""
def __init__(self, options, dataset):
model = choose_model(options)
self.move_model_to_gpu(model, options)
self.optimizer = GD(model.parameters(),
lr=options['lr'],
weight_decay=options['wd'])
self.num_epoch = options['num_epoch']
worker = LrdWorker(model, self.optimizer, options)
super(FedAvg5Trainer, self).__init__(options, dataset, worker=worker)
def train(self):
print('>>> Select {} clients per round \n'.format(
self.clients_per_round))
# Fetch latest flat model parameter
self.latest_model_params = self.worker.get_flat_model_params().detach()
for round_i in range(self.num_round):
# Test latest model on train data
self.test_latest_model_on_traindata(round_i)
self.test_latest_model_on_evaldata(round_i)
# Choose K clients prop to data size
selected_clients = self.select_clients(seed=round_i)
# Solve minimization locally
solns, stats = self.local_train(round_i, selected_clients)
# Track communication cost
self.metrics.extend_commu_stats(round_i, stats)
# Update latest model
self.latest_model_params = self.aggregate(solns)
self.optimizer.inverse_prop_decay_learning_rate(round_i)
# Test final model on train data
self.test_latest_model_on_traindata(self.num_round)
self.test_latest_model_on_evaldata(self.num_round)
# Save tracked information
self.metrics.write()
def aggregate(self, solns):
averaged_solution = torch.zeros_like(self.latest_model_params)
accum_sample_num = 0
for num_sample, local_solution in solns:
accum_sample_num += num_sample
averaged_solution += num_sample * local_solution
averaged_solution /= self.all_train_data_num
averaged_solution += (1 - accum_sample_num / self.all_train_data_num
) * self.latest_model_params
return averaged_solution.detach()
class FedAvgTrainer(BaseTrainer):
def __init__(self, options, dataset):
model = choose_model(options)
self.move_model_to_gpu(model, options)
self.optimizer = GD(model.parameters(),
lr=options['lr'],
weight_decay=options['wd'])
super(FedAvgTrainer, self).__init__(options, dataset, model,
self.optimizer)
def train(self):
print('>>> Select {} clients per round \n'.format(
self.clients_per_round))
# Fetch latest flat model parameter
self.latest_model = self.worker.get_flat_model_params().detach()
for round_i in range(self.num_round):
# Test latest model on train data
self.test_latest_model_on_traindata(round_i)
self.test_latest_model_on_evaldata(round_i)
# Choose K clients prop to data size
selected_clients = self.select_clients(seed=round_i)
# Solve minimization locally
solns, stats = self.local_train(round_i, selected_clients)
# Track communication cost
self.metrics.extend_commu_stats(round_i, stats)
# Update latest model
self.latest_model = self.aggregate(solns)
self.optimizer.inverse_prop_decay_learning_rate(round_i)
# Test final model on train data
self.test_latest_model_on_traindata(self.num_round)
self.test_latest_model_on_evaldata(self.num_round)
# Save tracked information
self.metrics.write()
class FedAvg9Trainer(BaseTrainer):
"""
Only Transformed II
"""
def __init__(self, options, dataset):
model = choose_model(options)
self.move_model_to_gpu(model, options)
self.optimizer = GD(model.parameters(),
lr=options['lr'],
weight_decay=options['wd'])
self.num_epoch = options['num_epoch']
worker = LrAdjustWorker(model, self.optimizer, options)
super(FedAvg9Trainer, self).__init__(options, dataset, worker=worker)
def train(self):
print('>>> Select {} clients per round \n'.format(
self.clients_per_round))
# Fetch latest flat model parameter
self.latest_model = self.worker.get_flat_model_params().detach()
for round_i in range(self.num_round):
# Test latest model on train data
self.test_latest_model_on_traindata(round_i)
self.test_latest_model_on_evaldata(round_i)
# Choose K clients prop to data size
selected_clients = self.select_clients(seed=round_i)
# Solve minimization locally
solns, stats = self.local_train(round_i, selected_clients)
# Track communication cost
self.metrics.extend_commu_stats(round_i, stats)
# Update latest model
self.latest_model = self.aggregate(solns)
self.optimizer.inverse_prop_decay_learning_rate(round_i)
# Test final model on train data
self.test_latest_model_on_traindata(self.num_round)
self.test_latest_model_on_evaldata(self.num_round)
# Save tracked information
self.metrics.write()
def aggregate(self, solns, **kwargs):
averaged_solution = torch.zeros_like(self.latest_model)
# averaged_solution = np.zeros(self.latest_model.shape)
assert self.simple_average
for num_sample, local_solution in solns:
averaged_solution += local_solution
averaged_solution /= self.clients_per_round
# averaged_solution = from_numpy(averaged_solution, self.gpu)
return averaged_solution.detach()
def local_train(self, round_i, selected_clients, **kwargs):
solns = [] # Buffer for receiving client solutions
stats = [] # Buffer for receiving client communication costs
for i, c in enumerate(selected_clients, start=1):
# Communicate the latest model
c.set_flat_model_params(self.latest_model)
# Solve minimization locally
m = len(c.train_data) / self.all_train_data_num * 100
soln, stat = c.local_train(multiplier=m)
if self.print_result:
print("Round: {:>2d} | CID: {: >3d} ({:>2d}/{:>2d})| "
"Param: norm {:>.4f} ({:>.4f}->{:>.4f})| "
"Loss {:>.4f} | Acc {:>5.2f}% | Time: {:>.2f}s".format(
round_i, c.cid, i, self.clients_per_round,
stat['norm'], stat['min'], stat['max'], stat['loss'],
stat['acc'] * 100, stat['time']))
# Add solutions and stats
solns.append(soln)
stats.append(stat)
return solns, stats
class FedAvg4Trainer(BaseTrainer):
"""
Scheme I and Scheme II, based on the flag of self.simple_average
"""
def __init__(self, options, dataset):
model = choose_model(options)
self.move_model_to_gpu(model, options)
self.optimizer = GD(model.parameters(),
lr=options['lr'],
weight_decay=options['wd'])
self.num_epoch = options['num_epoch']
worker = LrdWorker(model, self.optimizer, options)
super(FedAvg4Trainer, self).__init__(options, dataset, worker=worker)
self.prob = self.compute_prob()
def train(self):
print('>>> Select {} clients per round \n'.format(
self.clients_per_round))
# Fetch latest flat model parameter
self.latest_model = self.worker.get_flat_model_params().detach()
for round_i in range(self.num_round):
# Test latest model on train data
self.test_latest_model_on_traindata(round_i)
self.test_latest_model_on_evaldata(round_i)
# Choose K clients prop to data size
if self.simple_average:
selected_clients, repeated_times = self.select_clients_with_prob(
seed=round_i)
else:
selected_clients = self.select_clients(seed=round_i)
repeated_times = None
# Solve minimization locally
solns, stats = self.local_train(round_i, selected_clients)
# Track communication cost
self.metrics.extend_commu_stats(round_i, stats)
# Update latest model
self.latest_model = self.aggregate(solns,
repeated_times=repeated_times)
self.optimizer.inverse_prop_decay_learning_rate(round_i)
# Test final model on train data
self.test_latest_model_on_traindata(self.num_round)
self.test_latest_model_on_evaldata(self.num_round)
# Save tracked information
self.metrics.write()
def compute_prob(self):
probs = []
for c in self.clients:
probs.append(len(c.train_data))
return np.array(probs) / sum(probs)
def select_clients_with_prob(self, seed=1):
num_clients = min(self.clients_per_round, len(self.clients))
np.random.seed(seed)
index = np.random.choice(len(self.clients), num_clients, p=self.prob)
index = sorted(index.tolist())
select_clients = []
select_index = []
repeated_times = []
for i in index:
if i not in select_index:
select_clients.append(self.clients[i])
select_index.append(i)
repeated_times.append(1)
else:
repeated_times[-1] += 1
return select_clients, repeated_times
def aggregate(self, solns, **kwargs):
averaged_solution = torch.zeros_like(self.latest_model)
# averaged_solution = np.zeros(self.latest_model.shape)
if self.simple_average:
repeated_times = kwargs['repeated_times']
assert len(solns) == len(repeated_times)
for i, (num_sample, local_solution) in enumerate(solns):
averaged_solution += local_solution * repeated_times[i]
averaged_solution /= self.clients_per_round
else:
for num_sample, local_solution in solns:
averaged_solution += num_sample * local_solution
averaged_solution /= self.all_train_data_num
averaged_solution *= (100 / self.clients_per_round)
# averaged_solution = from_numpy(averaged_solution, self.gpu)
return averaged_solution.detach()