def validate(self,
                 epoch=0,
                 is_test=False,
                 run_str='',
                 net=None,
                 data_loader=None,
                 no_logs=False):
        if net is None:
            net = self.net
        if data_loader is None:
            if is_test:
                data_loader = self.run_config.test_loader
            else:
                data_loader = self.run_config.valid_loader

        net.eval()

        losses = DistributedMetric('val_loss')
        metric_dict = self.get_metric_dict()

        with torch.no_grad():
            with tqdm(total=len(data_loader),
                      desc='Validate Epoch #{} {}'.format(epoch + 1, run_str),
                      disable=no_logs or not self.is_root) as t:
                for i, (images, labels) in enumerate(data_loader):
                    images, labels = images.cuda(), labels.cuda()
                    # compute output
                    output = net(images)
                    loss = self.test_criterion(output, labels)
                    # measure accuracy and record loss
                    losses.update(loss, images.size(0))
                    self.update_metric(metric_dict, output, labels)
                    t.set_postfix({
                        'loss':
                        losses.avg.item(),
                        **self.get_metric_vals(metric_dict, return_dict=True),
                        'img_size':
                        images.size(2),
                    })
                    t.update(1)
        return losses.avg.item(), self.get_metric_vals(metric_dict)
def train_one_epoch(run_manager, args, epoch, warmup_epochs=0, warmup_lr=0):
    dynamic_net = run_manager.network
    distributed = isinstance(run_manager, DistributedRunManager)

    # switch to train mode
    dynamic_net.train()
    if distributed:
        run_manager.run_config.train_loader.sampler.set_epoch(epoch)
    MyRandomResizedCrop.EPOCH = epoch

    nBatch = len(run_manager.run_config.train_loader)

    data_time = AverageMeter()
    losses = DistributedMetric('train_loss') if distributed else AverageMeter()
    metric_dict = run_manager.get_metric_dict()

    with tqdm(total=nBatch,
              desc='Train Epoch #{}'.format(epoch + 1),
              disable=distributed and not run_manager.is_root) as t:
        end = time.time()
        for i, (images,
                labels) in enumerate(run_manager.run_config.train_loader):
            MyRandomResizedCrop.BATCH = i
            data_time.update(time.time() - end)
            if epoch < warmup_epochs:
                new_lr = run_manager.run_config.warmup_adjust_learning_rate(
                    run_manager.optimizer,
                    warmup_epochs * nBatch,
                    nBatch,
                    epoch,
                    i,
                    warmup_lr,
                )
            else:
                new_lr = run_manager.run_config.adjust_learning_rate(
                    run_manager.optimizer, epoch - warmup_epochs, i, nBatch)

            images, labels = images.cuda(), labels.cuda()
            target = labels

            # soft target
            if args.kd_ratio > 0:
                args.teacher_model.train()
                with torch.no_grad():
                    soft_logits = args.teacher_model(images).detach()
                    soft_label = F.softmax(soft_logits, dim=1)

            # clean gradients
            dynamic_net.zero_grad()

            loss_of_subnets = []
            # compute output
            subnet_str = ''
            for _ in range(args.dynamic_batch_size):
                # set random seed before sampling
                subnet_seed = int('%d%.3d%.3d' % (epoch * nBatch + i, _, 0))
                random.seed(subnet_seed)
                subnet_settings = dynamic_net.sample_active_subnet()
                subnet_str += '%d: ' % _ + ','.join([
                    '%s_%s' %
                    (key, '%.1f' %
                     subset_mean(val, 0) if isinstance(val, list) else val)
                    for key, val in subnet_settings.items()
                ]) + ' || '

                output = run_manager.net(images)
                if args.kd_ratio == 0:
                    loss = run_manager.train_criterion(output, labels)
                    loss_type = 'ce'
                else:
                    if args.kd_type == 'ce':
                        kd_loss = cross_entropy_loss_with_soft_target(
                            output, soft_label)
                    else:
                        kd_loss = F.mse_loss(output, soft_logits)
                    loss = args.kd_ratio * kd_loss + run_manager.train_criterion(
                        output, labels)
                    loss_type = '%.1fkd-%s & ce' % (args.kd_ratio,
                                                    args.kd_type)

                # measure accuracy and record loss
                loss_of_subnets.append(loss)
                run_manager.update_metric(metric_dict, output, target)

                loss.backward()
            run_manager.optimizer.step()

            losses.update(list_mean(loss_of_subnets), images.size(0))

            t.set_postfix({
                'loss':
                losses.avg.item(),
                **run_manager.get_metric_vals(metric_dict, return_dict=True),
                'R':
                images.size(2),
                'lr':
                new_lr,
                'loss_type':
                loss_type,
                'seed':
                str(subnet_seed),
                'str':
                subnet_str,
                'data_time':
                data_time.avg,
            })
            t.update(1)
            end = time.time()
    return losses.avg.item(), run_manager.get_metric_vals(metric_dict)
    def train_one_epoch(self, args, epoch, warmup_epochs=5, warmup_lr=0):
        self.net.train()
        self.run_config.train_loader.sampler.set_epoch(
            epoch)  # required by distributed sampler
        MyRandomResizedCrop.EPOCH = epoch  # required by elastic resolution

        nBatch = len(self.run_config.train_loader)

        losses = DistributedMetric('train_loss')
        metric_dict = self.get_metric_dict()
        data_time = AverageMeter()

        with tqdm(total=nBatch,
                  desc='Train Epoch #{}'.format(epoch + 1),
                  disable=not self.is_root) as t:
            end = time.time()
            for i, (images, labels) in enumerate(self.run_config.train_loader):
                MyRandomResizedCrop.BATCH = i
                data_time.update(time.time() - end)
                if epoch < warmup_epochs:
                    new_lr = self.run_config.warmup_adjust_learning_rate(
                        self.optimizer,
                        warmup_epochs * nBatch,
                        nBatch,
                        epoch,
                        i,
                        warmup_lr,
                    )
                else:
                    new_lr = self.run_config.adjust_learning_rate(
                        self.optimizer, epoch - warmup_epochs, i, nBatch)

                images, labels = images.cuda(), labels.cuda()
                target = labels
                if isinstance(self.run_config.mixup_alpha, float):
                    # transform data
                    random.seed(int('%d%.3d' % (i, epoch)))
                    lam = random.betavariate(self.run_config.mixup_alpha,
                                             self.run_config.mixup_alpha)
                    images = mix_images(images, lam)
                    labels = mix_labels(
                        labels, lam, self.run_config.data_provider.n_classes,
                        self.run_config.label_smoothing)

                # soft target
                if args.teacher_model is not None:
                    args.teacher_model.train()
                    with torch.no_grad():
                        soft_logits = args.teacher_model(images).detach()
                        soft_label = F.softmax(soft_logits, dim=1)

                # compute output
                output = self.net(images)

                if args.teacher_model is None:
                    loss = self.train_criterion(output, labels)
                    loss_type = 'ce'
                else:
                    if args.kd_type == 'ce':
                        kd_loss = cross_entropy_loss_with_soft_target(
                            output, soft_label)
                    else:
                        kd_loss = F.mse_loss(output, soft_logits)
                    loss = args.kd_ratio * kd_loss + self.train_criterion(
                        output, labels)
                    loss_type = '%.1fkd+ce' % args.kd_ratio

                # update
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

                # measure accuracy and record loss
                losses.update(loss, images.size(0))
                self.update_metric(metric_dict, output, target)

                t.set_postfix({
                    'loss':
                    losses.avg.item(),
                    **self.get_metric_vals(metric_dict, return_dict=True),
                    'img_size':
                    images.size(2),
                    'lr':
                    new_lr,
                    'loss_type':
                    loss_type,
                    'data_time':
                    data_time.avg,
                })
                t.update(1)
                end = time.time()

        return losses.avg.item(), self.get_metric_vals(metric_dict)
 def get_metric_dict(self):
     return {
         'top1': DistributedMetric('top1'),
         'top5': DistributedMetric('top5'),
     }