77 lines
2.5 KiB
Python
77 lines
2.5 KiB
Python
|
|
||
|
|
import models, torch, copy
|
||
|
|
class Client(object):
|
||
|
|
|
||
|
|
def __init__(self, conf, model, train_dataset, id = -1):
|
||
|
|
|
||
|
|
self.conf = conf
|
||
|
|
|
||
|
|
self.local_model = models.get_model(self.conf["model_name"])
|
||
|
|
|
||
|
|
self.client_id = id
|
||
|
|
|
||
|
|
self.train_dataset = train_dataset
|
||
|
|
|
||
|
|
all_range = list(range(len(self.train_dataset)))
|
||
|
|
data_len = int(len(self.train_dataset) / self.conf['no_models'])
|
||
|
|
train_indices = all_range[id * data_len: (id + 1) * data_len]
|
||
|
|
|
||
|
|
self.train_loader = torch.utils.data.DataLoader(self.train_dataset, batch_size=conf["batch_size"],
|
||
|
|
sampler=torch.utils.data.sampler.SubsetRandomSampler(train_indices))
|
||
|
|
|
||
|
|
|
||
|
|
def local_train(self, model):
|
||
|
|
|
||
|
|
for name, param in model.state_dict().items():
|
||
|
|
self.local_model.state_dict()[name].copy_(param.clone())
|
||
|
|
|
||
|
|
#print("\n\nlocal model train ... ... ")
|
||
|
|
#for name, layer in self.local_model.named_parameters():
|
||
|
|
# print(name, "->", torch.mean(layer.data))
|
||
|
|
|
||
|
|
#print("\n\n")
|
||
|
|
optimizer = torch.optim.SGD(self.local_model.parameters(), lr=self.conf['lr'],
|
||
|
|
momentum=self.conf['momentum'])
|
||
|
|
|
||
|
|
|
||
|
|
self.local_model.train()
|
||
|
|
for e in range(self.conf["local_epochs"]):
|
||
|
|
|
||
|
|
for batch_id, batch in enumerate(self.train_loader):
|
||
|
|
data, target = batch
|
||
|
|
#for name, layer in self.local_model.named_parameters():
|
||
|
|
# print(torch.mean(self.local_model.state_dict()[name].data))
|
||
|
|
#print("\n\n")
|
||
|
|
if torch.cuda.is_available():
|
||
|
|
data = data.cuda()
|
||
|
|
target = target.cuda()
|
||
|
|
|
||
|
|
optimizer.zero_grad()
|
||
|
|
output = self.local_model(data)
|
||
|
|
loss = torch.nn.functional.cross_entropy(output, target)
|
||
|
|
loss.backward()
|
||
|
|
|
||
|
|
optimizer.step()
|
||
|
|
|
||
|
|
#for name, layer in self.local_model.named_parameters():
|
||
|
|
# print(torch.mean(self.local_model.state_dict()[name].data))
|
||
|
|
#print("\n\n")
|
||
|
|
if self.conf["dp"]:
|
||
|
|
model_norm = models.model_norm(model, self.local_model)
|
||
|
|
|
||
|
|
norm_scale = min(1, self.conf['C'] / (model_norm))
|
||
|
|
#print(model_norm, norm_scale)
|
||
|
|
for name, layer in self.local_model.named_parameters():
|
||
|
|
clipped_difference = norm_scale * (layer.data - model.state_dict()[name])
|
||
|
|
layer.data.copy_(model.state_dict()[name] + clipped_difference)
|
||
|
|
|
||
|
|
print("Epoch %d done." % e)
|
||
|
|
diff = dict()
|
||
|
|
for name, data in self.local_model.state_dict().items():
|
||
|
|
diff[name] = (data - model.state_dict()[name])
|
||
|
|
|
||
|
|
#print("\n\nfinishing local model training ... ... ")
|
||
|
|
#for name, layer in self.local_model.named_parameters():
|
||
|
|
# print(name, "->", torch.mean(layer.data))
|
||
|
|
return diff
|
||
|
|
|