import os
import pickle
import torchvision as tv
import numpy as np
from PIL import Image
from torch.utils.data import Dataset
import random
import time
from IPython import embed
def parse_args():
parser = argparse.ArgumentParser(description='command for the first train')
parser.add_argument('--lr', type=float, default=0.1, help='learning rate')
parser.add_argument('--batch_size', type=int, default=128, help='#images in each mini-batch')
parser.add_argument('--test_batch_size', type=int, default=100, help='#images in each mini-batch')
parser.add_argument('--cuda_dev', type=int, default=0, help='GPU to select')
parser.add_argument('--epoch', type=int, default=200, help='training epoches')
parser.add_argument('--num_classes', type=int, default=10, help='Number of in-distribution classes')
parser.add_argument('--wd', type=float, default=1e-4, help='weight decay')
parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
parser.add_argument('--noise_type', default='asymmetric', help='symmetric or asymmetric')
parser.add_argument('--train_root', default='./data', help='root for train data')
parser.add_argument('--noise_ratio', type=float, default=0.4, help='percent of noise')
parser.add_argument('--out', type=str, default='./data/model_data', help='Directory of the output')
parser.add_argument('--alpha', type=float, default=1.0, help='Beta distribution parameter for mixup')
parser.add_argument('--download', type=bool, default=False, help='download dataset')
parser.add_argument('--network', type=str, default='PR18', help='Network architecture')
parser.add_argument('--seed_initialization', type=int, default=1, help='random seed (default: 1)')
parser.add_argument('--seed_dataset', type=int, default=42, help='random seed (default: 1)')
parser.add_argument('--M', action='append', type=int, default=[], help="Milestones for the LR sheduler")
parser.add_argument('--experiment_name', type=str, default = 'Proof',help='name of the experiment (for the output files)')
parser.add_argument('--method', type=str, default='MOIT', help='MOIT')
parser.add_argument('--dataset', type=str, default='CIFAR-10', help='CIFAR-10, CIFAR-100')
parser.add_argument('--initial_epoch', type=int, default=1, help="Star training at initial_epoch")
parser.add_argument('--DA', type=str, default="complex", help='Choose simple or complex data augmentation')
parser.add_argument('--low_dim', type=int, default=128, help='Size of contrastive learning embedding')
parser.add_argument('--mix_labels', type=int, default=0, help='1: Interpolate two input images and "interpolate" labels')
parser.add_argument('--batch_t', default=0.1, type=float, help='Contrastive learning temperature')
parser.add_argument('--aprox', type=int, default=1, help='Warm-up epochs')
parser.add_argument('--headType', type=str, default="Linear", help='Linear, NonLinear')
parser.add_argument('--xbm_use', type=int, default=1, help='1: Use xbm')
parser.add_argument('--xbm_begin', type=int, default=1, help='Epoch to begin using memory')
parser.add_argument('--xbm_per_class', type=int, default=20, help='Num of samples per class to store in the memory. Memory size = xbm_per_class*num_classes')
parser.add_argument('--startLabelCorrection', type=int, default=9999, help='Epoch to start label correction')
parser.add_argument('--k_val', type=int, default=5, help='k for k-nn correction')
parser.add_argument('--use_cleanLabels', type=int, default=0, help='Train the classifier with clean labels')
parser.add_argument('--PredictiveCorrection', type=int, default=0, help='Enable predictive label correction')
parser.add_argument('--balance_crit', type=str, default="none", help='None, max, min. median')
parser.add_argument('--discrepancy_corrected', type=int, default=1, help='Use corrected label for discrepancy measure')
parser.add_argument('--validation_exp', type=int, default=0, help='Using clean train subset for validation')
parser.add_argument('--val_samples', type=int, default=5000, help='Number of samples used for validation')
args = parser.parse_args()
return args
def get_dataset(args, transform_train, transform_test):
# prepare datasets
if args.validation_exp == 1:
#################################### Train set #############################################
temp_dataset = Cifar10Train(args, train=True, transform=transform_train, target_transform=transform_test, download=args.download)
train_indexes, val_indexes = train_val_split(args, temp_dataset.targets)
cifar_train = Cifar10Train(args, train=True, transform=transform_train, target_transform=transform_test, sample_indexes=train_indexes)
#################################### Noise corruption ######################################
if args.noise_type == "symmetric":
cifar_train.random_in_noise()
elif args.noise_type == "asymmetric":
cifar_train.real_in_noise()
else:
print('No noise')
cifar_train.labelsNoisyOriginal = cifar_train.targets.copy()
#################################### Test set #############################################
testset = Cifar10Train(args, train=True, transform=transform_test, sample_indexes=val_indexes)
###########################################################################################
else:
#################################### Train set #############################################
cifar_train = Cifar10Train(args, train=True, transform=transform_train, target_transform=transform_test, download=args.download)
#################################### Noise corruption ######################################
if args.noise_type == 'symmetric':
cifar_train.random_in_noise()
elif args.noise_type == 'asymmetric':
cifar_train.real_in_noise()
else:
print ('No noise')
cifar_train.labelsNoisyOriginal = cifar_train.targets.copy()
#################################### Test set #############################################
testset = tv.datasets.CIFAR10(root=args.train_root, train=False, download=False, transform=transform_test)
###########################################################################################
return cifar_train, testset, cifar_train.clean_labels, cifar_train.noisy_labels, cifar_train.noisy_indexes, cifar_train.labelsNoisyOriginal
def train_val_split(args, train_val):
np.random.seed(args.seed_dataset)
train_val = np.array(train_val)
train_indexes = []
val_indexes = []
val_num = int(args.val_samples / args.num_classes)
for id in range(args.num_classes):
indexes = np.where(train_val == id)[0]
np.random.shuffle(indexes)
val_indexes.extend(indexes[:val_num])
train_indexes.extend(indexes[val_num:])
np.random.shuffle(train_indexes)
np.random.shuffle(val_indexes)
return train_indexes, val_indexes
class Cifar10Train(tv.datasets.CIFAR10):
def __init__(self, args, train=True, transform=None, target_transform=None, sample_indexes=None, download=False):
super(Cifar10Train, self).__init__(args.train_root, train=train, transform=transform,
target_transform=target_transform, download=download)
self.root = os.path.expanduser(args.train_root)
self.transform = transform
self.target_transform = target_transform
self.train = train # Training set or validation set
self.args = args
if sample_indexes is not None:
self.data = self.data[sample_indexes]
self.targets = list(np.asarray(self.targets)[sample_indexes])
self.num_classes = self.args.num_classes
self.in_index = []
self.out_index = []
self.noisy_indexes = []
self.clean_indexes = []
self.clean_labels = []
self.noisy_labels = []
self.out_data = []
self.out_labels = []
self.soft_labels = []
self.labelsNoisyOriginal = []
self._num = []
self._count = 1
self.prediction = []
self.confusion_matrix_in = np.array([])
self.confusion_matrix_out = np.array([])
self.labeled_idx = []
self.unlabeled_idx = []
# From in ou split function:
self.soft_labels = np.zeros((len(self.targets), self.num_classes), dtype=np.float32)
self._num = int(len(self.targets) * self.args.noise_ratio)
################# Random in-distribution noise #########################
def random_in_noise(self):
# to be more equal, every category can be processed separately
np.random.seed(self.args.seed_dataset)
idxes = np.random.permutation(len(self.targets))
clean_labels = np.copy(self.targets)
noisy_indexes = idxes[0:self._num]
clean_indexes = idxes[self._num:]
for i in range(len(idxes)):
if i < self._num:
self.soft_labels[idxes[i]][self.targets[idxes[i]]] = 0 ## Remove soft-label created during label mapping
# targets[idxes[i]] -> another category
label_sym = np.random.randint(self.num_classes, dtype=np.int32)
while(label_sym==self.targets[idxes[i]]):#To exclude the original label
label_sym = np.random.randint(self.num_classes, dtype=np.int32)
self.targets[idxes[i]] = label_sym
self.soft_labels[idxes[i]][self.targets[idxes[i]]] = 1
self.targets = np.asarray(self.targets, dtype=np.long)
self.noisy_labels = np.copy(self.targets)
self.noisy_indexes = noisy_indexes
self.clean_labels = clean_labels
self.clean_indexes = clean_indexes
self.confusion_matrix_in = (np.ones((self.args.num_classes, self.args.num_classes)) - np.identity(self.args.num_classes))\
*(self.args.noise_ratio/(self.num_classes -1)) + \
np.identity(self.args.num_classes)*(1 - self.args.noise_ratio)
##########################################################################
################# Real in-distribution noise #########################
def real_in_noise(self):
# to be more equal, every category can be processed separately
np.random.seed(self.args.seed_dataset)
##### Create te confusion matrix #####
self.confusion_matrix_in = np.identity(self.args.num_classes)
# truck -> automobile
self.confusion_matrix_in[9, 9] = 1 - self.args.noise_ratio
self.confusion_matrix_in[9, 1] = self.args.noise_ratio
# bird -> airplane
self.confusion_matrix_in[2, 2] = 1 - self.args.noise_ratio
self.confusion_matrix_in[2, 0] = self.args.noise_ratio
# cat -> dog
self.confusion_matrix_in[3, 3] = 1 - self.args.noise_ratio
self.confusion_matrix_in[3, 5] = self.args.noise_ratio
# dog -> cat
self.confusion_matrix_in[5, 5] = 1 - self.args.noise_ratio
self.confusion_matrix_in[5, 3] = self.args.noise_ratio
# deer -> horse
self.confusion_matrix_in[4, 4] = 1 - self.args.noise_ratio
self.confusion_matrix_in[4, 7] = self.args.noise_ratio
idxes = np.random.permutation(len(self.targets))
clean_labels = np.copy(self.targets)
for i in range(len(idxes)):
self.soft_labels[idxes[i]][self.targets[idxes[i]]] = 0 ## Remove soft-label created during label mapping
current_label = self.targets[idxes[i]]
if self._num > 0:
# current_label = self.targets[idxes[i]]
conf_vec = self.confusion_matrix_in[current_label,:]
label_sym = np.random.choice(np.arange(0, self.num_classes), p=conf_vec.transpose())
self.targets[idxes[i]] = label_sym
else:
label_sym = current_label
self.soft_labels[idxes[i]][self.targets[idxes[i]]] = 1
if label_sym == current_label:
self.clean_indexes.append(idxes[i])
else:
self.noisy_indexes.append(idxes[i])
self.targets = np.asarray(self.targets, dtype=np.long)
self.clean_indexes = np.asarray(self.clean_indexes, dtype=np.long)
self.noisy_indexes = np.asarray(self.noisy_indexes, dtype=np.long)
self.noisy_labels = self.targets
self.clean_labels = clean_labels
def __getitem__(self, index):
if len(self.data) > self.args.val_samples:
img, labels, soft_labels, noisy_labels, clean_labels = self.data[index], self.targets[index], self.soft_labels[
index], self.labelsNoisyOriginal[index], self.clean_labels[index]
# doing this so that it is consistent with all other datasets.
img = Image.fromarray(img)
img_noDA = self.target_transform(img)
img1 = self.transform(img)
if self.args.method == "MOIT":
if self.train is True:
img2 = self.transform(img)
else:
img2 = 0
return img1, img2, img_noDA, labels, soft_labels, index, noisy_labels, clean_labels
else:
return img1, img_noDA, labels, soft_labels, index, noisy_labels, clean_labels
else:
img, labels = self.data[index], self.targets[index]
# doing this so that it is consistent with all other datasets.
img = Image.fromarray(img)
img = self.transform(img)
return img, labels
创建噪声标签的工具类
并且自定义了Dataset,从而能够得到想要的返回
return img1, img2, img_noDA, labels, soft_labels, index, noisy_labels, clean_labels