def main():
global best_acc
start_epoch = args.start_epoch # start from epoch 0 or last checkpoint epoch
if not os.path.isdir(args.out):
mkdir_p(args.out)
# Data
print(f'==> Preparing open-set nosiy cifar10')
transform_train = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
trainset = dataset.get_cifar10('./data',
args,
train=True,
download=True,
transform=transform_train)
testset = torchvision.datasets.CIFAR10('./data',
train=False,
transform=transform_test)
trainloader = data.DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=0)
testloader = data.DataLoader(testset,
batch_size=args.batch_size,
shuffle=False,
num_workers=0)
print("==> creating resnet34")
model = models.resnet34(num_classes=11)
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
val_criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=args.pre_lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# 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'])
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([
'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.',
'Valid Acc.'
])
# Label update stage
for epoch in range(start_epoch, args.pre_epochs):
# adjust_learning_rate(optimizer, epoch + 1)
print('\nEpoch: [%d | %d] LR: %f' %
(epoch + 1, args.epochs, state['pre_lr']))
train_loss, train_acc = train(trainloader, model, optimizer, epoch,
use_cuda)
val_loss, val_acc = validate(testloader, model, val_criterion, epoch,
use_cuda)
# append logger file
logger.append(
[state['pre_lr'], train_loss, val_loss, train_acc, 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(),
'acc': val_acc,
'best_val_acc': best_acc,
'optimizer': optimizer.state_dict(),
}, is_best)
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# Fine-tune stage
for epoch in range(start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch + 1)
print('\nEpoch: [%d | %d] LR: %f' %
(epoch + 1, args.epochs, state['lr']))
train_loss, train_acc = train_soft(trainloader, model, optimizer,
epoch, use_cuda)
val_loss, val_acc = validate(testloader, model, val_criterion, epoch,
use_cuda)
# append logger file
logger.append([state['lr'], train_loss, val_loss, train_acc, 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(),
'acc': val_acc,
'best_val_acc': best_acc,
'optimizer': optimizer.state_dict(),
}, is_best)
logger.close()
logger.plot()
savefig(os.path.join(args.out, 'log.eps'))
print(f'Best val acc: {best_acc}')
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
start_epoch = args.start_epoch # start from epoch 0 or last checkpoint epoch
if not os.path.isdir(args.out):
mkdir_p(args.out)
# Data
print(
f'==> Preparing {"asymmetric" if args.asym else "symmetric"} nosiy cifar10'
)
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=8),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
transform_val = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
trainset, valset = dataset.get_cifar10('./data',
args,
train=True,
download=True,
transform_train=transform_train,
transform_val=transform_val)
trainloader = data.DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
valloader = data.DataLoader(valset,
batch_size=args.batch_size,
shuffle=False,
num_workers=4)
# Model
print("==> creating preact_resnet")
model = models.PreActResNet34()
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
val_criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# 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'])
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([
'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.',
'Valid Acc.'
])
# 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_acc = train(trainloader, model, optimizer, epoch,
use_cuda)
val_loss, val_acc = validate(valloader, model, val_criterion, epoch,
use_cuda)
# append logger file
logger.append([state['lr'], train_loss, val_loss, train_acc, 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(),
'acc': val_acc,
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
}, is_best)
logger.close()
logger.plot()
savefig(os.path.join(args.out, 'log.eps'))
print('Best acc:')
print(best_acc)
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)
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():
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:]))
def main():
global best_acc
if not os.path.isdir(args.out):
mkdir_p(args.out)
# Data
dataset.validate_dataset(args.dataset)
print(f'==> Preparing %s' % args.dataset)
transform_train, transform_val = dataset.get_transforms(
args.dataset, args.resolution)
if args.dataset == 'C10':
train_labeled_set, train_unlabeled_set, val_set, test_set = dataset.get_cifar10(
args.data_root,
args.n_labeled,
transform_train=transform_train,
transform_val=transform_val)
elif args.dataset == 'STL10':
if args.n_labeled != 5000:
raise ValueError("For STL10 the only supported n_labeled is 5000")
train_labeled_set, train_unlabeled_set, val_set, test_set = dataset.get_stl10(
args.data_root,
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=True,
num_workers=0,
drop_last=True)
if test_set is not None:
test_loader = data.DataLoader(test_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=0)
else:
test_loader = None
# Model
print("==> creating %s" % which_model.__name__)
def create_model(ema=False):
model = which_model(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!'
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 = []
# 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(
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')
if test_loader is not None:
test_loss, test_acc = validate(test_loader,
ema_model,
criterion,
epoch,
use_cuda,
mode='Test Stats ')
else:
test_loss, test_acc = [-1, -1]
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
])
# 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:]))
