from __future__ import print_function import sys import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F import torch.backends.cudnn as cudnn import random import os import argparse import numpy as np from PreResNet import * from sklearn.mixture import GaussianMixture import dataloader_cifar as dataloader import matplotlib.pyplot as plt import copy import seaborn as sns # from sklearn.mixture import GaussianMixture from sklearn.cluster import KMeans from sklearn.cluster import Birch import matplotlib parser = argparse.ArgumentParser(description='PyTorch CIFAR Training') parser.add_argument('--batch_size', default=128, type=int, help='train batchsize') parser.add_argument('--lr', '--learning_rate', default=0.02, type=float, help='initial learning rate') parser.add_argument('--noise_mode', default='asym') parser.add_argument('--alpha', default=4, type=float, help='parameter for Beta') parser.add_argument('--lambda_u', default=150, type=float, help='weight for unsupervised loss') parser.add_argument('--p_threshold', default=0.5, type=float, help='clean probability threshold') parser.add_argument('--T', default=0.5, type=float, help='sharpening temperature') parser.add_argument('--num_epochs', default=300, type=int) parser.add_argument('--r', default=0.3, type=float, help='noise ratio') parser.add_argument('--id', default='') parser.add_argument('--seed', default=123) parser.add_argument('--gpuid', default=0, type=int) parser.add_argument('--num_class', default=100, type=int) # parser.add_argument('--data_path', default='./data/cifar-10-batches-py', type=str, help='path to dataset') # parser.add_argument('--dataset', default='cifar10', type=str) parser.add_argument('--data_path', default='./data/cifar-100-python', type=str, help='path to dataset') parser.add_argument('--dataset', default='cifar100', type=str) args = parser.parse_args() torch.cuda.set_device(args.gpuid) random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) mse = torch.nn.MSELoss(reduction='none').cuda() # Training def train(epoch, net, net2, optimizer, labeled_trainloader, unlabeled_trainloader, mask=None, f_G=None, new_y=None): net.train() net2.eval() # fix one network and train the other unlabeled_train_iter = iter(unlabeled_trainloader) num_iter = (len(labeled_trainloader.dataset) // args.batch_size) + 1 mse_total = 0 for batch_idx, (inputs_x, inputs_x2, labels_x, w_x) in enumerate(labeled_trainloader): try: inputs_u, inputs_u2 = unlabeled_train_iter.__next__() except: unlabeled_train_iter = iter(unlabeled_trainloader) inputs_u, inputs_u2 = unlabeled_train_iter.__next__() batch_size = inputs_x.size(0) # Transform label to one-hot，转为0-1矩阵 labels_x = torch.zeros(batch_size, args.num_class).scatter_(1, labels_x.view(-1, 1), 1) w_x = w_x.view(-1, 1).type(torch.FloatTensor) inputs_x, inputs_x2, labels_x, w_x = inputs_x.cuda(), inputs_x2.cuda(), labels_x.cuda(), w_x.cuda() inputs_u, inputs_u2 = inputs_u.cuda(), inputs_u2.cuda() with torch.no_grad(): # label co-guessing of unlabeled samples outputs_u11, feat_u11 = net(inputs_u, feat_out=True) outputs_u12, feat_u12 = net(inputs_u2, feat_out=True) outputs_u21, feat_u21 = net2(inputs_u, feat_out=True) outputs_u22, feat_u22 = net2(inputs_u2, feat_out=True) # 取average of 所有网络的输出，作者利用了所谓的augmentation pu = (torch.softmax(outputs_u11, dim=1) + torch.softmax(outputs_u12, dim=1) + torch.softmax(outputs_u21, dim=1) + torch.softmax(outputs_u22, dim=1)) / 4 ptu = pu ** (1 / args.T) # temparature sharpening # Algorithm 1 中的shapen(qb,T) targets_u = ptu / ptu.sum(dim=1, keepdim=True) # normalize targets_u = targets_u.detach() # label refinement of labeled samples outputs_x, feat_x1 = net(inputs_x, feat_out=True) outputs_x2, feat_x2 = net(inputs_x2, feat_out=True) # 取labeled的输出平均值 px = (torch.softmax(outputs_x, dim=1) + torch.softmax(outputs_x2, dim=1)) / 2 # 公式(3)(4)退火 px = w_x * labels_x + (1 - w_x) * px ptx = px ** (1 / args.T) # temparature sharpening targets_x = ptx / ptx.sum(dim=1, keepdim=True) # normalize targets_x = targets_x.detach() # aaa = torch.argmax(labels_x, dim=1) # mse_loss = torch.sum(mse((feat_x1+feat_x2)/2, f_G[aaa]), 1) # mse_total = (mse_total + torch.sum(mse_loss) / len(mse_loss))/2 # mixmatch l = np.random.beta(args.alpha, args.alpha) # 促使X'更加靠近labeled sample而不是无监督样本 l = max(l, 1 - l) all_inputs = torch.cat([inputs_x, inputs_x2, inputs_u, inputs_u2], dim=0) all_targets = torch.cat([targets_x, targets_x, targets_u, targets_u], dim=0) # 随机输出mini batch的序号，来mixup idx = torch.randperm(all_inputs.size(0)) input_a, input_b = all_inputs, all_inputs[idx] target_a, target_b = all_targets, all_targets[idx] # 利用mix但是促使模型更偏向于label而不是UNlabel mixed_input = l * input_a + (1 - l) * input_b mixed_target = l * target_a + (1 - l) * target_b logits = net(mixed_input) # 输出被排列成两部分，input_x、Input_u logits_x = logits[:batch_size * 2] logits_u = logits[batch_size * 2:] # 利用公式(9)-(10)计算损失函数，其中lamb是所谓的warm up Lx, Lu, lamb = criterion(logits_x, mixed_target[:batch_size * 2], logits_u, mixed_target[batch_size * 2:], epoch + batch_idx / num_iter, warm_up) # regularization prior = torch.ones(args.num_class) / args.num_class prior = prior.cuda() pred_mean = torch.softmax(logits, dim=1).mean(0) # 一般来说会省略固定的prior部分，只取last term # lambR=1 penalty = torch.sum(prior * torch.log(prior / pred_mean)) # lamb是通过warm和current epoch比较得出的百分数，意味着随着epoch进行，Lu所占比重会逐渐增加 # 前期需要保持标准CE损失，但是实际还有penalty # loss = Lx + lamb * Lu + penalty loss = Lx + penalty + lamb * Lu # compute gradient and do SGD step optimizer.zero_grad() loss.backward() optimizer.step() if batch_idx % 200 == 0: sys.stdout.write('\r') sys.stdout.write( '%s:%.1f-%s | Epoch [%3d/%3d] Iter[%3d/%3d]\t Labeled loss: %.2f Unlabeled loss: %.2f\n' % (args.dataset, args.r, args.noise_mode, epoch, args.num_epochs, batch_idx + 1, num_iter, Lx.item(), Lu.item())) sys.stdout.flush() # print('\r mse loss:%.4f\n' % mse_total, end='end', flush=True) # print('\r mse loss:%.4f\n' % mse_total, end='end', flush=True) def mixup_criterion(pred, y_a, y_b, lam): c = F.log_softmax(pred, 1) return lam * F.cross_entropy(c, y_a) + (1 - lam) * F.cross_entropy(c, y_b) soft_mix_warm = False def warmup(epoch, net, optimizer, dataloader): net.train() num_iter = (len(dataloader.dataset) // dataloader.batch_size) + 1 for batch_idx, (inputs, labels, path) in enumerate(dataloader): optimizer.zero_grad() l = np.random.beta(args.alpha, args.alpha) # 促使X'更加靠近labeled sample而不是无监督样本 l = max(l, 1 - l) idx = torch.randperm(inputs.size(0)) targets = torch.zeros(inputs.size(0), args.num_class).scatter_(1, labels.view(-1, 1), 1).cuda() targets = torch.clamp(targets, 1e-4, 1.) inputs, labels = inputs.cuda(), labels.cuda() if soft_mix_warm: input_a, input_b = inputs, inputs[idx] target_a, target_b = targets, targets[idx] labels_a, labels_b = labels, labels[idx] # 利用mix但是促使模型更偏向于label而不是UNlabel mixed_input = l * input_a + (1 - l) * input_b mixed_target = l * target_a + (1 - l) * target_b outputs = net(mixed_input) loss = mixup_criterion(outputs, labels_a, labels_b, l) L = loss else: outputs = net(inputs) loss = CEloss(outputs, labels) if args.noise_mode == 'asym': # penalize confident prediction for asymmetric noise penalty = conf_penalty(outputs) L = loss + penalty elif args.noise_mode == 'sym': L = loss L.backward() optimizer.step() if batch_idx % 200 == 0: sys.stdout.write('\r') sys.stdout.write('%s:%.1f-%s | Epoch [%3d/%3d] Iter[%3d/%3d]\t CE-loss: %.4f' % (args.dataset, args.r, args.noise_mode, epoch, args.num_epochs, batch_idx + 1, num_iter, loss.item())) sys.stdout.flush() def test(epoch, net1, net2, best_acc, w_glob=None): if w_glob is None: net1.eval() net2.eval() correct = 0 total = 0 with torch.no_grad(): for batch_idx, (inputs, targets) in enumerate(test_loader): inputs, targets = inputs.cuda(), targets.cuda() outputs1 = net1(inputs) outputs2 = net2(inputs) outputs = outputs1 + outputs2 _, predicted = torch.max(outputs, 1) total += targets.size(0) correct += predicted.eq(targets).cpu().sum().item() acc = 100. * correct / total if best_acc < acc: best_acc = acc print("\n| Ensemble network Test Epoch #%d\t Accuracy: %.2f, best_acc: %.2f%%\n" % (epoch, acc, best_acc)) test_log.write('ensemble_Epoch:%d Accuracy:%.2f, best_acc: %.2f\n' % (epoch, acc, best_acc)) test_log.flush() else: net1_w_bak = net1.state_dict() net1.load_state_dict(w_glob) net1.eval() correct = 0 total = 0 with torch.no_grad(): for batch_idx, (inputs, targets) in enumerate(test_loader): inputs, targets = inputs.cuda(), targets.cuda() outputs1 = net1(inputs) _, predicted = torch.max(outputs1, 1) total += targets.size(0) correct += predicted.eq(targets).cpu().sum().item() acc = 100. * correct / total if best_acc < acc: best_acc = acc print("\n| Global network Test Epoch #%d\t Accuracy: %.2f, best_acc: %.2f%%\n" % (epoch, acc, best_acc)) test_log.write('global_Epoch:%d Accuracy:%.2f, best_acc: %.2f\n' % (epoch, acc, best_acc)) test_log.flush() # 恢复权重 net1.load_state_dict(net1_w_bak) return best_acc feat_dim = 512 #是否可以加个全连接改成128 sim = torch.nn.CosineSimilarity(dim=1) loss_func = torch.nn.CrossEntropyLoss(reduction='none') def get_small_loss_samples(y_pred, y_true, forget_rate): loss = loss_func(y_pred, y_true) ind_sorted = np.argsort(loss.data.cpu()).cuda() loss_sorted = loss[ind_sorted] remember_rate = 1 - forget_rate num_remember = int(remember_rate * len(loss_sorted)) ind_update = ind_sorted[:num_remember] return ind_update def get_small_loss_by_loss_list(loss_list, forget_rate, eval_loader): remember_rate = 1 - forget_rate idx_list = [] for i in range(10): class_idx = np.where(np.array(eval_loader.dataset.noise_label)[:] == i)[0] # class_idx = torch.from_numpy(class_idx).cuda() loss_per_class = loss_list[class_idx] #取对应target的loss num_remember = int(remember_rate * len(loss_per_class)) ind_sorted = np.argsort(loss_per_class.data.cpu()) ind_update = ind_sorted[:num_remember].tolist() idx_list.append(ind_update) return idx_list def eval_train(model, all_loss): model.eval() losses = torch.zeros(50000) f_G = torch.zeros(args.num_class, feat_dim).cuda() f_all = torch.zeros(50000, feat_dim).cuda() n_labels = torch.zeros(args.num_class, 1).cuda() y_k_tilde = torch.zeros(50000) mask = np.zeros(50000) with torch.no_grad(): for batch_idx, (inputs, targets, index) in enumerate(eval_loader): inputs, targets = inputs.cuda(), targets.cuda() outputs, feat = model(inputs, feat_out=True) loss = CE(outputs, targets) _, predicted = torch.max(outputs, 1) for b in range(inputs.size(0)): losses[index[b]] = loss[b] f_G[predicted[b]] += feat[b] n_labels[predicted[b]] += 1 f_all[index] = feat assert torch.sum(n_labels) == 50000 for i in range(len(n_labels)): if n_labels[i] == 0: n_labels[i] = 1 f_G = torch.div(f_G, n_labels) f_G = F.normalize(f_G, dim=1) f_all = F.normalize(f_all, dim=1) temp = f_G.t() sim_all = torch.mm(f_all, temp) # .cpu().numpy() y_k_tilde = torch.argmax(sim_all.cpu(), dim=1) with torch.no_grad(): for batch_idx, (inputs, targets, index) in enumerate(eval_loader): for i in range(len(index)): if y_k_tilde[index[i]] == targets[i]: mask[index[i]] = 1 losses = (losses - losses.min()) / (losses.max() - losses.min()) all_loss.append(losses) if args.r == 0.9: # average loss over last 5 epochs to improve convergence stability history = torch.stack(all_loss) input_loss = history[-5:].mean(0) input_loss = input_loss.reshape(-1, 1) else: input_loss = losses.reshape(-1, 1) # fit a two-component GMM to the loss # 参数如下： # n_components 聚类数量，max_iter 最大迭代次数，tol 阈值低于停止，reg_covar 协方差矩阵对角线上非负正则化参数，接近0即可 gmm = GaussianMixture(n_components=2, max_iter=10, tol=1e-2, reg_covar=5e-4) gmm.fit(input_loss) prob = gmm.predict_proba(input_loss) prob = prob[:, gmm.means_.argmin()] return prob, all_loss, losses.numpy(), mask, f_G def mix_data_lab(x, y, alpha=1.0): '''Returns mixed inputs, pairs of targets, and lambda''' if alpha > 0: lam = np.random.beta(alpha, alpha) else: lam = 1 batch_size = x.size()[0] index = torch.randperm(batch_size).cuda() lam = max(lam, 1 - lam) mixed_x = lam * x + (1 - lam) * x[index, :] y_a, y_b = y, y[index] return mixed_x, y_a, y_b, index, lam def linear_rampup(current, warm_up, rampup_length=16): # 线性warm_up，对sym噪声使用标准CE训练一段时间 # 实际warm up epoch是warm_up+rampup_length current = np.clip((current - warm_up) / rampup_length, 0.0, 1.0) re_val = args.lambda_u * float(current) # print(" current warm up parameters:", current) # print("return parameters:", re_val) return re_val class SemiLoss(object): def __call__(self, outputs_x, targets_x, outputs_u, targets_u, epoch, warm_up): probs_u = torch.softmax(outputs_u, dim=1) # 利用mixup后的交叉熵，px输出*log(px_model) Lx = -torch.mean(torch.sum(F.log_softmax(outputs_x, dim=1) * targets_x, dim=1)) # 而UNlabel则是均方误差，p_u输出-pu_model Lu = torch.mean((probs_u - targets_u) ** 2) return Lx, Lu, linear_rampup(epoch, warm_up) class NegEntropy(object): def __call__(self, outputs): probs = torch.softmax(outputs, dim=1) return torch.mean(torch.sum(probs.log() * probs, dim=1)) def create_model(): # 其实是pre-resnet18，使用的是pre-resnet block model = ResNet18(num_classes=args.num_class) model = model.cuda() return model def plotHistogram(model_1_loss, model_2_loss, noise_index, clean_index, epoch, round, noise_rate): title = 'Epoch-' + str(epoch)+':' fig = plt.figure() plt.subplot(121) gmm = GaussianMixture(n_components=2, max_iter=20, tol=1e-2, random_state=0, reg_covar=5e-4) model_1_loss = np.reshape(model_1_loss, (-1, 1)) gmm.fit(model_1_loss) # fit the loss # plot resulting fit x_range = np.linspace(0, 1, 1000) pdf = np.exp(gmm.score_samples(x_range.reshape(-1, 1))) responsibilities = gmm.predict_proba(x_range.reshape(-1, 1)) pdf_individual = responsibilities * pdf[:, np.newaxis] plt.hist(np.array(model_1_loss[noise_index]), density=True, bins=100, alpha=0.5,histtype='bar', color='red', label='Noisy subset') plt.hist(np.array(model_1_loss[clean_index]), density=True, bins=100, alpha=0.5,histtype='bar', color='blue', label='Clean subset') plt.plot(x_range, pdf, '-k', label='Mixture') plt.plot(x_range, pdf_individual, '--', label='Component') plt.legend(loc='upper right', prop={'size': 12}) plt.xlabel('Normalized loss') plt.ylabel('Estimated pdf') plt.title(title+'Model_1') plt.subplot(122) gmm = GaussianMixture(n_components=2, max_iter=20, tol=1e-2, random_state=0, reg_covar=5e-4) model_2_loss = np.reshape(model_2_loss, (-1, 1)) gmm.fit(model_2_loss) # fit the loss # plot resulting fit x_range = np.linspace(0, 1, 1000) pdf = np.exp(gmm.score_samples(x_range.reshape(-1, 1))) responsibilities = gmm.predict_proba(x_range.reshape(-1, 1)) pdf_individual = responsibilities * pdf[:, np.newaxis] plt.hist(np.array(model_2_loss[noise_index]), density=True, bins=100, alpha=0.5,histtype='bar', color='red', label='Noisy subset') plt.hist(np.array(model_2_loss[clean_index]), density=True, bins=100, alpha=0.5,histtype='bar', color='blue', label='Clean subset') plt.plot(x_range, pdf, '-k', label='Mixture') plt.plot(x_range, pdf_individual, '--', label='Component') plt.legend(loc='upper right', prop={'size': 12}) plt.xlabel('Normalized loss') plt.ylabel('Estimated pdf') plt.title(title+'Model_2') print('\nlogging histogram...') title = 'cifar10_' + str(args.noise_mode) + '_moit_double_' + str(noise_rate) plt.savefig(os.path.join('./figure_his/', 'two_model_{}_{}_{}_{}.{}'.format(epoch, round, title, int(soft_mix_warm), ".tif")), dpi=300) # plt.show() plt.close() def loss_dist_plot(loss, noisy_index, clean_index, epoch, rou=None, g_file=True, model_name='', loss2=None): """ plot the loss distribution :param loss: the list contains the loss per sample :param noisy_index: contains the indices of real noisy label :param clean_index: contains the indices of real clean label :param filename: the generated pdf file name :param title: the figure title :param g_file: whether to generate the pdf figure file :return: None """ if loss2 is None: filename = 'one_model_'+str(args.dataset)+'_'+str(args.noise_mode)+'_'+str(args.r)+'_epoch='+str(epoch) if rou is None: title = 'Epoch-'+str(epoch) + ': ' + str(args.dataset)+' '+str(args.r*100)+'%-'+str(args.noise_mode) else: title = 'Epoch-' + str(epoch) + ' ' +'Round-'+str(rou)+ ': ' + str(args.dataset) + ' ' + str(int(args.r * 100)) + '%-' + str(args.noise_mode) if type(loss) is not np.ndarray: loss= loss.numpy() sns.set(style='whitegrid') gmm = GaussianMixture(n_components=2, max_iter=20, tol=1e-2, random_state=0, reg_covar=5e-4) loss = np.reshape(loss, (-1, 1)) gmm.fit(loss) # fit the loss # plot resulting fit x_range = np.linspace(0, 1, 1000) pdf = np.exp(gmm.score_samples(x_range.reshape(-1, 1))) responsibilities = gmm.predict_proba(x_range.reshape(-1, 1)) pdf_individual = responsibilities * pdf[:, np.newaxis] # sns.distplot(loss[noisy_index], color="red", rug=False,kde=False, label="incorrect", # hist_kws={"color": "r", "alpha": 0.5}) # sns.distplot(loss[clean_index], color="skyblue", rug=False,kde=False, label="correct", # hist_kws={"color": "b", "alpha": 0.5}) plt.hist(np.array(loss[noisy_index]), density=True, bins=100, histtype='bar', alpha=0.5, color='red', label='Noisy subset') plt.hist(np.array(loss[clean_index]), density=True, bins=100, histtype='bar', alpha=0.5, color='blue', label='Clean subset') plt.plot(x_range, pdf, '-k', label='Mixture') plt.plot(x_range, pdf_individual, '--', label='Component') # plt.plot(x_range, pdf_individual[:][1], '--', color='blue', label='Component 1') plt.title(title, fontsize=20) plt.xlabel('Normalized loss', fontsize=24) plt.ylabel('Estimated pdf', fontsize=24) plt.tick_params(labelsize=24) plt.legend(loc='upper right', prop={'size': 12}) # plt.tight_layout() if g_file: plt.savefig('./figure_his/{0}.tif'.format(filename+model_name), bbox_inches='tight', dpi=300) #plt.show() plt.close() else: filename = 'noise_'+str(args.dataset) + '_' + str(args.noise_mode) + '_' + str(args.r) + '_epoch=' + str(epoch) if rou is None: title = 'Epoch-' + str(epoch) + ': ' + str(args.dataset) + ' ' + str(args.r * 100) + '%-' + str( args.noise_mode) else: title = 'Epoch-' + str(epoch) + ' ' + 'Round-' + str(rou) + ': ' + str(args.dataset) + ' ' + str( args.r * 100) + '%-' + str(args.noise_mode) if type(loss) is not np.ndarray: loss = loss.numpy() if type(loss2) is not np.ndarray: loss2 = loss2.numpy() fig = plt.figure() plt.subplot(121) sns.set(style='whitegrid') sns.distplot(loss[noisy_index], color="red", rug=False, kde=False, label="incorrect", hist_kws={"color": "r", "alpha": 0.5}) sns.distplot(loss[clean_index], color="skyblue", rug=False, kde=False, label="correct", hist_kws={"color": "b", "alpha": 0.5}) plt.title('Model_1', fontsize=32) plt.xlabel('Normalized loss', fontsize=32) plt.ylabel('Sample number', fontsize=32) plt.tick_params(labelsize=32) plt.legend(loc='upper right', prop={'size': 24}) plt.subplot(122) sns.set(style='whitegrid') sns.distplot(loss2[noisy_index], color="red", rug=False, kde=False, label="incorrect", hist_kws={"color": "r", "alpha": 0.5}) sns.distplot(loss2[clean_index], color="skyblue", rug=False, kde=False, label="correct", hist_kws={"color": "b", "alpha": 0.5}) plt.title('Model_2', fontsize=32) plt.xlabel('Normalized loss', fontsize=32) plt.ylabel('Sample number', fontsize=32) plt.tick_params(labelsize=32) plt.legend(loc='upper right', prop={'size': 24}) # plt.tight_layout() if g_file: plt.savefig('./figure_his/{0}.tif'.format(filename + model_name), bbox_inches='tight', dpi=300) # plt.show() plt.close() def loss_dist_plot_real(loss, epoch, rou=None, g_file=True, model_name=''): """ plot the loss distribution :param loss: the list contains the loss per sample :param noisy_index: contains the indices of real noisy label :param clean_index: contains the indices of real clean label :param filename: the generated pdf file name :param title: the figure title :param g_file: whether to generate the pdf figure file :return: None """ filename = str(args.dataset) + '_' + str(args.noise_mode) + '_' + str(args.r) + '_epoch=' + str(epoch) if rou is None: title = 'Epoch-' + str(epoch) + ': ' + str(args.dataset) + ' ' + str(args.r * 100) + '%-' + str(args.noise_mode) else: title = 'Epoch-' + str(epoch) + ' ' + 'Round-' + str(rou) + ': ' + str(args.dataset) + ' ' + str(args.r * 100) + '%-' + str(args.noise_mode) if type(loss) is not np.ndarray: loss= loss.numpy() sns.set(style='whitegrid') gmm = GaussianMixture(n_components=2, max_iter=20, tol=1e-2, random_state=0, reg_covar=5e-4) loss = np.reshape(loss, (-1, 1)) gmm.fit(loss) # fit the loss # plot resulting fit x_range = np.linspace(0, 1, 1000) pdf = np.exp(gmm.score_samples(x_range.reshape(-1, 1))) responsibilities = gmm.predict_proba(x_range.reshape(-1, 1)) pdf_individual = responsibilities * pdf[:, np.newaxis] plt.hist(loss, bins=60, density=True, histtype='bar', alpha=0.3) plt.plot(x_range, pdf, '-k', label='Mixture') plt.plot(x_range, pdf_individual, '--', label='Component') plt.legend() # plt.tight_layout() plt.title(title, fontsize=32) plt.xlabel('Normalized loss', fontsize=32) plt.ylabel('Estimated PDF', fontsize=32) plt.tick_params(labelsize=32) plt.legend(loc='upper right', prop={'size': 22}) if g_file: plt.savefig('./figure_his/{0}.tif'.format(filename+model_name), bbox_inches='tight', dpi=300) #plt.show() plt.close() def FedAvg(w): w_avg = copy.deepcopy(w[0]) for k in w_avg.keys(): for i in range(1, len(w)): w_avg[k] += w[i][k] # 只考虑iid noise的话，每个client训练样本数一样，所以不用做nk/n w_avg[k] = torch.div(w_avg[k], len(w)) return w_avg if os.path.exists('checkpoint') == False: os.mkdir('checkpoint') print("新建日志文件夹") stats_log = open('./checkpoint/single_%s_%.1f_%s_%d' % (args.dataset, args.r, args.noise_mode, int(soft_mix_warm)) + '_stats.txt', 'w') test_log = open('./checkpoint/single_%s_%.1f_%s_%d' % (args.dataset, args.r, args.noise_mode, int(soft_mix_warm)) + '_acc.txt', 'w') warm_up = 10 dmix_epoch = 150 args.num_epochs = dmix_epoch + 150 # 第6页提及的warm up的epoch if args.dataset == 'cifar10': warm_up = 10 dmix_epoch = 150 args.num_epochs = dmix_epoch + 50 elif args.dataset == 'cifar100': warm_up = 30 dmix_epoch = 150 args.num_epochs = dmix_epoch + 50 loader = dataloader.cifar_dataloader(args.dataset, r=args.r, noise_mode=args.noise_mode, batch_size=args.batch_size, num_workers=0, root_dir=args.data_path, log=stats_log, noise_file='%s/%.1f_%s.json' % (args.data_path, args.r, args.noise_mode)) print('| Building net') net1 = create_model() net2 = create_model() cudnn.benchmark = True criterion = SemiLoss() optimizer1 = optim.SGD(net1.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4) optimizer2 = optim.SGD(net2.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4) CE = nn.CrossEntropyLoss(reduction='none') CEloss = nn.CrossEntropyLoss() if args.noise_mode == 'asym': # 本文第一个问题，对于非对称和对称需要不同措施，这很不适用 # 其次本文在不同步骤中噪声数据处理措施很凌乱 conf_penalty = NegEntropy() all_loss = [[], []] # save the history of losses from two networks local_round = 5 first = True balance_crit = 'median' exp_path = './checkpoint/single_%s_%.1f_%s_double_m2_' % (args.dataset, args.r, args.noise_mode) save_clean_idx = exp_path + "clean_idx.npy" boot_loader = None w_glob = None if args.r == 0.9: args.p_threshold = 0.6 best_en_acc = 0. best_gl_acc = 0. resume_epoch = 0 if resume_epoch > 0: snapLast = exp_path + str(resume_epoch-1) + "_global_model.pth" global_state = torch.load(snapLast) # 先更新还是后跟新 w_glob = global_state net1.load_state_dict(global_state) net2.load_state_dict(global_state) for epoch in range(resume_epoch, args.num_epochs + 1): test_loader = loader.run('test') eval_loader = loader.run('eval_train') lr = args.lr if epoch >= dmix_epoch: lr /= 10 for param_group in optimizer1.param_groups: param_group['lr'] = lr for param_group in optimizer2.param_groups: param_group['lr'] = lr noise_ind, clean_ind = eval_loader.dataset.if_noise() print(len(np.where(np.array(eval_loader.dataset.noise_label) != np.array(eval_loader.dataset.clean_label))[0]) / len(eval_loader.dataset.clean_label)) local_weights = [] if epoch < warm_up: # 考虑warm up时是否需要merge warmup_trainloader = loader.run('warmup') print('Warmup Net1') warmup(epoch, net1, optimizer1, warmup_trainloader) print('\nWarmup Net2') warmup(epoch, net2, optimizer2, warmup_trainloader) if epoch == (warm_up-1): snapLast = exp_path+str(epoch) + "_1_model.pth" torch.save(net1.state_dict(), snapLast) snapLast = exp_path+str(epoch) + "_2_model.pth" torch.save(net1.state_dict(), snapLast) local_weights.append(net1.state_dict()) local_weights.append(net2.state_dict()) w_glob = FedAvg(local_weights) else: if epoch != warm_up: net1.load_state_dict(w_glob) net2.load_state_dict(w_glob) for rou in range(local_round): prob1, all_loss[0], loss1, mask1, f_G1 = eval_train(net1, all_loss[0]) prob2, all_loss[1], loss2, mask2, f_G2 = eval_train(net2, all_loss[1]) # 加载完global后第一次评估 if rou == 0: # plotHistogram(np.array(loss1), np.array(loss2), noise_ind, clean_ind, epoch, rou, args.r) loss_dist_plot(loss1, noise_ind, clean_ind, epoch, model_name='model_1') # loss_dist_plot_real(loss1, epoch, model_name='model_1') if rou == local_round-1: plotHistogram(np.array(loss1), np.array(loss2), noise_ind, clean_ind, epoch, rou, args.r) # pred1 = (prob1 > args.p_threshold) & (mask1 != 0) # pred2 = (prob2 > args.p_threshold) & (mask2 != 0) pred1 = (prob1 > args.p_threshold) pred2 = (prob2 > args.p_threshold) non_zero_idx = pred1.nonzero()[0].tolist() aaa = len(non_zero_idx) if balance_crit == "max" or balance_crit == "min" or balance_crit == "median": num_clean_per_class = np.zeros(args.num_class) target_label = np.array(eval_loader.dataset.noise_label)[non_zero_idx] for i in range(args.num_class): idx_class = np.where(target_label == i)[0] num_clean_per_class[i] = len(idx_class) if balance_crit == "median": num_samples2select_class = np.median(num_clean_per_class) for i in range(args.num_class): idx_class = np.where(np.array(eval_loader.dataset.noise_label) == i)[0] cur_num = num_clean_per_class[i] idx_class2 = non_zero_idx if num_samples2select_class > cur_num: remian_idx = list(set(idx_class.tolist()) - set(idx_class2)) idx = list(range(len(remian_idx))) random.shuffle(idx) num_app = int(num_samples2select_class - cur_num) idx = idx[:num_app] for j in idx: non_zero_idx.append(remian_idx[j]) non_zero_idx = np.array(non_zero_idx).reshape(-1, ) bbb = len(non_zero_idx) num_per_class2 = [] for i in range(max(eval_loader.dataset.noise_label)): temp = np.where(np.array(eval_loader.dataset.noise_label)[non_zero_idx.tolist()] == i)[0] num_per_class2.append(len(temp)) print('\npred1 appended num per class:', num_per_class2, aaa, bbb) idx_per_class = np.zeros_like(pred1).astype(bool) for i in non_zero_idx: idx_per_class[i] = True pred1 = idx_per_class non_aaa = pred1.nonzero()[0].tolist() assert len(non_aaa) == len(non_zero_idx) non_zero_idx2 = pred2.nonzero()[0].tolist() aaa = len(non_zero_idx2) if balance_crit == "max" or balance_crit == "min" or balance_crit == "median": num_clean_per_class = np.zeros(args.num_class) target_label = np.array(eval_loader.dataset.noise_label)[non_zero_idx2] for i in range(args.num_class): idx_class = np.where(target_label == i)[0] num_clean_per_class[i] = len(idx_class) if balance_crit == "median": num_samples2select_class = np.median(num_clean_per_class) for i in range(args.num_class): idx_class = np.where(np.array(eval_loader.dataset.noise_label) == i)[0] cur_num = num_clean_per_class[i] idx_class2 = non_zero_idx2 if num_samples2select_class > cur_num: remian_idx = list(set(idx_class.tolist()) - set(idx_class2)) idx = list(range(len(remian_idx))) random.shuffle(idx) num_app = int(num_samples2select_class - cur_num) idx = idx[:num_app] for j in idx: non_zero_idx2.append(remian_idx[j]) non_zero_idx2 = np.array(non_zero_idx2).reshape(-1, ) bbb = len(non_zero_idx2) num_per_class2 = [] for i in range(max(eval_loader.dataset.noise_label)): temp = np.where(np.array(eval_loader.dataset.noise_label)[non_zero_idx2.tolist()] == i)[0] num_per_class2.append(len(temp)) print('\npred2 appended num per class:', num_per_class2, aaa, bbb) idx_per_class2 = np.zeros_like(pred2).astype(bool) for i in non_zero_idx2: idx_per_class2[i] = True pred2 = idx_per_class2 non_aaa = pred2.nonzero()[0].tolist() assert len(non_aaa) == len(non_zero_idx2) correct_num = len(pred1.nonzero()[0]) eval_loader.dataset.if_noise(pred1) eval_loader.dataset.if_noise(pred2) print(f'round={rou}/{local_round}, dmix selection, Train Net1') # prob2就是先验概率wi,通过GMM拟合出来的，大于阈值就认为是clean，否则noisy labeled_trainloader, unlabeled_trainloader = loader.run('train', pred2, prob2) # co-divide train(epoch, net1, net2, optimizer1, labeled_trainloader, unlabeled_trainloader) # train net1 print(f'\nround={rou}/{local_round}, dmix selection, Train Net2') labeled_trainloader, unlabeled_trainloader = loader.run('train', pred1, prob1) # co-divide train(epoch, net2, net1, optimizer2, labeled_trainloader, unlabeled_trainloader) # train net2 local_weights.append(net1.state_dict()) local_weights.append(net2.state_dict()) w_glob = FedAvg(local_weights) if epoch % 5 == 0: snapLast = exp_path + str(epoch) + "_global_model.pth" torch.save(w_glob, snapLast) best_en_acc = test(epoch, net1, net2, best_en_acc) best_gl_acc= test(epoch, net1, net2, best_gl_acc, w_glob)