def train(epoch, train_loader, model, classifier, criterion, optimizer, opt):
"""
one epoch training
"""
model.eval()
classifier.train()
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
end = time.time()
for idx, (x, y) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
x = x.cuda(non_blocking=True)
y = y.cuda(non_blocking=True)
# ===================forward=====================
with torch.no_grad():
feat = model(x, opt.layer)
output = classifier(feat)
loss = criterion(output, y)
acc1, acc5 = accuracy(output, y, topk=(1, 5))
losses.update(loss.item(), x.size(0))
top1.update(acc1[0], x.size(0))
top5.update(acc5[0], x.size(0))
# ===================backward=====================
optimizer.zero_grad()
loss.backward()
optimizer.step()
# ===================meters=====================
batch_time.update(time.time() - end)
end = time.time()
# print info
if opt.local_rank == 0:
if idx % opt.print_freq == 0:
lr = optimizer.param_groups[0]['lr']
print(f'Epoch: [{epoch}][{idx}/{len(train_loader)}]\t'
f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
f'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
f'Lr {lr:.3f} \t'
f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
f'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
f'Acc@5 {top5.val:.3f} ({top5.avg:.3f})')
return top1.avg, top5.avg, losses.avg
def validate(val_loader, model, classifier, criterion, args):
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to evaluate mode
model.eval()
classifier.eval()
with torch.no_grad():
end = time.time()
for idx, (x, y) in enumerate(val_loader):
x = x.cuda(non_blocking=True)
y = y.cuda(non_blocking=True)
# compute output
feat = model(x, args.layer)
output = classifier(feat)
loss = criterion(output, y)
# measure accuracy and record loss
acc1, acc5 = accuracy(output, y, topk=(1, 5))
losses.update(loss.item(), x.size(0))
top1.update(acc1[0], x.size(0))
top5.update(acc5[0], x.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if idx % args.print_freq == 0:
print(f'Test: [{idx}/{len(val_loader)}]\t'
f'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
f'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t'
f'Acc@5 {top5.val:.3f} ({top5.avg:.3f})')
print(' * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'.format(top1=top1,
top5=top5))
return top1.avg, top5.avg, losses.avg
def train_moco(epoch, train_loader, model, model_ema, contrast, criterion,
optimizer, args):
"""
one epoch training for moco
"""
model.train()
set_bn_train(model_ema)
batch_time = AverageMeter()
data_time = AverageMeter()
loss_meter = AverageMeter()
prob_meter = AverageMeter()
end = time.time()
for idx, (
inputs,
_,
) in enumerate(train_loader):
data_time.update(time.time() - end)
bsz = inputs.size(0)
# forward
x1, x2 = torch.split(inputs, [3, 3], dim=1)
x1.contiguous()
x2.contiguous()
x1 = x1.cuda(non_blocking=True)
x2 = x2.cuda(non_blocking=True)
feat_q = model(x1)
with torch.no_grad():
x2_shuffled, backward_inds = DistributedShufle.forward_shuffle(
x2, epoch)
feat_k = model_ema(x2_shuffled)
feat_k_all, feat_k = DistributedShufle.backward_shuffle(
feat_k, backward_inds, return_local=True)
out = contrast(feat_q, feat_k, feat_k_all)
loss = criterion(out)
prob = F.softmax(out, dim=1)[:, 0].mean()
# backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
moment_update(model, model_ema, args.alpha)
# update meters
loss_meter.update(loss.item(), bsz)
prob_meter.update(prob.item(), bsz)
batch_time.update(time.time() - end)
end = time.time()
# print info
if args.local_rank == 0 and idx % args.print_freq == 0:
print(f'Train: [{epoch}][{idx}/{len(train_loader)}]\t'
f'T {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
f'DT {data_time.val:.3f} ({data_time.avg:.3f})\t'
f'loss {loss_meter.val:.3f} ({loss_meter.avg:.3f})\t'
f'prob {prob_meter.val:.3f} ({prob_meter.avg:.3f})')
return loss_meter.avg, prob_meter.avg