Add files via upload

PreResNet.py

@@ -0,0 +1,184 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch.autograd import Variable
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(in_planes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion*planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion*planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class PreActBlock(nn.Module):
+    '''Pre-activation version of the BasicBlock.'''
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(PreActBlock, self).__init__()
+        self.bn1 = nn.BatchNorm2d(in_planes)
+        self.conv1 = conv3x3(in_planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion*planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(x))
+        shortcut = self.shortcut(out)
+        out = self.conv1(out)
+        out = self.conv2(F.relu(self.bn2(out)))
+        out += shortcut
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(self.expansion*planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion*planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion*planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = F.relu(self.bn2(self.conv2(out)))
+        out = self.bn3(self.conv3(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+
+class PreActBottleneck(nn.Module):
+    '''Pre-activation version of the original Bottleneck module.'''
+    expansion = 4
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(PreActBottleneck, self).__init__()
+        self.bn1 = nn.BatchNorm2d(in_planes)
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, self.expansion*planes, kernel_size=1, bias=False)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion*planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(x))
+        shortcut = self.shortcut(out)
+        out = self.conv1(out)
+        out = self.conv2(F.relu(self.bn2(out)))
+        out = self.conv3(F.relu(self.bn3(out)))
+        out += shortcut
+        return out
+
+
+class ResNet(nn.Module):
+    def __init__(self, block, num_blocks, num_classes=10):
+        super(ResNet, self).__init__()
+        self.in_planes = 64
+
+        self.conv1 = conv3x3(3,64)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
+        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
+        self.linear = nn.Linear(512*block.expansion, num_classes)
+
+    def _make_layer(self, block, planes, num_blocks, stride):
+        strides = [stride] + [1]*(num_blocks-1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.in_planes, planes, stride))
+            self.in_planes = planes * block.expansion
+        return nn.Sequential(*layers)
+
+    def forward(self, x, lin=0, lout=5, feat_out=False):
+        out = x
+        if lin < 1 and lout > -1:
+            out = self.conv1(out)
+            out = self.bn1(out)
+            out = F.relu(out)
+        if lin < 2 and lout > 0:
+            out = self.layer1(out)
+        if lin < 3 and lout > 1:
+            out = self.layer2(out)
+        if lin < 4 and lout > 2:
+            out = self.layer3(out)
+        if lin < 5 and lout > 3:
+            out = self.layer4(out)
+        if lout > 4:
+            out = F.avg_pool2d(out, 4)
+            feat = out.view(out.size(0), -1)
+            out = self.linear(feat)
+        if feat_out:
+            return out, feat
+        else:
+            return out
+
+
+def ResNet18(num_classes=10):
+    return ResNet(PreActBlock, [2,2,2,2], num_classes=num_classes)
+
+def ResNet34(num_classes=10):
+    return ResNet(BasicBlock, [3,4,6,3], num_classes=num_classes)
+
+def ResNet50(num_classes=10):
+    return ResNet(Bottleneck, [3,4,6,3], num_classes=num_classes)
+
+def ResNet101(num_classes=10):
+    return ResNet(Bottleneck, [3,4,23,3], num_classes=num_classes)
+
+def ResNet152(num_classes=10):
+    return ResNet(Bottleneck, [3,8,36,3], num_classes=num_classes)
+
+
+def test():
+    net = ResNet18()
+    y = net(Variable(torch.randn(1,3,32,32)))
+    print(y.size())

Train_animal10N.py

@@ -0,0 +1,487 @@
+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 torchvision
+import torchvision.models as models
+from models.CNN import CNN
+import random
+import os
+import argparse
+import numpy as np
+import dataloader_animal10N as animal_dataloader
+from sklearn.mixture import GaussianMixture
+import copy
+
+parser = argparse.ArgumentParser(description='PyTorch Clothing1M Training')
+parser.add_argument('--batch_size', default=128, type=int, help='train batchsize')
+parser.add_argument('--lr', '--learning_rate', default=0.01, type=float, help='initial learning rate')
+parser.add_argument('--alpha', default=4, type=float, help='parameter for Beta')
+parser.add_argument('--lambda_u', default=0, 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('--id', default='animal10N')
+# parser.add_argument('--data_path', default='E:/Dataset_All/clothing1M/images', type=str, help='path to dataset')
+parser.add_argument('--data_path', default='C:/Users/Administrator/Desktop/DatasetAll/Animal-10N', type=str, help='path to dataset')
+parser.add_argument('--seed', default=123)
+parser.add_argument('--gpuid', default=0, type=int)
+parser.add_argument('--num_class', default=10, type=int)
+# parser.add_argument('--num_batches', default=1000, type=int)
+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)
+
+
+# Training
+def train(epoch, net, net2, optimizer, labeled_trainloader, unlabeled_trainloader):
+    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
+    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
+        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 = net(inputs_u)
+            outputs_u12 = net(inputs_u2)
+            outputs_u21 = net2(inputs_u)
+            outputs_u22 = net2(inputs_u2)
+
+            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
+
+            targets_u = ptu / ptu.sum(dim=1, keepdim=True)  # normalize
+            targets_u = targets_u.detach()
+
+            # label refinement of labeled samples
+            outputs_x = net(inputs_x)
+            outputs_x2 = net(inputs_x2)
+
+            px = (torch.softmax(outputs_x, dim=1) + torch.softmax(outputs_x2, dim=1)) / 2
+            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()
+
+            # mixmatch
+        l = np.random.beta(args.alpha, args.alpha)
+        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)
+
+        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]
+
+        mixed_input = l * input_a[:batch_size * 2] + (1 - l) * input_b[:batch_size * 2]
+        mixed_target = l * target_a[:batch_size * 2] + (1 - l) * target_b[:batch_size * 2]
+
+        logits = net(mixed_input)
+
+        Lx = -torch.mean(torch.sum(F.log_softmax(logits, dim=1) * mixed_target, dim=1))
+
+        # regularization
+        prior = torch.ones(args.num_class) / args.num_class
+        prior = prior.cuda()
+        pred_mean = torch.softmax(logits, dim=1).mean(0)
+        penalty = torch.sum(prior * torch.log(prior / pred_mean))
+
+        loss = Lx + penalty
+
+        # compute gradient and do SGD step
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        sys.stdout.write('\r')
+        sys.stdout.write('Animal10N | Epoch [%3d/%3d] Iter[%3d/%3d]\t  Labeled loss: %.4f '
+                         % (epoch, args.num_epochs, batch_idx + 1, num_iter, Lx.item()))
+        sys.stdout.flush()
+
+
+def warmup(net, optimizer, dataloader):
+    net.train()
+    num_batches = 50000/args.batch_size
+    for batch_idx, (inputs, labels, path) in enumerate(dataloader):
+        inputs, labels = inputs.cuda(), labels.cuda()
+        optimizer.zero_grad()
+        outputs = net(inputs)
+        loss = CEloss(outputs, labels)
+
+        penalty = conf_penalty(outputs)
+        L = loss + penalty
+        L.backward()
+        optimizer.step()
+
+        sys.stdout.write('\r')
+        sys.stdout.write('|Warm-up: Iter[%3d/%3d]\t CE-loss: %.4f  Conf-Penalty: %.4f'
+                         % (2*(batch_idx + 1), num_batches, loss.item(), penalty.item()))
+        sys.stdout.flush()
+
+
+def val(net, val_loader, best_acc, w_glob=None):
+    net.eval()
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for batch_idx, (inputs, targets) in enumerate(val_loader):
+            inputs, targets = inputs.cuda(), targets.cuda()
+            outputs = net(inputs)
+            _, predicted = torch.max(outputs, 1)
+
+            total += targets.size(0)
+            correct += predicted.eq(targets).cpu().sum().item()
+    acc = 100. * correct / total
+    print("\n| Validation\t Net%d  Acc: %.2f%%" % (k, acc))
+    if acc > best_acc[k - 1]:
+        best_acc[k - 1] = acc
+        print('| Saving Best Net%d ...' % k)
+        save_point = './checkpoint/%s_net%d.pth.tar' % (args.id, k)
+        torch.save(net.state_dict(), save_point)
+    return acc
+
+
+def test(epoch, net1, net2, test_loader, best_acc, w_glob=None):
+    if w_glob is None:
+        net1.eval()
+        net2.eval()
+        correct = 0
+        correct2 = 0
+        correct1 = 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)
+                _, predicted1 = torch.max(outputs1, 1)
+                _, predicted2 = torch.max(outputs2, 1)
+
+                total += targets.size(0)
+                correct += predicted.eq(targets).cpu().sum().item()
+                correct1 += predicted1.eq(targets).cpu().sum().item()
+                correct2 += predicted2.eq(targets).cpu().sum().item()
+        acc = 100. * correct / total
+        acc1 = 100. * correct / total
+        acc2 = 100. * correct / total
+        if best_acc < acc:
+            best_acc = acc
+        print(
+            "\n| Ensemble network Test Epoch #%d\t Accuracy: %.2f, Accuracy1: %.2f, Accuracy2: %.2f, best_acc: %.2f%%\n" % (
+            epoch, acc, acc1, acc2, best_acc))
+        log.write('ensemble_Epoch:%d   Accuracy:%.2f, Accuracy1: %.2f, Accuracy2: %.2f, best_acc: %.2f\n' % (
+        epoch, acc, acc1, acc2, best_acc))
+        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))
+        log.write('global_Epoch:%d   Accuracy:%.2f, best_acc: %.2f\n' % (epoch, acc, best_acc))
+        log.flush()
+        #   恢复权重
+        net1.load_state_dict(net1_w_bak)
+    return best_acc
+
+
+def eval_train(epoch, model):
+    model.eval()
+    num_samples = eval_loader.dataset.__len__()
+    losses = torch.zeros(num_samples)
+    paths = []
+    n = 0
+    with torch.no_grad():
+        for batch_idx, (inputs, targets, path) in enumerate(eval_loader):
+            inputs, targets = inputs.cuda(), targets.cuda()
+            outputs = model(inputs)
+            loss = CE(outputs, targets)
+            for b in range(inputs.size(0)):
+                losses[n] = loss[b]
+                paths.append(path[b])
+                n += 1
+            sys.stdout.write('\r')
+            sys.stdout.write('| Evaluating loss Iter %3d\t' % (batch_idx))
+            sys.stdout.flush()
+
+    losses = (losses - losses.min()) / (losses.max() - losses.min())
+    losses = losses.reshape(-1, 1)
+    gmm = GaussianMixture(n_components=2, max_iter=10, reg_covar=5e-4, tol=1e-2)
+    gmm.fit(losses)
+    prob = gmm.predict_proba(losses)
+    prob = prob[:, gmm.means_.argmin()]
+    return prob, paths
+
+
+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():
+    use_cnn = False
+    if use_cnn:
+        model = CNN()
+        model = model.cuda()
+    else:
+        model = models.vgg19_bn(pretrained=False)
+        model.classifier._modules['6'] = nn.Linear(4096, 10)
+        model = model.cuda()
+    return model
+
+
+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
+
+
+log = open('./checkpoint/%s.txt' % args.id, 'w')
+log.flush()
+
+loader = animal_dataloader.animal_dataloader(root=args.data_path, batch_size=args.batch_size, num_workers=0)
+
+print('| Building net')
+net1 = create_model()
+net2 = create_model()
+cudnn.benchmark = True
+
+optimizer1 = optim.SGD(net1.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-3)
+optimizer2 = optim.SGD(net2.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-3)
+
+CE = nn.CrossEntropyLoss(reduction='none')
+CEloss = nn.CrossEntropyLoss()
+conf_penalty = NegEntropy()
+
+local_round = 5
+balance_crit = 'median'  # 'median'
+exp_path = './checkpoint/c2mt_animal10N'
+
+boot_loader = None
+w_glob = None
+best_en_acc = 0.
+best_gl_acc = 0.
+resume_epoch = 0
+warm_up = 10
+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)
+
+# if True:
+#     snapLast = exp_path + "0_1_model.pth"
+#     global_state = torch.load(snapLast)
+#     net1.load_state_dict(global_state)
+#     snapLast = exp_path + "0_2_model.pth"
+#     global_state = torch.load(snapLast)
+#     net2.load_state_dict(global_state)
+#     test_loader = loader.run('test')
+#     best_en_acc = test(0, net1, net2, test_loader, best_en_acc)
+
+for epoch in range(resume_epoch, args.num_epochs + 1):
+    lr = args.lr
+    if 100 <= epoch < 150:
+        lr /= 10
+    elif epoch >= 150:
+        lr /= 10
+    # if 15 <= epoch:
+    #     lr /= 2
+    for param_group in optimizer1.param_groups:
+        param_group['lr'] = lr
+    for param_group in optimizer2.param_groups:
+        param_group['lr'] = lr
+
+    local_weights = []
+    if epoch < warm_up:  # warm up
+        train_loader = loader.run('warmup')
+        print('Warmup Net1')
+        warmup(net1, optimizer1, train_loader)
+        train_loader = loader.run('warmup')
+        print('\nWarmup Net2')
+        warmup(net2, optimizer2, train_loader)
+        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):
+            print('\n==== net 1 evaluate next epoch training data loss ====')
+            eval_loader = loader.run('eval_train')  # evaluate training data loss for next epoch
+            prob1, paths1 = eval_train(epoch, net1)
+            print('\n==== net 2 evaluate next epoch training data loss ====')
+            eval_loader = loader.run('eval_train')
+            prob2, paths2 = eval_train(epoch, net2)
+
+            pred1 = (prob1 > args.p_threshold)  # divide dataset
+            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)
+                ppp = np.array(paths1)[non_zero_idx].tolist()
+                target_label = np.array([eval_loader.dataset.train_labels[it] for it in ppp])
+                # target_label = np.array(eval_loader.dataset.train_labels[paths1])[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 == "max":
+                    num_samples2select_class = np.max(num_clean_per_class)
+                elif balance_crit == "min":
+                    num_samples2select_class = np.min(num_clean_per_class)
+                elif 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.train_labels[it] for it in paths1]) == 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(10):
+                temp = \
+                np.where(np.array([eval_loader.dataset.train_labels[it] for it in paths1])[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)
+                ppp = np.array(paths2)[non_zero_idx].tolist()
+                target_label = np.array([eval_loader.dataset.train_labels[it] for it in ppp])
+                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 == "max":
+                    num_samples2select_class = np.max(num_clean_per_class)
+                elif balance_crit == "min":
+                    num_samples2select_class = np.min(num_clean_per_class)
+                elif 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.train_labels[it] for it in paths1]) == 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(10):
+                temp = np.where(
+                    np.array([eval_loader.dataset.train_labels[it] for it in paths1])[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)
+
+            print(f'round={rou}/{local_round}, dmix selection, Train Net1')
+            labeled_trainloader, unlabeled_trainloader = loader.run('train', pred2, prob2, paths=paths2)  # 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, paths=paths1)  # co-divide
+            train(epoch, net2, net1, optimizer2, labeled_trainloader, unlabeled_trainloader)  # train net2
+
+            test_loader = loader.run('test')
+            if rou != local_round-1:
+                best_en_acc = test(epoch, net1, net2, test_loader, best_en_acc)
+            # best_gl_acc = test(epoch, net1, net2, test_loader, best_gl_acc, w_glob=w_glob)
+
+        print(f'c2m, get global network\n')
+        local_weights.append(net1.state_dict())
+        local_weights.append(net2.state_dict())
+        w_glob = FedAvg(local_weights)
+        if epoch % 1 == 0:
+            snapLast = exp_path + str(epoch) + "_global_model.pth"
+            torch.save(w_glob, snapLast)
+
+    test_loader = loader.run('test')
+    best_en_acc = test(epoch, net1, net2, test_loader, best_en_acc)
+    best_gl_acc = test(epoch, net1, net2, test_loader, best_gl_acc, w_glob=w_glob)
+

dataloader_animal10N.py

@@ -0,0 +1,200 @@
+from torch.utils.data import Dataset, DataLoader
+import torchvision.transforms as transforms
+import random
+import numpy as np
+from PIL import Image
+import json
+import torch
+import os
+import matplotlib
+
+def unpickle(file):
+    fo = open(file, 'rb').read()
+    size = 64 * 64 * 3 + 1
+    for i in range(50000):
+        arr = np.fromstring(fo[i * size:(i + 1) * size], dtype=np.uint8)
+        lab = np.identity(10)[arr[0]]
+        img = arr[1:].reshape((3, 64, 64)).transpose((1, 2, 0))
+    return img, lab
+
+class animal_dataset(Dataset):
+    def __init__(self, root, transform, mode, pred=[], path=[], probability=[], num_class=10):
+
+        self.root = root
+        self.transform = transform
+        self.mode = mode
+
+        self.train_dir = root + '/training/'
+        self.test_dir = root + '/testing/'
+        train_imgs = os.listdir(self.train_dir)
+        test_imgs = os.listdir(self.test_dir)
+        self.test_data = []
+        self.test_labels = []
+        noise_file1 = './training_batch.json'
+        noise_file2 = './testing_batch.json'
+        if mode == 'test':
+            if os.path.exists(noise_file2):
+                dict = json.load(open(noise_file2, "r"))
+                self.test_labels = dict['data']
+                self.test_data = dict['label']
+            else:
+                for img in test_imgs:
+                    self.test_data.append(self.test_dir+img)
+                    self.test_labels.append(int(img[0]))
+                dicts = {}
+                dicts['data'] = self.test_data
+                dicts['label'] = self.test_labels
+                # json.dump(dicts, open(noise_file2, "w"))
+        else:
+            if os.path.exists(noise_file1):
+                dict = json.load(open(noise_file1, "r"))
+                train_data = dict['data']
+                train_labels = dict['label']
+            else:
+                train_data = []
+                train_labels = {}
+                for img in train_imgs:
+                    img_path = self.train_dir+img
+                    train_data.append(img_path)
+                    train_labels[img_path] = (int(img[0]))
+                dicts = {}
+                dicts['data'] = train_data
+                dicts['label'] = train_labels
+                # json.dump(dicts, open(noise_file1, "w"))
+            if self.mode == "all":
+                self.train_data = train_data
+                self.train_labels = train_labels
+            elif self.mode == "labeled":
+                pred_idx = pred.nonzero()[0]
+                train_img = path
+                self.train_data = [train_img[i] for i in pred_idx]
+                self.probability = probability[pred_idx]
+                # self.train_labels = train_labels[pred_idx]
+                self.train_labels = train_labels
+                print("%s data has a size of %d" % (self.mode, len(self.train_data)))
+            elif self.mode == "unlabeled":
+                pred_idx = (1 - pred).nonzero()[0]
+                train_img = path
+                self.train_data = [train_img[i] for i in pred_idx]
+                self.probability = probability[pred_idx]
+                # self.train_labels = train_labels[pred_idx]
+                print("%s data has a size of %d" % (self.mode, len(self.train_data)))
+                self.train_labels = train_labels
+
+    def __getitem__(self, index):
+        if self.mode == 'labeled':
+            img_path = self.train_data[index]
+            target = self.train_labels[img_path]
+            prob = self.probability[index]
+            image = Image.open(img_path).convert('RGB')
+            img1 = self.transform(image)
+            img2 = self.transform(image)
+            return img1, img2, target, prob
+        elif self.mode == 'unlabeled':
+            img_path = self.train_data[index]
+            image = Image.open(img_path).convert('RGB')
+            img1 = self.transform(image)
+            img2 = self.transform(image)
+            return img1, img2
+        elif self.mode == 'all':
+            img_path = self.train_data[index]
+            target = self.train_labels[img_path]
+            image = Image.open(img_path).convert('RGB')
+            img = self.transform(image)
+            return img, target,img_path
+        elif self.mode == 'test':
+            img_path = self.test_data[index]
+            target = self.test_labels[index]
+            image = Image.open(img_path).convert('RGB')
+            img = self.transform(image)
+            return img, target
+
+    def __len__(self):
+        if self.mode == 'test':
+            return len(self.test_data)
+        else:
+            return len(self.train_data)
+
+
+class animal_dataloader():
+    def __init__(self, root='E:/2_Dataset_All/Animal-10N', batch_size=32, num_workers=0):
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.root = root
+
+        self.transform_train = transforms.Compose([
+            transforms.Resize(64),
+            transforms.RandomCrop(64),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            transforms.Normalize((0.6959, 0.6537, 0.6371), (0.3113, 0.3192, 0.3214)),
+        ])
+        self.transform_test = transforms.Compose([
+            # transforms.Resize(64),
+            # transforms.CenterCrop(64),
+            transforms.ToTensor(),
+            transforms.Normalize((0.6959, 0.6537, 0.6371), (0.3113, 0.3192, 0.3214)),
+        ])
+
+    def run(self, mode, pred=[], prob=[], paths=[]):
+        if mode == 'warmup':
+            warmup_dataset = animal_dataset(self.root, transform=self.transform_train, mode='all')
+            warmup_loader = DataLoader(
+                dataset=warmup_dataset,
+                batch_size=self.batch_size * 2,
+                shuffle=True,
+                num_workers=self.num_workers,
+                pin_memory=True)
+            return warmup_loader
+        elif mode == 'train':
+            labeled_dataset = animal_dataset(self.root, transform=self.transform_train, mode='labeled', pred=pred, path=paths,
+                                             probability=prob)
+            labeled_loader = DataLoader(
+                dataset=labeled_dataset,
+                batch_size=self.batch_size,
+                shuffle=True,
+                num_workers=self.num_workers,
+                pin_memory=True)
+            unlabeled_dataset = animal_dataset(self.root, transform=self.transform_train, mode='unlabeled', pred=pred,path=paths,
+                                               probability=prob)
+            unlabeled_loader = DataLoader(
+                dataset=unlabeled_dataset,
+                batch_size=int(self.batch_size),
+                shuffle=True,
+                num_workers=self.num_workers,
+                pin_memory=True)
+            return labeled_loader, unlabeled_loader
+        elif mode == 'eval_train':
+            eval_dataset = animal_dataset(self.root, transform=self.transform_test, mode='all')
+            eval_loader = DataLoader(
+                dataset=eval_dataset,
+                batch_size=self.batch_size,
+                shuffle=False,
+                num_workers=self.num_workers,
+                pin_memory=True)
+            return eval_loader
+        elif mode == 'test':
+            test_dataset = animal_dataset(self.root, transform=self.transform_test, mode='test')
+            test_loader = DataLoader(
+                dataset=test_dataset,
+                batch_size=1000,
+                shuffle=False,
+                num_workers=self.num_workers,
+                pin_memory=True)
+            return test_loader
+
+# if __name__ == '__main__':
+#     loader = animal_dataloader()
+#     train_loader = loader.run('warmup')
+#     import matplotlib.pyplot as plt
+#     for batch_idx, (inputs, labels, idx, img_path) in enumerate(train_loader):
+#         print(img_path[0])
+#         plt.figure(dpi=300)
+#         # plt.imshow(inputs[0])
+#         plt.imshow(inputs[0].reshape(64, 64, 3))
+#         plt.show()
+#         plt.close()
+#         print(inputs.shape())
+#         print(idx)
+#         print(labels, len(labels))
+#     # print(train_loader.dataset.__len__())

dataloader_cifar.py

@@ -0,0 +1,276 @@
+from torch.utils.data import Dataset, DataLoader
+import torchvision.transforms as transforms
+import random
+import numpy as np
+from PIL import Image
+import json
+import os
+import torch
+from torchnet.meter import AUCMeter
+
+            
+def unpickle(file):
+    import _pickle as cPickle
+    with open(file, 'rb') as fo:
+        dict = cPickle.load(fo, encoding='latin1')
+    return dict
+
+class cifar_dataset(Dataset): 
+    def __init__(self, dataset, r, noise_mode, root_dir, transform, mode, noise_file='', pred=[], probability=[], log='', clean_idx=[], test_form = None):
+        
+        self.r = r # noise ratio
+        self.transform = transform
+        self.test_form = test_form
+        self.mode = mode  
+        self.transition = {0:0,2:0,4:7,7:7,1:1,9:1,3:5,5:3,6:6,8:8} # class transition for asymmetric noise
+        self.noise_file = noise_file
+     
+        if self.mode=='test':
+            if dataset=='cifar10':                
+                test_dic = unpickle('%s/test_batch'%root_dir)
+                self.test_data = test_dic['data']
+                self.test_data = self.test_data.reshape((10000, 3, 32, 32))
+                self.test_data = self.test_data.transpose((0, 2, 3, 1))  
+                self.test_label = test_dic['labels']
+            elif dataset=='cifar100':
+                test_dic = unpickle('%s/test'%root_dir)
+                self.test_data = test_dic['data']
+                self.test_data = self.test_data.reshape((10000, 3, 32, 32))
+                self.test_data = self.test_data.transpose((0, 2, 3, 1))  
+                self.test_label = test_dic['fine_labels']                            
+        else:    
+            train_data=[]
+            train_label=[]
+            if dataset=='cifar10': 
+                for n in range(1,6):
+                    dpath = '%s/data_batch_%d'%(root_dir,n)
+                    data_dic = unpickle(dpath)
+                    train_data.append(data_dic['data'])
+                    train_label = train_label+data_dic['labels']
+                train_data = np.concatenate(train_data)
+            elif dataset=='cifar100':    
+                train_dic = unpickle('%s/train'%root_dir)
+                train_data = train_dic['data']
+                train_label = train_dic['fine_labels']
+            train_data = train_data.reshape((50000, 3, 32, 32))
+            train_data = train_data.transpose((0, 2, 3, 1))
+
+            self.clean_label = np.array(train_label)
+
+            if os.path.exists(noise_file):
+                noise_label = json.load(open(noise_file,"r"))
+            else:    #inject noise   
+                noise_label = []
+                idx = list(range(50000))
+                random.shuffle(idx)
+                num_noise = int(self.r*50000)            
+                noise_idx = idx[:num_noise]
+                for i in range(50000):
+                    if i in noise_idx:
+                        if noise_mode=='sym':
+                            if dataset=='cifar10': 
+                                noiselabel = random.randint(0,9)
+                            elif dataset=='cifar100':    
+                                noiselabel = random.randint(0,99)
+                            noise_label.append(noiselabel)
+                        elif noise_mode=='asym':   
+                            noiselabel = self.transition[train_label[i]]
+                            noise_label.append(noiselabel)                    
+                    else:    
+                        noise_label.append(train_label[i])   
+                print("save noisy labels to %s ..."%noise_file)        
+                json.dump(noise_label,open(noise_file,"w"))       
+            
+            if self.mode == 'all':
+                self.train_data = train_data
+                self.noise_label = np.array(noise_label).astype(np.int64)
+            else:                   
+                if self.mode == "labeled":
+                    pred_idx = pred.nonzero()[0]
+                    self.probability = [probability[i] for i in pred_idx]   
+                    
+                    clean = (np.array(noise_label)==np.array(train_label))                                                       
+                    auc_meter = AUCMeter()
+                    auc_meter.reset()
+                    auc_meter.add(probability,clean)        
+                    auc,_,_ = auc_meter.value()
+                    clean_index = np.where(np.array(noise_label)[pred_idx.tolist()] == np.array(self.clean_label)[pred_idx.tolist()])[0]
+
+                    num_per_class = []
+                    for i in range(max(noise_label)):
+                        temp = np.where(np.array(noise_label)[clean_index.tolist()] == i)[0]
+                        num_per_class.append(len(temp))
+                    num_per_class2 = []
+                    for i in range(max(noise_label)):
+                        temp = np.where(np.array(noise_label)[pred_idx.tolist()] == i)[0]
+                        num_per_class2.append(len(temp))
+                    print('clean num per class:', num_per_class, num_per_class2)
+
+                    log.write('Numer of labeled samples:%d   AUC:%.3f   corrected clean num:%d, uncorrected noisy num:%d\n'
+                              % (pred.sum(), auc, len(clean_index), len(pred_idx) - len(clean_index)))
+                    log.flush()
+                    
+                elif self.mode == "unlabeled":
+                    pred_idx = (1-pred).nonzero()[0]
+                    noise_index = np.where(np.array(noise_label)[pred_idx.tolist()] != np.array(self.clean_label)[pred_idx.tolist()])[0]
+                    log.write('Numer of unlabeled samples:%d   corrected noisy num:%d, uncorrected clean num:%d\n'
+                              % (pred.sum(), len(noise_index), len(pred_idx) - len(noise_index)))
+                    log.flush()
+                elif self.mode == 'boost':
+                    pred_idx = clean_idx
+                
+                self.train_data = train_data[pred_idx]
+                self.noise_label = [noise_label[i] for i in pred_idx]                          
+                print("%s data has a size of %d"%(self.mode,len(self.noise_label)))
+
+    def if_noise(self, pred=None):
+        if pred is None:
+            noise_index = np.where(self.noise_label[:] != self.clean_label[:])[0]
+            clean_index = np.where(self.noise_label[:] == self.clean_label[:])[0]
+            return noise_index, clean_index
+        else:
+            pred_idx1 = pred.nonzero()[0].tolist()
+            clean_index = np.where(np.array(self.noise_label)[pred_idx1] == np.array(self.clean_label)[pred_idx1])[0]
+            pred_idx = (1 - pred).nonzero()[0].tolist()
+            noise_index = np.where(np.array(self.noise_label)[pred_idx] != np.array(self.clean_label)[pred_idx])[0]
+            print(
+                f'选择的非mask样本中正确选取的干净标签数量{len(clean_index)}, 不正确选取的非干净数量{len(pred_idx1) - len(clean_index)}.\t '
+                f'选择的mask样本中正确选取的不干净标签数量{len(noise_index)}, 不正确选取的干净数量{len(pred_idx) - len(noise_index)}')
+            return len(clean_index), (len(pred_idx1) - len(clean_index)), len(noise_index), len(pred_idx) - len(
+                noise_index)
+    def print_noise_rate(self, new_y):
+        temp_y = np.array(new_y.reshape(1, -1).squeeze())
+        clean_index = np.where(temp_y[:] == np.array(self.clean_label)[:])
+        print(f'clean rate is: {len(clean_index[0]) / len(self.clean_label)}')
+
+    def load_train_label(self, new_y):
+        temp_y = np.array(new_y.reshape(1, -1).squeeze()).astype(np.int64)
+        self.noise_label[:] = np.array(temp_y)[:]
+        if os.path.exists(self.noise_file):
+            result = os.path.splitext(self.noise_file)
+            noise_file_temp = result[0]+'_old'+result[1]
+            if not os.path.exists(noise_file_temp):
+                os.rename(self.noise_file, noise_file_temp)
+        #   覆盖原来的noise_file
+        json.dump(self.noise_label.tolist(), open(self.noise_file, "w"))
+
+    def __getitem__(self, index):
+        if self.mode=='labeled':
+            img, target, prob = self.train_data[index], self.noise_label[index], self.probability[index]
+            img = Image.fromarray(img)
+            img1 = self.transform(img) 
+            img2 = self.transform(img) 
+            return img1, img2, target, prob            
+        elif self.mode=='unlabeled':
+            img = self.train_data[index]
+            img = Image.fromarray(img)
+            img1 = self.transform(img) 
+            img2 = self.transform(img) 
+            return img1, img2
+        elif self.mode=='all':
+            img, target = self.train_data[index], self.noise_label[index]
+            img = Image.fromarray(img)
+            img = self.transform(img)            
+            return img, target, index        
+        elif self.mode=='test':
+            img, target = self.test_data[index], self.test_label[index]
+            img = Image.fromarray(img)
+            img = self.transform(img)            
+            return img, target
+        elif self.mode=='boost':
+            img, target = self.train_data[index], self.noise_label[index]
+            img = Image.fromarray(img)
+            img_no_da = self.test_form(img)
+            img = self.transform(img)
+            return img, img_no_da, target, index
+           
+    def __len__(self):
+        if self.mode!='test':
+            return len(self.train_data)
+        else:
+            return len(self.test_data)         
+        
+        
+class cifar_dataloader():  
+    def __init__(self, dataset, r, noise_mode, batch_size, num_workers, root_dir, log, noise_file=''):
+        self.dataset = dataset
+        self.r = r
+        self.noise_mode = noise_mode
+        self.batch_size = batch_size
+        self.num_workers = num_workers
+        self.root_dir = root_dir
+        self.log = log
+        self.noise_file = noise_file
+        if self.dataset=='cifar10':
+            self.transform_train = transforms.Compose([
+                    transforms.RandomCrop(32, padding=4),
+                    transforms.RandomHorizontalFlip(),
+                    transforms.ToTensor(),
+                    transforms.Normalize((0.4914, 0.4822, 0.4465),(0.2023, 0.1994, 0.2010)),
+                ]) 
+            self.transform_test = transforms.Compose([
+                    transforms.ToTensor(),
+                    transforms.Normalize((0.4914, 0.4822, 0.4465),(0.2023, 0.1994, 0.2010)),
+                ])    
+        elif self.dataset=='cifar100':    
+            self.transform_train = transforms.Compose([
+                    transforms.RandomCrop(32, padding=4),
+                    transforms.RandomHorizontalFlip(),
+                    transforms.ToTensor(),
+                    transforms.Normalize((0.507, 0.487, 0.441), (0.267, 0.256, 0.276)),
+                ]) 
+            self.transform_test = transforms.Compose([
+                    transforms.ToTensor(),
+                    transforms.Normalize((0.507, 0.487, 0.441), (0.267, 0.256, 0.276)),
+                ])   
+    def run(self,mode,pred=[],prob=[], clean_idx=[]):
+        if mode=='warmup':
+            all_dataset = cifar_dataset(dataset=self.dataset, noise_mode=self.noise_mode, r=self.r, root_dir=self.root_dir, transform=self.transform_train, mode="all",noise_file=self.noise_file)                
+            trainloader = DataLoader(
+                dataset=all_dataset, 
+                batch_size=self.batch_size*2,
+                shuffle=True,
+                num_workers=self.num_workers)             
+            return trainloader
+                                     
+        elif mode=='train':
+            labeled_dataset = cifar_dataset(dataset=self.dataset, noise_mode=self.noise_mode, r=self.r, root_dir=self.root_dir, transform=self.transform_train, mode="labeled", noise_file=self.noise_file, pred=pred, probability=prob,log=self.log)              
+            labeled_trainloader = DataLoader(
+                dataset=labeled_dataset, 
+                batch_size=self.batch_size,
+                shuffle=True,
+                num_workers=self.num_workers)   
+            
+            unlabeled_dataset = cifar_dataset(dataset=self.dataset, noise_mode=self.noise_mode, r=self.r, root_dir=self.root_dir, transform=self.transform_train, mode="unlabeled", noise_file=self.noise_file, pred=pred, log=self.log)
+            unlabeled_trainloader = DataLoader(
+                dataset=unlabeled_dataset, 
+                batch_size=self.batch_size,
+                shuffle=True,
+                num_workers=self.num_workers)     
+            return labeled_trainloader, unlabeled_trainloader
+        
+        elif mode=='test':
+            test_dataset = cifar_dataset(dataset=self.dataset, noise_mode=self.noise_mode, r=self.r, root_dir=self.root_dir, transform=self.transform_test, mode='test')      
+            test_loader = DataLoader(
+                dataset=test_dataset, 
+                batch_size=self.batch_size,
+                shuffle=False,
+                num_workers=self.num_workers)          
+            return test_loader
+        
+        elif mode=='eval_train':
+            eval_dataset = cifar_dataset(dataset=self.dataset, noise_mode=self.noise_mode, r=self.r, root_dir=self.root_dir, transform=self.transform_test, mode='all', noise_file=self.noise_file)      
+            eval_loader = DataLoader(
+                dataset=eval_dataset, 
+                batch_size=self.batch_size,
+                shuffle=False,
+                num_workers=self.num_workers)          
+            return eval_loader
+        elif mode=='boost':
+            eval_dataset = cifar_dataset(dataset=self.dataset, noise_mode=self.noise_mode, r=self.r, root_dir=self.root_dir, transform=self.transform_train, mode=mode, noise_file=self.noise_file, clean_idx=clean_idx, test_form=self.transform_test)
+            eval_loader = DataLoader(
+                dataset=eval_dataset,
+                batch_size=self.batch_size,
+                shuffle=False,
+                num_workers=self.num_workers)
+            return eval_loader

models/CNN.py

@@ -0,0 +1,193 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.init as init
+import torch.nn.functional as F
+import torch.utils.model_zoo as model_zoo
+
+class HiddenLayer(nn.Module):
+    def __init__(self, input_size, output_size):
+        super(HiddenLayer, self).__init__()
+        self.fc = nn.Linear(input_size, output_size)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        return self.relu(self.fc(x))
+
+
+class VNet(nn.Module):
+    def __init__(self, hidden_size=100, num_layers=1):
+        super(VNet, self).__init__()
+        self.first_hidden_layer = HiddenLayer(1, hidden_size)
+        self.rest_hidden_layers = nn.Sequential(*[HiddenLayer(hidden_size, hidden_size) for _ in range(num_layers - 1)])
+        self.output_layer = nn.Linear(hidden_size, 1)
+
+    def forward(self, x):
+        x = self.first_hidden_layer(x)
+        x = self.rest_hidden_layers(x)
+        x = self.output_layer(x)
+        return torch.sigmoid(x)
+
+
+class CNN(nn.Module):
+    def __init__(self, input_channel=3, n_outputs=10, dropout_rate=0.25):
+        self.dropout_rate = dropout_rate
+        super(CNN, self).__init__()
+
+        #block1
+        self.conv1 = nn.Conv2d(input_channel, 128, kernel_size=3, stride=1, padding=1)
+        self.bn1=nn.BatchNorm2d(128)
+        self.conv2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
+        self.bn2=nn.BatchNorm2d(128)
+        self.conv3 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
+        self.bn3=nn.BatchNorm2d(128)
+
+        #block2
+        self.conv4 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
+        self.bn4=nn.BatchNorm2d(256)
+        self.conv5 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
+        self.bn5=nn.BatchNorm2d(256)
+        self.conv6 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
+        self.bn6=nn.BatchNorm2d(256)
+
+        #block3
+        self.conv7 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=0)
+        self.bn7=nn.BatchNorm2d(512)
+        self.conv8 = nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=0)
+        self.bn8=nn.BatchNorm2d(256)
+        self.conv9 = nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=0)
+        self.bn9=nn.BatchNorm2d(128)
+
+        self.pool = nn.MaxPool2d(2, 2)
+        self.avgpool = nn.AvgPool2d(kernel_size=2)
+        
+        self.fc=nn.Linear(128,n_outputs)
+        
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        
+        #block1
+        x=F.leaky_relu(self.bn1(self.conv1(x)), negative_slope=0.01)
+        x=F.leaky_relu(self.bn2(self.conv2(x)), negative_slope=0.01)
+        x=F.leaky_relu(self.bn3(self.conv3(x)), negative_slope=0.01)
+        x=self.pool(x)
+        x=F.dropout2d(x, p=self.dropout_rate)
+
+        #block2
+        x=F.leaky_relu(self.bn4(self.conv4(x)), negative_slope=0.01)
+        x=F.leaky_relu(self.bn5(self.conv5(x)), negative_slope=0.01)
+        x=F.leaky_relu(self.bn6(self.conv6(x)), negative_slope=0.01)
+        x=self.pool(x)
+        x=F.dropout2d(x, p=self.dropout_rate)
+
+        #block3
+        x=F.leaky_relu(self.bn7(self.conv7(x)), negative_slope=0.01)
+        x=F.leaky_relu(self.bn8(self.conv8(x)), negative_slope=0.01)
+        x=F.leaky_relu(self.bn9(self.conv9(x)), negative_slope=0.01)
+        x=self.avgpool(x)
+
+        x = x.view(x.size(0), x.size(1))
+        x=self.fc(x)
+        return x
+    
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=14):
+        self.inplanes = 64
+        super(ResNet, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AvgPool2d(7, stride=1)
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        return x
+
+def resnet18(pretrained=False, **kwargs):
+    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
+    return model

models/InceptionResNetV2.py

@@ -0,0 +1,345 @@
+from __future__ import print_function, division, absolute_import
+import torch
+import torch.nn as nn
+import os
+import sys
+
+
+class BasicConv2d(nn.Module):
+
+    def __init__(self, in_planes, out_planes, kernel_size, stride, padding=0):
+        super(BasicConv2d, self).__init__()
+        self.conv = nn.Conv2d(in_planes, out_planes,
+                              kernel_size=kernel_size, stride=stride,
+                              padding=padding, bias=False) # verify bias false
+        self.bn = nn.BatchNorm2d(out_planes,
+                                 eps=0.001, # value found in tensorflow
+                                 momentum=0.1, # default pytorch value
+                                 affine=True)
+        self.relu = nn.ReLU(inplace=False)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+
+class Mixed_5b(nn.Module):
+
+    def __init__(self):
+        super(Mixed_5b, self).__init__()
+
+        self.branch0 = BasicConv2d(192, 96, kernel_size=1, stride=1)
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(192, 48, kernel_size=1, stride=1),
+            BasicConv2d(48, 64, kernel_size=5, stride=1, padding=2)
+        )
+
+        self.branch2 = nn.Sequential(
+            BasicConv2d(192, 64, kernel_size=1, stride=1),
+            BasicConv2d(64, 96, kernel_size=3, stride=1, padding=1),
+            BasicConv2d(96, 96, kernel_size=3, stride=1, padding=1)
+        )
+
+        self.branch3 = nn.Sequential(
+            nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False),
+            BasicConv2d(192, 64, kernel_size=1, stride=1)
+        )
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        x3 = self.branch3(x)
+        out = torch.cat((x0, x1, x2, x3), 1)
+        return out
+
+
+class Block35(nn.Module):
+
+    def __init__(self, scale=1.0):
+        super(Block35, self).__init__()
+
+        self.scale = scale
+
+        self.branch0 = BasicConv2d(320, 32, kernel_size=1, stride=1)
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(320, 32, kernel_size=1, stride=1),
+            BasicConv2d(32, 32, kernel_size=3, stride=1, padding=1)
+        )
+
+        self.branch2 = nn.Sequential(
+            BasicConv2d(320, 32, kernel_size=1, stride=1),
+            BasicConv2d(32, 48, kernel_size=3, stride=1, padding=1),
+            BasicConv2d(48, 64, kernel_size=3, stride=1, padding=1)
+        )
+
+        self.conv2d = nn.Conv2d(128, 320, kernel_size=1, stride=1)
+        self.relu = nn.ReLU(inplace=False)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        out = torch.cat((x0, x1, x2), 1)
+        out = self.conv2d(out)
+        out = out * self.scale + x
+        out = self.relu(out)
+        return out
+
+
+class Mixed_6a(nn.Module):
+
+    def __init__(self):
+        super(Mixed_6a, self).__init__()
+
+        self.branch0 = BasicConv2d(320, 384, kernel_size=3, stride=2)
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(320, 256, kernel_size=1, stride=1),
+            BasicConv2d(256, 256, kernel_size=3, stride=1, padding=1),
+            BasicConv2d(256, 384, kernel_size=3, stride=2)
+        )
+
+        self.branch2 = nn.MaxPool2d(3, stride=2)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        out = torch.cat((x0, x1, x2), 1)
+        return out
+
+
+class Block17(nn.Module):
+
+    def __init__(self, scale=1.0):
+        super(Block17, self).__init__()
+
+        self.scale = scale
+
+        self.branch0 = BasicConv2d(1088, 192, kernel_size=1, stride=1)
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(1088, 128, kernel_size=1, stride=1),
+            BasicConv2d(128, 160, kernel_size=(1,7), stride=1, padding=(0,3)),
+            BasicConv2d(160, 192, kernel_size=(7,1), stride=1, padding=(3,0))
+        )
+
+        self.conv2d = nn.Conv2d(384, 1088, kernel_size=1, stride=1)
+        self.relu = nn.ReLU(inplace=False)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        out = torch.cat((x0, x1), 1)
+        out = self.conv2d(out)
+        out = out * self.scale + x
+        out = self.relu(out)
+        return out
+
+
+class Mixed_7a(nn.Module):
+
+    def __init__(self):
+        super(Mixed_7a, self).__init__()
+
+        self.branch0 = nn.Sequential(
+            BasicConv2d(1088, 256, kernel_size=1, stride=1),
+            BasicConv2d(256, 384, kernel_size=3, stride=2)
+        )
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(1088, 256, kernel_size=1, stride=1),
+            BasicConv2d(256, 288, kernel_size=3, stride=2)
+        )
+
+        self.branch2 = nn.Sequential(
+            BasicConv2d(1088, 256, kernel_size=1, stride=1),
+            BasicConv2d(256, 288, kernel_size=3, stride=1, padding=1),
+            BasicConv2d(288, 320, kernel_size=3, stride=2)
+        )
+
+        self.branch3 = nn.MaxPool2d(3, stride=2)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        x2 = self.branch2(x)
+        x3 = self.branch3(x)
+        out = torch.cat((x0, x1, x2, x3), 1)
+        return out
+
+
+class Block8(nn.Module):
+
+    def __init__(self, scale=1.0, noReLU=False):
+        super(Block8, self).__init__()
+
+        self.scale = scale
+        self.noReLU = noReLU
+
+        self.branch0 = BasicConv2d(2080, 192, kernel_size=1, stride=1)
+
+        self.branch1 = nn.Sequential(
+            BasicConv2d(2080, 192, kernel_size=1, stride=1),
+            BasicConv2d(192, 224, kernel_size=(1,3), stride=1, padding=(0,1)),
+            BasicConv2d(224, 256, kernel_size=(3,1), stride=1, padding=(1,0))
+        )
+
+        self.conv2d = nn.Conv2d(448, 2080, kernel_size=1, stride=1)
+        if not self.noReLU:
+            self.relu = nn.ReLU(inplace=False)
+
+    def forward(self, x):
+        x0 = self.branch0(x)
+        x1 = self.branch1(x)
+        out = torch.cat((x0, x1), 1)
+        out = self.conv2d(out)
+        out = out * self.scale + x
+        if not self.noReLU:
+            out = self.relu(out)
+        return out
+
+
+class InceptionResNetV2(nn.Module):
+
+    def __init__(self, num_classes=50):
+        super(InceptionResNetV2, self).__init__()
+        # Special attributs
+        self.num_classes = num_classes
+        self.input_space = None
+        self.input_size = (299, 299, 3)
+        self.mean = None
+        self.std = None
+        # Modules
+        self.conv2d_1a = BasicConv2d(3, 32, kernel_size=3, stride=2)
+        self.conv2d_2a = BasicConv2d(32, 32, kernel_size=3, stride=1)
+        self.conv2d_2b = BasicConv2d(32, 64, kernel_size=3, stride=1, padding=1)
+        self.maxpool_3a = nn.MaxPool2d(3, stride=2)
+        self.conv2d_3b = BasicConv2d(64, 80, kernel_size=1, stride=1)
+        self.conv2d_4a = BasicConv2d(80, 192, kernel_size=3, stride=1)
+        self.maxpool_5a = nn.MaxPool2d(3, stride=2)
+        self.mixed_5b = Mixed_5b()
+        self.repeat = nn.Sequential(
+            Block35(scale=0.17),
+            Block35(scale=0.17),
+            Block35(scale=0.17),
+            Block35(scale=0.17),
+            Block35(scale=0.17),
+            Block35(scale=0.17),
+            Block35(scale=0.17),
+            Block35(scale=0.17),
+            Block35(scale=0.17),
+            Block35(scale=0.17)
+        )
+        self.mixed_6a = Mixed_6a()
+        self.repeat_1 = nn.Sequential(
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10),
+            Block17(scale=0.10)
+        )
+        self.mixed_7a = Mixed_7a()
+        self.repeat_2 = nn.Sequential(
+            Block8(scale=0.20),
+            Block8(scale=0.20),
+            Block8(scale=0.20),
+            Block8(scale=0.20),
+            Block8(scale=0.20),
+            Block8(scale=0.20),
+            Block8(scale=0.20),
+            Block8(scale=0.20),
+            Block8(scale=0.20)
+        )
+        self.block8 = Block8(noReLU=True)
+        self.conv2d_7b = BasicConv2d(2080, 1536, kernel_size=1, stride=1)
+        self.avgpool_1a = nn.AdaptiveAvgPool2d((1, 1))#nn.AvgPool2d(8, count_include_pad=False)
+        self.last_linear = nn.Linear(1536, num_classes)
+        
+        self.branch = self._make_branch(320, 1536, 3)
+        self.branch1 = self._make_branch(1088, 1536, 3)
+        self.branch2 = self._make_branch(2080, 1536, 3)
+        
+    def _make_branch(self, channel_in, channel_out, kernel_size):
+        middle_channel = channel_out // 4
+        return nn.Sequential(
+            nn.Conv2d(channel_in, middle_channel, kernel_size=1, stride=1),
+            nn.BatchNorm2d(middle_channel),
+            nn.ReLU(),
+            
+            nn.Conv2d(middle_channel, middle_channel, kernel_size=kernel_size, stride=kernel_size),
+            nn.BatchNorm2d(middle_channel),
+            nn.ReLU(),
+            
+            nn.Conv2d(middle_channel, channel_out, kernel_size=1, stride=1),
+            nn.BatchNorm2d(channel_out),
+            nn.ReLU(),
+            
+            nn.AdaptiveAvgPool2d((1,1)),
+            nn.Flatten(),
+            nn.Linear(channel_out, self.num_classes)
+        )
+        
+
+    def features(self, input):
+        x = self.conv2d_1a(input)
+        x = self.conv2d_2a(x)
+        x = self.conv2d_2b(x)
+        x = self.maxpool_3a(x)
+        x = self.conv2d_3b(x)
+        x = self.conv2d_4a(x)
+        x = self.maxpool_5a(x)
+        x = self.mixed_5b(x)
+        x = self.repeat(x)
+        x1 = self.branch(x)
+        
+        x = self.mixed_6a(x)
+        x = self.repeat_1(x)
+        x2 = self.branch1(x)
+        
+        x = self.mixed_7a(x)
+        x = self.repeat_2(x)
+        x3 = self.branch2(x)
+        
+        x = self.block8(x)
+        x = self.conv2d_7b(x)
+        return x, x1, x2, x3
+
+    def logits(self, features):
+        x = self.avgpool_1a(features)
+        x = x.view(x.size(0), -1)
+        out = self.last_linear(x)
+        return out
+
+
+    def forward(self, input):
+        x, x1, x2, x3, = self.features(input)
+        out = self.logits(x)
+        return {'outputs': [out, x1, x2, x3]}
+
+
+def test():
+    net = InceptionResNetV2().cuda()
+    y = net(torch.randn(1,3,227,227).cuda())
+    print(y.size())
+#test()

