def main(): global args, switched args = parser.parse_args() print(args) criterion = get_criterion() if args.seed is not None: random.seed(args.seed) torch.manual_seed(args.seed) cudnn.deterministic = True if args.dataset == "mnist": (trainX, trainY), (testX, testY) = get_mnist() _trainY, _testY = binarize_mnist_class(trainY, testY) dataset_train1_clean = MNIST_Dataset_FixSample(1000, 60000, trainX, _trainY, testX, _testY, split='train', ids=[], increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, top = args.top1, type="clean", seed = args.seed) # clean dataset初始化为空 dataset_train1_noisy = MNIST_Dataset_FixSample(1000, 60000, trainX, _trainY, testX, _testY, split='train', increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, top = args.top1, type="noisy", seed = args.seed) dataset_train1_noisy.copy(dataset_train1_clean) # 和clean dataset使用相同的随机顺序 dataset_train1_noisy.reset_ids() # 让初始化的noisy dataset使用全部数据 dataset_test = MNIST_Dataset_FixSample(1000, 60000, trainX, _trainY, testX, _testY, split='test', increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, type="clean") dataset_train2_noisy = MNIST_Dataset_FixSample(1000, 60000, trainX, _trainY, testX, _testY, split='train', increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, top = args.top2, type="noisy", seed = args.seed) dataset_train2_clean = MNIST_Dataset_FixSample(1000, 60000, trainX, _trainY, testX, _testY, split='train', ids=[], increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, top = args.top2, type="clean", seed = args.seed) dataset_train2_noisy.copy(dataset_train1_noisy) dataset_train2_noisy.reset_ids() dataset_train2_clean.copy(dataset_train1_clean) #dataset_train2_clean.set_ids([]) assert np.all(dataset_train1_clean.X == dataset_train1_noisy.X) assert np.all(dataset_train2_clean.X == dataset_train1_noisy.X) assert np.all(dataset_train2_noisy.X == dataset_train1_noisy.X) elif args.dataset == 'cifar': data_transforms = { 'train': transforms.Compose([ transforms.ToPILImage(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]), 'val': transforms.Compose([ transforms.ToPILImage(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) } (trainX, trainY), (testX, testY) = get_cifar() _trainY, _testY = binarize_cifar_class(trainY, testY) dataset_train1_clean = CIFAR_Dataset(1000, 50000, trainX, _trainY, testX, _testY, split='train', ids=[], increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, top = args.top1, transform = data_transforms['train'], type="clean", seed = args.seed) # clean dataset初始化为空 dataset_train1_noisy = CIFAR_Dataset(1000, 50000, trainX, _trainY, testX, _testY, split='train', increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, top = args.top1, transform = data_transforms['train'], type="noisy", seed = args.seed) dataset_train1_noisy.copy(dataset_train1_clean) # 和clean dataset使用相同的随机顺序 dataset_train1_noisy.reset_ids() # 让初始化的noisy dataset使用全部数据 dataset_test = CIFAR_Dataset(1000, 50000, trainX, _trainY, testX, _testY, split='test', increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, transform = data_transforms['val'], type="clean", seed = args.seed) dataset_train2_noisy = CIFAR_Dataset(1000, 50000, trainX, _trainY, testX, _testY, split='train', increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, transform = data_transforms['train'], top = args.top2, type="noisy", seed = args.seed) dataset_train2_clean = CIFAR_Dataset(1000, 50000, trainX, _trainY, testX, _testY, split='train', ids=[], increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, transform = data_transforms['train'], top = args.top2, type="clean",) dataset_train2_noisy.copy(dataset_train1_noisy) dataset_train2_noisy.reset_ids() dataset_train2_clean.copy(dataset_train1_clean) #dataset_train2_clean.set_ids([]) assert np.all(dataset_train1_clean.X == dataset_train1_noisy.X) assert np.all(dataset_train2_clean.X == dataset_train1_noisy.X) assert np.all(dataset_train2_noisy.X == dataset_train1_noisy.X) assert np.all(dataset_train1_clean.Y == dataset_train1_noisy.Y) assert np.all(dataset_train2_clean.Y == dataset_train1_noisy.Y) assert np.all(dataset_train2_noisy.Y == dataset_train1_noisy.Y) assert np.all(dataset_train1_clean.T == dataset_train1_noisy.T) assert np.all(dataset_train2_clean.T == dataset_train1_noisy.T) assert np.all(dataset_train2_noisy.T == dataset_train1_noisy.T) criterion.update_p(0.4) dataloader_train1_clean = None dataloader_train1_noisy = DataLoader(dataset_train1_noisy, batch_size=args.batch_size, num_workers=args.workers, shuffle=False, pin_memory=True) dataloader_train2_clean = None dataloader_train2_noisy = DataLoader(dataset_train2_noisy, batch_size=args.batch_size, num_workers=args.workers, shuffle=False, pin_memory=True) dataloader_test = DataLoader(dataset_test, batch_size=args.batch_size, num_workers=args.workers, shuffle=False, pin_memory=True) consistency_criterion = losses.softmax_mse_loss if args.dataset == 'mnist': model1 = create_model() model2 = create_model() ema_model1 = create_model(ema = True) ema_model2 = create_model(ema = True) elif args.dataset == 'cifar': model1 = create_cifar_model() model2 = create_cifar_model() ema_model1 = create_cifar_model(ema = True) ema_model2 = create_cifar_model(ema = True) if args.gpu is not None: model1 = model1.cuda(args.gpu) model2 = model2.cuda(args.gpu) ema_model1 = ema_model1.cuda(args.gpu) ema_model2 = ema_model2.cuda(args.gpu) else: model1 = model1.cuda(args.gpu) model2 = model2.cuda(args.gpu) ema_model1 = ema_model1.cuda(args.gpu) ema_model2 = ema_model2.cuda(args.gpu) optimizer1 = torch.optim.Adam(model1.parameters(), lr=args.lr, weight_decay=args.weight_decay ) optimizer2 = torch.optim.Adam(model2.parameters(), lr=args.lr, weight_decay=args.weight_decay ) stats_ = stats(args.modeldir, 0) scheduler1 = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer1, args.epochs) scheduler2 = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer2, args.epochs) if args.evaluation: print("Evaluation mode!") best_acc1 = 0 best_acc2 = 0 best_acc3 = 0 best_acc4 = 0 best_acc = 0 best_model1 = copy.deepcopy(model1) best_model2 = copy.deepcopy(model2) for epoch in range(args.epochs): print("Self paced status: {}".format(check_self_paced(epoch))) print("Mean Teacher status: {}".format(check_mean_teacher(epoch))) if check_mean_teacher(epoch) and not check_mean_teacher(epoch - 1) and not switched: ema_model1.load_state_dict(best_model1.state_dict()) ema_model2.load_state_dict(best_model2.state_dict()) switched = True print("SWITCHED!") trainPacc, trainNacc, trainPNacc = train(dataloader_train1_clean, dataloader_train1_noisy, dataloader_train2_clean, dataloader_train2_noisy, model1, model2, ema_model1, ema_model2, criterion, consistency_criterion, optimizer1, scheduler1, optimizer2, scheduler2, epoch) valPacc, valNacc, valPNacc1, valPNacc2, valPNacc3, valPNacc4 = validate(dataloader_test, model1, model2, ema_model1, ema_model2, epoch) #print(valPacc, valNacc, valPNacc1, valPNacc2, valPNacc3) stats_._update(trainPacc, trainNacc, trainPNacc, valPacc, valNacc, valPNacc1) if valPNacc1 > best_acc1 and not check_mean_teacher(epoch): best_model1 = copy.deepcopy(model1) if valPNacc3 > best_acc3 and not check_mean_teacher(epoch): best_model2 = copy.deepcopy(model2) best_acc1 = max(valPNacc1, best_acc1) best_acc2 = max(valPNacc2, best_acc2) best_acc3 = max(valPNacc3, best_acc3) best_acc4 = max(valPNacc4, best_acc4) all_accuracy = [valPNacc1, valPNacc2, valPNacc3, valPNacc4] models = [model1, model2, ema_model1, ema_model2] if (check_self_paced(epoch)) and (epoch - args.self_paced_start) % args.self_paced_frequency == 0: dataloader_train1_clean, dataloader_train1_noisy, dataloader_train2_clean, dataloader_train2_noisy = update_dataset(model1, model2, ema_model1, ema_model2, dataset_train1_clean, dataset_train1_noisy, dataset_train2_clean, dataset_train2_noisy, epoch) plot_curve(stats_, args.modeldir, 'model', True) if (max(all_accuracy) > best_acc): torch.save({ 'epoch': epoch + 1, 'state_dict': models[all_accuracy.index(max(all_accuracy))].state_dict(), 'best_prec1': best_acc1, }, 'model_best.pth.tar') best_acc = max(all_accuracy) dataset_train1_noisy.shuffle() dataset_train2_noisy.shuffle() dataloader_train1_noisy = DataLoader(dataset_train1_noisy, batch_size=args.batch_size, num_workers=args.workers, shuffle=False, pin_memory=True) dataloader_train2_noisy = DataLoader(dataset_train2_noisy, batch_size=args.batch_size, num_workers=args.workers, shuffle=False, pin_memory=True) print(best_acc1) print(best_acc2) print(best_acc3) print(best_acc4)
def main(): global args, switched, single_epoch_steps, step args = parser.parse_args() criterion = get_criterion() torch.cuda.set_device(int(args.gpu)) cudnn.benchmark = True if args.dataset == "mnist": (trainX, trainY), (testX, testY) = get_mnist() _trainY, _testY = binarize_mnist_class(trainY, testY) dataset_train_clean = MNIST_Dataset(1000, 60000, trainX, _trainY, testX, _testY, split='train', ids=[], increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, top = args.top, type="clean", seed = args.seed) # clean dataset初始化为空 dataset_train_noisy = MNIST_Dataset(1000, 60000, trainX, _trainY, testX, _testY, split='train', increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, top = args.top, type="noisy", seed = args.seed) dataset_train_noisy.copy(dataset_train_clean) # 和clean dataset使用相同的随机顺序 dataset_train_noisy.reset_ids() # 让初始化的noisy dataset使用全部数据 dataset_test = MNIST_Dataset(1000, 60000, trainX, _trainY, testX, _testY, split='test', increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, top = args.top, type="clean", seed = args.seed) elif args.dataset == 'cifar': data_transforms = { 'train': transforms.Compose([ transforms.ToPILImage(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]), 'val': transforms.Compose([ transforms.ToPILImage(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) } (trainX, trainY), (testX, testY) = get_cifar() _trainY, _testY = binarize_cifar_class(trainY, testY) dataset_train_clean = CIFAR_Dataset(1000, 50000, trainX, _trainY, testX, _testY, split='train', ids=[], increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, top = args.top, transform = data_transforms['train'], type="clean", seed = args.seed) # clean dataset初始化为空 dataset_train_noisy = CIFAR_Dataset(1000, 50000, trainX, _trainY, testX, _testY, split='train', increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, top = args.top, transform = data_transforms['train'], type="noisy", seed = args.seed) dataset_train_noisy.copy(dataset_train_clean) # 和clean dataset使用相同的随机顺序 dataset_train_noisy.reset_ids() # 让初始化的noisy dataset使用全部数据 dataset_test = CIFAR_Dataset(1000, 50000, trainX, _trainY, testX, _testY, split='test', increasing=args.increasing, replacement=args.replacement, mode=args.self_paced_type, top = args.top, transform = data_transforms['val'], type="clean", seed = args.seed) criterion.update_p(0.4) assert np.all(dataset_train_noisy.X == dataset_train_clean.X) assert np.all(dataset_train_noisy.Y == dataset_train_clean.Y) assert np.all(dataset_train_noisy.oids == dataset_train_clean.oids) assert np.all(dataset_train_noisy.T == dataset_train_clean.T) #step = args.ema_start * 2 + 1 if len(dataset_train_clean) > 0: dataloader_train_clean = DataLoader(dataset_train_clean, batch_size=args.batch_size, num_workers=args.workers, shuffle=True, pin_memory=True) else: dataloader_train_clean = None if len(dataset_train_noisy) > 0: dataloader_train_noisy = DataLoader(dataset_train_noisy, batch_size=args.batch_size, num_workers=args.workers, shuffle=False, pin_memory=True) else: dataloader_train_noisy = None if len(dataset_test): dataloader_test = DataLoader(dataset_test, batch_size=args.batch_size, num_workers=0, shuffle=False, pin_memory=True) else: dataloader_test = None single_epoch_steps = len(dataloader_train_noisy) + 1 print('Steps: {}'.format(single_epoch_steps)) consistency_criterion = losses.softmax_mse_loss if args.dataset == 'mnist': model = create_model() ema_model = create_model(ema = True) elif args.dataset == 'cifar': model = create_cifar_model() ema_model = create_cifar_model(ema = True) if args.gpu is not None: model = model.cuda() ema_model = ema_model.cuda() else: model = model.cuda() ema_model = ema_model.cuda() params_list = [{'params': model.parameters(), 'lr': args.lr},] #optimizer = torch.optim.Adam(params_list, lr=args.lr, # weight_decay=args.weight_decay #) optimizer = torch.optim.Adam(params_list, lr=args.lr, weight_decay=args.weight_decay ) stats_ = stats(args.modeldir, 0) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.epochs, eta_min = args.lr * 0.2) #scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[40, 80], gamma=0.6) if args.evaluation: print("Evaluation mode!") best_acc = 0 for epoch in range(args.warmup): print("Warming up {}/{}".format(epoch + 1, args.warmup)) trainPacc, trainNacc, trainPNacc = train(dataloader_train_clean, dataloader_train_noisy, model, ema_model, criterion, consistency_criterion, optimizer, scheduler, -1, warmup = True) valPacc, valNacc, valPNacc = validate(dataloader_test, model, ema_model, criterion, consistency_criterion, -1) dataset_train_noisy.shuffle() dataloader_train_noisy = DataLoader(dataset_train_noisy, batch_size=args.batch_size, num_workers=args.workers, shuffle=False, pin_memory=True) val = [] for epoch in range(args.epochs): print("Self paced status: {}".format(check_self_paced(epoch))) print("Mean teacher status: {}".format(check_mean_teacher(epoch))) print("Noisy status: {}".format(check_noisy(epoch))) if check_mean_teacher(epoch) and (not check_mean_teacher(epoch - 1)) and not switched: model.eval() ema_model.eval() ema_model.load_state_dict(model.state_dict()) switched = True print("SWITCHED!") validate(dataloader_test, model, ema_model, criterion, consistency_criterion, epoch) validate(dataloader_test, ema_model, model, criterion, consistency_criterion, epoch) model.train() ema_model.train() if epoch == 0: switched = False if (not check_mean_teacher(epoch)) and check_mean_teacher(epoch - 1) and not switched: model.eval() ema_model.eval() model.load_state_dict(ema_model.state_dict()) switched = True print("SWITCHED!") validate(dataloader_test, model, ema_model, criterion, consistency_criterion, epoch) validate(dataloader_test, ema_model, model, criterion, consistency_criterion, epoch) model.train() ema_model.train() trainPacc, trainNacc, trainPNacc = train(dataloader_train_clean, dataloader_train_noisy, model, ema_model, criterion, consistency_criterion, optimizer, scheduler, epoch, self_paced_pick = len(dataset_train_clean)) valPacc, valNacc, valPNacc = validate(dataloader_test, model, ema_model, criterion, consistency_criterion, epoch) val.append(valPNacc) #validate_2(dataloader_test, model, ema_model, criterion, consistency_criterion, epoch) stats_._update(trainPacc, trainNacc, trainPNacc, valPacc, valNacc, valPNacc) is_best = valPNacc > best_acc best_acc = max(valPNacc, best_acc) filename = [] filename.append(os.path.join(args.modeldir, 'checkpoint.pth.tar')) filename.append(os.path.join(args.modeldir, 'model_best.pth.tar')) if (check_self_paced(epoch)) and (epoch - args.self_paced_start) % args.self_paced_frequency == 0: dataloader_train_clean, dataloader_train_noisy = update_dataset(model, ema_model, dataset_train_clean, dataset_train_noisy, epoch) plot_curve(stats_, args.modeldir, 'model', True) save_checkpoint({ 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'best_prec1': best_acc, 'optimizer' : optimizer.state_dict(), }, is_best, filename) dataset_train_noisy.shuffle() #dataloader_train_clean = DataLoader(dataset_train_clean, batch_size=args.batch_size, num_workers=args.workers, shuffle=True, pin_memory=True) dataloader_train_noisy = DataLoader(dataset_train_noisy, batch_size=args.batch_size, num_workers=args.workers, shuffle=False, pin_memory=True) print(best_acc) print(val)
def main(): global args, switched args = parser.parse_args() print(args) if args.seed is not None: random.seed(args.seed) torch.manual_seed(args.seed) cudnn.deterministic = True if args.dataset == "mnist": (trainX, trainY), (testX, testY) = get_mnist() _trainY, _testY = binarize_mnist_class(trainY, testY) dataset_test = MNIST_Dataset_FixSample(1000, 60000, trainX, _trainY, testX, _testY, split='test', type="clean", seed=args.seed) elif args.dataset == 'cifar': data_transforms = { 'train': transforms.Compose([ transforms.ToPILImage(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]), 'val': transforms.Compose([ transforms.ToPILImage(), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) } (trainX, trainY), (testX, testY) = get_cifar() _trainY, _testY = binarize_cifar_class(trainY, testY) dataset_test = CIFAR_Dataset(1000, 50000, trainX, _trainY, testX, _testY, split='test', transform=data_transforms['val'], type="clean", seed=args.seed) dataloader_test = DataLoader(dataset_test, batch_size=1, num_workers=args.workers, shuffle=False, pin_memory=True) consistency_criterion = losses.softmax_mse_loss if args.dataset == 'mnist': model = create_model() elif args.dataset == 'cifar': model = create_cifar_model() if args.gpu is not None: model = model.cuda() else: model = model.cuda() print("Evaluation mode!") if args.model is None: raise RuntimeError("Please specify a model file.") else: state_dict = torch.load(args.model)['state_dict'] model.load_state_dict(state_dict) valPacc, valNacc, valPNacc = validate(dataloader_test, model)