global best_acc

if not os.path.isdir(args.out):

mkdir_p(args.out)

# Data

print(f'==> Preparing DL Dataset')

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)

normalize = transforms.Normalize(mean=[0.4837, 0.4531, 0.4015],

std=[0.2212, 0.2165, 0.2156])

# Inverse transformation: needed for plotting.

unnormalize = transforms.Normalize(

mean=[-0.4837/0.2212, -0.4531/0.2165, -0.4015/0.2156],

std=[1/0.2212, 1/0.2165, 1/0.2156]

)

train_transforms = transforms.Compose([

transforms.Resize((32, 32)),

transforms.RandomHorizontalFlip(),

transforms.RandomVerticalFlip(),

transforms.ColorJitter(hue=.1, saturation=.1, contrast=.1),

transforms.RandomRotation(20, resample=Image.BILINEAR),

transforms.GaussianBlur(7, sigma=(0.1, 1.0)),

transforms.ToTensor(),

normalize,

])

val_transforms = transforms.Compose([

transforms.Resize((32, 32)),

transforms.ToTensor(),

normalize,

])

train_dataset, unlabeled_dataset, val_dataset = dataset.get_dl_dataset(transform_train=train_transforms, transform_val=val_transforms)

# train_dataset = CustomDataset(root="dataset", split="unlabeled", transform=unlabeled_transforms)

labeled_trainloader = data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)

unlabeled_trainloader = data.DataLoader(unlabeled_dataset, batch_size=args.batch_size, shuffle=True)

val_loader = data.DataLoader(val_dataset, 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=args.num_classes)

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'

title = 'dldataset-800'

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 = []

# 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')

# 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/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:')

batch_time = AverageMeter()

data_time = AverageMeter()

losses = AverageMeter()

losses_x = AverageMeter()

losses_u = AverageMeter()

ws = AverageMeter()

end = time.time()

# bar = Bar('Training', max=args.train_iteration)

labeled_train_iter = iter(labeled_trainloader)

unlabeled_train_iter = iter(unlabeled_trainloader)

model.train()

for batch_idx in range(args.train_iteration):

try:

inputs_x, targets_x = labeled_train_iter.next()

except:

labeled_train_iter = iter(labeled_trainloader)

inputs_x, targets_x = labeled_train_iter.next()

try:

(inputs_u, inputs_u2), _ = unlabeled_train_iter.next()

except:

unlabeled_train_iter = iter(unlabeled_trainloader)

(inputs_u, inputs_u2), _ = unlabeled_train_iter.next()

# measure data loading time

data_time.update(time.time() - end)

batch_size = inputs_x.size(0)

# Transform label to one-hot

targets_x = torch.zeros(batch_size, args.num_classes).scatter_(1, targets_x.view(-1,1).long(), 1)

if use_cuda:

inputs_x, targets_x = inputs_x.cuda(), targets_x.cuda(non_blocking=True)

inputs_u = inputs_u.cuda()

inputs_u2 = inputs_u2.cuda()

with torch.no_grad():

# compute guessed labels of unlabel samples

outputs_u = model(inputs_u)

outputs_u2 = model(inputs_u2)

p = (torch.softmax(outputs_u, dim=1) + torch.softmax(outputs_u2, dim=1)) / 2

pt = p**(1/args.T)

targets_u = pt / pt.sum(dim=1, keepdim=True)

targets_u = targets_u.detach()

# mixup

all_inputs = torch.cat([inputs_x, inputs_u, inputs_u2], dim=0)

all_targets = torch.cat([targets_x, targets_u, targets_u], dim=0)

l = np.random.beta(args.alpha, args.alpha)

l = max(l, 1-l)

idx = torch.randperm(all_inputs.size(0))

input_a, input_b = all_inputs, all_inputs[idx]

target_a, target_b = all_targets, all_targets[idx]

mixed_input = l * input_a + (1 - l) * input_b

mixed_target = l * target_a + (1 - l) * target_b

# interleave labeled and unlabed samples between batches to get correct batchnorm calculation

mixed_input = list(torch.split(mixed_input, batch_size))

mixed_input = interleave(mixed_input, batch_size)

logits = [model(mixed_input[0])]

for input in mixed_input[1:]:

logits.append(model(input))

# put interleaved samples back

logits = interleave(logits, batch_size)

logits_x = logits[0]

logits_u = torch.cat(logits[1:], dim=0)

