def train(epoch, train_loader, model, classifier, criterion, optimizer,
scheduler, args):
"""
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):
x = x.cuda(non_blocking=True)
y = y.cuda(non_blocking=True)
# measure data loading time
data_time.update(time.time() - end)
# ===================forward=====================
with torch.no_grad():
feat = model(x, args.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()
if args.amp_opt_level != "O0":
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optimizer.step()
scheduler.step()
# ===================meters=====================
batch_time.update(time.time() - end)
end = time.time()
# print info
if idx % args.print_freq == 0:
logger.info(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 {optimizer.param_groups[0]["lr"]:.3f} \t'
f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
f'Acc@1 {top1.val:.3%} ({top1.avg:.3%})\t'
f'Acc@5 {top5.val:.3%} ({top5.avg:.3%})')
def validate(val_loader, model, criterion, args, indices_in_1k=None):
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
# switch to evaluate mode
model.eval()
with torch.no_grad():
end = time.time()
prefetcher = data_prefetcher(val_loader)
x, y = prefetcher.next()
idx = 0
while x is not None:
# compute output
output = model(x)
if indices_in_1k is not None:
output = output[:, indices_in_1k]
loss = criterion(output, y)
# measure accuracy and record loss
acc1, acc5 = accuracy(output, y, topk=(1, 5))
acc1 = reduce_tensor(acc1)
acc5 = reduce_tensor(acc5)
loss = reduce_tensor(loss)
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()
x, y = prefetcher.next()
if idx % args.print_freq == 0:
logger.info(
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:.3%} ({top1.avg:.3%})\t'
f'Acc@5 {top5.val:.3%} ({top5.avg:.3%})')
idx = idx + 1
logger.info(f' * Acc@1 {top1.avg:.3%} Acc@5 {top5.avg:.3%}')
return top1.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))
acc1 = reduce_tensor(acc1)
acc5 = reduce_tensor(acc5)
loss = reduce_tensor(loss)
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:
logger.info(
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:.3%} ({top1.avg:.3%})\t'
f'Acc@5 {top5.val:.3%} ({top5.avg:.3%})')
logger.info(f' * Acc@1 {top1.avg:.3%} Acc@5 {top5.avg:.3%}')
return top1.avg