text: 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 fix_seed(seed=888):
      np.random.seed(seed)
      random.seed(seed)
      torch.manual_seed(seed)
      torch.cuda.manual_seed_all(seed)
      torch.backends.cudnn.deterministic = True
      np.random.seed(seed)
              
  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='', teacher_idx=None, truncate_mode=None, refinement=None): 
          
          self.r = r # noise ratio
          self.transform = transform
          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
          
          # For distill test
          self.teacher_idx = teacher_idx
          self.truncate_mode = truncate_mode
          self.train_label = None
          self.refinement = refinement
       
          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.train_label = train_label
              
              if os.path.exists(noise_file):
                  noise_label = json.load(open(noise_file,"r"))
              else:    #inject noise   
                  fix_seed()
                  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 = noise_label
                  if self.truncate_mode == 'initial':
                      self.train_data = self.train_data[teacher_idx]
                      self.noise_label = [noise_label[i] for i in teacher_idx]
              else:                   
                  if self.mode == "labeled":
                      pred_idx = pred.nonzero()[0]
                      if self.truncate_mode == 'initial':
                          pred_idx = pred_idx.tolist()
                          teacher_idx = teacher_idx.tolist()
                          pred_idx = list(set(pred_idx) & set(teacher_idx))
                          pred_idx = torch.tensor(pred_idx)
                      
                      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()      
                      log.write('Numer of labeled samples:%d   AUC:%.3f\n'%(pred.sum(),auc))
                      log.flush()      
                      
                  elif self.mode == "unlabeled":
                      pred_idx = (1-pred).nonzero()[0]
                      if self.truncate_mode == 'initial':
                          whole_idx = list(range(50000))
                          pred_idx = pred_idx.tolist()
                          teacher_idx = teacher_idx.tolist()
                          tmp_set = set(whole_idx) - set(teacher_idx)
                          tmp_set = tmp_set | set(pred_idx)
                          pred_idx = torch.tensor(list(tmp_set))
                      
                  elif self.mode == "labeled_svd":
                      if self.refinement:
                          pred_idx = pred.nonzero()[0]
                          pred_idx_set = set(pred_idx.tolist())
                          teacher_idx_set = set(teacher_idx.tolist())
                          pred_idx = torch.tensor(list(pred_idx_set & teacher_idx_set))
                          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()               
                          log.write('Numer of labeled samples:%d   AUC:%.3f\n'%(pred.sum(),auc))
                          log.flush()
                      else:
                          pred_idx = teacher_idx
                          probability = torch.ones(50000,)
                          self.probability = [probability[i] for i in pred_idx]
  
                          log.write('Number of labeled samples (by svd) : %d' % teacher_idx.shape[0])
                  
                  elif self.mode == "unlabeled_svd":
                      if self.refinement:
                          clean_pred_idx = pred.nonzero()[0]
                          clean_pred_idx_set = set(clean_pred_idx.tolist())
                          teacher_idx_set = set(teacher_idx.tolist())
                          all_idx_set = set(range(50000))
                          pred_idx = torch.tensor(list(all_idx_set - (clean_pred_idx_set & teacher_idx_set)))                    
                      else:
                          pred_idx = torch.arange(0, 50000)
                          pred_idx_set = set(pred_idx.tolist()) - set(teacher_idx.tolist())
                          pred_idx = torch.tensor(list(pred_idx_set))
                  
                  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 __getitem__(self, index):
          if self.mode=='labeled' or self.mode=='labeled_svd':
              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' or self.mode=='unlabeled_svd':
              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
             
      def __len__(self):
          if self.mode!='test':
              return len(self.train_data)
          else:
              return len(self.test_data)
      
      def _print_statistics(self, teacher_idx):
          fn = 0
          tn = 0
          tp = 0
          fp = 0
          noisy = 0
          clean = 0
          for i in teacher_idx:
              if self.noise_label[i] == self.train_label[i]:
                  tn += 1
              else:
                  fn += 1
          for i in range(50000):
              if self.noise_label[i] == self.train_label[i]:
                  clean += 1
              else:
                  noisy += 1
          fp = clean - tn
          tp = noisy - fn
          print("precision : " + str(tp / (tp + fp)))
          print("recall : " + str(tp / (tp + fn)))
          print("specificity : " + str(tn / (tn + fp)))      
          
          
  class cifar_dataloader():  
      def __init__(self, dataset, r, noise_mode, batch_size, num_workers, root_dir, log, noise_file='', _teacher_idx=None, _truncate_mode=None):
          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
          self.teacher_idx = _teacher_idx
          self.truncate_mode = _truncate_mode
          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 print_statistics(self, teacher_idx):
          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,teacher_idx=self.teacher_idx,truncate_mode=self.truncate_mode)
          all_dataset._print_statistics(teacher_idx)
      
      def run(self,mode,pred=[],prob=[], teacher_idx=None, refinement=None):
          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,teacher_idx=self.teacher_idx,truncate_mode=self.truncate_mode)
              trainloader = DataLoader(
                  dataset=all_dataset, 
                  batch_size=self.batch_size*2,
                  shuffle=True,
                  num_workers=self.num_workers)             
              return trainloader
          
          elif mode=='train_svd':
              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_svd", noise_file=self.noise_file, pred=pred, probability=prob,log=self.log,teacher_idx=teacher_idx, refinement=refinement)
              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_svd", noise_file=self.noise_file, pred=pred, probability=prob,log=self.log,teacher_idx=teacher_idx, refinement=refinement)
              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=='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,teacher_idx=self.teacher_idx,truncate_mode=self.truncate_mode)              
              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,teacher_idx=self.teacher_idx,truncate_mode=self.truncate_mode)                    
              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        