def main(): global best_acc if not os.path.isdir(args.out): mkdir_p(args.out) # Data print(f'==> Preparing cifar10') transform_train = transforms.Compose([ dataset.RandomPadandCrop(32), dataset.RandomFlip(), dataset.ToTensor(), ]) transform_val = transforms.Compose([ dataset.ToTensor(), ]) train_labeled_set, train_unlabeled_set, val_set, test_set = dataset.get_cifar10('./data', args.n_labeled, transform_train=transform_train, transform_val=transform_val) labeled_trainloader = data.DataLoader(train_labeled_set, batch_size=args.batch_size, shuffle=True, num_workers=0, drop_last=True) unlabeled_trainloader = data.DataLoader(train_unlabeled_set, batch_size=args.batch_size, shuffle=True, num_workers=0, drop_last=True) val_loader = data.DataLoader(val_set, batch_size=args.batch_size, shuffle=False, num_workers=0) test_loader = data.DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=0) # Model print("==> creating WRN-28-2") def create_model(ema=False): model = models.WideResNet(num_classes=10) model = model.cuda() if ema: for param in model.parameters(): param.detach_() return model model = create_model() ema_model = create_model(ema=True) cudnn.benchmark = True print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters())/1000000.0)) train_criterion = SemiLoss() criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=args.lr) ema_optimizer= WeightEMA(model, ema_model, alpha=args.ema_decay) start_epoch = 0 # Resume title = 'noisy-cifar-10' if args.resume: # Load checkpoint. print('==> Resuming from checkpoint..') assert os.path.isfile(args.resume), 'Error: no checkpoint directory found!' args.out = os.path.dirname(args.resume) checkpoint = torch.load(args.resume) best_acc = checkpoint['best_acc'] start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) ema_model.load_state_dict(checkpoint['ema_state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger = Logger(os.path.join(args.out, 'log.txt'), title=title, resume=True) else: logger = Logger(os.path.join(args.out, 'log.txt'), title=title) logger.set_names(['Train Loss', 'Train Loss X', 'Train Loss U', 'Valid Loss', 'Valid Acc.', 'Test Loss', 'Test Acc.']) writer = SummaryWriter(args.out) step = 0 test_accs = [] no_improvement = 0 # Train and val for epoch in range(start_epoch, args.epochs): print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, args.epochs, state['lr'])) if args.sup_only: train_loss, train_loss_x, train_loss_u = train_sup(labeled_trainloader, model, optimizer, ema_optimizer, criterion, epoch, use_cuda) else: train_loss, train_loss_x, train_loss_u = train(labeled_trainloader, unlabeled_trainloader, model, optimizer, ema_optimizer, train_criterion, epoch, use_cuda) _, train_acc = validate(labeled_trainloader, ema_model, criterion, epoch, use_cuda, mode='Train Stats') val_loss, val_acc = validate(val_loader, ema_model, criterion, epoch, use_cuda, mode='Valid Stats') test_loss, test_acc = validate(test_loader, ema_model, criterion, epoch, use_cuda, mode='Test Stats ') step = args.val_iteration * (epoch + 1) writer.add_scalar('losses/train_loss', train_loss, step) writer.add_scalar('losses/valid_loss', val_loss, step) writer.add_scalar('losses/test_loss', test_loss, step) writer.add_scalar('accuracy/train_acc', train_acc, step) writer.add_scalar('accuracy/val_acc', val_acc, step) writer.add_scalar('accuracy/test_acc', test_acc, step) # append logger file logger.append([train_loss, train_loss_x, train_loss_u, val_loss, val_acc, test_loss, test_acc]) # early stopping if args.early_stopping: if no_improvement >= args.early_stopping: break if val_acc > best_acc: no_improvement = 0 else: no_improvement += 1 # save model is_best = val_acc > best_acc best_acc = max(val_acc, best_acc) save_checkpoint({ 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'ema_state_dict': ema_model.state_dict(), 'acc': val_acc, 'best_acc': best_acc, 'optimizer' : optimizer.state_dict(), }, is_best) test_accs.append(test_acc) logger.close() writer.close() print('Best acc:') print(best_acc) print('Mean acc:') print(np.mean(test_accs[-20:]))
def main(): global best_acc if not os.path.isdir(args.out): mkdir_p(args.out) import dataset.cifar10 as dataset print(f'==> Preparing cifar10') transform_train = transforms.Compose([ dataset.RandomPadandCrop(32), dataset.RandomFlip(), dataset.ToTensor(), ]) transform_val = transforms.Compose([ dataset.ToTensor(), ]) train_labeled_set, train_unlabeled_set, val_set = dataset.get_cifar10('./data', args.n_labeled, transform_train=transform_train, transform_val=transform_val) num_labl = 10 labeled_trainloader = data.DataLoader(train_labeled_set, batch_size=args.batch_size, shuffle=True, num_workers=0, drop_last=True) unlabeled_trainloader = data.DataLoader(train_unlabeled_set, batch_size=args.batch_size, shuffle=True, num_workers=0, drop_last=True) val_loader = data.DataLoader(val_set, batch_size=args.batch_size * 4, shuffle=False, num_workers=0) # Model print("==> creating WRN-28-2") def create_model(ema=False): model = models.WideResNet(num_classes=num_labl) model = model.cuda() if ema: for param in model.parameters(): param.detach_() return model model = create_model() ema_model = create_model(ema=True) cudnn.benchmark = True print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters())/1000000.0)) train_criterion = SemiLoss() criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=args.lr) if args.ALR: from ALR import ALR alr = ALR( d_X=lambda x, x_hat: torch.norm((x - x_hat).view(x.size(0), -1), p=2, dim=1, keepdim=True), d_Y=lambda y, y_hat: kl_divergence(logits_p=y, logits_q=y_hat), lambda_lp=args.alr_lambda, eps_min=0.1, eps_max=10, xi=10, ip=2, K=0 #1 ) ema_optimizer= WeightEMA(model, ema_model, alpha=args.ema_decay) start_epoch = 0 # Resume title = 'cifar-10' if args.resume: # Load checkpoint. print('==> Resuming from checkpoint..') assert os.path.isfile(args.resume), 'Error: no checkpoint directory found!' args.out = os.path.dirname(args.resume) checkpoint = torch.load(args.resume) best_acc = checkpoint['best_acc'] start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) ema_model.load_state_dict(checkpoint['ema_state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) logger = Logger(os.path.join(args.out, 'log.txt'), title=title, resume=True) else: logger = Logger(os.path.join(args.out, 'log.txt'), title=title) logger.set_names(['Train Loss', 'Train Loss X', 'Train Loss U', 'Valid Loss', 'Valid Acc.']) writer = SummaryWriter(args.out) step = 0 test_accs = [] train_acc = 0 ssize = 64 # Train and val for epoch in range(start_epoch, args.epochs): print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, args.epochs, state['lr'])) if args.ALR: train_loss, train_loss_x, train_loss_u = train(labeled_trainloader, unlabeled_trainloader, model, optimizer, ema_optimizer, train_criterion, epoch, use_cuda, alr, num_lab=num_labl) else: train_loss, train_loss_x, train_loss_u = train(labeled_trainloader, unlabeled_trainloader, model, optimizer, ema_optimizer, train_criterion, epoch, use_cuda, num_lab=num_labl) _, train_acc = validate(labeled_trainloader, ema_model, criterion, epoch, use_cuda, mode='Train Stats') val_loss, val_acc = validate(val_loader, ema_model, criterion, epoch, use_cuda, mode='Valid Stats') step = args.val_iteration * (epoch + 1) writer.add_scalar('losses/train_loss', train_loss, step) writer.add_scalar('losses/valid_loss', val_loss, step) writer.add_scalar('accuracy/train_acc', train_acc, step) writer.add_scalar('accuracy/val_acc', val_acc, step) # append logger file logger.append([train_loss, train_loss_x, train_loss_u, val_loss, val_acc]) # save model is_best = val_acc > best_acc best_acc = max(val_acc, best_acc) save_checkpoint({ 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'ema_state_dict': ema_model.state_dict(), 'acc': val_acc, 'best_acc': best_acc, 'optimizer' : optimizer.state_dict(), }, is_best) test_accs.append(val_acc) logger.close() writer.close() print('Best acc:') print(best_acc) print('Median (of last 20 epochs) acc:') print(np.median(test_accs[-20:]))
def main(create_model, black_model, output_dir, labeled_num): global best_acc if not os.path.isdir(output_dir): mkdir_p(output_dir) # Data print(f'==> Preparing cifar10') transform_train = transforms.Compose([ dataset.RandomPadandCrop(32), dataset.RandomFlip(), dataset.ToTensor(), ]) transform_val = transforms.Compose([ dataset.ToTensor(), ]) train_labeled_set, train_unlabeled_set, val_set, test_set = dataset.get_cifar10( '../data', labeled_num, transform_train=transform_train, transform_val=transform_val) labeled_trainloader = torch.utils.data.DataLoader(train_labeled_set, batch_size=10, shuffle=True, num_workers=0, drop_last=True) unlabeled_trainloader = torch.utils.data.DataLoader(train_unlabeled_set, batch_size=10, shuffle=True, num_workers=0, drop_last=True) val_loader = torch.utils.data.DataLoader(val_set, batch_size=10, shuffle=False, num_workers=0) test_loader = torch.utils.data.DataLoader(test_set, batch_size=10, shuffle=False, num_workers=0) model = create_model() ema_model = create_model(ema=True) cudnn.benchmark = True train_criterion = SemiLoss() criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.002) ema_optimizer = WeightEMA(model, ema_model, alpha=0.999) start_epoch = 0 # Resume title = 'noisy-cifar-10' logger = Logger(os.path.join(output_dir, 'log.txt'), title=title) logger.set_names([ 'Train Loss', 'Train Loss X', 'Train Loss U', 'Valid Loss', 'Valid Acc.', 'Test Loss', 'Test Acc.' ]) writer = SummaryWriter(output_dir) step = 0 test_accs = [] black_model.eval() meancifar, stdcifar = torch.Tensor([0.471, 0.448, 0.408 ]), torch.Tensor([0.234, 0.239, 0.242]) transformcifar = torchvision.transforms.Compose([ torchvision.transforms.ToTensor(), torchvision.transforms.Normalize(meancifar, stdcifar) ]) data = torchvision.datasets.CIFAR10(root="../data", train=False, download=False, transform=transformcifar) loadercifar = torch.utils.data.DataLoader(data, batch_size=32, shuffle=False) # Train and val for epoch in range(start_epoch, 1000): train_loss, train_loss_x, train_loss_u = train( labeled_trainloader, unlabeled_trainloader, model, optimizer, ema_optimizer, train_criterion, epoch, use_cuda, black_model) _, train_acc = validate(labeled_trainloader, ema_model, criterion, epoch, use_cuda, mode='Train Stats') val_loss, val_acc = validate(val_loader, ema_model, criterion, epoch, use_cuda, mode='Valid Stats') test_loss, test_acc = validate(test_loader, ema_model, criterion, epoch, use_cuda, mode='Test Stats ') step = 1024 * (epoch + 1) writer.add_scalar('losses/train_loss', train_loss, step) writer.add_scalar('losses/valid_loss', val_loss, step) writer.add_scalar('losses/test_loss', test_loss, step) writer.add_scalar('accuracy/train_acc', train_acc, step) writer.add_scalar('accuracy/val_acc', val_acc, step) writer.add_scalar('accuracy/test_acc', test_acc, step) # append logger file logger.append([ train_loss, train_loss_x, train_loss_u, val_loss, val_acc, test_loss, test_acc ]) # save model is_best = val_acc > best_acc best_acc = max(val_acc, best_acc) save_checkpoint( { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'ema_state_dict': ema_model.state_dict(), 'acc': val_acc, 'best_acc': best_acc, 'optimizer': optimizer.state_dict(), }, is_best, output_dir) test_accs.append(test_acc) logger.close() writer.close() print('Best acc:') print(best_acc) print('Mean acc:') print(np.mean(test_accs[-20:]))
def main(): global best_acc if not os.path.isdir(args.out): mkdir_p(args.out) # Data print(f'==> Preparing cifar10') transform_train = transforms.Compose([ dataset.RandomPadandCrop(32), dataset.RandomFlip(), dataset.ToTensor(), ]) transform_val = transforms.Compose([ dataset.ToTensor(), ]) train_labeled_set, train_unlabeled_set, val_set, test_set = dataset.get_cifar10( '/home/fengchan/stor/dataset/original-data/cifar10', args.n_labeled, transform_train=transform_train, transform_val=transform_val) labeled_trainloader = data.DataLoader(train_labeled_set, batch_size=args.batch_size, shuffle=True, num_workers=0, drop_last=True) unlabeled_trainloader = data.DataLoader(train_unlabeled_set, batch_size=args.batch_size, shuffle=True, num_workers=0, drop_last=True) val_loader = data.DataLoader(val_set, batch_size=args.batch_size, shuffle=False, num_workers=0) test_loader = data.DataLoader(test_set, batch_size=args.batch_size, shuffle=False, num_workers=0) # Model print("==> creating WRN-28-2") def create_model(ema=False): model = models.WideResNet(num_classes=num_classes) model = model.cuda() if ema: for param in model.parameters(): param.detach_() return model data_shape = [3, 32, 32] regularization_fns, regularization_coeffs = create_regularization_fns(args) def create_cnf(): # generate cnf # cnf = create_cnf_model_1(args, data_shape, regularization_fns=None) # cnf = create_cnf_model(args, data_shape, regularization_fns=regularization_fns) cnf = create_nf_model(args, data_shape, regularization_fns=None) cnf = cnf.cuda() if use_cuda else cnf return cnf model = create_model() ema_model = create_model(ema=True) cnf = create_cnf() if args.spectral_norm: add_spectral_norm(cnf, logger) set_cnf_options(args, cnf) cudnn.benchmark = True print(' Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0)) train_criterion = SemiLoss() criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=args.lr) #CNF cnf_optimizer = optim.Adam(cnf.parameters(), lr=args.lr, weight_decay=args.weight_decay) ema_optimizer = WeightEMA(model, ema_model, alpha=args.ema_decay) start_epoch = 0 # Resume #generate prior means = generate_gaussian_means(num_classes, data_shape, seed=num_classes) title = 'noisy-cifar-10' if args.resume: # Load checkpoint. print('==> Resuming from checkpoint..') assert os.path.isfile( args.resume), 'Error: no checkpoint directory found!' args.out = os.path.dirname(args.resume) checkpoint = torch.load(args.resume) best_acc = checkpoint['best_acc'] start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) ema_model.load_state_dict(checkpoint['ema_state_dict']) cnf.load_state_dict(checkpoint['cnf_state_dict']) means = checkpoint['means'] cnf_optimizer.load_state_dict(checkpoint['cnf_optimizer']) optimizer.load_state_dict(checkpoint['optimizer']) logger = Logger(os.path.join(args.out, 'log.txt'), title=title, resume=True) else: logger = Logger(os.path.join(args.out, 'log.txt'), title=title) logger.set_names([ 'Train Loss', 'Train Loss X', 'Train Loss U', 'Train loss NLL X', 'Train loss NLL U', 'Train loss mixed X', 'Valid Loss', 'Valid Acc.', 'Test Loss', 'Test Acc.' ]) means = means.cuda() if use_cuda else means prior = SSLGaussMixture(means, device='cuda' if use_cuda else 'cpu') writer = SummaryWriter(args.out) step = 0 test_accs = [] # Train and val for epoch in range(start_epoch, args.epochs): print('\nEpoch: [%d | %d] LR: %f' % (epoch + 1, args.epochs, state['lr'])) train_loss, train_loss_x, train_loss_u, train_loss_nll_x, train_loss_nll_u, train_loss_mixed_x = train( labeled_trainloader, unlabeled_trainloader, model, cnf, prior, cnf_optimizer, optimizer, ema_optimizer, train_criterion, epoch, use_cuda) _, train_acc = validate(labeled_trainloader, ema_model, criterion, epoch, use_cuda, mode='Train Stats') val_loss, val_acc = validate(val_loader, ema_model, criterion, epoch, use_cuda, mode='Valid Stats') test_loss, test_acc = validate(test_loader, ema_model, criterion, epoch, use_cuda, mode='Test Stats ') step = args.train_iteration * (epoch + 1) writer.add_scalar('losses/train_loss', train_loss, step) writer.add_scalar('losses/train_loss_nll_x', train_loss_nll_x, step) writer.add_scalar('losses/train_loss_nll_u', train_loss_nll_u, step) writer.add_scalar('losses/train_loss_mixed_x', train_loss_mixed_x, step) writer.add_scalar('losses/train_loss_nll_x', train_loss_nll_x, step) writer.add_scalar('losses/valid_loss', val_loss, step) writer.add_scalar('losses/test_loss', test_loss, step) writer.add_scalar('accuracy/train_acc', train_acc, step) writer.add_scalar('accuracy/val_acc', val_acc, step) writer.add_scalar('accuracy/test_acc', test_acc, step) # append logger file logger.append([ train_loss, train_loss_x, train_loss_u, train_loss_nll_x, train_loss_nll_u, train_loss_mixed_x, val_loss, val_acc, test_loss, test_acc ]) # save model is_best = val_acc > best_acc best_acc = max(val_acc, best_acc) save_checkpoint( { 'epoch': epoch + 1, 'state_dict': model.state_dict(), 'cnf_state_dict': cnf.state_dict(), 'means': means, 'ema_state_dict': ema_model.state_dict(), 'acc': val_acc, 'best_acc': best_acc, 'optimizer': optimizer.state_dict(), 'cnf_optimizer': cnf_optimizer.state_dict(), }, is_best) test_accs.append(test_acc) logger.close() writer.close() print('Best acc:') print(best_acc) print('Mean acc:') print(np.mean(test_accs[-20:]))