ablation.py

# -*- coding:utf-8 -*-
import os
import torch
import torch.nn.functional as F
from torch.autograd import Variable
from data.datasets import input_dataset
from models import *
import argparse
import copy
import wandb
import logging

parser = argparse.ArgumentParser()
parser.add_argument('--lr', type=float, default=0.1)
parser.add_argument('--noise_type', type=str, help='clean, aggre, worst, rand1, rand2, rand3, clean100, noisy100',
                    default='clean')
parser.add_argument('--noise_path', type=str, help='path of CIFAR-10_human.pt', default=None)
parser.add_argument('--dataset', type=str, help=' cifar10 or cifar100', default='cifar10')
parser.add_argument('--n_epoch', type=int, default=100)
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--print_freq', type=int, default=50)
parser.add_argument('--num_workers', type=int, default=4, help='how many subprocesses to use for data loading')
parser.add_argument('--is_human', action='store_true', default=False)
##########################################################################################################################################################################################
######################################################################################################################################################################## New Code Starting
##########################################################################################################################################################################################
parser.add_argument('--lam', default=1, type=float)
parser.add_argument('--method', default='CT', type=str)


##########################################################################################################################################################################################
########################################################################################################################################################################## New Code Ending
##########################################################################################################################################################################################

# Adjust learning rate and for SGD Optimizer
def adjust_learning_rate(optimizer, epoch, alpha_plan):
    for param_group in optimizer.param_groups:
        param_group['lr'] = alpha_plan[epoch]
    print('lr = ', alpha_plan[epoch])