Lx, Lu, w = criterion(logits_x, mixed_target[:batch_size], logits_u, mixed_target[batch_size:], epoch+batch_idx/args.train_iteration)

loss = Lx + w * Lu

# record loss

losses.update(loss.item(), inputs_x.size(0))

losses_x.update(Lx.item(), inputs_x.size(0))

losses_u.update(Lu.item(), inputs_x.size(0))

ws.update(w, inputs_x.size(0))

# compute gradient and do SGD step

optimizer.zero_grad()

loss.backward()

optimizer.step()

ema_optimizer.step()

# measure elapsed time

batch_time.update(time.time() - end)

end = time.time()

# plot progress

# bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | Loss_x: {loss_x:.4f} | Loss_u: {loss_u:.4f} | W: {w:.4f}'.format(

# batch=batch_idx + 1,

# size=args.train_iteration,

# data=data_time.avg,

# bt=batch_time.avg,

# total=bar.elapsed_td,

# eta=bar.eta_td,

# loss=losses.avg,

# loss_x=losses_x.avg,

# loss_u=losses_u.avg,

# w=ws.avg,

# )

# bar.next()

# bar.finish()

batch_time = AverageMeter()

data_time = AverageMeter()

losses = AverageMeter()

top1 = AverageMeter()

top5 = AverageMeter()

# switch to evaluate mode

model.eval()

end = time.time()

# bar = Bar(f'{mode}', max=len(valloader))

with torch.no_grad():

for batch_idx, (inputs, targets) in enumerate(valloader):

# measure data loading time

data_time.update(time.time() - end)

if use_cuda:

inputs, targets = inputs.cuda(), targets.cuda(non_blocking=True)

# compute output

outputs = model(inputs)

loss = criterion(outputs, targets)

# measure accuracy and record loss

prec1, prec5 = accuracy(outputs, targets, topk=(1, 5))

losses.update(loss.item(), inputs.size(0))

top1.update(prec1.item(), inputs.size(0))

top5.update(prec5.item(), inputs.size(0))

# measure elapsed time

batch_time.update(time.time() - end)

end = time.time()

# plot progress

# bar.suffix = '({batch}/{size}) Data: {data:.3f}s | Batch: {bt:.3f}s | Total: {total:} | ETA: {eta:} | Loss: {loss:.4f} | top1: {top1: .4f} | top5: {top5: .4f}'.format(

# batch=batch_idx + 1,

# size=len(valloader),

# data=data_time.avg,

# bt=batch_time.avg,

# total=bar.elapsed_td,

# eta=bar.eta_td,

# loss=losses.avg,

# top1=top1.avg,

# top5=top5.avg,

# )

# bar.next()

# bar.finish()

filepath = os.path.join(checkpoint, filename)

torch.save(state, filepath)

if is_best:

if rampup_length == 0:

return 1.0

else:

current = np.clip(current / rampup_length, 0.0, 1.0)

def __call__(self, outputs_x, targets_x, outputs_u, targets_u, epoch):

probs_u = torch.softmax(outputs_u, dim=1)

Lx = -torch.mean(torch.sum(F.log_softmax(outputs_x, dim=1) * targets_x, dim=1))

Lu = torch.mean((probs_u - targets_u)**2)

def __init__(self, model, ema_model, alpha=0.999):

self.model = model

self.ema_model = ema_model

self.alpha = alpha

self.params = list(model.state_dict().values())

self.ema_params = list(ema_model.state_dict().values())

self.wd = 0.02 * args.lr

for param, ema_param in zip(self.params, self.ema_params):

param.data.copy_(ema_param.data)

def step(self):

one_minus_alpha = 1.0 - self.alpha

for param, ema_param in zip(self.params, self.ema_params):

if ema_param.dtype==torch.float32:

ema_param.mul_(self.alpha)

ema_param.add_(param * one_minus_alpha)

# customized weight decay

groups = [batch // (nu + 1)] * (nu + 1)

for x in range(batch - sum(groups)):

groups[-x - 1] += 1

offsets = [0]

for g in groups:

offsets.append(offsets[-1] + g)

assert offsets[-1] == batch

nu = len(xy) - 1

offsets = interleave_offsets(batch, nu)

xy = [[v[offsets[p]:offsets[p + 1]] for p in range(nu + 1)] for v in xy]

for i in range(1, nu + 1):

xy[0][i], xy[i][i] = xy[i][i], xy[0][i]