def FEMNIST_iid(dataset,num_users):
args = args_parser()
if os.path.exists(args.private_dataset_index):
print('private_dataset_index exist ~ ')
f = open(args.private_dataset_index,'r')
temp = f.read()
dict_users =eval(temp)
if args.testBestCondition:
setvolume = set()
for key in dict_users.keys():
setvolume = setvolume.union(dict_users[key])
for key in dict_users.keys():
dict_users[key] = setvolume
return dict_users
else:
print('create private_dataset_index ~ ')
# num_item = int(len(dataset) / num_users)
num_class_item = 3 # 用于达到作者目的！
dict_users = {}
for i in range(num_users):
dict_users[i]=set() # create the set dictionary
for i in range(num_users):
for key in dict.keys():
temp = set(np.random.choice(dict[key], num_class_item, replace=False)) # 样本无重复值
dict_users[i] = dict_users[i] | temp
dict[key] = list(set(dict[key]) - temp)
f = open(args.private_dataset_index,'w')
f.write(str(dict_users))
f.close()
return dict_users
.format(n, model_type_list[n], epoch + 1,
batch_idx * len(images),
len(trainloader.dataset),
100. * batch_idx / len(trainloader),
loss.item()))
train_batch_loss[n].append(loss.item())
torch.save(
model.state_dict(),
'Src/CollaborativeModel/LocalModel{}Type{}.pkl'.format(
n, model_type_list[n]))
for index in range(len(train_batch_loss)):
loss_avg = sum(train_batch_loss[index]) / len(
train_batch_loss[index])
train_loss[index].append(loss_avg)
plt.figure()
for index in range(len(train_loss)):
plt.plot(range(len(train_loss[index])), train_loss[index])
plt.title('collaborative_train_losses')
plt.xlabel('epoches')
plt.ylabel('Train loss')
plt.savefig('Src/Figure/collaborative_train_losses.png')
plt.show()
print('End Public Training')
from option import args_parser
if __name__ == '__main__':
args = args_parser()
collaborative_private_model_mnist_train(args)