models/ResNet_Imagenet.py

@@ -0,0 +1,289 @@
+import torch
+import torch.nn as nn
+
+try:
+    from torch.hub import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+    'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
+    'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
+    'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
+    'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+def branchBottleNeck(channel_in, channel_out, kernel_size):
+    middle_channel = channel_out//4
+    return nn.Sequential(
+        nn.Conv2d(channel_in, middle_channel, kernel_size=1, stride=1),
+        nn.BatchNorm2d(middle_channel),
+        nn.ReLU(),
+        
+        nn.Conv2d(middle_channel, middle_channel, kernel_size=kernel_size, stride=kernel_size),
+        nn.BatchNorm2d(middle_channel),
+        nn.ReLU(),
+        
+        nn.Conv2d(middle_channel, channel_out, kernel_size=1, stride=1),
+        nn.BatchNorm2d(channel_out),
+        nn.ReLU(),
+        )
+
+class LambdaLayer(nn.Module):
+    def __init__(self, lambd):
+        super(LambdaLayer, self).__init__()
+        self.lambd = lambd
+
+    def forward(self, x):
+        return self.lambd(x)
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(Bottleneck, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
+                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
+                 norm_layer=None):
+        super(ResNet, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+        self.num_classes = num_classes
+
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
+                                       dilate=replace_stride_with_dilation[2])
+        
+        self.branch1 = self._make_branch(64*block.expansion, 512*block.expansion, kernel_size=8)
+        self.branch2 = self._make_branch(128*block.expansion, 512*block.expansion, kernel_size=4)
+        self.branch3 = self._make_branch(256*block.expansion, 512*block.expansion, kernel_size=2)
+        self.branch4 = self._make_branch(512*block.expansion, 512*block.expansion, kernel_size=1)
+        
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
+                            self.base_width, previous_dilation, norm_layer))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=self.groups,
+                                base_width=self.base_width, dilation=self.dilation,
+                                norm_layer=norm_layer))
+
+        return nn.Sequential(*layers)
+    
+    def _make_branch(self, channel_in, channel_out, kernel_size):
+        middle_channel = channel_out // 4
+        return nn.Sequential(
+            nn.Conv2d(channel_in, middle_channel, kernel_size=1, stride=1),
+            nn.BatchNorm2d(middle_channel),
+            nn.ReLU(),
+
+            nn.Conv2d(middle_channel, middle_channel, kernel_size=kernel_size, stride=kernel_size),
+            nn.BatchNorm2d(middle_channel),
+            nn.ReLU(),
+
+            nn.Conv2d(middle_channel, channel_out, kernel_size=1, stride=1),
+            nn.BatchNorm2d(channel_out),
+            nn.ReLU(),
+            
+            nn.AdaptiveAvgPool2d((1,1)),
+            nn.Flatten(),
+            nn.Linear(channel_out, self.num_classes)
+        )
+
+    def _forward_impl(self, x):
+        # See note [TorchScript super()]
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x1 = self.branch1(x)
+
+        x = self.layer2(x)
+        x2 = self.branch2(x)
+
+        x = self.layer3(x)
+        x3 = self.branch3(x)
+
+        x = self.layer4(x)
+        x = self.avgpool(x)
+        final_fea = x
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+        
+        return {'outputs': [x, x1, x2, x3]}
+
+    def forward(self, x):
+        return self._forward_impl(x)
+    
+def sdresnet50(num_classes=14, pretrained=True):
+    if pretrained:
+        model = ResNet(Bottleneck, [3,4,6,3], num_classes=14)
+        num_ftrs = model.fc.in_features
+        model.fc = nn.Linear(num_ftrs, 1000)
+        state_dict = load_state_dict_from_url(model_urls['resnet50'], progress=True)
+        model.load_state_dict(state_dict, strict=False)
+
+        num_ftrs = model.fc.in_features
+        model.fc = nn.Linear(num_ftrs, num_classes)
+    else:
+        model = ResNet(Bottleneck, [3,4,6,3], num_classes=50)
+    return model

