def validate_general(val_loader, model, infer_model_fn, cuda=False): from examples.classification.main import AverageMeter, accuracy top1 = AverageMeter() top5 = AverageMeter() for i, (input_, target) in enumerate(val_loader): # compute output output = infer_model_fn(model, input_) if cuda: target = target.cuda(None, non_blocking=True) # measure accuracy and record loss acc1, acc5 = accuracy(output, target, topk=(1, 5)) top1.update(acc1, input_.size(0)) top5.update(acc5, input_.size(0)) if i % 10 == 0: print('IE Test : [{0}/{1}]\t' 'Acc@1 {top1.val:.3f} ({top1.avg:.3f})\t' 'Acc@5 {top5.val:.3f} ({top5.avg:.3f})'.format( i, len(val_loader), top1=top1, top5=top5)) print(' * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'.format(top1=top1, top5=top5)) return top1.avg, top5.avg
def train_epoch_bin(train_loader, batch_multiplier, model, criterion, optimizer, optimizer_scheduler: BinarizationOptimizerScheduler, kd_loss_calculator: KDLossCalculator, compression_ctrl, epoch, config, is_inception=False): batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() kd_losses_meter = AverageMeter() criterion_losses = AverageMeter() top1 = AverageMeter() top5 = AverageMeter() compression_scheduler = compression_ctrl.scheduler # switch to train mode model.train() end = time.time() for i, (input_, target) in enumerate(train_loader): # measure data loading time data_time.update(time.time() - end) input_ = input_.to(config.device) target = target.to(config.device) # compute output if is_inception: # From https://discuss.pytorch.org/t/how-to-optimize-inception-model-with-auxiliary-classifiers/7958 output, aux_outputs = model(input_) loss1 = criterion(output, target) loss2 = criterion(aux_outputs, target) criterion_loss = loss1 + 0.4 * loss2 else: output = model(input_) criterion_loss = criterion(output, target) # compute KD loss kd_loss = kd_loss_calculator.loss(input_, output) loss = criterion_loss + kd_loss # measure accuracy and record loss acc1, acc5 = accuracy(output, target, topk=(1, 5)) losses.update(loss.item(), input_.size(0)) comp_loss_val = kd_loss.item() kd_losses_meter.update(comp_loss_val, input_.size(0)) criterion_losses.update(criterion_loss.item(), input_.size(0)) top1.update(acc1, input_.size(0)) top1.update(acc1, input_.size(0)) top5.update(acc5, input_.size(0)) # compute gradient and do SGD step if i % batch_multiplier == 0: optimizer.zero_grad() loss.backward() optimizer.step() else: loss.backward() compression_scheduler.step() optimizer_scheduler.step(float(i) / len(train_loader)) # measure elapsed time batch_time.update(time.time() - end) end = time.time() if i % config.print_freq == 0: logger.info( '{rank}: ' 'Epoch: [{0}][{1}/{2}] ' 'Lr: {3:.3} ' 'Wd: {4:.3} ' 'Time: {batch_time.val:.3f} ({batch_time.avg:.3f}) ' 'Data: {data_time.val:.3f} ({data_time.avg:.3f}) ' 'CE_loss: {ce_loss.val:.4f} ({ce_loss.avg:.4f}) ' 'KD_loss: {kd_loss.val:.4f} ({kd_loss.avg:.4f}) ' 'Loss: {loss.val:.4f} ({loss.avg:.4f}) ' 'Acc@1: {top1.val:.3f} ({top1.avg:.3f}) ' 'Acc@5: {top5.val:.3f} ({top5.avg:.3f})'.format( epoch, i, len(train_loader), get_lr(optimizer), get_wd(optimizer), batch_time=batch_time, data_time=data_time, ce_loss=criterion_losses, kd_loss=kd_losses_meter, loss=losses, top1=top1, top5=top5, rank='{}:'.format(config.rank) if config.multiprocessing_distributed else '' )) if is_main_process(): global_step = len(train_loader) * epoch config.tb.add_scalar("train/learning_rate", get_lr(optimizer), i + global_step) config.tb.add_scalar("train/criterion_loss", criterion_losses.avg, i + global_step) config.tb.add_scalar("train/kd_loss", kd_losses_meter.avg, i + global_step) config.tb.add_scalar("train/loss", losses.avg, i + global_step) config.tb.add_scalar("train/top1", top1.avg, i + global_step) config.tb.add_scalar("train/top5", top5.avg, i + global_step) for stat_name, stat_value in compression_ctrl.statistics().items(): if isinstance(stat_value, (int, float)): config.tb.add_scalar('train/statistics/{}'.format(stat_name), stat_value, i + global_step)