def accuracy(logit, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    output = F.softmax(logit, dim=1)
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.reshape(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res


# Train the Model
def train(epoch, train_loader, model, optimizer):
    train_total = 0
    train_correct = 0

    for i, (images, _, _, labels, indexes) in enumerate(train_loader):
        ind = indexes.cpu().numpy().transpose()
        batch_size = len(ind)

        images = Variable(images).cuda()
        labels = Variable(labels).cuda()

        # Forward + Backward + Optimize
        logits = model(images)

        prec, _ = accuracy(logits, labels, topk=(1, 5))
        # prec = 0.0
        train_total += 1
        train_correct += prec
        loss = F.cross_entropy(logits, labels, reduce=True)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if (i + 1) % args.print_freq == 0:
            print('Epoch [%d/%d], Iter [%d/%d] Training Accuracy: %.4F, Loss: %.4f'
                  % (epoch + 1, args.n_epoch, i + 1, len(train_dataset) // batch_size, prec, loss.data))

    train_acc = float(train_correct) / float(train_total)
    return train_acc


##########################################################################################################################################################################################
######################################################################################################################################################################## New Code Starting
##########################################################################################################################################################################################
consistency_criterion = nn.KLDivLoss(reduction='batchmean').cuda()


# Train the Model
def train_ours5(epoch, train_loader, model1, optimizer1, args):
    train_total = 0
    train_correct = 0
    EMA1 = copy.deepcopy(model1)
    for i, (images, images_s1, images_s2, images_s3, labels, indexes) in enumerate(train_loader):
        ind = indexes.cpu().numpy().transpose()
        batch_size = len(ind)

        images = Variable(images).cuda()
        images_s1 = Variable(images_s1).cuda()
        images_s2 = Variable(images_s2).cuda()
        images_s3 = Variable(images_s3).cuda()
        labels = Variable(labels).cuda()

        # Forward + Backward + Optimize
        logits1 = model1(images)
        prec, _ = accuracy(logits1, labels, topk=(1, 5))
        # prec = 0.0
        train_total += 1
        train_correct += prec

        loss_ce1 = F.cross_entropy(logits1, labels, reduce=True)

        # Consistency Loss

        logits_s1 = model1(images_s1)
        logits_s2 = model1(images_s2)
        logits_s3 = model1(images_s3)

        outputs = torch.pow(torch.softmax(logits1, dim=-1).detach(), 1 / (1 + 3)) \
                  * torch.pow(torch.softmax(logits_s1, dim=-1).detach(), 1 / (1 + 3)) \
                  * torch.pow(torch.softmax(logits_s2, dim=-1).detach(), 1 / (1 + 3)) \
                  * torch.pow(torch.softmax(logits_s3, dim=-1).detach(), 1 / (1 + 3))  # outputs = torch.softmax(logits, dim=-1).detach() #outputs = torch.softmax((logits + logits_s1 + logits_s2) / 3, dim=-1).detach()

        log_outputs_s1 = torch.log_softmax(logits_s1, dim=-1)
        log_outputs_s2 = torch.log_softmax(logits_s2, dim=-1)
        log_outputs_s3 = torch.log_softmax(logits_s3, dim=-1)

        consist_loss1 = consistency_criterion(log_outputs_s1, outputs)
        consist_loss2 = consistency_criterion(log_outputs_s2, outputs)
        consist_loss3 = consistency_criterion(log_outputs_s3, outputs)

        log_outputs1 = torch.log_softmax(logits1, dim=-1)

        consist_loss01 = consistency_criterion(log_outputs1, outputs)

        lam = min((epoch / 100) * args.lam, args.lam)

        loss1 = (1 - lam) * loss_ce1 + lam * (consist_loss01 + consist_loss1)

        optimizer1.zero_grad()
        loss1.backward()
        optimizer1.step()

        loss = loss1

        if (i + 1) % args.print_freq == 0:
            print('Epoch [%d/%d], Iter [%d/%d] Training Accuracy: %.4F, Loss: %.4f, Lamda: %.4f'
                  % (epoch + 1, args.n_epoch, i + 1, len(train_dataset) // batch_size, prec, loss.data, lam))

    train_acc = float(train_correct) / float(train_total)
    return train_acc


##########################################################################################################################################################################################
########################################################################################################################################################################## New Code Ending
##########################################################################################################################################################################################

# Evaluate the Model
def evaluate(test_loader, model):
    model.eval()  # Change model to 'eval' mode.
    print('previous_best', best_acc_)
    correct = 0
    total = 0
    for images, _, _, _, labels, _ in test_loader:
        images = Variable(images).cuda()
        logits = model(images)
        outputs = F.softmax(logits, dim=1)
        _, pred = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (pred.cpu() == labels).sum()
    acc = 100 * float(correct) / float(total)

    return acc


# Evaluate the Model
def evaluate_ensemble(test_loader, model1):
    model1.eval()  # Change model to 'eval' mode.

    print('previous_best', best_acc_)
    correct = 0
    total = 0
    for images, _, _, _, labels, _ in test_loader:
        images = Variable(images).cuda()
        logits1 = model1(images)
        logits = logits1
        outputs = F.softmax(logits, dim=1)
        _, pred = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (pred.cpu() == labels).sum()
    acc = 100 * float(correct) / float(total)

    return acc


#####################################main code ################################################
args = parser.parse_args()
################################
os.environ["WANDB_API_KEY"] = "1038e371f2c76d1a383f9c4757c3de55f142fa93"
os.environ["WANDB_MODE"] = "dryrun"
wandb.login()
wandb.init(project="CT", entity="wu-dongdong", config=args, name=args.method)
logging.basicConfig(format='[%(asctime)s] - %(message)s',
                    datefmt='%Y/%m/%d %H:%M:%S',
                    level=logging.DEBUG,
                    handlers=[
                        logging.StreamHandler()
                    ])
#########################
# Seed
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)

# Hyper Parameters
batch_size = 128
learning_rate = args.lr
noise_type_map = {'clean': 'clean_label', 'worst': 'worse_label', 'aggre': 'aggre_label', 'rand1': 'random_label1',
                  'rand2': 'random_label2', 'rand3': 'random_label3', 'clean100': 'clean_label',
                  'noisy100': 'noisy_label'}
args.noise_type = noise_type_map[args.noise_type]
# load dataset
if args.noise_path is None:
    if args.dataset == 'cifar10':
        args.noise_path = './data/CIFAR-10_human.pt'
    elif args.dataset == 'cifar100':
        args.noise_path = './data/CIFAR-100_human.pt'
    else:
        raise NameError(f'Undefined dataset {args.dataset}')

train_dataset, test_dataset, num_classes, num_training_samples = input_dataset(args.dataset, args.noise_type,
                                                                               args.noise_path, args.is_human)
noise_prior = train_dataset.noise_prior
noise_or_not = train_dataset.noise_or_not
print('train_labels:', len(train_dataset.train_labels), train_dataset.train_labels[:10])
# load model


print('building model1...')
model1 = ResNet34(num_classes)
print('building model1 done')
optimizer1 = torch.optim.SGD(model1.parameters(), lr=learning_rate, weight_decay=0.0005, momentum=0.9)


train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=128,
                                           num_workers=args.num_workers,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=64,
                                          num_workers=args.num_workers,
                                          shuffle=False)
##########################################################################################################################################################################################
######################################################################################################################################################################## New Code Starting
##########################################################################################################################################################################################
alpha_plan = [0.1] * 60 + [0.05] * 15 + [0.01] * 15 + [0.001] * 10
##########################################################################################################################################################################################
########################################################################################################################################################################## New Code Ending
##########################################################################################################################################################################################
model1.cuda()

epoch = 0
train_acc = 0
best_acc_ = 0

# training
noise_prior_cur = noise_prior
for epoch in range(args.n_epoch):
    # train models
    print(f'epoch {epoch}')
    adjust_learning_rate(optimizer1, epoch, alpha_plan)
    model1.train()
    ##########################################################################################################################################################################################
    ######################################################################################################################################################################## New Code Starting
    ##########################################################################################################################################################################################
    train_acc1 = train_ours5(epoch, train_loader, model1, optimizer1, args)

    # evaluate models
    test_acc1 = evaluate(test_loader=test_loader, model=model1)
    print('test acc on test images is ', test_acc1)

    # evaluate all models
    test_acc_ensemble = evaluate_ensemble(test_loader, model1)
    # print('******test acc on test images is ', test_acc_ensemble, '*******')
    # save results
    logging.info('test acc on test images is {}'.format(test_acc_ensemble))
    best_acc_ = max(best_acc_, test_acc_ensemble)
    wandb.log({"test_acc": test_acc_ensemble, "epoch": epoch})
    wandb.run.summary["best_acc"] = best_acc_

    ##########################################################################################################################################################################################
    ########################################################################################################################################################################## New Code Ending
    ##########################################################################################################################################################################################