models/ResNet_cifar.py

@@ -0,0 +1,688 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+from torch.nn import Parameter
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+    'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
+    'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
+    'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
+    'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
+}
+
+def conv3x3(in_planes, out_planes, stride=1):
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3,
+                     stride=stride, padding=1, bias=False)
+
+def conv1x1(in_planes, planes, stride=1):
+    return nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride, bias=False)
+
+def branchBottleNeck(channel_in, channel_out, kernel_size):
+    middle_channel = channel_out//4
+    return nn.Sequential(
+        nn.Conv2d(channel_in, middle_channel, kernel_size=1, stride=1),
+        nn.BatchNorm2d(middle_channel),
+        nn.ReLU(),
+        
+        nn.Conv2d(middle_channel, middle_channel, kernel_size=kernel_size, stride=kernel_size),
+        nn.BatchNorm2d(middle_channel),
+        nn.ReLU(),
+        
+        nn.Conv2d(middle_channel, channel_out, kernel_size=1, stride=1),
+        nn.BatchNorm2d(channel_out),
+        nn.ReLU(),
+        )
+
+def branchMLP(channel_in, channel_out):
+    middle_channel = channel_out//4
+    return nn.Sequential(
+            conv1x1(channel_in, channel_in, stride=8),
+            nn.BatchNorm2d(512 * block.expansion),
+            nn.ReLU(),
+        )
+
+def invertedBottleNeck(channel_in, channel_out, kernel_size):
+    middle_channel = channel_out * 2
+    return nn.Sequential(
+        nn.Conv2d(channel_in, middle_channel, kernel_size=1, stride=1),
+        nn.BatchNorm2d(middle_channel),
+        nn.ReLU(),
+        
+        nn.Conv2d(middle_channel, middle_channel, kernel_size=kernel_size, stride=kernel_size),
+        nn.BatchNorm2d(middle_channel),
+        nn.ReLU(),
+        
+        nn.Conv2d(middle_channel, channel_out, kernel_size=1, stride=1),
+        nn.BatchNorm2d(channel_out),
+        nn.ReLU(),
+        )
+
+class BatchNorm2dMul(nn.Module):
+    def __init__(self, num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True):
+        super(BatchNorm2dMul, self).__init__()
+        self.bn = nn.BatchNorm2d(num_features, eps=eps, momentum=momentum, affine=False, track_running_stats=track_running_stats)
+        self.gamma = nn.Parameter(torch.ones(num_features))
+        self.beta = nn.Parameter(torch.zeros(num_features))
+        self.affine = affine
+
+    def forward(self, x):
+        bn_out = self.bn(x)
+        if self.affine:
+            out = self.gamma[None, :, None, None] * bn_out + self.beta[None, :, None, None]
+        return out, bn_out
+    
+def _weights_init(m):
+    classname = m.__class__.__name__
+    if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d):
+        init.kaiming_normal_(m.weight)
+
+class LambdaLayer(nn.Module):
+    def __init__(self, lambd):
+        super(LambdaLayer, self).__init__()
+        self.lambd = lambd
+
+    def forward(self, x):
+        return self.lambd(x)
+
+class BasicBlock_s(nn.Module):
+    expansion = 1
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(BasicBlock_s, self).__init__()
+        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.shortcut = nn.Sequential()
+        if stride != 1 or in_planes != self.expansion*planes:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(self.expansion*planes)
+            )
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+class BasicBlock(nn.Module):
+    expansion = 1
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = BatchNorm2dMul(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = BatchNorm2dMul(planes)
+        self.downsample = downsample
+        self.stride = stride
+    
+    def forward(self, x):
+        bn_outputs = []
+        
+        residual = x
+        output = self.conv1(x)
+        output, bn_out = self.bn1(output)
+        bn_outputs.append(bn_out)
+        output = self.relu(output)
+
+        output = self.conv2(output)
+        output, bn_out = self.bn2(output)
+        bn_outputs.append(bn_out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        
+        output += residual
+        output = self.relu(output)
+        return output, bn_outputs
+
+class BottleneckBlock(nn.Module):
+    expansion = 4
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BottleneckBlock, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.conv3 = conv1x1(planes, planes*self.expansion)
+        self.bn3 = nn.BatchNorm2d(planes*self.expansion)
+
+        self.downsample = downsample
+        self.stride = stride
+    
+    def forward(self, x):
+        residual = x
+
+        output = self.conv1(x)
+        output = self.bn1(output)
+        output = self.relu(output)
+
+        output = self.conv2(output)
+        output = self.bn2(output)
+        output = self.relu(output)
+
+        output = self.conv3(output)
+        output = self.bn3(output)
+        
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        output += residual
+        output = self.relu(output)
+
+        return output
+    
+class LayerBlock(nn.Module):
+    def __init__(self, block, inplanes, planes, num_blocks, stride):
+        super(LayerBlock, self).__init__()
+        downsample = None
+        if stride !=1 or inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+        layer = []
+        layer.append(block(inplanes, planes, stride=stride, downsample=downsample))
+        inplanes = planes * block.expansion
+        for i in range(1, num_blocks):
+            layer.append(block(inplanes, planes))
+        self.layers = nn.Sequential(*layer)
+        
+    def forward(self, x):
+        bn_outputs = []
+        for layer in self.layers:
+            x, bn_output = layer(x)
+            bn_outputs.extend(bn_output)
+        return x, bn_outputs
+    
+class SDResNet(nn.Module):
+    """
+    Resnet model
+    
+    Args:
+        block (class): block type, BasicBlock or BottlenetckBlock
+        layers (int list): layer num in each block
+        num_classes (int): class num
+    """
+    
+    def __init__(self, block, layers, num_classes=10, position_all=True):
+        super(SDResNet, self).__init__()
+        
+        self.position_all = position_all
+        
+        self.inplanes = 64
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        # self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        
+        self.layer1 = LayerBlock(block, 64, 64, layers[0], stride=1)
+        self.layer2 = LayerBlock(block, 64, 128, layers[1], stride=2)
+        self.layer3 = LayerBlock(block, 128, 256, layers[2], stride=2)
+        self.layer4 = LayerBlock(block, 256, 512, layers[3], stride=2)
+
+        self.downsample1_1 = nn.Sequential(
+                            conv1x1(64 * block.expansion, 512 * block.expansion, stride=8),
+                            nn.BatchNorm2d(512 * block.expansion),
+                            nn.ReLU(),
+        )
+        self.bottleneck1_1 = branchBottleNeck(64 * block.expansion, 512 * block.expansion, kernel_size=8)
+        self.avgpool1 = nn.AdaptiveAvgPool2d((1,1))
+        self.middle_fc1 = nn.Linear(512 * block.expansion, num_classes)
+
+
+        self.downsample2_1 = nn.Sequential(
+                            conv1x1(128 * block.expansion, 512 * block.expansion, stride=4),
+                            nn.BatchNorm2d(512 * block.expansion),
+            )
+        self.bottleneck2_1 = branchBottleNeck(128 * block.expansion, 512 * block.expansion, kernel_size=4)
+        self.avgpool2 = nn.AdaptiveAvgPool2d((1,1))
+        self.middle_fc2 = nn.Linear(512 * block.expansion, num_classes)
+
+
+        self.downsample3_1 = nn.Sequential(
+                            conv1x1(256 * block.expansion, 512 * block.expansion, stride=2),
+                            nn.BatchNorm2d(512 * block.expansion),
+        )
+        self.bottleneck3_1 = branchBottleNeck(256 * block.expansion, 512 * block.expansion, kernel_size=2)
+        self.avgpool3 = nn.AdaptiveAvgPool2d((1,1))
+        self.middle_fc3 = nn.Linear(512 * block.expansion, num_classes)
+
+        self.avgpool = nn.AdaptiveAvgPool2d((1,1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+        
+        self.apply(_weights_init)
+        
+    def _make_layer(self, block, planes, layers, stride=1):
+        """A block with 'layers' layers
+        Args:
+            block (class): block type
+            planes (int): output channels = planes * expansion
+            layers (int): layer num in the block
+            stride (int): the first layer stride in the block
+        """
+        downsample = None
+        if stride !=1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+        layer = []
+        layer.append(block(self.inplanes, planes, stride=stride, downsample=downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, layers):
+            layer.append(block(self.inplanes, planes))
+        
+        return nn.Sequential(*layer)
+    
+    def forward(self, x, feat_out=False):
+        all_bn_outputs = []
+        
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        # x = self.maxpool(x)
+
+        x, bn_outputs = self.layer1(x)
+        all_bn_outputs.extend(bn_outputs)
+        middle_output1 = self.bottleneck1_1(x)
+        middle_output1 = self.avgpool1(middle_output1)
+        middle1_fea = middle_output1
+        middle_output1 = torch.flatten(middle_output1, 1)
+        middle_output1 = self.middle_fc1(middle_output1)
+
+        x, bn_outputs = self.layer2(x)
+        all_bn_outputs.extend(bn_outputs)
+        middle_output2 = self.bottleneck2_1(x)
+        middle_output2 = self.avgpool2(middle_output2)
+        middle2_fea = middle_output2
+        middle_output2 = torch.flatten(middle_output2, 1)
+        middle_output2 = self.middle_fc2(middle_output2)
+
+        x, bn_outputs = self.layer3(x)
+        all_bn_outputs.extend(bn_outputs)
+        middle_output3 = self.bottleneck3_1(x)
+        middle_output3 = self.avgpool3(middle_output3)
+        middle3_fea = middle_output3
+        middle_output3 = torch.flatten(middle_output3, 1)
+        middle_output3 = self.middle_fc3(middle_output3)
+
+        x, bn_outputs = self.layer4(x)
+        all_bn_outputs.extend(bn_outputs)
+        x = self.avgpool(x)
+        final_fea = x
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        if self.position_all and feat_out:
+            return {'outputs': [x, middle_output1, middle_output2, middle_output3],
+                'features': [final_fea, middle1_fea, middle2_fea, middle3_fea],
+                'bn_outputs': all_bn_outputs}
+        else:
+            return x
+        
+class SDResNet_mlp(nn.Module):
+    """
+    Resnet model
+    
+    Args:
+        block (class): block type, BasicBlock or BottlenetckBlock
+        layers (int list): layer num in each block
+        num_classes (int): class num
+    """
+    
+    def __init__(self, block, layers, num_classes=10, position_all=True):
+        super(SDResNet_mlp, self).__init__()
+        
+        self.position_all = position_all
+        
+        self.inplanes = 64
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        # self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        
+        self.layer1 = LayerBlock(block, 64, 64, layers[0], stride=1)
+        self.layer2 = LayerBlock(block, 64, 128, layers[1], stride=2)
+        self.layer3 = LayerBlock(block, 128, 256, layers[2], stride=2)
+        self.layer4 = LayerBlock(block, 256, 512, layers[3], stride=2)
+
+        self.downsample1_1 = nn.Sequential(
+            conv1x1(64 * block.expansion, 512 * block.expansion),
+            nn.BatchNorm2d(512 * block.expansion),
+            nn.ReLU(),
+        )
+        self.avgpool1 = nn.AdaptiveAvgPool2d((1,1))
+        self.middle_fc1 = nn.Linear(512 * block.expansion, num_classes)
+
+
+        self.downsample2_1 = nn.Sequential(
+            conv1x1(128 * block.expansion, 512 * block.expansion),
+            nn.BatchNorm2d(512 * block.expansion),
+            nn.ReLU()
+            )
+        self.avgpool2 = nn.AdaptiveAvgPool2d((1,1))
+        self.middle_fc2 = nn.Linear(512 * block.expansion, num_classes)
+
+
+        self.downsample3_1 = nn.Sequential(
+            conv1x1(256 * block.expansion, 512 * block.expansion),
+            nn.BatchNorm2d(512 * block.expansion),
+            nn.ReLU()
+        )
+        self.avgpool3 = nn.AdaptiveAvgPool2d((1,1))
+        self.middle_fc3 = nn.Linear(512 * block.expansion, num_classes)
+
+        self.avgpool = nn.AdaptiveAvgPool2d((1,1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+        
+        self.apply(_weights_init)
+        
+    def _make_layer(self, block, planes, layers, stride=1):
+        """A block with 'layers' layers
+        Args:
+            block (class): block type
+            planes (int): output channels = planes * expansion
+            layers (int): layer num in the block
+            stride (int): the first layer stride in the block
+        """
+        downsample = None
+        if stride !=1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+        layer = []
+        layer.append(block(self.inplanes, planes, stride=stride, downsample=downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, layers):
+            layer.append(block(self.inplanes, planes))
+        
+        return nn.Sequential(*layer)
+    
+    def forward(self, x):
+        all_bn_outputs = []
+        
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x, bn_outputs = self.layer1(x)
+        all_bn_outputs.extend(bn_outputs)
+#         middle_output1 = self.downsample1_1(x)
+#         middle_output1 = self.avgpool1(middle_output1)
+#         middle1_fea = middle_output1
+#         middle_output1 = torch.flatten(middle_output1, 1)
+#         middle_output1 = self.middle_fc1(middle_output1)
+
+        x, bn_outputs = self.layer2(x)
+        all_bn_outputs.extend(bn_outputs)
+#         middle_output2 = self.downsample2_1(x)
+#         middle_output2 = self.avgpool2(middle_output2)
+#         middle2_fea = middle_output2
+#         middle_output2 = torch.flatten(middle_output2, 1)
+#         middle_output2 = self.middle_fc2(middle_output2)
+
+        x, bn_outputs = self.layer3(x)
+        all_bn_outputs.extend(bn_outputs)
+#         middle_output3 = self.downsample3_1(x)
+#         middle_output3 = self.avgpool3(middle_output3)
+#         middle3_fea = middle_output3
+#         middle_output3 = torch.flatten(middle_output3, 1)
+#         middle_output3 = self.middle_fc3(middle_output3)
+
+        x, bn_outputs = self.layer4(x)
+        all_bn_outputs.extend(bn_outputs)
+        x = self.avgpool(x)
+        final_fea = x
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        if self.position_all:
+            return {'outputs': [x, middle_output1, middle_output2, middle_output3],
+                'bn_outputs': all_bn_outputs}
+        else:
+            return {'outputs': [x, x],
+                    'bn_outputs': all_bn_outputs}
+        
+class SDResNet_residual(nn.Module):
+    """
+    Resnet model
+    
+    Args:
+        block (class): block type, BasicBlock or BottlenetckBlock
+        layers (int list): layer num in each block
+        num_classes (int): class num
+    """
+    
+    def __init__(self, block, layers, num_classes=10, position_all=True):
+        super(SDResNet_residual, self).__init__()
+        
+        self.position_all = position_all
+        
+        self.inplanes = 64
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        # self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        
+        self.layer1 = LayerBlock(block, 64, 64, layers[0], stride=1)
+        self.layer2 = LayerBlock(block, 64, 128, layers[1], stride=2)
+        self.layer3 = LayerBlock(block, 128, 256, layers[2], stride=2)
+        self.layer4 = LayerBlock(block, 256, 512, layers[3], stride=2)
+        
+        self.bottleneck1_1 = LayerBlock(block, 64, 512, 1, stride=8)
+#         branchBottleNeck(64 * block.expansion, 512 * block.expansion, kernel_size=8)
+        self.avgpool1 = nn.AdaptiveAvgPool2d((1,1))
+        self.middle_fc1 = nn.Linear(512 * block.expansion, num_classes)
+
+        self.bottleneck2_1 = LayerBlock(block, 128, 512, 1, stride=4)
+#         branchBottleNeck(128 * block.expansion, 512 * block.expansion, kernel_size=4)
+        self.avgpool2 = nn.AdaptiveAvgPool2d((1,1))
+        self.middle_fc2 = nn.Linear(512 * block.expansion, num_classes)
+
+
+#         self.downsample3_1 = nn.Sequential(
+#                             conv1x1(256 * block.expansion, 512 * block.expansion, stride=2),
+#                             nn.BatchNorm2d(512 * block.expansion),
+#         )
+        self.bottleneck3_1 = LayerBlock(block, 256, 512, 1, stride=2)
+        self.avgpool3 = nn.AdaptiveAvgPool2d((1,1))
+        self.middle_fc3 = nn.Linear(512 * block.expansion, num_classes)
+
+        self.avgpool = nn.AdaptiveAvgPool2d((1,1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+        
+        self.apply(_weights_init)
+        
+    def _make_layer(self, block, planes, layers, stride=1):
+        """A block with 'layers' layers
+        Args:
+            block (class): block type
+            planes (int): output channels = planes * expansion
+            layers (int): layer num in the block
+            stride (int): the first layer stride in the block
+        """
+        downsample = None
+        if stride !=1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+        layer = []
+        layer.append(block(self.inplanes, planes, stride=stride, downsample=downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, layers):
+            layer.append(block(self.inplanes, planes))
+        
+        return nn.Sequential(*layer)
+    
+    def forward(self, x):
+        all_bn_outputs = []
+        
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        # x = self.maxpool(x)
+
+        x, bn_outputs = self.layer1(x)
+        all_bn_outputs.extend(bn_outputs)
+        middle_output1, _ = self.bottleneck1_1(x)
+        middle_output1 = self.avgpool1(middle_output1)
+        middle1_fea = middle_output1
+        middle_output1 = torch.flatten(middle_output1, 1)
+        middle_output1 = self.middle_fc1(middle_output1)
+
+        x, bn_outputs = self.layer2(x)
+        all_bn_outputs.extend(bn_outputs)
+        middle_output2, _ = self.bottleneck2_1(x)
+        middle_output2 = self.avgpool2(middle_output2)
+        middle2_fea = middle_output2
+        middle_output2 = torch.flatten(middle_output2, 1)
+        middle_output2 = self.middle_fc2(middle_output2)
+
+        x, bn_outputs = self.layer3(x)
+        all_bn_outputs.extend(bn_outputs)
+        middle_output3, _ = self.bottleneck3_1(x)
+        middle_output3 = self.avgpool3(middle_output3)
+        middle3_fea = middle_output3
+        middle_output3 = torch.flatten(middle_output3, 1)
+        middle_output3 = self.middle_fc3(middle_output3)
+
+        x, bn_outputs = self.layer4(x)
+        all_bn_outputs.extend(bn_outputs)
+        x = self.avgpool(x)
+        final_fea = x
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        if self.position_all:
+            return {'outputs': [x, middle_output1, middle_output2, middle_output3],
+                'features': [final_fea, middle1_fea, middle2_fea, middle3_fea],
+                'bn_outputs': all_bn_outputs}
+        else:
+            return {'outputs': [x, middle_output3],
+                    'features': [final_fea, middle1_fea, middle2_fea, middle3_fea],
+                    'bn_outputs': all_bn_outputs}
+    
+class SDResNet_s(nn.Module):
+    """
+    Resnet model small
+    
+    Args:
+        block (class): block type, BasicBlock or BottlenetckBlock
+        layers (int list): layer num in each block
+        num_classes (int): class num
+    """
+    
+    def __init__(self, block, layers, num_classes=10):
+        super(SDResNet_s, self).__init__()
+        
+        self.inplanes = 16
+        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        # self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        self.layer1 = self._make_layer(block, 16, layers[0])
+        self.layer2 = self._make_layer(block, 32, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 64, layers[2], stride=2)
+
+        self.downsample1_1 = nn.Sequential(
+                            conv1x1(16 * block.expansion, 64 * block.expansion, stride=4),
+                            nn.BatchNorm2d(64 * block.expansion),
+        )
+        self.bottleneck1_1 = branchBottleNeck(16 * block.expansion, 64 * block.expansion, kernel_size=4)
+        self.avgpool1 = nn.AdaptiveAvgPool2d((1,1))
+        self.middle_fc1 = nn.Linear(64 * block.expansion, num_classes)
+
+
+        self.downsample2_1 = nn.Sequential(
+                            conv1x1(32 * block.expansion, 64 * block.expansion, stride=2),
+                            nn.BatchNorm2d(64 * block.expansion),
+            )
+        self.bottleneck2_1 = branchBottleNeck(32 * block.expansion, 64 * block.expansion, kernel_size=2)
+        self.avgpool2 = nn.AdaptiveAvgPool2d((1,1))
+        self.middle_fc2 = nn.Linear(64 * block.expansion, num_classes)
+
+        self.avgpool = nn.AdaptiveAvgPool2d((1,1))
+        self.fc = nn.Linear(64 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d): 
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+        
+    def _make_layer(self, block, planes, layers, stride=1):
+        """A block with 'layers' layers
+        Args:
+            block (class): block type
+            planes (int): output channels = planes * expansion
+            layers (int): layer num in the block
+            stride (int): the first layer stride in the block
+        """
+        strides = [stride] + [1]*(layers-1)
+        layers = []
+        for stride in strides:
+            layers.append(block(self.inplanes, planes, stride))
+            self.inplanes = planes * block.expansion
+
+        return nn.Sequential(*layers)
+    
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x = self.layer1(x)
+        middle_output1 = self.bottleneck1_1(x)
+        middle_output1 = self.avgpool1(middle_output1)
+        middle1_fea = middle_output1
+        middle_output1 = torch.flatten(middle_output1, 1)
+        middle_output1 = self.middle_fc1(middle_output1)
+
+        x = self.layer2(x)
+        middle_output2 = self.bottleneck2_1(x)
+        middle_output2 = self.avgpool2(middle_output2)
+        middle2_fea = middle_output2
+        middle_output2 = torch.flatten(middle_output2, 1)
+        middle_output2 = self.middle_fc2(middle_output2)
+
+        x = self.layer3(x)
+        x = self.avgpool(x)
+        final_fea = x
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return {'outputs': [x, middle_output1, middle_output2],
+                'features': [final_fea, middle1_fea, middle2_fea]}
+    
+def sdresnet18(num_classes=10, position_all=True):
+    return SDResNet(BasicBlock, [2,2,2,2], num_classes=num_classes, position_all=position_all)
+
+def sdresnet34(num_classes=10, position_all=True):
+    return SDResNet(BasicBlock, [3,4,6,3], num_classes=num_classes, position_all=position_all)
+
+def sdresnet34_mlp(num_classes=10, position_all=True):
+    return SDResNet_mlp(BasicBlock, [3,4,6,3], num_classes=num_classes, position_all=position_all)
+
+def sdresnet34_residual(num_classes=10, position_all=True):
+    return SDResNet_residual(BasicBlock, [3,4,6,3], num_classes=num_classes, position_all=position_all)
+
+def sdresnet32(num_classes=10):
+    return SDResNet_s(BasicBlock_s, [5,5,5], num_classes=num_classes)

models/__init__.py

@@ -0,0 +1,4 @@
+from .ResNet_Imagenet import sdresnet50
+from .ResNet_cifar import sdresnet18, sdresnet34, sdresnet32
+from .ResNet_cifar import sdresnet34_mlp, sdresnet34_residual
+from .InceptionResNetV2 import InceptionResNetV2

requirements.txt

@@ -0,0 +1,7 @@
+numpy
+pillow
+matplotlib
+scikit-learn
+scipy
+torch
+torchvision

train_cifar_c2mt.py

@@ -0,0 +1,807 @@
+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)
+
+
+