net_glob.train()
# copy weights
w_init = copy.deepcopy(net_glob.state_dict())
local_acc_final = []
total_acc_final = []
local_acc = np.zeros([args.num_users, args.epochs])
total_acc = np.zeros([args.num_users, args.epochs])
# training
for idx in range(args.num_users):
# print(w_init)
net_glob.load_state_dict(w_init)
optimizer = optim.Adam(net_glob.parameters())
train_loader = DataLoader(DatasetSplit(dataset_train, dict_users[idx]),
batch_size=64,
shuffle=True)
image_trainset_weight = np.zeros(10)
for label in np.array(dataset_train.targets)[dict_users[idx]]:
image_trainset_weight[label] += 1
image_trainset_weight = image_trainset_weight / image_trainset_weight.sum(
)
list_loss = []
net_glob.train()
for epoch in range(args.epochs):
batch_loss = []
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(args.device), target.to(args.device)
optimizer.zero_grad()
output = net_glob(data)
best_loss = None
best_acc = None
best_epoch = None
lr = args.lr
results = []
criterion = nn.CrossEntropyLoss()
for user, net_local in enumerate(net_local_list):
model_save_path = os.path.join(base_dir,
'local/model_user{}.pt'.format(user))
net_best = None
best_acc = None
ldr_train = DataLoader(DatasetSplit(dataset_train,
dict_users_train[user]),
batch_size=args.local_bs,
shuffle=True)
optimizer = torch.optim.SGD(net_local.parameters(),
lr=lr,
momentum=0.5)
for iter in range(args.epochs):
for batch_idx, (images, labels) in enumerate(ldr_train):
images, labels = images.to(args.device), labels.to(args.device)
net_local.zero_grad()
log_probs = net_local(images)
loss = criterion(log_probs, labels)
loss.backward()
optimizer.step()
best_loss_valid = loss_valid
w_best = copy.deepcopy(w_glob)
loss_avg = sum(loss_locals) / len(loss_locals)
print('Round {:3d}, Average loss {:.3f}, acc_valid {:.2f}%'.format(
iter, loss_avg, acc_valid))
loss_train.append(loss_avg)
print("\n Begin test")
net_glob.load_state_dict(w_best)
net_glob.eval()
users_labeled = set()
for i in range(len(dict_users_labeled)):
users_labeled = users_labeled | dict_users_labeled[i]
users_unlabeled = set()
for i in range(len(dict_users_labeled)):
users_unlabeled = users_unlabeled | (dict_users[i] -
dict_users_labeled[i])
dataset_train_labeled = DatasetSplit(dataset=dataset_train,
idxs=users_labeled,
pseudo_label=pseudo_label)
acc_train_labeled, loss_train_test_labeled = test_img(
net_glob, dataset_train_labeled, args)
print("labeled Training accuracy: {:.2f}%".format(acc_train_labeled))
acc_test, loss_test = test_img(net_glob, dataset_test, args)
print("Testing accuracy: {:.2f}% \n\n".format(acc_test))
net_glob.load_state_dict(w_glob)
net_glob.eval()
acc_valid, loss_valid = test_img(net_glob, dataset_valid, args)
if loss_valid <= best_loss_valid:
best_loss_valid = loss_valid
w_best = copy.deepcopy(w_glob)
loss_avg = sum(loss_locals) / len(loss_locals)
print('Round {:3d}, Average loss {:.3f}, acc_valid {:.2f}%'.format(
iter, loss_avg, acc_valid))
loss_train.append(loss_avg)
print('\n Begin Pseduo Label')
for batch_idx in range(len(dataset_train)):
if batch_idx % 10000 == 0:
print("Pseduo Label: {:3d}".format(batch_idx))
try_data = DataLoader(DatasetSplit(dataset=dataset_train,
idxs=[batch_idx],
pseudo_label=pseudo_label),
batch_size=1,
shuffle=True)
for i, (img, label) in enumerate(try_data):
if (pseudo_label[batch_idx] == -1):
img, label = img.to(args.device), label.to(args.device)
log_probs = net_glob(img)
y_pred = log_probs.data.max(1, keepdim=True)[1]
pseudo_label[batch_idx] = y_pred.cpu().detach().numpy()
print("\n Begin mid test")
net_glob.eval()
users_labeled = set()
for i in range(len(dict_users_labeled)):
users_labeled = users_labeled | dict_users_labeled[i]
users_unlabeled = set()
def create_shared_dataset(valid_ds, size):
data_loader = DataLoader(valid_ds, batch_size=1)
size_final_dataset = size
each_label_size = size_final_dataset/10
final_test_dataset_idx = np.zeros(size_final_dataset, dtype = int)
# final_test_dataset_idx = {i: np.array([], dtype='int64') for i in range(10)}
print(len(valid_ds))
counter = np.zeros(10)
counter_arr = 0
for idx, (data, target) in enumerate(data_loader):
# print(idx, ' ', ' ', target, ' ')
if target == 0:
if counter[0] < each_label_size:
counter[0] += 1
final_test_dataset_idx[counter_arr] = int(idx)
counter_arr += 1
elif target == 1:
if counter[1] < each_label_size:
counter[1] += 1
final_test_dataset_idx[counter_arr] = int(idx)
counter_arr += 1
elif target == 2:
if counter[2] < each_label_size:
counter[2] += 1
final_test_dataset_idx[counter_arr] = int(idx)
counter_arr += 1
elif target == 3:
if counter[3] < each_label_size:
counter[3] += 1
final_test_dataset_idx[counter_arr] = int(idx)
counter_arr += 1
elif target == 4:
if counter[4] < each_label_size:
counter[4] += 1
final_test_dataset_idx[counter_arr] = int(idx)
counter_arr += 1
elif target == 5:
if counter[5] < each_label_size:
counter[5] += 1
final_test_dataset_idx[counter_arr] = int(idx)
counter_arr += 1
elif target == 6:
if counter[6] < each_label_size:
counter[6] += 1
final_test_dataset_idx[counter_arr] = int(idx)
counter_arr += 1
elif target == 7:
if counter[7] < each_label_size:
counter[7] += 1
final_test_dataset_idx[counter_arr] = int(idx)
counter_arr += 1
elif target == 8:
if counter[8] < each_label_size:
counter[8] += 1
final_test_dataset_idx[counter_arr] = int(idx)
counter_arr += 1
elif target == 9:
if counter[9] < each_label_size:
counter[9] += 1
final_test_dataset_idx[counter_arr] = int(idx)
counter_arr += 1
ds = DatasetSplit(valid_ds, final_test_dataset_idx)
return ds
net_glob = gate_vgg16(args=args).to(args.device)
else:
exit('Error: unrecognized model')
image, target = next(iter(test_loader))
writer.add_graph(net_glob, image.to(args.device))
gate_epochs = 200
local_acc = np.zeros([args.num_users, args.epochs + gate_epochs + 1])
total_acc = np.zeros([args.num_users, args.epochs + gate_epochs + 1])
local_acc2 = np.zeros([args.num_users, args.epochs + gate_epochs + 1])
total_acc2 = np.zeros([args.num_users, args.epochs + gate_epochs + 1])
for user_num in range(len(dict_users)):
# user data
user_train = DatasetSplit(dataset_train, dict_users[user_num])
np.random.shuffle(dict_users[user_num])
cut_point = len(dict_users[user_num]) // 4
train_loader = DataLoader(DatasetSplit(
dataset_train, dict_users[user_num][cut_point:]),
batch_size=64,
shuffle=True)
gate_loader = DataLoader(DatasetSplit(
dataset_train, dict_users[user_num][:cut_point]),
batch_size=64,
shuffle=True)
class_weight = np.zeros(10)
for image, label in user_train:
class_weight[label] += 1
transform=trans_svhn_strong)
dataset_train_weak = datasets.SVHN('../data/svhn',
split='train',
download=True,
transform=trans_svhn_weak)
dataset_test = datasets.SVHN('../data/svhn',
split='test',
download=True,
transform=transform_svhn_test)
else:
exit('Error: unrecognized dataset')
a = np.arange(len(dataset_test))
np.random.shuffle(a)
dataset_valid = DatasetSplit(dataset=dataset_test, idxs=a[:1000])
dataset_test = DatasetSplit(dataset=dataset_test, idxs=a[1000:])
if args.iid == 'noniid_ssl' and args.dataset == 'cifar':
dict_users, dict_users_labeled, pseudo_label = noniid_ssl(
dataset_train_weak, args.num_users, args.label_rate)
else:
dict_users, dict_users_labeled, pseudo_label = sample(
dataset_train_weak, args.num_users, args.label_rate, args.iid)
if args.dataset == 'cifar':
net_glob = CNNCifar(args=args).to(args.device)
net_glob_helper_1 = CNNCifar(args=args).to(args.device)
net_glob_helper_2 = CNNCifar(args=args).to(args.device)
net_glob_valid = CNNCifar(args=args).to(args.device)
print(type(dataset_train))
data_weight = len(dataset_train) / args.num_users / 100
if args.random_idx:
idx = random.randint(0, args.num_users - 1)
else:
idx = args.idx
local = LocalUpdate(args=args, dataset=dataset_train, idxs=dict_users[idx])
w, loss = local.train(net=copy.deepcopy(net_local))
print(loss)
net_local.load_state_dict(w)
#Here let's just define the trained portion of train_set for finding acccuracy
acc_train, loss_train = test_img(
net_local, DatasetSplit(dataset_train, dict_users[idx]), args)
#acc_train, loss_train = test_img(net_local, dataset_train, args)
acc_test, loss_test = test_img(net_local, dataset_test, args)
print("Training accuracy: {:.2f}".format(acc_train))
print("Testing accuracy: {:.2f}".format(acc_test))
#w_locals.append(copy.deepcopy(w))
#loss_locals.append(copy.deepcopy(loss))
#w['epoch']=0
if args.new:
checkpoint = {'data_weight': data_weight, 'state_dict': w}
else:
checkpoint = {'data_weight': data_weight, 'state_dict': w}