def run_epoch(model, loader, loss_fn, optimizer, desc_default='', epoch=0, writer=None, verbose=1, scheduler=None): tqdm_disable = bool(os.environ.get('TASK_NAME', '')) # KakaoBrain Environment if verbose: loader = tqdm(loader, disable=tqdm_disable) loader.set_description('[%s %04d/%04d]' % (desc_default, epoch, C.get()['epoch'])) metrics = Accumulator() cnt = 0 total_steps = len(loader) steps = 0 for data, label in loader: steps += 1 data, label = data.cuda(), label.cuda() if optimizer: optimizer.zero_grad() preds = model(data) loss = loss_fn(preds, label) if optimizer: loss.backward() if getattr(optimizer, "synchronize", None): optimizer.synchronize() # for horovod if C.get()['optimizer'].get('clip', 5) > 0: nn.utils.clip_grad_norm_(model.parameters(), C.get()['optimizer'].get('clip', 5)) optimizer.step() top1, top5 = accuracy(preds, label, (1, 5)) metrics.add_dict({ 'loss': loss.item() * len(data), 'top1': top1.item() * len(data), 'top5': top5.item() * len(data), }) cnt += len(data) if verbose: postfix = metrics / cnt if optimizer: postfix['lr'] = optimizer.param_groups[0]['lr'] loader.set_postfix(postfix) if scheduler is not None: scheduler.step(epoch - 1 + float(steps) / total_steps) del preds, loss, top1, top5, data, label if tqdm_disable: if optimizer: logger.info('[%s %03d/%03d] %s lr=%.6f', desc_default, epoch, C.get()['epoch'], metrics / cnt, optimizer.param_groups[0]['lr']) else: logger.info('[%s %03d/%03d] %s', desc_default, epoch, C.get()['epoch'], metrics / cnt) metrics /= cnt if optimizer: metrics.metrics['lr'] = optimizer.param_groups[0]['lr'] if verbose: for key, value in metrics.items(): writer.add_scalar(key, value, epoch) return metrics
def get_affinity(aug, aff_bases, config, augment): C.get() C.get().conf = config # setup - provided augmentation rules C.get()['aug'] = aug load_paths = augment['load_paths'] cv_num = augment["cv_num"] aug_loaders = [] for cv_id in range(cv_num): _, tl, validloader, tl2 = get_dataloaders(C.get()['dataset'], C.get()['batch'], augment['dataroot'], augment['cv_ratio_test'], split_idx=cv_id) aug_loaders.append(validloader) del tl, tl2 loss_fn = torch.nn.CrossEntropyLoss(reduction='none') aug_accs = [] for cv_id, loader in enumerate(aug_loaders): # eval model = get_model(C.get()['model'], num_class(C.get()['dataset'])) ckpt = torch.load(load_paths[cv_id]) if 'model' in ckpt: model.load_state_dict(ckpt['model']) else: model.load_state_dict(ckpt) model.eval() metrics = Accumulator() for data, label in loader: data = data.cuda() label = label.cuda() pred = model(data) loss = loss_fn(pred, label) # (N) _, pred = pred.topk(1, 1, True, True) pred = pred.t() correct = pred.eq(label.view( 1, -1).expand_as(pred)).detach().cpu().numpy() # (1,N) metrics.add_dict({ 'minus_loss': -1 * np.sum(loss.detach().cpu().numpy()), 'correct': np.sum(correct), 'cnt': len(data) }) del loss, correct, pred, data, label aug_accs.append(metrics['correct'] / metrics['cnt']) del model affs = [] for aug_valid, clean_valid in zip(aug_accs, aff_bases): affs.append(aug_valid - clean_valid) return affs
def run_epoch(model, loader, loss_fn, optimizer, desc_default='', epoch=0, writer=None, verbose=1): if verbose: loader = tqdm(loader) if optimizer: curr_lr = optimizer.param_groups[0]['lr'] loader.set_description( '[%s %04d/%04d] lr=%.4f' % (desc_default, epoch, C.get()['epoch'], curr_lr)) else: loader.set_description('[%s %04d/%04d]' % (desc_default, epoch, C.get()['epoch'])) metrics = Accumulator() cnt = 0 for data, label in loader: data, label = data.cuda(), label.cuda() if optimizer: optimizer.zero_grad() preds = model(data) loss = loss_fn(preds, label) if optimizer: nn.utils.clip_grad_norm_(model.parameters(), 5) loss.backward() optimizer.step() top1, top5 = accuracy(preds, label, (1, 5)) metrics.add_dict({ 'loss': loss.item() * len(data), 'top1': top1.item() * len(data), 'top5': top5.item() * len(data), }) cnt += len(data) if verbose: loader.set_postfix(metrics / cnt) del preds, loss, top1, top5, data, label metrics /= cnt if optimizer: metrics.metrics['lr'] = optimizer.param_groups[0]['lr'] if verbose: for key, value in metrics.items(): writer.add_scalar(key, value, epoch) return metrics
def eval_tta3(config, augment, reporter): C.get() C.get().conf = config save_path = augment['save_path'] cv_id, gr_id = augment["cv_id"], augment["gr_id"] gr_ids = augment["gr_ids"] # setup - provided augmentation rules C.get()['aug'] = policy_decoder(augment, augment['num_policy'], augment['num_op']) # eval model = get_model(C.get()['model'], num_class(C.get()['dataset'])) ckpt = torch.load(save_path) if 'model' in ckpt: model.load_state_dict(ckpt['model']) else: model.load_state_dict(ckpt) del ckpt model.eval() loader = get_post_dataloader(C.get()["dataset"], C.get()['batch'], augment["dataroot"], augment['cv_ratio_test'], cv_id, gr_id, gr_ids) start_t = time.time() metrics = Accumulator() loss_fn = torch.nn.CrossEntropyLoss(reduction='none') for data, label in loader: data = data.cuda() label = label.cuda() pred = model(data) loss = loss_fn(pred, label) # (N) _, pred = pred.topk(1, 1, True, True) pred = pred.t() correct = pred.eq(label.view( 1, -1).expand_as(pred)).detach().cpu().numpy() # (1,N) metrics.add_dict({ 'loss': np.sum(loss.detach().cpu().numpy()), 'correct': np.sum(correct), 'cnt': len(data) }) del loss, correct, pred, data, label del model, loader metrics = metrics / 'cnt' gpu_secs = (time.time() - start_t) * torch.cuda.device_count() reporter(loss=metrics['loss'], top1_valid=metrics['correct'], elapsed_time=gpu_secs, done=True) return metrics['correct']
def eval_tta2(config, augment, reporter): C.get() C.get().conf = config cv_ratio_test, cv_id, save_path = augment['cv_ratio_test'], augment['cv_id'], augment['save_path'] gr_id = augment["gr_id"] num_repeat = 1 # setup - provided augmentation rules C.get()['aug'] = policy_decoder(augment, augment['num_policy'], augment['num_op']) # eval model = get_model(C.get()['model'], num_class(C.get()['dataset'])) ckpt = torch.load(save_path) if 'model' in ckpt: model.load_state_dict(ckpt['model']) else: model.load_state_dict(ckpt) model.eval() loaders = [] for i in range(num_repeat): _, tl, validloader, tl2 = get_dataloaders(C.get()['dataset'], C.get()['batch'], augment['dataroot'], cv_ratio_test, split_idx=cv_id, gr_assign=augment["gr_assign"], gr_id=gr_id) loaders.append(validloader) del tl, tl2 start_t = time.time() metrics = Accumulator() loss_fn = torch.nn.CrossEntropyLoss(reduction='none') for loader in loaders: for data, label in loader: data = data.cuda() label = label.cuda() pred = model(data) loss = loss_fn(pred, label) # (N) _, pred = pred.topk(1, 1, True, True) pred = pred.t() correct = pred.eq(label.view(1, -1).expand_as(pred)).detach().cpu().numpy() # (1,N) metrics.add_dict({ 'minus_loss': -1 * np.sum(loss.detach().cpu().numpy()), 'correct': np.sum(correct), 'cnt': len(data) }) del loss, correct, pred, data, label del model metrics = metrics / 'cnt' gpu_secs = (time.time() - start_t) * torch.cuda.device_count() reporter(minus_loss=metrics['minus_loss'], top1_valid=metrics['correct'], elapsed_time=gpu_secs, done=True) return metrics['correct']
def eval_tta(config, augment, reporter, num_class, get_model, get_dataloaders): C.get() C.get().conf = config cv_ratio_test, cv_fold, save_path = ( augment["cv_ratio_test"], augment["cv_fold"], augment["save_path"], ) # setup - provided augmentation rules C.get()["aug"] = policy_decoder(augment, augment["num_policy"], augment["num_op"]) # eval model = get_model(C.get()["model"], num_class(C.get()["dataset"])) ckpt = torch.load(save_path) if "model" in ckpt: model.load_state_dict(ckpt["model"]) else: model.load_state_dict(ckpt) model.eval() loaders = [] for _ in range(augment["num_policy"]): # TODO _, tl, validloader, tl2 = get_dataloaders( C.get()["dataset"], C.get()["batch"], augment["dataroot"], cv_ratio_test, split_idx=cv_fold, ) loaders.append(iter(validloader)) del tl, tl2 start_t = time.time() metrics = Accumulator() loss_fn = torch.nn.CrossEntropyLoss(reduction="none") try: while True: losses = [] corrects = [] for loader in loaders: data, label = next(loader) data = data.cuda() label = label.cuda() pred = model(data) loss = loss_fn(pred, label) losses.append(loss.detach().cpu().numpy()) _, pred = pred.topk(1, 1, True, True) pred = pred.t() correct = ( pred.eq(label.view(1, -1).expand_as(pred)).detach().cpu().numpy() ) corrects.append(correct) del loss, correct, pred, data, label losses = np.concatenate(losses) losses_min = np.min(losses, axis=0).squeeze() corrects = np.concatenate(corrects) corrects_max = np.max(corrects, axis=0).squeeze() metrics.add_dict( { "minus_loss": -1 * np.sum(losses_min), "correct": np.sum(corrects_max), "cnt": len(corrects_max), } ) del corrects, corrects_max except StopIteration: pass del model metrics = metrics / "cnt" gpu_secs = (time.time() - start_t) * torch.cuda.device_count() reporter( minus_loss=metrics["minus_loss"], top1_valid=metrics["correct"], elapsed_time=gpu_secs, done=True, ) return metrics["correct"]
def run_epoch(model, loader, loss_fn, optimizer, desc_default='', epoch=0, writer=None, verbose=1, scheduler=None, is_master=True, ema=None, wd=0.0, tqdm_disabled=False): if verbose: loader = tqdm(loader, disable=tqdm_disabled) loader.set_description('[%s %04d/%04d]' % (desc_default, epoch, C.get()['epoch'])) params_without_bn = [ params for name, params in model.named_parameters() if not ('_bn' in name or '.bn' in name) ] loss_ema = None metrics = Accumulator() cnt = 0 total_steps = len(loader) steps = 0 for data, label in loader: steps += 1 data, label = data.cuda(), label.cuda() if C.get().conf.get('mixup', 0.0) <= 0.0 or optimizer is None: preds = model(data) loss = loss_fn(preds, label) else: # mixup data, targets, shuffled_targets, lam = mixup( data, label, C.get()['mixup']) preds = model(data) loss = loss_fn(preds, targets, shuffled_targets, lam) del shuffled_targets, lam if optimizer: loss += wd * (1. / 2.) * sum( [torch.sum(p**2) for p in params_without_bn]) loss.backward() grad_clip = C.get()['optimizer'].get('clip', 5.0) if grad_clip > 0: nn.utils.clip_grad_norm_(model.parameters(), grad_clip) optimizer.step() optimizer.zero_grad() if ema is not None: ema(model, (epoch - 1) * total_steps + steps) top1, top5 = accuracy(preds, label, (1, 5)) metrics.add_dict({ 'loss': loss.item() * len(data), 'top1': top1.item() * len(data), 'top5': top5.item() * len(data), }) cnt += len(data) if loss_ema: loss_ema = loss_ema * 0.9 + loss.item() * 0.1 else: loss_ema = loss.item() if verbose: postfix = metrics / cnt if optimizer: postfix['lr'] = optimizer.param_groups[0]['lr'] postfix['loss_ema'] = loss_ema loader.set_postfix(postfix) if scheduler is not None: scheduler.step(epoch - 1 + float(steps) / total_steps) del preds, loss, top1, top5, data, label if tqdm_disabled and verbose: if optimizer: logger.info('[%s %03d/%03d] %s lr=%.6f', desc_default, epoch, C.get()['epoch'], metrics / cnt, optimizer.param_groups[0]['lr']) else: logger.info('[%s %03d/%03d] %s', desc_default, epoch, C.get()['epoch'], metrics / cnt) metrics /= cnt if optimizer: metrics.metrics['lr'] = optimizer.param_groups[0]['lr'] if verbose: for key, value in metrics.items(): writer.add_scalar(key, value, epoch) return metrics
def eval_tta(config, augment): C.get() C.get().conf = config cv_ratio_test, cv_fold, save_path = augment['cv_ratio_test'], augment[ 'cv_fold'], augment['save_path'] print(augment) # setup - provided augmentation rules C.get()['aug'] = policy_decoder(augment, augment['num_policy'], augment['num_op']) # eval model = get_model(C.get()['model'], num_class(C.get()['dataset'])) ckpt = torch.load(save_path) if 'model' in ckpt: model.load_state_dict(ckpt['model']) else: model.load_state_dict(ckpt) model.eval() loaders = [] for _ in range(augment['num_policy']): # TODO _, tl, validloader, tl2 = get_dataloaders(C.get()['dataset'], C.get()['batch'], augment['dataroot'], cv_ratio_test, split_idx=cv_fold) loaders.append(iter(validloader)) del tl, tl2 start_t = time.time() metrics = Accumulator() loss_fn = torch.nn.CrossEntropyLoss(reduction='none') try: while True: losses = [] corrects = [] for loader in loaders: data, label = next(loader) data = data.cuda() label = label.cuda() pred = model(data) loss = loss_fn(pred, label) losses.append(loss.detach().cpu().numpy()) _, pred = pred.topk(1, 1, True, True) pred = pred.t() correct = pred.eq(label.view( 1, -1).expand_as(pred)).detach().cpu().numpy() corrects.append(correct) del loss, correct, pred, data, label losses = np.concatenate(losses) losses_min = np.min(losses, axis=0).squeeze() corrects = np.concatenate(corrects) corrects_max = np.max(corrects, axis=0).squeeze() metrics.add_dict({ 'minus_loss': -1 * np.sum(losses_min), 'correct': np.sum(corrects_max), 'cnt': len(corrects_max) }) del corrects, corrects_max except StopIteration: pass del model metrics = metrics / 'cnt' gpu_secs = (time.time() - start_t) * torch.cuda.device_count() # reporter(minus_loss=metrics['minus_loss'], top1_valid=metrics['correct'], elapsed_time=gpu_secs, done=True) tune.track.log(minus_loss=metrics['minus_loss'], top1_valid=metrics['correct'], elapsed_time=gpu_secs, done=True) return metrics['correct']
def run_epoch(model, loader, loss_fn, optimizer, desc_default='', epoch=0, writer=None, verbose=1, scheduler=None): model_name = C.get()['model']['type'] alpha = C.get()['alpha'] skip_ratios = ListAverageMeter() tqdm_disable = bool(os.environ.get('TASK_NAME', '')) if verbose: loader = tqdm(loader, disable=tqdm_disable) loader.set_description('[%s %04d/%04d]' % (desc_default, epoch, C.get()['epoch'])) metrics = Accumulator() cnt = 0 total_steps = len(loader) steps = 0 for data, label in loader: steps += 1 data, label = data.cuda(), label.cuda() if optimizer: optimizer.zero_grad() if model_name == 'pyramid_skip': if desc_default == '*test': with torch.no_grad(): preds, masks, gprobs = model(data) skips = [mask.data.le(0.5).float().mean() for mask in masks] if skip_ratios.len != len(skips): skip_ratios.set_len(len(skips)) skip_ratios.update(skips, data.size(0)) else: preds, masks, gprobs = model(data) sparsity_loss = 0 for mask in masks: sparsity_loss += mask.mean() loss1 = loss_fn(preds, label) loss2 = alpha * sparsity_loss loss = loss1 + loss2 else: preds = model(data) loss = loss_fn(preds, label) if optimizer: loss.backward() if getattr(optimizer, "synchronize", None): optimizer.skip_synchronize() if C.get()['optimizer'].get('clip', 5) > 0: nn.utils.clip_grad_norm_(model.parameters(), C.get()['optimizer'].get('clip', 5)) optimizer.step() top1, top5 = accuracy(preds, label, (1, 5)) if model_name == 'pyramid_skip': metrics.add_dict({ 'loss1': loss1.item() * len(data), 'loss2': loss2.item() * len(data), 'top1': top1.item() * len(data), 'top5': top5.item() * len(data), }) else: metrics.add_dict({ 'loss': loss.item() * len(data), 'top1': top1.item() * len(data), 'top5': top5.item() * len(data), }) cnt += len(data) if verbose: postfix = metrics / cnt if optimizer: postfix['lr'] = optimizer.param_groups[0]['lr'] loader.set_postfix(postfix) # if scheduler is not None: # scheduler.step(epoch - 1 + float(steps) / total_steps) if model_name == 'pyramid_skip': del masks[:], gprobs[:] del preds, loss, top1, top5, data, label if model_name == 'pyramid_skip': if desc_default == '*test': skip_summaries = [] for idx in range(skip_ratios.len): skip_summaries.append(1 - skip_ratios.avg[idx]) cp = ((sum(skip_summaries) + 1) / (len(skip_summaries) + 1)) * 100 if tqdm_disable: logger.info('[%s %03d/%03d] %s', desc_default, epoch, C.get()['epoch'], metrics / cnt) metrics /= cnt if optimizer: metrics.metrics['lr'] = optimizer.param_groups[0]['lr'] if verbose: for key, value in metrics.items(): writer.add_scalar(key, value, epoch) if model_name == 'pyramid_skip': if desc_default == '*test': writer.add_scalar('Computation Percentage', cp, epoch) return metrics
def train_controller(controller, dataloaders, save_path, ctl_save_path): dataset = C.get()['test_dataset'] ctl_train_steps = 1500 ctl_num_aggre = 10 ctl_entropy_w = 1e-5 ctl_ema_weight = 0.95 metrics = Accumulator() cnt = 0 controller.train() test_ratio = 0. _, _, dataloader, _ = dataloaders # validloader optimizer = optim.SGD(controller.parameters(), lr=0.00035, momentum=0.9, weight_decay=0.0, nesterov=True) # optimizer = optim.Adam(controller.parameters(), lr = 0.00035) # create a model & a criterion model = get_model(C.get()['model'], num_class(dataset), local_rank=-1) criterion = CrossEntropyLabelSmooth(num_class(dataset), C.get().conf.get('lb_smooth', 0), reduction="batched_sum").cuda() # load model weights data = torch.load(save_path) key = 'model' if 'model' in data else 'state_dict' if 'epoch' not in data: model.load_state_dict(data) else: logger.info('checkpoint epoch@%d' % data['epoch']) if not isinstance(model, (DataParallel, DistributedDataParallel)): model.load_state_dict( {k.replace('module.', ''): v for k, v in data[key].items()}) else: model.load_state_dict({ k if 'module.' in k else 'module.' + k: v for k, v in data[key].items() }) del data model.eval() loader_iter = iter(dataloader) # [(image)->ToTensor->Normalize] baseline = None if os.path.isfile(ctl_save_path): logger.info('------Controller load------') checkpoint = torch.load(ctl_save_path) controller.load_state_dict(checkpoint['ctl_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) cnt = checkpoint['cnt'] mean_probs = checkpoint['mean_probs'] accs = checkpoint['accs'] metrics_dict = checkpoint['metrics'] metrics.metrics = metrics_dict init_step = checkpoint['step'] else: logger.info('------Train Controller from scratch------') mean_probs = [] accs = [] init_step = 0 for step in tqdm(range(init_step + 1, ctl_train_steps * ctl_num_aggre + 1)): try: inputs, labels = next(loader_iter) except: loader_iter = iter(dataloader) inputs, labels = next(loader_iter) batch_size = len(labels) inputs, labels = inputs.cuda(), labels.cuda() log_probs, entropys, sampled_policies = controller(inputs) # evaluate model with augmented validation dataset with torch.no_grad(): # compare Accuracy before/after augmentation # ori_preds = model(inputs) # ori_top1, ori_top5 = accuracy(ori_preds, labels, (1, 5)) batch_policies = batch_policy_decoder( sampled_policies ) # (list:list:list:tuple) [batch, num_policy, n_op, 3] aug_inputs, applied_policy = augment_data(inputs, batch_policies) aug_inputs = aug_inputs.cuda() # assert type(aug_inputs) == torch.Tensor, "Augmented Input Type Error: {}".format(type(aug_inputs)) preds = model(aug_inputs) model_losses = criterion(preds, labels) # (tensor)[batch] top1, top5 = accuracy(preds, labels, (1, 5)) # logger.info("Acc B/A Aug, {:.2f}->{:.2f}".format(ori_top1, top1)) # assert model_losses.shape == entropys.shape == log_probs.shape, \ # "[Size miss match] loss: {}, entropy: {}, log_prob: {}".format(model_losses.shape, entropys.shape, log_probs.shape) rewards = -model_losses + ctl_entropy_w * entropys # (tensor)[batch] if baseline is None: baseline = -model_losses.mean() # scalar tensor else: # assert baseline, "len(baseline): {}".format(len(baseline)) baseline = baseline - (1 - ctl_ema_weight) * ( baseline - rewards.mean().detach()) # baseline = 0. loss = -1 * (log_probs * (rewards - baseline)).mean() #scalar tensor # Average gradient over controller_num_aggregate samples loss = loss / ctl_num_aggre loss.backward(retain_graph=True) metrics.add_dict({ 'loss': loss.item() * batch_size, 'top1': top1.item() * batch_size, 'top5': top5.item() * batch_size, }) cnt += batch_size if (step + 1) % ctl_num_aggre == 0: torch.nn.utils.clip_grad_norm_(controller.parameters(), 5.0) optimizer.step() controller.zero_grad() # torch.cuda.empty_cache() logger.info('\n[Train Controller %03d/%03d] log_prob %02f, %s', step, ctl_train_steps*ctl_num_aggre, \ log_probs.mean().item(), metrics / cnt ) if step % 100 == 0 or step == ctl_train_steps * ctl_num_aggre: save_pic(inputs, aug_inputs, labels, applied_policy, batch_policies, step) ps = [] for pol in batch_policies: # (list:list:list:tuple) [batch, num_policy, n_op, 3] for ops in pol: for op in ops: p = op[1] ps.append(p) mean_prob = np.mean(ps) mean_probs.append(mean_prob) accs.append(top1.item()) print("Mean probability: {:.2f}".format(mean_prob)) torch.save( { 'step': step, 'ctl_state_dict': controller.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'metrics': dict(metrics.metrics), 'cnt': cnt, 'mean_probs': mean_probs, 'accs': accs }, ctl_save_path) return metrics, None #baseline.item()
def eval_tta(config, augment, reporter): C.get() C.get().conf = config save_path = augment['save_path'] cv_id, gr_id = augment["cv_id"], augment["gr_id"] gr_ids = augment["gr_ids"] # setup - provided augmentation rules C.get()['aug'] = policy_decoder(augment, augment['num_policy'], augment['num_op']) # eval model = get_model(C.get()['model'], num_class(C.get()['dataset'])) ckpt = torch.load(save_path) if 'model' in ckpt: model.load_state_dict(ckpt['model']) else: model.load_state_dict(ckpt) model.eval() loaders = [] for _ in range(augment['num_policy']): # TODO loader = get_post_dataloader(C.get()["dataset"], C.get()['batch'], augment["dataroot"], augment['cv_ratio_test'], cv_id, gr_id, gr_ids) loaders.append(iter(loader)) start_t = time.time() metrics = Accumulator() loss_fn = torch.nn.CrossEntropyLoss(reduction='none') try: while True: losses = [] corrects = [] for loader in loaders: data, label = next(loader) data = data.cuda() label = label.cuda() pred = model(data) loss = loss_fn(pred, label) losses.append(loss.detach().cpu().numpy().reshape(1, -1)) # (1,N) _, pred = pred.topk(1, 1, True, True) pred = pred.t() correct = pred.eq(label.view( 1, -1).expand_as(pred)).detach().cpu().numpy() # (1,N) corrects.append(correct) del loss, correct, pred, data, label losses = np.concatenate(losses) losses_min = np.min(losses, axis=0).squeeze() # (N,) corrects = np.concatenate(corrects) corrects_max = np.max(corrects, axis=0).squeeze() # (N,) metrics.add_dict({ 'loss': np.sum(losses_min), 'correct': np.sum(corrects_max), 'cnt': corrects_max.size }) del corrects, corrects_max except StopIteration: pass del model metrics = metrics / 'cnt' gpu_secs = (time.time() - start_t) * torch.cuda.device_count() reporter(loss=metrics['loss'], top1_valid=metrics['correct'], elapsed_time=gpu_secs, done=True) return metrics['correct']