예제 #1
0
def step(args, split, epoch, loader, model, optimizer = None, M = None, f = None, tag = None):
  losses, mpjpe, mpjpe_r = AverageMeter(), AverageMeter(), AverageMeter()
  viewLosses, shapeLosses, supLosses = AverageMeter(), AverageMeter(), AverageMeter()
  
  if split == 'train':
    model.train()
  else:
    model.eval()
  bar = Bar('{}'.format(ref.category), max=len(loader))
  
  nViews = loader.dataset.nViews
  for i, (input, target, meta) in enumerate(loader):
    input_var = torch.autograd.Variable(input)
    target_var = torch.autograd.Variable(target)
    output = model(input_var)
    loss = ShapeConsistencyCriterion(nViews, supWeight = 1, unSupWeight = args.shapeWeight, M = M)(output, target_var, torch.autograd.Variable(meta))

    if split == 'test':
      for j in range(input.numpy().shape[0]):
        img = (input.numpy()[j] * 255).transpose(1, 2, 0).astype(np.uint8)
        cv2.imwrite('{}/img_{}/{}.png'.format(args.save_path, tag, i * input.numpy().shape[0] + j), img)
        gt = target.cpu().numpy()[j]
        pred = (output.data).cpu().numpy()[j]
        vis = meta.cpu().numpy()[j][5:]
        for t in range(ref.J):
          f.write('{} {} {} '.format(pred[t * 3], pred[t * 3 + 1], pred[t * 3 + 2]))
        f.write('\n')
        for t in range(ref.J):
          f.write('{} {} {} '.format(gt[t, 0], gt[t, 1], gt[t, 2]))
        f.write('\n')
        if args.saveVis:
          for t in range(ref.J):
            f.write('{} 0 0 '.format(vis[t]))
          f.write('\n')

    mpjpe_this = accuracy(output.data, target, meta)
    mpjpe_r_this = accuracy_dis(output.data, target, meta)
    shapeLoss = shapeConsistency(output.data, meta, nViews, M, split = split)

    losses.update(loss.data[0], input.size(0))
    shapeLosses.update(shapeLoss, input.size(0))
    mpjpe.update(mpjpe_this, input.size(0))
    mpjpe_r.update(mpjpe_r_this, input.size(0))
    
    if split == 'train':
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()
    
    Bar.suffix = '{split:10}: [{0:2}][{1:3}/{2:3}] | Total: {total:} | ETA: {eta:} | Loss {loss.avg:.6f} | shapeLoss {shapeLoss.avg:.6f} | AE {mpjpe.avg:.6f} | ShapeDis {mpjpe_r.avg:.6f}'.format(epoch, i, len(loader), total=bar.elapsed_td, eta=bar.eta_td, loss=losses, mpjpe=mpjpe, split = split, shapeLoss = shapeLosses, mpjpe_r = mpjpe_r)
    bar.next()
      
  bar.finish()
  return mpjpe.avg, losses.avg, shapeLosses.avg
예제 #2
0
    def run_epoch(self, phase, epoch, data_loader):
        model_with_loss = self.model_with_loss
        if len(self.opt.gpus) > 1:
            if self.opt.weight_strategy == 'GRADNORM':
                model_with_loss.module.loss.optimizer = get_loss_optimizer(model=model_with_loss.module.loss.loss_model,
                                                                           opt=self.opt)
            model_with_loss.module.loss.update_weight(epoch)
        else:
            if self.opt.weight_strategy == 'GRADNORM':
                model_with_loss.loss.optimizer = get_loss_optimizer(model=model_with_loss.loss.loss_model, opt=self.opt)
            model_with_loss.loss.update_weight(epoch)
        if phase == 'train':
            model_with_loss.train()
        else:
            if len(self.opt.gpus) > 1:
                model_with_loss = self.model_with_loss.module
            model_with_loss.eval()
            torch.cuda.empty_cache()

        opt = self.opt
        results = {}
        data_time, batch_time = AverageMeter(), AverageMeter()
        avg_loss_stats = {l: AverageMeter() for l in self.loss_stats} #\
                          # if l == 'tot' or opt.weights[l] > 0}
        num_iters = len(data_loader) if opt.num_iters < 0 else opt.num_iters
        bar = Bar('{}/{}'.format(opt.task, opt.exp_id), max=num_iters)
        end = time.time()
        for iter_id, batch in enumerate(data_loader):
            if iter_id >= num_iters:
                break
            data_time.update(time.time() - end)

            for k in batch:
                if k != 'meta':
                    batch[k] = batch[k].to(device=opt.device, non_blocking=True)
            output, loss, loss_stats = model_with_loss(batch, epoch)
            loss = loss.mean()
            if phase == 'train':
                self.optimizer.zero_grad()
                if opt.weight_strategy == 'GRADNORM':
                    loss.backward(retain_graph=True)
                    if len(self.opt.gpus) > 1:
                        # model_with_loss.module.loss.loss_model.update_weight(model_with_loss.module.model, model_with_loss.module.loss.optimizer, loss_stats)
                        # torch.sum(loss_stats['update']).backward()
                        model_with_loss.module.loss.optimizer.zero_grad()
                        temp_grad = torch.autograd.grad(torch.sum(loss_stats['update']),
                                                        model_with_loss.module.loss.loss_model.weight)[0]
                        grad_norm = torch.norm(temp_grad.data, 1)
                        print(grad_norm)
                        if grad_norm > opt.gradnorm_thred:
                            temp_grad = torch.zeros_like(temp_grad)
                        model_with_loss.module.loss.loss_model.weight.grad = temp_grad
                        model_with_loss.module.loss.optimizer.step()
                    else:
                        # model_with_loss.loss.loss_model.update_weight(model_with_loss.model,
                        #                                               model_with_loss.loss.optimizer, loss_stats)
                        model_with_loss.loss.optimizer.zero_grad()
                        temp_grad = torch.autograd.grad(loss_stats['update'],
                                                        model_with_loss.loss.loss_model.
                                                        weight)[0]
                        grad_norm = torch.norm(temp_grad.data, 1)
                        if grad_norm > opt.gradnorm_thred:
                            temp_grad = torch.zeros_like(temp_grad)
                        model_with_loss.loss.loss_model.weight.grad = temp_grad
                        model_with_loss.loss.optimizer.step()
                else:
                    # torch.autograd.grad(loss, model_with_loss.model.padnet.parameters())
                    loss.backward()
                self.optimizer.step()
                if opt.weight_strategy == 'UNCER':
                    if len(self.opt.gpus) > 1:
                        model_with_loss.module.loss.optimizer.step()
                        model_with_loss.module.loss.optimizer.zero_grad()
                        print(model_with_loss.module.loss.group_weight)
                    else:
                        model_with_loss.loss.optimizer.step()
                        model_with_loss.loss.optimizer.zero_grad()
                        print(model_with_loss.loss.group_weight)
            batch_time.update(time.time() - end)
            end = time.time()

            Bar.suffix = '{phase}: [{0}][{1}/{2}]|Tot: {total:} |ETA: {eta:} '.format(
                epoch, iter_id, num_iters, phase=phase,
                total=bar.elapsed_td, eta=bar.eta_td)
            for l in avg_loss_stats:
                avg_loss_stats[l].update(
                    loss_stats[l].mean().item(), batch['image'].size(0))
                Bar.suffix = Bar.suffix + '|{} {:.4f} '.format(l, avg_loss_stats[l].avg)
            Bar.suffix = Bar.suffix + '|Data {dt.val:.3f}s({dt.avg:.3f}s) ' \
                                      '|Net {bt.avg:.3f}s'.format(dt=data_time, bt=batch_time)
            if opt.print_iter > 0:  # If not using progress bar
                if iter_id % opt.print_iter == 0:
                    print('{}/{}| {}'.format(opt.task, opt.exp_id, Bar.suffix))
            else:
                bar.next()

            if opt.debug > 0:
                self.debug(batch, output, iter_id, dataset=data_loader.dataset)

            del output, loss, loss_stats

        bar.finish()
        ret = {k: v.avg for k, v in avg_loss_stats.items()}
        ret['time'] = bar.elapsed_td.total_seconds() / 60.
        if len(self.opt.gpus) > 1:
            model_with_loss.module.loss.update_loss(epoch, ret)
        else:
            model_with_loss.loss.update_loss(epoch, ret)
        return ret, results
예제 #3
0
    def run_epoch(self, phase, epoch, data_loader):
        model_with_loss = self.model_with_loss
        if phase == 'train':
            model_with_loss.train()
        else:
            if len(self.opt.gpus) > 1:
                model_with_loss = self.model_with_loss.module
            model_with_loss.eval()
            torch.cuda.empty_cache()

        opt = self.opt
        results = {}
        data_time, batch_time = AverageMeter(), AverageMeter()
        avg_loss_stats = {l: AverageMeter() for l in self.loss_stats}
        num_iters = len(data_loader) if opt.num_iters < 0 else opt.num_iters
        bar = Bar('{}/{}'.format(opt.task, opt.exp_id), max=num_iters)
        end = time.time()
        for iter_id, batch in enumerate(data_loader):
            if iter_id >= num_iters:
                break
            data_time.update(time.time() - end)

            for k in batch:
                if k != 'meta':
                    batch[k] = batch[k].to(device=opt.device,
                                           non_blocking=True)
            output, loss, loss_stats = model_with_loss(batch)
            loss = loss.mean()
            if phase == 'train':
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()
            batch_time.update(time.time() - end)
            end = time.time()

            Bar.suffix = '{phase}: [{0}][{1}/{2}]|Tot: {total:} |ETA: {eta:} '.format(
                epoch,
                iter_id,
                num_iters,
                phase=phase,
                total=bar.elapsed_td,
                eta=bar.eta_td)
            for l in avg_loss_stats:
                avg_loss_stats[l].update(loss_stats[l].mean().item(),
                                         batch['input'].size(0))
                Bar.suffix = Bar.suffix + '|{} {:.4f} '.format(
                    l, avg_loss_stats[l].avg)
            if not opt.hide_data_time:
                Bar.suffix = Bar.suffix + '|Data {dt.val:.3f}s({dt.avg:.3f}s) ' \
                  '|Net {bt.avg:.3f}s'.format(dt=data_time, bt=batch_time)
            if opt.print_iter > 0:
                if iter_id % opt.print_iter == 0:
                    print('{}/{}| {}'.format(opt.task, opt.exp_id, Bar.suffix))
            else:
                bar.next()

            if opt.debug > 0:
                self.debug(batch, output, iter_id)

            if opt.test:
                self.save_result(output, batch, results)
            del output, loss, loss_stats

        bar.finish()
        ret = {k: v.avg for k, v in avg_loss_stats.items()}
        ret['time'] = bar.elapsed_td.total_seconds() / 60.
        return ret, results
예제 #4
0
파일: train.py 프로젝트: ggsDing/SARSeg
def validate(val_loader, net, criterion, curr_epoch):
    # the following code is written assuming that batch size is 1
    net.eval()
    if args['gpu']:
        torch.cuda.empty_cache()
    start = time.time()

    val_loss = AverageMeter()
    acc_meter = AverageMeter()
    fwIoU_meter = AverageMeter()

    for vi, (imgs, labels) in enumerate(val_loader):
        imgs = imgs.float()
        labels = labels.long()

        if args['gpu']:
            imgs = imgs.cuda().float()
            labels = labels.cuda().long()

        with torch.no_grad():
            outputs, aux = net(imgs)
            loss = criterion(outputs, labels)
        val_loss.update(loss.cpu().detach().numpy())

        outputs = outputs.cpu().detach()
        labels = labels.cpu().detach().numpy()
        _, preds = torch.max(outputs, dim=1)
        preds = preds.numpy()
        for (pred, label) in zip(preds, labels):
            acc, valid_sum = accuracy(pred, label)
            fwiou = FWIoU(pred.squeeze(), label.squeeze(), ignore_zero=True)
            acc_meter.update(acc)
            fwIoU_meter.update(fwiou)

        if curr_epoch % args['predict_step'] == 0 and vi == 0:
            pred_color = RS.Index2Color(preds[0])
            io.imsave(os.path.join(args['pred_dir'], NET_NAME + '.png'),
                      pred_color)
            print('Prediction saved!')

    curr_time = time.time() - start
    print('%.1fs Val loss: %.2f, Accuracy: %.2f, fwIoU: %.2f' %
          (curr_time, val_loss.average(), acc_meter.average() * 100,
           fwIoU_meter.average() * 100))

    writer.add_scalar('val_loss', val_loss.average(), curr_epoch)
    writer.add_scalar('val_Accuracy', acc_meter.average(), curr_epoch)
    writer.add_scalar('val_fwIoU', fwIoU_meter.average(), curr_epoch)

    return acc_meter.avg, fwIoU_meter.avg, val_loss.avg
예제 #5
0
파일: test.py 프로젝트: INSFAN/FreeGazeNet
def validate(args, val_dataloader, model, auxiliarynet, epoch):

    batch_time = AverageMeter('Time', ':6.3f')
    data_time = AverageMeter('Data', ':6.3f')
    losses = AverageMeter('Loss', ':.4e')
    error = AverageMeter('error', ':6.2f')

    progress = ProgressMeter(len(val_dataloader),
                             batch_time,
                             data_time,
                             losses,
                             error,
                             prefix="Val Epoch: [{}]".format(epoch))

    model.eval()
    # auxiliarynet.eval()
    end = time.time()
    with torch.no_grad():
        end = time.time()
        for i, (patch, gaze_norm_g, head_norm,
                rot_vec_norm) in enumerate(val_dataloader):
            # measure data loading time
            data_time.update(time.time() - end)
            patch = patch.to(args.device)
            gaze_norm_g = gaze_norm_g.to(args.device)

            head_norm = head_norm.to(args.device)

            rot_vec_norm = rot_vec_norm.to(args.device)

            # model = model.to(args.device)
            gaze_pred, _ = model(patch)
            # hp_pred = auxiliarynet(features)

            head_norm = 10 * head_norm
            gaze_norm_g = 10 * gaze_norm_g
            # loss = criterion(gaze_norm_g, head_norm, gaze_pred[:,0:2], gaze_pred[:,2:4])

            angle_error = mean_angle_error(
                gaze_pred.cpu().detach().numpy() / 10,
                gaze_norm_g.cpu().detach().numpy() / 10,
                rot_vec_norm.cpu().detach().numpy())
            # losses.update(loss.item())
            error.update(angle_error)

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if (i + 1) % args.print_freq == 0:
                progress.print(i + 1)
                # img = patch.cpu().detach().numpy()[0].deepcopy()
                # to_visualize = draw_gaze(img[0], (0.25 * img.shape[1], 0.25 * img.shape[1]), gaze_pred,
                # gaze_norm_g, length=80.0, thickness=1)
                # cv2.imshow('vis', to_visualize)
                # cv2.waitKey(1)

    return losses.get_avg(), error.get_avg()
예제 #6
0
def validate(args, val_loader, model, epoch, win_feats5, win_fusion, viz,
             global_step):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    feats5_losses = AverageMeter()
    fusion_losses = AverageMeter()
    total_losses = AverageMeter()

    # switch to train mode
    model.eval()
    torch.no_grad()

    end = time.time()
    for i, (img, target) in enumerate(val_loader):
        # measure data loading time
        data_time.update(time.time() - end)

        # Input for Image CNN.
        img_var = utils.check_gpu(0, img)  # BS X 3 X H X W
        target_var = utils.check_gpu(0, target)  # BS X H X W X NUM_CLASSES

        bs = img.size()[0]

        score_feats5, fused_feats = model(
            img_var)  # BS X NUM_CLASSES X 472 X 472

        feats5_loss = WeightedMultiLabelSigmoidLoss(score_feats5, target_var)
        fused_feats_loss = WeightedMultiLabelSigmoidLoss(
            fused_feats, target_var)
        loss = feats5_loss + fused_feats_loss

        feats5_losses.update(feats5_loss.data, bs)
        fusion_losses.update(fused_feats_loss.data, bs)
        total_losses.update(loss.data, bs)

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        # clear memory
        del img_var
        del target_var
        del score_feats5
        del fused_feats_loss
        del feats5_loss
        torch.cuda.empty_cache()

        if (i % args.print_freq == 0):
            print("\n\n")
            print('Epoch: [{0}][{1}/{2}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Total Loss {total_loss.val:.4f} ({total_loss.avg:.4f})\n'.
                  format(epoch,
                         i,
                         len(val_loader),
                         batch_time=batch_time,
                         data_time=data_time,
                         total_loss=total_losses))

    #viz.line(win=win_feats5, name='val_feats5', update='append', X=np.array([global_step]), Y=np.array([feats5_losses.avg]))
    #viz.line(win=win_fusion, name='val_fusion', update='append', X=np.array([global_step]), Y=np.array([fusion_losses.avg]))

    return fusion_losses.avg
예제 #7
0
def train(train_loader, miml, encoder, decoder, criterion, encoder_optimizer,
          decoder_optimizer, epoch, writer):
    """
    Performs one epoch's training.

    :param train_loader: DataLoader for training data
    :param encoder: encoder model
    :param decoder: decoder model
    :param criterion: loss layer
    :param encoder_optimizer: optimizer to update encoder's weights (if fine-tuning)
    :param decoder_optimizer: optimizer to update decoder's weights
    :param epoch: epoch number
    """
    encoder.train()
    decoder.train()  # train mode (dropout and batchnorm is used)
    miml.train()
    total_step = len(train_loader)
    batch_time = AverageMeter()  # forward prop. + back prop. time
    data_time = AverageMeter()  # data loading time
    losses = AverageMeter()  # loss (per word decoded)
    top5accs = AverageMeter()  # top5 accuracy

    start = time.time()

    # Batches
    for i, (imgs, caps, caplens) in enumerate(train_loader):
        data_time.update(time.time() - start)

        # Move to GPU, if available
        imgs = imgs.to(device)
        caps = caps.to(device)
        caplens = caplens.to(device)

        # Forward prop.
        # attrs = miml(imgs)
        # imgs = encoder(imgs)
        scores, caps_sorted, decode_lengths, alphas, sort_ind = decoder(
            miml(imgs), encoder(imgs), caps, caplens)

        # Since we decoded starting with <start>, the targets are all words after <start>, up to <end>
        targets = caps_sorted[:, 1:]

        # Remove timesteps that we didn't decode at, or are pads
        # pack_padded_sequence is an easy trick to do this
        # torch在计算时会自动除去pad,这样不带pad计算不影响精度
        scores, _ = pack_padded_sequence(scores,
                                         decode_lengths,
                                         batch_first=True)
        targets, _ = pack_padded_sequence(targets,
                                          decode_lengths,
                                          batch_first=True)

        # Calculate loss
        loss = criterion(scores, targets)
        loss += alpha_c * ((1. - alphas.sum(dim=1))**2).mean()
        # Back prop.
        decoder_optimizer.zero_grad()
        if encoder_optimizer is not None:
            encoder_optimizer.zero_grad()

        loss.backward()

        # Clip gradients
        if grad_clip is not None:
            clip_gradient(decoder_optimizer, grad_clip)
            if encoder_optimizer is not None:
                clip_gradient(encoder_optimizer, grad_clip)
        # Update weights
        if encoder_optimizer is not None:
            encoder_optimizer.step()
        decoder_optimizer.step()

        # Keep track of metrics
        top5 = accuracy(scores, targets, 5)
        losses.update(loss.item(), sum(decode_lengths))
        top5accs.update(top5, sum(decode_lengths))
        batch_time.update(time.time() - start)

        start = time.time()

        # Print status
        if i % print_freq == 0:
            writer.add_scalars('train', {
                'loss': loss.item(),
                'mAp': top5accs.val
            }, epoch * total_step + i)
            print('Epoch: [{0}][{1}/{2}]\t'
                  'Batch Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data Load Time {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Top-5 Accuracy {top5.val:.3f} ({top5.avg:.3f})'.format(
                      epoch,
                      i,
                      len(train_loader),
                      batch_time=batch_time,
                      data_time=data_time,
                      loss=losses,
                      top5=top5accs))
예제 #8
0
def main():
    # Arguments
    parser = argparse.ArgumentParser(
        description=
        'High Quality Monocular Depth Estimation via Transfer Learning')
    parser.add_argument('--epochs',
                        default=20,
                        type=int,
                        help='number of total epochs to run')
    parser.add_argument('--lr',
                        '--learning-rate',
                        default=0.0001,
                        type=float,
                        help='initial learning rate')
    parser.add_argument('--bs', default=4, type=int, help='batch size')
    args = parser.parse_args()

    # Create model
    model = Model().cuda()
    print('Model created.')

    # Training parameters
    optimizer = torch.optim.Adam(model.parameters(), args.lr)
    batch_size = args.bs
    prefix = 'densenet_' + str(batch_size)

    # Load data
    train_loader, test_loader = getTrainingTestingData(batch_size=batch_size)

    # Logging
    writer = SummaryWriter(comment='{}-lr{}-e{}-bs{}'.format(
        prefix, args.lr, args.epochs, args.bs),
                           flush_secs=30)

    # Loss
    l1_criterion = nn.L1Loss()

    # Start training...
    for epoch in range(args.epochs):
        batch_time = AverageMeter()
        losses = AverageMeter()
        N = len(train_loader)

        # Switch to train mode
        model.train()

        end = time.time()

        for i, sample_batched in enumerate(train_loader):
            optimizer.zero_grad()

            # Prepare sample and target
            image = torch.autograd.Variable(sample_batched['image'].cuda())
            depth = torch.autograd.Variable(
                sample_batched['depth'].cuda(non_blocking=True))

            # Normalize depth
            depth_n = DepthNorm(depth)

            # Predict
            output = model(image)

            # Compute the loss
            l_depth = l1_criterion(output, depth_n)
            l_ssim = torch.clamp(
                (1 - ssim(output, depth_n, val_range=1000.0 / 10.0)) * 0.5, 0,
                1)

            loss = (1.0 * l_ssim) + (0.1 * l_depth)

            # Update step
            losses.update(loss.data.item(), image.size(0))
            loss.backward()
            optimizer.step()

            # Measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()
            eta = str(datetime.timedelta(seconds=int(batch_time.val *
                                                     (N - i))))

            # Log progress
            niter = epoch * N + i
            if i % 5 == 0:
                # Print to console
                print('Epoch: [{0}][{1}/{2}]\t'
                      'Time {batch_time.val:.3f} ({batch_time.sum:.3f})\t'
                      'ETA {eta}\t'
                      'Loss {loss.val:.4f} ({loss.avg:.4f})'.format(
                          epoch,
                          i,
                          N,
                          batch_time=batch_time,
                          loss=losses,
                          eta=eta))

                # Log to tensorboard
                writer.add_scalar('Train/Loss', losses.val, niter)

            if i % 300 == 0:
                LogProgress(model, writer, test_loader, niter)

        # Record epoch's intermediate results
        LogProgress(model, writer, test_loader, niter)
        writer.add_scalar('Train/Loss.avg', losses.avg, epoch)
예제 #9
0
    def update_network(self, **kwargs):
        stop = False
        self.train_data['source']['iterator'] = iter(
            self.train_data['source']['loader'])
        self.train_data['target']['iterator'] = iter(
            self.train_data['target']['loader'])
        self.iters_per_epoch = len(self.train_data['target']['loader'])
        iters_counter_within_epoch = 0
        data_time = AverageMeter()
        batch_time = AverageMeter()
        classifier_loss = AverageMeter()
        feature_extractor_loss = AverageMeter()
        prec1_fs = AverageMeter()
        prec1_ft = AverageMeter()
        self.feature_extractor.train()
        self.classifier.train()
        end = time.time()
        if self.opt.TRAIN.PROCESS_COUNTER == 'epoch':
            lam = 2 / (1 + math.exp(
                -1 * 10 * self.epoch / self.opt.TRAIN.MAX_EPOCH)) - 1
            self.update_lr()
            print('value of lam is: %3f' % (lam))
        while not stop:
            if self.opt.TRAIN.PROCESS_COUNTER == 'iteration':
                lam = 2 / (1 + math.exp(
                    -1 * 10 * self.iters /
                    (self.opt.TRAIN.MAX_EPOCH * self.iters_per_epoch))) - 1
                print('value of lam is: %3f' % (lam))
                self.update_lr()
            source_data, source_gt = self.get_samples('source')
            target_data, _ = self.get_samples('target')
            source_data = to_cuda(source_data)
            source_gt = to_cuda(source_gt)
            target_data = to_cuda(target_data)
            data_time.update(time.time() - end)

            feature_source = self.feature_extractor(source_data)
            output_source = self.classifier(feature_source)
            feature_target = self.feature_extractor(target_data)
            output_target = self.classifier(feature_target)

            loss_task_fs = self.CELoss(output_source[:, :self.num_classes],
                                       source_gt)
            loss_task_ft = self.CELoss(output_source[:, self.num_classes:],
                                       source_gt)
            loss_discrim_source = self.SourceDiscrimLoss(output_source)
            loss_discrim_target = self.TargetDiscrimLoss(output_target)
            loss_summary_classifier = loss_task_fs + loss_task_ft + loss_discrim_source + loss_discrim_target

            source_gt_for_ft_in_fst = source_gt + self.num_classes
            loss_confusion_source = 0.5 * self.CELoss(
                output_source, source_gt) + 0.5 * self.CELoss(
                    output_source, source_gt_for_ft_in_fst)
            loss_confusion_target = 0.5 * self.SourceDiscrimLoss(
                output_target) + 0.5 * self.TargetDiscrimLoss(output_target)
            loss_em = self.ConcatenatedEMLoss(output_target)
            loss_summary_feature_extractor = loss_confusion_source + lam * (
                loss_confusion_target + loss_em)

            self.optimizer_classifier.zero_grad()
            loss_summary_classifier.backward(retain_graph=True)
            self.optimizer_classifier.step()

            self.optimizer_feature_extractor.zero_grad()
            loss_summary_feature_extractor.backward()
            self.optimizer_feature_extractor.step()

            classifier_loss.update(loss_summary_classifier,
                                   source_data.size()[0])
            feature_extractor_loss.update(loss_summary_feature_extractor,
                                          source_data.size()[0])
            prec1_fs.update(
                accuracy(output_source[:, :self.num_classes], source_gt),
                source_data.size()[0])
            prec1_ft.update(
                accuracy(output_source[:, self.num_classes:], source_gt),
                source_data.size()[0])

            print("  Train:epoch: %d:[%d/%d], LossCla: %3f, LossFeat: %3f, AccFs: %3f, AccFt: %3f" % \
                  (self.epoch, iters_counter_within_epoch, self.iters_per_epoch, classifier_loss.avg, feature_extractor_loss.avg, prec1_fs.avg, prec1_ft.avg))

            batch_time.update(time.time() - end)
            end = time.time()
            self.iters += 1
            iters_counter_within_epoch += 1
            if iters_counter_within_epoch >= self.iters_per_epoch:
                log = open(os.path.join(self.opt.SAVE_DIR, 'log.txt'), 'a')
                log.write("\n")
                log.write("  Train:epoch: %d:[%d/%d], LossCla: %3f, LossFeat: %3f, AccFs: %3f, AccFt: %3f" % \
                  (self.epoch, iters_counter_within_epoch, self.iters_per_epoch, classifier_loss.avg, feature_extractor_loss.avg, prec1_fs.avg, prec1_ft.avg))
                log.close()
                stop = True
예제 #10
0
    def test(self):
        self.feature_extractor.eval()
        self.classifier.eval()
        prec1_fs = AverageMeter()
        prec1_ft = AverageMeter()
        counter_all_fs = torch.FloatTensor(
            self.opt.DATASET.NUM_CLASSES).fill_(0)
        counter_all_ft = torch.FloatTensor(
            self.opt.DATASET.NUM_CLASSES).fill_(0)
        counter_acc_fs = torch.FloatTensor(
            self.opt.DATASET.NUM_CLASSES).fill_(0)
        counter_acc_ft = torch.FloatTensor(
            self.opt.DATASET.NUM_CLASSES).fill_(0)

        for i, (input, target) in enumerate(self.test_data['loader']):
            input, target = to_cuda(input), to_cuda(target)
            with torch.no_grad():
                feature_test = self.feature_extractor(input)
                output_test = self.classifier(feature_test)

            if self.opt.EVAL_METRIC == 'accu':
                prec1_fs_iter = accuracy(output_test[:, :self.num_classes],
                                         target)
                prec1_ft_iter = accuracy(output_test[:, self.num_classes:],
                                         target)
                prec1_fs.update(prec1_fs_iter, input.size(0))
                prec1_ft.update(prec1_ft_iter, input.size(0))
                if i % self.opt.PRINT_STEP == 0:
                    print("  Test:epoch: %d:[%d/%d], AccFs: %3f, AccFt: %3f" % \
                          (self.epoch, i, len(self.test_data['loader']), prec1_fs.avg, prec1_ft.avg))
            elif self.opt.EVAL_METRIC == 'accu_mean':
                prec1_ft_iter = accuracy(output_test[:, self.num_classes:],
                                         target)
                prec1_ft.update(prec1_ft_iter, input.size(0))
                counter_all_fs, counter_acc_fs = accuracy_for_each_class(
                    output_test[:, :self.num_classes], target, counter_all_fs,
                    counter_acc_fs)
                counter_all_ft, counter_acc_ft = accuracy_for_each_class(
                    output_test[:, self.num_classes:], target, counter_all_ft,
                    counter_acc_ft)
                if i % self.opt.PRINT_STEP == 0:
                    print("  Test:epoch: %d:[%d/%d], Task: %3f" % \
                          (self.epoch, i, len(self.test_data['loader']), prec1_ft.avg))
            else:
                raise NotImplementedError
        acc_for_each_class_fs = counter_acc_fs / counter_all_fs
        acc_for_each_class_ft = counter_acc_ft / counter_all_ft
        log = open(os.path.join(self.opt.SAVE_DIR, 'log.txt'), 'a')
        log.write("\n")
        if self.opt.EVAL_METRIC == 'accu':
            log.write(
                "                                                          Test:epoch: %d, AccFs: %3f, AccFt: %3f" % \
                (self.epoch, prec1_fs.avg, prec1_ft.avg))
            log.close()
            return max(prec1_fs.avg, prec1_ft.avg)
        elif self.opt.EVAL_METRIC == 'accu_mean':
            log.write(
                "                                            Test:epoch: %d, AccFs: %3f, AccFt: %3f" % \
                (self.epoch,acc_for_each_class_fs.mean(), acc_for_each_class_ft.mean()))
            log.write(
                "\nClass-wise Acc of Ft:")  ## based on the task classifier.
            for i in range(self.opt.DATASET.NUM_CLASSES):
                if i == 0:
                    log.write("%dst: %3f" % (i + 1, acc_for_each_class_ft[i]))
                elif i == 1:
                    log.write(",  %dnd: %3f" %
                              (i + 1, acc_for_each_class_ft[i]))
                elif i == 2:
                    log.write(", %drd: %3f" %
                              (i + 1, acc_for_each_class_ft[i]))
                else:
                    log.write(", %dth: %3f" %
                              (i + 1, acc_for_each_class_ft[i]))
            log.close()
            return max(acc_for_each_class_ft.mean(),
                       acc_for_each_class_fs.mean())
예제 #11
0
def test_gallery(net, dataloader, output_dir, thresh=0.):
    """test gallery images"""

    with open('config.yml', 'r') as f:
        config = yaml.load(f)

    num_images = len(dataloader.dataset)
    all_boxes = []
    all_features = []
    end = time.time()
    time_cost = AverageMeter()
    net.eval()

    for i, data in enumerate(dataloader):
        with torch.no_grad():
            im, (orig_shape, im_info) = data
            im = im.to(device)
            im_info = im_info.numpy().squeeze(0)
            orig_shape = [x.item() for x in orig_shape]

            scores, bbox_pred, rois, features = net.forward(im, None, im_info)

        boxes = rois[:, 1:5] / im_info[2]
        scores = np.reshape(scores, [scores.shape[0], -1])
        bbox_pred = np.reshape(bbox_pred, [bbox_pred.shape[0], -1])
        if config['test_bbox_reg']:
            # Apply bounding-box regression deltas
            box_deltas = bbox_pred
            pred_boxes = bbox_transform_inv(
                torch.from_numpy(boxes), torch.from_numpy(box_deltas)).numpy()
            pred_boxes = clip_boxes(pred_boxes, orig_shape)
        else:
            # Simply repeat the boxes, once for each class
            pred_boxes = np.tile(boxes, (1, scores.shape[1]))

        boxes = pred_boxes

        # skip j = 0, because it's the background class
        j = 1
        inds = np.where(scores[:, j] > thresh)[0]
        cls_scores = scores[inds, j]
        cls_boxes = boxes[inds, j * 4:(j + 1) * 4]
        cls_dets = np.hstack(
            (cls_boxes, cls_scores[:, np.newaxis])).astype(np.float32,
                                                           copy=False)
        keep = nms(torch.from_numpy(cls_dets),
                   config['test_nms']).numpy() if cls_dets.size > 0 else []
        cls_dets = cls_dets[keep, :]
        all_boxes.append(cls_dets)
        all_features.append(features[inds][keep])

        time_cost.update(time.time() - end)
        end = time.time()
        print('im_detect: {:d}/{:d} {:.3f}s'.format(i + 1, num_images,
                                                    time_cost.avg))

    det_file = os.path.join(output_dir, 'gboxes.pkl')
    with open(det_file, 'wb') as f:
        pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL)

    feature_file = os.path.join(output_dir, 'gfeatures.pkl')
    with open(feature_file, 'wb') as f:
        pickle.dump(all_features, f, pickle.HIGHEST_PROTOCOL)

    return all_boxes, all_features
예제 #12
0
def scan_train(train_loader,
               model,
               criterion,
               optimizer,
               epoch,
               update_cluster_head_only=False):
    """ 
    Train w/ SCAN-Loss
    """
    total_losses = AverageMeter('Total Loss', ':.4e')
    consistency_losses = AverageMeter('Consistency Loss', ':.4e')
    entropy_losses = AverageMeter('Entropy', ':.4e')
    progress = ProgressMeter(
        len(train_loader), [total_losses, consistency_losses, entropy_losses],
        prefix="Epoch: [{}]".format(epoch))

    if update_cluster_head_only:
        model.eval()  # No need to update BN
    else:
        model.train()  # Update BN

    for i, batch in enumerate(train_loader):
        # Forward pass
        anchors = batch['anchor'].cuda(non_blocking=True)
        neighbors = batch['neighbor'].cuda(non_blocking=True)

        if update_cluster_head_only:  # Only calculate gradient for backprop of linear layer
            with torch.no_grad():
                anchors_features = model(anchors, forward_pass='******')
                neighbors_features = model(neighbors, forward_pass='******')
            anchors_output = model(anchors_features, forward_pass='******')
            neighbors_output = model(neighbors_features, forward_pass='******')

        else:  # Calculate gradient for backprop of complete network
            anchors_output = model(anchors)
            neighbors_output = model(neighbors)

        # Loss for every head
        total_loss, consistency_loss, entropy_loss = [], [], []
        for anchors_output_subhead, neighbors_output_subhead in zip(
                anchors_output, neighbors_output):
            total_loss_, consistency_loss_, entropy_loss_ = criterion(
                anchors_output_subhead, neighbors_output_subhead)
            total_loss.append(total_loss_)
            consistency_loss.append(consistency_loss_)
            entropy_loss.append(entropy_loss_)

        # Register the mean loss and backprop the total loss to cover all subheads
        total_losses.update(np.mean([v.item() for v in total_loss]))
        consistency_losses.update(np.mean([v.item()
                                           for v in consistency_loss]))
        entropy_losses.update(np.mean([v.item() for v in entropy_loss]))

        total_loss = torch.sum(torch.stack(total_loss, dim=0))

        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()

        if i % 25 == 0:
            progress.display(i)
예제 #13
0
def test(opt, model, loader):
    """
    Validate the model at the current state

    Args:
        opt (Namspace): training options
        model (LaneNet): a LaneNet model
        criterion: a CrossEntropyLoss criterion
        loader: val data loader

    Returns:
        The average loss value on val data

    """

    model.eval()

    run_time = AverageMeter()
    end = time.time()
    palette = get_palette(opt.dataset, opt.num_classes + 1)

    total_inter, total_union, total_correct, total_label = 0, 0, 0, 0
    pixAcc = 0
    IoU = 0
    mIoU = 0
    pbar = tqdm(loader)
    with torch.no_grad():
        for data in pbar:
            images, labels, sizes, image_names, org_images = data
            sizes = sizes[0].numpy()

            images = Variable(images)
            N_, C_, H_, W_ = images.shape

            if torch.cuda.is_available():
                images = images.cuda()

            preds = model(images)
            preds = gather(preds, 0, dim=0)
            if 'dsn' in opt.model_type:
                preds = preds[-1]

            preds = F.upsample(input=preds,
                               size=(H_, W_),
                               mode='bilinear',
                               align_corners=True)
            output = preds.cpu().data.numpy().transpose(0, 2, 3, 1)

            seg_pred = np.asarray(np.argmax(output, axis=3), dtype=np.uint8)

            # store images
            if opt.store_output:
                for i in range(N_):
                    output_im = Image.fromarray(seg_pred[i])
                    output_im.putpalette(palette)
                    output_file = os.path.join(opt.output_dir,
                                               image_names[i] + '.png')
                    output_im.save(output_file)

                    src_img = org_images[i].data.numpy()

                    if opt.dataset == 'cityscapes':
                        drivable_img = np.where(seg_pred[i] == 0, 0,
                                                19).astype(np.uint8)
                        drivable_img = Image.fromarray(drivable_img)
                        drivable_img.putpalette(palette)
                        drivable_img = np.array(
                            drivable_img.convert('RGB')).astype(src_img.dtype)
                        #overlay_img = cv2.addWeighted(src_img, 1.0, drivable_img, 1.0, 0)
                        src_img[drivable_img > 0] = 0
                    else:
                        drivable_img = seg_pred[i]
                        drivable_img = Image.fromarray(drivable_img)
                        drivable_img.putpalette(palette)
                        drivable_img = np.array(
                            drivable_img.convert('RGB')).astype(src_img.dtype)

                    overlay_img = cv2.add(src_img, drivable_img)
                    output_file = os.path.join(
                        opt.output_dir, image_names[i] + '_drivable.png')
                    cv2.imwrite(output_file,
                                cv2.cvtColor(overlay_img, cv2.COLOR_RGB2BGR))

            if len(labels) > 0:
                labels[labels == opt.ignore_label] = -1  #
                correct, labeled = batch_pix_accuracy(preds.data, labels)
                inter, union = batch_intersection_union(
                    preds.data, labels, opt.num_classes)
                total_correct += correct
                total_label += labeled
                total_inter += inter
                total_union += union
                pixAcc = 1.0 * total_correct / (np.spacing(1) + total_label)
                IoU = 1.0 * total_inter / (np.spacing(1) + total_union)
                mIoU = IoU.mean()

            # measure speed test
            run_time.update(time.time() - end)
            nd = time.time()
            fps = N_ / run_time.avg
            pbar.set_description(
                'Average run time: {:.3f} fps, pixAcc={:.6f}, mIoU={:.6f}'.
                format(fps, pixAcc, mIoU))

    print({'meanIU': mIoU, 'IU_array': IoU})
    return mIoU
예제 #14
0
def train_fixed(starting_epoch, data_loader, fixed_cnn, criterion, fixed_cnn_optmizer, fixed_cnn_scheduler, utilHelper):
    global GLOBAL_STEP, reduction_arc, cell_arc, EVAL_BEST_ACC, EVAL_STEP, TRAIN_HISTORY

    for epoch in tqdm(range(starting_epoch, args.epoch)):
        logger.info("=" * 20 + "Training" + "=" * 20)
        fixed_cnn.train()
        train_loss_meters = AverageMeter()
        train_acc_meters = AverageMeter()
        train_acc5_meters = AverageMeter()

        for images, labels in tqdm(data_loader["train_dataset"]):
            images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True)
            start = time.time()
            fixed_cnn.zero_grad()
            fixed_cnn_optmizer.zero_grad()
            pred = fixed_cnn(images)
            loss = criterion(pred, labels)
            loss.backward()
            grad_norm = torch.nn.utils.clip_grad_norm_(fixed_cnn.parameters(), args.grad_bound)
            fixed_cnn_optmizer.step()
            acc, acc5 = accuracy(pred, labels, topk=(1, 5))

            train_loss_meters.update(loss.item(), labels.shape[0])
            train_acc_meters.update(acc.item(), labels.shape[0])
            train_acc5_meters.update(acc5.item(), labels.shape[0])

            end = time.time()

            GLOBAL_STEP += 1
            if GLOBAL_STEP % args.train_log_every == 0:
                lr = fixed_cnn_optmizer.param_groups[0]['lr']
                display = log_display(epoch=epoch,
                                      global_step=GLOBAL_STEP,
                                      time_elapse=end-start,
                                      loss=loss.item(),
                                      loss_avg=train_loss_meters.avg,
                                      acc=acc.item(),
                                      acc_top1_avg=train_acc_meters.avg,
                                      acc_top5_avg=train_acc5_meters.avg,
                                      lr=lr,
                                      gn=grad_norm)
                logger.info(display)
        if fixed_cnn_scheduler is not None:
            fixed_cnn_scheduler.step()
        logger.info("="*20 + "Eval" + "="*20)
        curr_acc, _ = model_eval(epoch, fixed_cnn, data_loader)
        logger.info("curr_acc\t%.4f" % curr_acc)
        logger.info("BEST_ACC\t%.4f" % EVAL_BEST_ACC)
        payload = '=' * 10 + '\n'
        payload = payload + ("curr_acc: %.4f\n best_acc: %.4f\n" % (curr_acc, EVAL_BEST_ACC))
        EVAL_BEST_ACC = max(curr_acc, EVAL_BEST_ACC)
        TRAIN_HISTORY["train_loss"].append(train_loss_meters.avg)
        TRAIN_HISTORY["train_acc"].append(train_acc_meters.avg)
        TRAIN_HISTORY["test_acc"].append(curr_acc)
        TRAIN_HISTORY["test_acc_best"] = [EVAL_BEST_ACC]
        with open(args.checkpoint_path + args.version + '.pickle', 'wb') as handle:
            pickle.dump(TRAIN_HISTORY, handle, protocol=pickle.HIGHEST_PROTOCOL)
        logger.info("Saved!\n")
    return
예제 #15
0
def model_eval(epoch, fixed_cnn, data_loader, dataset_type='test_dataset'):
    global EVAL_STEP
    fixed_cnn.eval()
    valid_loss_meters = AverageMeter()
    valid_acc_meters = AverageMeter()
    valid_acc5_meters = AverageMeter()
    ce_loss = torch.nn.CrossEntropyLoss()

    for images, labels in tqdm(data_loader[dataset_type]):
        start = time.time()
        images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True)
        with torch.no_grad():
            pred = fixed_cnn(images)
            loss = ce_loss(pred, labels)
            acc, acc5 = accuracy(pred, labels, topk=(1, 5))

        valid_loss_meters.update(loss.item(), labels.shape[0])
        valid_acc_meters.update(acc.item(), labels.shape[0])
        valid_acc5_meters.update(acc5.item(), labels.shape[0])
        end = time.time()

        EVAL_STEP += 1
        if EVAL_STEP % args.train_log_every == 0:
            display = log_display(epoch=epoch,
                                  global_step=GLOBAL_STEP,
                                  time_elapse=end-start,
                                  loss=loss.item(),
                                  test_loss_avg=valid_loss_meters.avg,
                                  acc=acc.item(),
                                  test_acc_avg=valid_acc_meters.avg,
                                  test_acc_top5_avg=valid_acc5_meters.avg)
            logger.info(display)
    display = log_display(epoch=epoch,
                          global_step=GLOBAL_STEP,
                          time_elapse=end-start,
                          loss=loss.item(),
                          test_loss_avg=valid_loss_meters.avg,
                          acc=acc.item(),
                          test_acc_avg=valid_acc_meters.avg,
                          test_acc_top5_avg=valid_acc5_meters.avg)
    logger.info(display)
    return valid_acc_meters.avg, valid_acc5_meters.avg
예제 #16
0
def do_train(cfg, model, data_loader, loss_factory, optimizer, epoch,
             output_dir, tb_log_dir, writer_dict):
    logger = logging.getLogger("Training")

    batch_time = AverageMeter()
    data_time = AverageMeter()

    heatmaps_loss_meter = [AverageMeter() for _ in range(cfg.LOSS.NUM_STAGES)]
    offset_loss_meter = [AverageMeter() for _ in range(cfg.LOSS.NUM_STAGES)]

    model.train()

    end = time.time()
    for i, (images, heatmaps, masks, offsets,
            weights) in enumerate(data_loader):
        data_time.update(time.time() - end)

        heatmaps = [
            list(map(lambda x: x.cuda(non_blocking=True), heatmap))
            for heatmap in heatmaps
        ]
        masks = [
            list(map(lambda x: x.cuda(non_blocking=True), mask))
            for mask in masks
        ]
        offsets = [
            list(map(lambda x: x.cuda(non_blocking=True), offset))
            for offset in offsets
        ]
        offset_weights = [
            list(map(lambda x: x.cuda(non_blocking=True), weight))
            for weight in weights
        ]

        ####################################################################
        if cfg.LOSS.HEATMAP_MIDDLE_LOSS:
            outputs, poffsets, middle_output = model(images)
            heatmaps_losses, offset_losses, middle_losses = \
                loss_factory(outputs, poffsets, heatmaps,
                             masks, offsets, offset_weights, middle_output)
        else:
            outputs, poffsets = model(images)
            heatmaps_losses, offset_losses = \
                loss_factory(outputs, poffsets, heatmaps,
                             masks, offsets, offset_weights)
        ####################################################################

        loss = 0
        for idx in range(cfg.LOSS.NUM_STAGES):
            if heatmaps_losses[idx] is not None:
                heatmaps_loss = heatmaps_losses[idx].mean(dim=0)
                heatmaps_loss_meter[idx].update(heatmaps_loss.item(),
                                                images.size(0))
                loss = loss + heatmaps_loss
            if offset_losses[idx] is not None:
                offset_loss = offset_losses[idx]
                offset_loss_meter[idx].update(offset_loss.item(),
                                              images.size(0))
                loss = loss + offset_loss

        ########################################################################
        if cfg.LOSS.HEATMAP_MIDDLE_LOSS:
            if middle_losses is not None:
                loss = loss + middle_losses.mean(dim=0)
        ########################################################################

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        batch_time.update(time.time() - end)
        end = time.time()

        if i % cfg.PRINT_FREQ == 0 and cfg.RANK == 0:
            msg = 'Epoch: [{0}][{1}/{2}]\t' \
                  'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s)\t' \
                  'Speed: {speed:.1f} samples/s\t' \
                  'Data: {data_time.val:.3f}s ({data_time.avg:.3f}s)\t' \
                  '{heatmaps_loss}{offset_loss}'.format(
                      epoch, i, len(data_loader),
                      batch_time=batch_time,
                      speed=images.size(0)/batch_time.val,
                      data_time=data_time,
                      heatmaps_loss=_get_loss_info(
                          heatmaps_loss_meter, 'heatmaps'),
                      offset_loss=_get_loss_info(offset_loss_meter, 'offset')
                  )
            logger.info(msg)

            writer = writer_dict['writer']
            global_steps = writer_dict['train_global_steps']
            for idx in range(cfg.LOSS.NUM_STAGES):
                writer.add_scalar('train_stage{}_heatmaps_loss'.format(i),
                                  heatmaps_loss_meter[idx].val, global_steps)
                writer.add_scalar('train_stage{}_offset_loss'.format(idx),
                                  offset_loss_meter[idx].val, global_steps)
            # 每经过PRINT_FREQ个batch,train_global_steps加1
            writer_dict['train_global_steps'] = global_steps + 1

            prefix = '{}_{}'.format(os.path.join(output_dir, 'train'), i)
            for scale_idx in range(len(cfg.DATASET.OUTPUT_SIZE)):
                prefix_scale = prefix + '_output_{}'.format(
                    cfg.DATASET.OUTPUT_SIZE[scale_idx])
                save_debug_images(cfg, images, heatmaps[scale_idx],
                                  masks[scale_idx], outputs[scale_idx],
                                  prefix_scale)
예제 #17
0
def adaptive_quantization(model, val_loader, checkpoint, quantizable_ind,
                          quantize_bit_choice):
    module_list = list(model.modules())

    model.load_state_dict(checkpoint)  # restore weight first
    org_acc = validate(val_loader, model)
    # 1. Calculate t_i
    # calculate the mean value of adversarial noise for the dataset, notice that we do not add Softmax to network

    mean_adv_noise_meter = AverageMeter()
    with torch.no_grad():
        for input_w, target in val_loader:
            input_var = torch.autograd.Variable(input_w).cuda()
            output = model(input_var)
            top2, inds = torch.topk(output, 2, dim=1)
            mean_adv_noise = torch.mean((top2[:, 0] - top2[:, 1])**2 / 2.)
            mean_adv_noise_meter.update(mean_adv_noise.cpu().data[0],
                                        output.size(0))
    mean_adv_noise_dset = mean_adv_noise_meter.avg
    print('Mean adversarial noise for the dataset is: {:.4f}'.format(
        mean_adv_noise_dset))

    d_acc = 10.  # choose 10% for delta_acc. Does not matter
    t_i_list = []
    for ind in quantizable_ind:
        layer = module_list[ind]
        assert hasattr(layer, 'weight')
        r_W_i_pi = torch.rand(layer.weight.data.size()).cuda(
        ) - 0.5  # re-normalize to [-0.5, 0.5]
        k = k_min = 1e-5
        k_max = 1e1
        # get initial acc'
        model.load_state_dict(checkpoint)  # restore weight first
        # assert validate(val_loader, model) == org_acc  # removed to accelerate
        layer.weight.data += k * r_W_i_pi
        new_acc = validate(val_loader, model)
        while not np.abs(org_acc - new_acc - d_acc) < 0.1:
            if org_acc - new_acc < d_acc:
                k_min = k
            else:
                k_max = k
            k = np.sqrt(k_min * k_max).item()
            model.load_state_dict(checkpoint)  # restore weight first
            layer.weight.data += k * r_W_i_pi
            new_acc = validate(val_loader, model)
            print('Layer {} current acc degradation: {:.3f}'.format(
                ind, org_acc - new_acc))
        mean_r_z_i = AverageMeter()
        with torch.no_grad():
            for i, (input_d, target) in enumerate(val_loader):
                input_var = torch.autograd.Variable(input_d).cuda()
                # compute output
                model.load_state_dict(checkpoint)  # restore weight first
                output1 = model(input_var)
                layer.weight.data += k * r_W_i_pi
                output2 = model(input_var)
                norm_r = torch.norm(output1 - output2, p=2, dim=1)**2
                mean_r_z_i.update(
                    torch.mean(norm_r).cpu().data[0], output1.size(0))
        t_i = mean_r_z_i.avg / mean_adv_noise_dset
        print('==> t_i for layer {}: {}'.format(ind, t_i))
        t_i_list.append(t_i)
    print('t_i_list: ')
    print(t_i_list)
    # t_i_list = [113.314645623, 108.02437323, 111.228006385, 109.585273767, 115.362011096, 111.136186881, 114.150737539,
    #             106.789374298, 135.436417323, 118.175965146, 136.776404035, 162.089406771, 224.905060191, 340.589419784,
    #             904.878690392, 256.250864841]

    # 2. Calculate p_i
    fix_b_i = 6
    p_i_list = []
    for i, ind in enumerate(quantizable_ind):
        model.load_state_dict(checkpoint)  # restore weight first
        centroids, labels = k_means_torch_compact(module_list[ind].weight.data,
                                                  2**fix_b_i,
                                                  init='linear',
                                                  use_gpu=False)
        w_q = reconstruct_weight_from_k_means_result(centroids, labels)
        del centroids, labels
        mean_r_z_i = AverageMeter()
        with torch.no_grad():
            for input_d, target in val_loader:
                input_var = torch.autograd.Variable(input_d).cuda()
                # compute output
                model.load_state_dict(checkpoint)  # restore weight first
                output1 = model(input_var)
                del module_list[ind].weight
                module_list[ind].weight = nn.Parameter(w_q.float())
                output2 = model(input_var)
                norm_r = torch.norm(output1 - output2, p=2, dim=1)**2
                mean_r_z_i.update(
                    torch.mean(norm_r).cpu().data[0], output1.size(0))
        del w_q
        p_i = mean_r_z_i.avg  # / np.exp(-np.log(4) * fix_b_i)
        print('==> p_i for layer {}: {}'.format(ind, p_i))
        p_i_list.append(p_i)

    # 3. Calculate b_i
    b1 = 8
    assert len(p_i_list) == len(t_i_list) == len(quantizable_ind)
    layer1_size = np.prod(module_list[quantizable_ind[0]].weight.data.size())
    bits_list = [b1]
    print('==> Layer1 size: {}, bits: {}'.format(layer1_size, b1))
    for i, ind in enumerate(quantizable_ind):
        if i == 0:
            continue
        this_size = np.prod(module_list[ind].weight.data.size())
        b_i = b1 + (1 / np.log(4)) * np.log(
            p_i_list[i] * t_i_list[0] * layer1_size /
            (p_i_list[0] * t_i_list[i] * this_size))
        print('Optimal bits for layer {}: {}'.format(ind, b_i))
        bits_list.append(b_i)

    print('Final result: {}'.format(bits_list))
예제 #18
0
def train(**kwargs):
    opt.parse(kwargs)
    if not os.path.exists(opt.save_folder):
        os.mkdir(opt.save_folder)
    tb_logger = SummaryWriter(opt.save_folder)
    logger = create_logger('global_logger', opt.save_folder + '/log.txt')
    batch_time = AverageMeter(10)
    data_time = AverageMeter(10)
    losses = AverageMeter(10)
    loss_meter = meter.AverageValueMeter()

    train_sets = []
    for data_txt in opt.train_txt:
        data_root, gt_root, list_file = data_txt.split(' ')
        train_sets.append(
            OCRDataset(data_root, gt_root, list_file, opt.input_size, 'train',
                       opt.chars_list, opt.max_seq))
    train_data = ConcatDataset(train_sets)
    train_loader = DataLoader(train_data,
                              batch_size=opt.batch_size,
                              shuffle=True,
                              num_workers=opt.num_works)

    valid_sets = []
    for valid_txt in opt.valid_txt:
        data_root, gt_root, list_file = valid_txt.split(' ')
        valid_sets.append(
            OCRDataset(data_root, gt_root, list_file, opt.input_size, 'valid',
                       opt.chars_list, opt.max_seq))
    valid_data = ConcatDataset(valid_sets)
    valid_loader = DataLoader(valid_data,
                              batch_size=opt.batch_size,
                              shuffle=False,
                              num_workers=opt.num_works)

    model = getattr(models, opt.model)(opt.basenet,
                                       opt.input_size,
                                       opt.max_seq,
                                       opt.num_classes,
                                       mode='train',
                                       attn=opt.attn)

    if opt.load_model_path is not None:
        load_state(model, opt.load_model_path, 'cuda:%d' % opt.gpus[0])
    if len(opt.gpus) > 1:
        model = torch.nn.DataParallel(model, device_ids=opt.gpus)
    model = gpu(model, opt)

    if len(opt.gpus) > 1:
        optimizer = torch.optim.Adam(model.module.parameters(),
                                     lr=opt.lr,
                                     betas=opt.betas,
                                     weight_decay=opt.weight_decay)
    else:
        optimizer = torch.optim.Adam(model.parameters(),
                                     lr=opt.lr,
                                     betas=opt.betas,
                                     weight_decay=opt.weight_decay)

    curr_step = 0
    total_step = int(len(train_data) / opt.batch_size * opt.epoches)
    best_val_error = 1e10
    previous_loss = 1e10
    # warmup
    warmup_epoches = opt.epoches // 10
    warmup_rate = math.pow(100, 1 / warmup_epoches)

    for epoch in range(opt.epoches):
        model.train()
        end = time.time()
        # loss_meter.reset()
        for i, (imgs, gt_chars_seg, gt_order_seg,
                gt_pos_seg) in enumerate(train_loader):
            # measure data loading time
            data_time.update(time.time() - end)
            # zero the parameter gradients
            optimizer.zero_grad()

            imgs = gpu(imgs, opt)
            gt_chars_seg = gpu(gt_chars_seg, opt)
            gt_order_seg = gpu(gt_order_seg, opt)
            gt_pos_seg = gpu(gt_pos_seg, opt)

            chars_seg, ord_seg, pos_seg = model(imgs)
            loss = get_loss(chars_seg, ord_seg, pos_seg, gt_chars_seg,
                            gt_order_seg, gt_pos_seg, opt)

            loss.backward()
            optimizer.step()
            losses.update(loss.item())
            loss_meter.add(loss.item())
            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            curr_step += 1
            current_lr = optimizer.param_groups[0]['lr']
            if curr_step % opt.print_freq == 0:
                tb_logger.add_scalar('loss_train', losses.avg, curr_step)
                tb_logger.add_scalar('lr', current_lr, curr_step)
                logger.info(
                    'Iter: [{0}/{1}]\t'
                    'Epoch: {2}\t'
                    'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                    'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                    'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                    'LR {lr:.4f}'.format(curr_step,
                                         total_step,
                                         epoch,
                                         batch_time=batch_time,
                                         data_time=data_time,
                                         loss=losses,
                                         lr=current_lr))

        # val
        model.eval()
        val_error = val(model, valid_loader, opt)
        logger.info('Mean error: {0}\t'.format(val_error))
        if not tb_logger is None:
            tb_logger.add_scalar('error_val', val_error, curr_step)
        if val_error < best_val_error:
            best_val_error = val_error
            if len(opt.gpus) > 1:
                torch.save(model.module.state_dict(),
                           os.path.join(opt.save_folder, "best_val_error.pth"))
            else:
                torch.save(model.state_dict(),
                           os.path.join(opt.save_folder, "best_val_error.pth"))
        # warmup
        if epoch < warmup_epoches:
            for param_group in optimizer.param_groups:
                param_group["lr"] *= warmup_rate
                # decay lr if loss no longer decrease
        else:
            if opt.lr_immediate_decay and loss_meter.value(
            )[0] > previous_loss:
                for param_group in optimizer.param_groups:
                    param_group["lr"] *= opt.lr_decay

            if epoch == int(opt.epoches * 0.6) or epoch == int(
                    opt.epoches * 0.9):
                for param_group in optimizer.param_groups:
                    param_group["lr"] *= opt.lr_decay
            previous_loss = loss_meter.value()[0]
    # save last pth
    if len(opt.gpus) > 1:
        torch.save(model.module.state_dict(),
                   os.path.join(opt.save_folder, "last.pth"))
    else:
        torch.save(model.state_dict(), os.path.join(opt.save_folder,
                                                    "last.pth"))
예제 #19
0
def validate(val_loader, model):
    from utils import accuracy
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    criterion = nn.CrossEntropyLoss().cuda()
    # switch to evaluate mode
    model.eval()
    with torch.no_grad():
        for i, (input, target) in enumerate(val_loader):
            target = target.cuda(async=True)
            input_var = torch.autograd.Variable(input).cuda()
            target_var = torch.autograd.Variable(target).cuda()

            # compute output
            output = model(input_var)
            loss = criterion(output, target_var)
            # measure accuracy and record loss
            prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
            losses.update(loss.data[0], input.size(0))
            top1.update(prec1[0], input.size(0))
            top5.update(prec5[0], input.size(0))

    return top1.avg
예제 #20
0
def validate(val_loader, miml, encoder, decoder, criterion, epoch, writer):
    """
    Performs one epoch's validation.

    :param val_loader: DataLoader for validation data.
    :param encoder: encoder model
    :param decoder: decoder model
    :param criterion: loss layer
    :return: BLEU-4 score
    """
    miml.eval()
    encoder.eval()
    decoder.eval()  # eval mode (no dropout or batchnorm)
    total_step = len(val_loader)
    batch_time = AverageMeter()
    losses = AverageMeter()
    top5accs = AverageMeter()

    start = time.time()

    references = list(
    )  # references (true captions) for calculating BLEU-4 score
    hypotheses = list()  # hypotheses (predictions)

    with torch.no_grad():

        # Batches
        for i, (imgs, caps, caplens, allcaps) in enumerate(val_loader):

            # Move to device, if available
            imgs = imgs.to(device)
            caps = caps.to(device)
            caplens = caplens.to(device)

            # Forward prop.

            attrs = miml(imgs)
            imgs = encoder(imgs)
            scores, caps_sorted, decode_lengths, alphas, sort_ind = decoder(
                attrs, imgs, caps, caplens)

            # Since we decoded starting with <start>, the targets are all words after <start>, up to <end>
            targets = caps_sorted[:, 1:]

            # Remove timesteps that we didn't decode at, or are pads
            # pack_padded_sequence is an easy trick to do this
            scores_copy = scores.clone()
            scores, _ = pack_padded_sequence(scores,
                                             decode_lengths,
                                             batch_first=True)
            targets, _ = pack_padded_sequence(targets,
                                              decode_lengths,
                                              batch_first=True)

            # Calculate loss
            loss = criterion(scores, targets)
            loss += alpha_c * ((1. - alphas.sum(dim=1))**2).mean()
            # Keep track of metrics
            losses.update(loss.item(), sum(decode_lengths))
            top5 = accuracy(scores, targets, 5)
            top5accs.update(top5, sum(decode_lengths))
            batch_time.update(time.time() - start)

            start = time.time()

            if i % print_freq == 0:
                writer.add_scalars('val', {
                    'loss': loss.item(),
                    'mAp': top5accs.val
                }, epoch * total_step + i)
                print(
                    'Validation: [{0}/{1}]\t'
                    'Batch Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                    'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                    'Top-5 Accuracy {top5.val:.3f} ({top5.avg:.3f})\t'.format(
                        i,
                        len(val_loader),
                        batch_time=batch_time,
                        loss=losses,
                        top5=top5accs))

            # Store references (true captions), and hypothesis (prediction) for each image
            # If for n images, we have n hypotheses, and references a, b, c... for each image, we need -
            # references = [[ref1a, ref1b, ref1c], [ref2a, ref2b], ...], hypotheses = [hyp1, hyp2, ...]

            # References
            # because images were sorted in the decoder
            allcaps = allcaps[sort_ind]
            for j in range(allcaps.shape[0]):
                img_caps = allcaps[j].tolist()
                img_captions = list(
                    map(
                        lambda c: [
                            w for w in c if w not in
                            {word_map['<start>'], word_map['<pad>']}
                        ], img_caps))  # remove <start> and pads
                references.append(img_captions)

            # Hypotheses
            _, preds = torch.max(scores_copy, dim=2)
            preds = preds.tolist()
            temp_preds = list()
            for j, p in enumerate(preds):
                temp_preds.append(preds[j][:decode_lengths[j]])  # remove pads
            preds = temp_preds
            hypotheses.extend(preds)

            assert len(references) == len(hypotheses)

        # Calculate BLEU-4 scores
        weights = (1.0 / 1.0, )
        bleu1 = corpus_bleu(references, hypotheses, weights)

        weights = (
            1.0 / 2.0,
            1.0 / 2.0,
        )
        bleu2 = corpus_bleu(references, hypotheses, weights)

        weights = (
            1.0 / 3.0,
            1.0 / 3.0,
            1.0 / 3.0,
        )
        bleu3 = corpus_bleu(references, hypotheses, weights)
        bleu4 = corpus_bleu(references, hypotheses)
        writer.add_scalars('Bleu', {
            'Bleu1': bleu1,
            'Bleu2': bleu2,
            'Bleu3': bleu3,
            'Bleu4': bleu4
        }, epoch)
        print(
            '\n * LOSS - {loss.avg:.3f}, TOP-5 ACCURACY - {top5.avg:.3f}, BLEU-4 - {bleu}\n'
            .format(loss=losses, top5=top5accs, bleu=bleu4))

    return bleu4
예제 #21
0
def test(model, test_dataset, test_loader, output_filename, result_dir=None):

    load_batch_size = test_loader.batch_size
    num_batches = len(test_loader)

    model.eval()

    fw_time_meter = AverageMeter()

    det_results = {}
    gt_results = {}

    for i, (data_dicts, datas) in enumerate(test_loader):
        torch.cuda.synchronize()
        tic = time.time()

        if data_dicts is None:
            continue
        predictions = predict(model, data_dicts)

        torch.cuda.synchronize()
        fw_time_meter.update((time.time() - tic))
        print('%d/%d %.3f' % (i, num_batches, fw_time_meter.val))

        # datas = test_dataset.get_frustum_data(i)
        # print(f'Datas len: {len(datas)}')
        # print(f'Predictions len: {len(predictions)}')

        for data, pred in zip(datas, predictions):

            if data is None:
                continue

            data_idx = data['idx_i']
            class_type = data['type_i']
            box2d = data['box2d_i']
            box3d = data['box3d_i']
            box3d_sizes = data['box3d_size_i']
            ry_gt = data['heading_i']
            box_pos_gt = data['obj_pos_i']

            if data_idx not in det_results:
                det_results[data_idx] = {}
                gt_results[data_idx] = {}

            if class_type not in det_results[data_idx]:
                det_results[data_idx][class_type] = []
                gt_results[data_idx][class_type] = []

            x1, y1, x2, y2 = box2d
            l_gt, w_gt, h_gt = box3d_sizes
            tx_gt, ty_gt, tz_gt = box_pos_gt

            output_gt = [
                x1, y1, x2, y2, tx_gt, ty_gt, tz_gt, h_gt, w_gt, l_gt, ry_gt
            ]
            gt_results[data_idx][class_type].append(output_gt)

            # print('****************')
            # print(tx_gt, ty_gt, tz_gt, h_gt, w_gt, l_gt, ry_gt)
            # print('================')
            for n in range(len(pred)):
                h, w, l, tx, ty, tz, ry, score = pred[n]
                output = [x1, y1, x2, y2, tx, ty, tz, h, w, l, ry, score]
                # print(tx, ty, tz, h, w, l, ry)
                # output = [x1, y1, x2,  y2, tx, ty, tz, h, w, l, ry,1]
                # print(score)
                det_results[data_idx][class_type].append(output)
            # print('****************')

    num_images = len(det_results)

    logging.info('Average time:')
    logging.info('batch:%0.3f' % fw_time_meter.avg)
    logging.info('avg_per_object:%0.3f' %
                 (fw_time_meter.avg / load_batch_size))
    logging.info('avg_per_image:%.3f' %
                 (fw_time_meter.avg * len(test_loader) / num_images))

    return gt_results, det_results
예제 #22
0
def train(args, train_loader, model, optimizer, epoch, curr_lr, win_feats5,
          win_fusion, viz, global_step, accumulation_steps):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    feats5_losses = AverageMeter()
    fusion_losses = AverageMeter()
    total_losses = AverageMeter()

    # switch to eval mode to make BN unchanged.
    model.train()
    optimizer.zero_grad()

    end = time.time()
    for i, (img, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)

        # Input for Image CNN.
        img_var = utils.check_gpu(0, img)  # BS X 3 X H X W
        target_var = utils.check_gpu(0, target)  # BS X H X W X NUM_CLASSES

        bs = img.size()[0] * accumulation_steps

        score_feats5, fused_feats = model(
            img_var)  # BS X NUM_CLASSES X 472 X 472

        feats5_loss = WeightedMultiLabelSigmoidLoss(score_feats5, target_var)
        fused_feats_loss = WeightedMultiLabelSigmoidLoss(
            fused_feats, target_var)
        loss = feats5_loss + fused_feats_loss

        loss.backward()

        # clear memory
        del img_var
        del target_var
        del score_feats5
        torch.cuda.empty_cache()

        # increase batch size by factor of accumulation steps (Gradient accumulation) for training with limited memory
        if (i + 1) % accumulation_steps == 0:
            feats5_losses.update(feats5_loss.data, bs)
            fusion_losses.update(fused_feats_loss.data, bs)
            total_losses.update(loss.data, bs)

            # Only plot the fused feats loss.
            trn_feats5_loss = feats5_loss.clone().cpu().data.numpy()
            trn_fusion_loss = fused_feats_loss.clone().cpu().data.numpy()
            viz.line(win=win_feats5,
                     name='train_feats5',
                     update='append',
                     X=np.array([global_step]),
                     Y=np.array([trn_feats5_loss]))
            viz.line(win=win_fusion,
                     name='train_fusion',
                     update='append',
                     X=np.array([global_step]),
                     Y=np.array([trn_fusion_loss]))

            optimizer.step()
            optimizer.zero_grad()
            global_step += 1

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if ((i + 1) % args.print_freq == 0):
                print("\n\n")
                print(
                    'Epoch: [{0}][{1}/{2}]\t'
                    'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                    'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                    'Total Loss {total_loss.val:.11f} ({total_loss.avg:.11f})\n'
                    'lr {learning_rate:.10f}\t'.format(
                        epoch,
                        int((i + 1) / accumulation_steps),
                        int(len(train_loader) / accumulation_steps),
                        batch_time=batch_time,
                        data_time=data_time,
                        total_loss=total_losses,
                        learning_rate=curr_lr))

    del feats5_loss
    del fused_feats_loss
    del feats5_losses
    del fusion_losses
    del total_losses
    torch.cuda.empty_cache()
    return global_step
def train_net(param, model, train_data, valid_data, plot=False,device='cuda'):
    # 初始化参数
    model_name      = param['model_name']
    epochs          = param['epochs']
    batch_size      = param['batch_size']
    lr              = param['lr']
    gamma           = param['gamma']
    step_size       = param['step_size']
    momentum        = param['momentum']
    weight_decay    = param['weight_decay']

    disp_inter      = param['disp_inter']
    save_inter      = param['save_inter']
    min_inter       = param['min_inter']
    iter_inter      = param['iter_inter']

    save_log_dir    = param['save_log_dir']
    save_ckpt_dir   = param['save_ckpt_dir']
    load_ckpt_dir   = param['load_ckpt_dir']

    #
    scaler = GradScaler() 

    # 网络参数
    train_data_size = train_data.__len__()
    valid_data_size = valid_data.__len__()
    c, y, x = train_data.__getitem__(0)['image'].shape
    train_loader = DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True, num_workers=1)
    valid_loader = DataLoader(dataset=valid_data, batch_size=batch_size, shuffle=False, num_workers=1)
    optimizer = optim.AdamW(model.parameters(), lr=3e-4 ,weight_decay=weight_decay)
    #optimizer = optim.SGD(model.parameters(), lr=1e-2, momentum=momentum, weight_decay=weight_decay)
    #scheduler = StepLR(optimizer, step_size=step_size, gamma=gamma)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=3, T_mult=2, eta_min=1e-5, last_epoch=-1)
    #criterion = nn.CrossEntropyLoss(reduction='mean').to(device)
    DiceLoss_fn=DiceLoss(mode='multiclass')
    SoftCrossEntropy_fn=SoftCrossEntropyLoss(smooth_factor=0.1)
    criterion = L.JointLoss(first=DiceLoss_fn, second=SoftCrossEntropy_fn,
                              first_weight=0.5, second_weight=0.5).cuda()
    logger = inial_logger(os.path.join(save_log_dir, time.strftime("%m-%d %H:%M:%S", time.localtime()) +'_'+model_name+ '.log'))

    # 主循环
    train_loss_total_epochs, valid_loss_total_epochs, epoch_lr = [], [], []
    train_loader_size = train_loader.__len__()
    valid_loader_size = valid_loader.__len__()
    best_iou = 0
    best_epoch=0
    best_mode = copy.deepcopy(model)
    epoch_start = 0
    if load_ckpt_dir is not None:
        ckpt = torch.load(load_ckpt_dir)
        epoch_start = ckpt['epoch']
        model.load_state_dict(ckpt['state_dict'])
        optimizer.load_state_dict(ckpt['optimizer'])

    logger.info('Total Epoch:{} Image_size:({}, {}) Training num:{}  Validation num:{}'.format(epochs, x, y, train_data_size, valid_data_size))
    #
    for epoch in range(epoch_start, epochs):
        epoch_start = time.time()
        # 训练阶段
        model.train()
        train_epoch_loss = AverageMeter()
        train_iter_loss = AverageMeter()
        for batch_idx, batch_samples in enumerate(train_loader):
            data, target = batch_samples['image'], batch_samples['label']
            data, target = Variable(data.to(device)), Variable(target.to(device))
            with autocast(): #need pytorch>1.6
                pred = model(data)
                loss = criterion(pred, target)
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad()
            scheduler.step(epoch + batch_idx / train_loader_size) 
            image_loss = loss.item()
            train_epoch_loss.update(image_loss)
            train_iter_loss.update(image_loss)
            if batch_idx % iter_inter == 0:
                spend_time = time.time() - epoch_start
                logger.info('[train] epoch:{} iter:{}/{} {:.2f}% lr:{:.6f} loss:{:.6f} ETA:{}min'.format(
                    epoch, batch_idx, train_loader_size, batch_idx/train_loader_size*100,
                    optimizer.param_groups[-1]['lr'],
                    train_iter_loss.avg,spend_time / (batch_idx+1) * train_loader_size // 60 - spend_time // 60))
                train_iter_loss.reset()

        # 验证阶段
        model.eval()
        valid_epoch_loss = AverageMeter()
        valid_iter_loss = AverageMeter()
        iou=IOUMetric(10)
        with torch.no_grad():
            for batch_idx, batch_samples in enumerate(valid_loader):
                data, target = batch_samples['image'], batch_samples['label']
                data, target = Variable(data.to(device)), Variable(target.to(device))
                pred = model(data)
                loss = criterion(pred, target)
                pred=pred.cpu().data.numpy()
                pred= np.argmax(pred,axis=1)
                iou.add_batch(pred,target.cpu().data.numpy())
                #
                image_loss = loss.item()
                valid_epoch_loss.update(image_loss)
                valid_iter_loss.update(image_loss)
                # if batch_idx % iter_inter == 0:
                #     logger.info('[val] epoch:{} iter:{}/{} {:.2f}% loss:{:.6f}'.format(
                #         epoch, batch_idx, valid_loader_size, batch_idx / valid_loader_size * 100, valid_iter_loss.avg))
            val_loss=valid_iter_loss.avg
            acc, acc_cls, iu, mean_iu, fwavacc=iou.evaluate()
            logger.info('[val] epoch:{} miou:{:.2f}'.format(epoch,mean_iu))
                

        # 保存loss、lr
        train_loss_total_epochs.append(train_epoch_loss.avg)
        valid_loss_total_epochs.append(valid_epoch_loss.avg)
        epoch_lr.append(optimizer.param_groups[0]['lr'])
        # 保存模型
        if epoch % save_inter == 0 and epoch > min_inter:
            state = {'epoch': epoch, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()}
            filename = os.path.join(save_ckpt_dir, 'checkpoint-epoch{}.pth'.format(epoch))
            torch.save(state, filename)  # pytorch1.6会压缩模型,低版本无法加载
        # 保存最优模型
        if mean_iu > best_iou:  # train_loss_per_epoch valid_loss_per_epoch
            state = {'epoch': epoch, 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict()}
            filename = os.path.join(save_ckpt_dir, 'checkpoint-best.pth')
            torch.save(state, filename)
            best_iou = mean_iu
            best_mode = copy.deepcopy(model)
            logger.info('[save] Best Model saved at epoch:{} ============================='.format(epoch))
        #scheduler.step()
        # 显示loss
    # 训练loss曲线
    if plot:
        x = [i for i in range(epochs)]
        fig = plt.figure(figsize=(12, 4))
        ax = fig.add_subplot(1, 2, 1)
        ax.plot(x, smooth(train_loss_total_epochs, 0.6), label='train loss')
        ax.plot(x, smooth(valid_loss_total_epochs, 0.6), label='val loss')
        ax.set_xlabel('Epoch', fontsize=15)
        ax.set_ylabel('CrossEntropy', fontsize=15)
        ax.set_title('train curve', fontsize=15)
        ax.grid(True)
        plt.legend(loc='upper right', fontsize=15)
        ax = fig.add_subplot(1, 2, 2)
        ax.plot(x, epoch_lr,  label='Learning Rate')
        ax.set_xlabel('Epoch', fontsize=15)
        ax.set_ylabel('Learning Rate', fontsize=15)
        ax.set_title('lr curve', fontsize=15)
        ax.grid(True)
        plt.legend(loc='upper right', fontsize=15)
        plt.show()
            
    return best_mode, model
예제 #24
0
def test(model, test_dataset, test_loader, output_filename, result_dir=None):

    load_batch_size = test_loader.batch_size
    num_batches = len(test_loader)

    model.eval()

    fw_time_meter = AverageMeter()

    det_results = {}

    for i, data_dicts in enumerate(test_loader):

        point_clouds = data_dicts['point_cloud']
        rot_angles = data_dicts['rot_angle']
        # optional
        ref_centers = data_dicts.get('ref_center')
        rgb_probs = data_dicts.get('rgb_prob')

        # from ground truth box detection
        if rgb_probs is None:
            rgb_probs = torch.ones_like(rot_angles)

        # not belong to refinement stage
        if ref_centers is None:
            ref_centers = torch.zeros((point_clouds.shape[0], 3))

        batch_size = point_clouds.shape[0]
        rot_angles = rot_angles.view(-1)
        rgb_probs = rgb_probs.view(-1)

        if 'box3d_center' in data_dicts:
            data_dicts.pop('box3d_center')

        data_dicts_var = {
            key: value.cuda()
            for key, value in data_dicts.items()
        }

        torch.cuda.synchronize()
        tic = time.time()
        with torch.no_grad():
            outputs = model(data_dicts_var)

        cls_probs, center_preds, heading_preds, size_preds = outputs

        torch.cuda.synchronize()
        fw_time_meter.update((time.time() - tic))

        num_pred = cls_probs.shape[1]
        print('%d/%d %.3f' % (i, num_batches, fw_time_meter.val))

        cls_probs = cls_probs.data.cpu().numpy()
        center_preds = center_preds.data.cpu().numpy()
        heading_preds = heading_preds.data.cpu().numpy()
        size_preds = size_preds.data.cpu().numpy()

        rgb_probs = rgb_probs.numpy()
        rot_angles = rot_angles.numpy()
        ref_centers = ref_centers.numpy()

        for b in range(batch_size):

            if cfg.TEST.METHOD == 'nms':
                fg_idx = (cls_probs[b, :, 0] < cls_probs[b, :, 1]).nonzero()[0]
                if fg_idx.size == 0:
                    fg_idx = np.argmax(cls_probs[b, :, 1])
                    fg_idx = np.array([fg_idx])
            else:
                fg_idx = np.argmax(cls_probs[b, :, 1])
                fg_idx = np.array([fg_idx])

            num_pred = len(fg_idx)

            single_centers = center_preds[b, fg_idx]
            single_headings = heading_preds[b, fg_idx]
            single_sizes = size_preds[b, fg_idx]
            single_scores = cls_probs[b, fg_idx, 1] + rgb_probs[b]

            data_idx = test_dataset.id_list[load_batch_size * i + b]
            class_type = test_dataset.type_list[load_batch_size * i + b]
            box2d = test_dataset.box2d_list[load_batch_size * i + b]
            rot_angle = rot_angles[b]
            ref_center = ref_centers[b]

            if data_idx not in det_results:
                det_results[data_idx] = {}

            if class_type not in det_results[data_idx]:
                det_results[data_idx][class_type] = []

            for n in range(num_pred):
                x1, y1, x2, y2 = box2d
                score = single_scores[n]
                h, w, l, tx, ty, tz, ry = from_prediction_to_label_format(
                    single_centers[n], single_headings[n], single_sizes[n],
                    rot_angle, ref_center)
                output = [x1, y1, x2, y2, tx, ty, tz, h, w, l, ry, score]
                det_results[data_idx][class_type].append(output)

    num_images = len(det_results)

    logging.info('Average time:')
    logging.info('batch:%0.3f' % fw_time_meter.avg)
    logging.info('avg_per_object:%0.3f' %
                 (fw_time_meter.avg / load_batch_size))
    logging.info('avg_per_image:%.3f' %
                 (fw_time_meter.avg * len(test_loader) / num_images))

    # Write detection results for KITTI evaluation

    if cfg.TEST.METHOD == 'nms':
        write_detection_results_nms(result_dir, det_results)
    else:
        write_detection_results(result_dir, det_results)

    output_dir = os.path.join(result_dir, 'data')

    if 'test' not in cfg.TEST.DATASET:
        evaluate_py_wrapper(result_dir)
        # evaluate_cuda_wrapper(output_dir, cfg.TEST.DATASET)
    else:
        logger.info('results file save in  {}'.format(result_dir))
        os.system('cd %s && zip -q -r ../results.zip *' % (result_dir))
예제 #25
0
def train_val(model, args):

    train_dir = args.train_dir
    val_dir = args.val_dir

    config = Config(args.config)
    cudnn.benchmark = True

    # train
    train_loader = torch.utils.data.DataLoader(lsp_lspet_data.LSP_Data(
        'lspet', train_dir, 8,
        Mytransforms.Compose([
            Mytransforms.RandomResized(),
            Mytransforms.RandomRotate(40),
            Mytransforms.RandomCrop(368),
            Mytransforms.RandomHorizontalFlip(),
        ])),
                                               batch_size=config.batch_size,
                                               shuffle=True,
                                               num_workers=config.workers,
                                               pin_memory=True)
    # val
    if args.val_dir is not None and config.test_interval != 0:
        # val
        val_loader = torch.utils.data.DataLoader(lsp_lspet_data.LSP_Data(
            'lsp', val_dir, 8,
            Mytransforms.Compose([
                Mytransforms.TestResized(368),
            ])),
                                                 batch_size=config.batch_size,
                                                 shuffle=True,
                                                 num_workers=config.workers,
                                                 pin_memory=True)

    if args.gpu[0] < 0:
        criterion = nn.MSELoss()
    else:
        criterion = nn.MSELoss().cuda()

    params, multiple = get_parameters(model, config, True)
    # params, multiple = get_parameters(model, config, False)

    optimizer = torch.optim.SGD(params,
                                config.base_lr,
                                momentum=config.momentum,
                                weight_decay=config.weight_decay)

    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    losses_list = [AverageMeter() for i in range(6)]
    end = time.time()
    iters = config.start_iters
    best_model = config.best_model

    heat_weight = 46 * 46 * 15 / 1.0

    while iters < config.max_iter:

        for i, (input, heatmap, centermap) in enumerate(train_loader):

            learning_rate = adjust_learning_rate(
                optimizer,
                iters,
                config.base_lr,
                policy=config.lr_policy,
                policy_parameter=config.policy_parameter,
                multiple=multiple)
            data_time.update(time.time() - end)

            if args.gpu[0] >= 0:
                heatmap = heatmap.cuda(async=True)
                centermap = centermap.cuda(async=True)

            input_var = torch.autograd.Variable(input)
            heatmap_var = torch.autograd.Variable(heatmap)
            centermap_var = torch.autograd.Variable(centermap)

            heat1, heat2, heat3, heat4, heat5, heat6 = model(
                input_var, centermap_var)

            loss1 = criterion(heat1, heatmap_var) * heat_weight
            loss2 = criterion(heat2, heatmap_var) * heat_weight
            loss3 = criterion(heat3, heatmap_var) * heat_weight
            loss4 = criterion(heat4, heatmap_var) * heat_weight
            loss5 = criterion(heat5, heatmap_var) * heat_weight
            loss6 = criterion(heat6, heatmap_var) * heat_weight

            loss = loss1 + loss2 + loss3 + loss4 + loss5 + loss6
            #print(input.size(0).item())
            losses.update(loss.item(), input.size(0))
            for cnt, l in enumerate([loss1, loss2, loss3, loss4, loss5,
                                     loss6]):
                losses_list[cnt].update(l.item(), input.size(0))

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            batch_time.update(time.time() - end)
            end = time.time()

            iters += 1
            if iters % config.display == 0:
                print(
                    'Train Iteration: {0}\t'
                    'Time {batch_time.sum:.3f}s / {1}iters, ({batch_time.avg:.3f})\t'
                    'Data load {data_time.sum:.3f}s / {1}iters, ({data_time.avg:3f})\n'
                    'Learning rate = {2}\n'
                    'Loss = {loss.val:.8f} (ave = {loss.avg:.8f})\n'.format(
                        iters,
                        config.display,
                        learning_rate,
                        batch_time=batch_time,
                        data_time=data_time,
                        loss=losses))
                for cnt in range(0, 6):
                    print(
                        'Loss{0} = {loss1.val:.8f} (ave = {loss1.avg:.8f})\t'.
                        format(cnt + 1, loss1=losses_list[cnt]))

                print(
                    time.strftime(
                        '%Y-%m-%d %H:%M:%S -----------------------------------------------------------------------------------------------------------------\n',
                        time.localtime()))

                batch_time.reset()
                data_time.reset()
                losses.reset()
                for cnt in range(6):
                    losses_list[cnt].reset()

            save_checkpoint({
                'iter': iters,
                'state_dict': model.state_dict(),
            }, 0, args.model_name)

            # val
            if args.val_dir is not None and config.test_interval != 0 and iters % config.test_interval == 0:

                model.eval()
                for j, (input, heatmap, centermap) in enumerate(val_loader):
                    if args.cuda[0] >= 0:
                        heatmap = heatmap.cuda(async=True)
                        centermap = centermap.cuda(async=True)

                    input_var = torch.autograd.Variable(input)
                    heatmap_var = torch.autograd.Variable(heatmap)
                    centermap_var = torch.autograd.Variable(centermap)

                    heat1, heat2, heat3, heat4, heat5, heat6 = model(
                        input_var, centermap_var)

                    loss1 = criterion(heat1, heatmap_var) * heat_weight
                    loss2 = criterion(heat2, heatmap_var) * heat_weight
                    loss3 = criterion(heat3, heatmap_var) * heat_weight
                    loss4 = criterion(heat4, heatmap_var) * heat_weight
                    loss5 = criterion(heat5, heatmap_var) * heat_weight
                    loss6 = criterion(heat6, heatmap_var) * heat_weight

                    loss = loss1 + loss2 + loss3 + loss4 + loss5 + loss6
                    losses.update(loss.data[0], input.size(0))
                    for cnt, l in enumerate(
                        [loss1, loss2, loss3, loss4, loss5, loss6]):
                        losses_list[cnt].update(l.data[0], input.size(0))

                    batch_time.update(time.time() - end)
                    end = time.time()
                    is_best = losses.avg < best_model
                    best_model = min(best_model, losses.avg)
                    save_checkpoint(
                        {
                            'iter': iters,
                            'state_dict': model.state_dict(),
                        }, is_best, args.model_name)

                    if j % config.display == 0:
                        print(
                            'Test Iteration: {0}\t'
                            'Time {batch_time.sum:.3f}s / {1}iters, ({batch_time.avg:.3f})\t'
                            'Data load {data_time.sum:.3f}s / {1}iters, ({data_time.avg:3f})\n'
                            'Loss = {loss.val:.8f} (ave = {loss.avg:.8f})\n'.
                            format(j,
                                   config.display,
                                   batch_time=batch_time,
                                   data_time=data_time,
                                   loss=losses))
                        for cnt in range(0, 6):
                            print(
                                'Loss{0} = {loss1.val:.8f} (ave = {loss1.avg:.8f})\t'
                                .format(cnt + 1, loss1=losses_list[cnt]))

                        print(
                            time.strftime(
                                '%Y-%m-%d %H:%M:%S -----------------------------------------------------------------------------------------------------------------\n',
                                time.localtime()))
                        batch_time.reset()
                        losses.reset()
                        for cnt in range(6):
                            losses_list[cnt].reset()

                model.train()
예제 #26
0
파일: train.py 프로젝트: ggsDing/SARSeg
def train(train_loader, net, criterion, optimizer, scheduler, args,
          val_loader):
    bestaccT = 0
    bestfwiou = 0.5
    bestaccV = 0.0
    bestloss = 1
    begin_time = time.time()
    all_iters = float(len(train_loader) * args['epochs'])
    curr_epoch = 0
    while True:
        if args['gpu']: torch.cuda.empty_cache()
        net.train()
        start = time.time()
        acc_meter = AverageMeter()
        train_main_loss = AverageMeter()
        train_aux_loss = AverageMeter()

        curr_iter = curr_epoch * len(train_loader)
        for i, (imgs, labels) in enumerate(train_loader):
            running_iter = curr_iter + i + 1
            adjust_lr(optimizer, running_iter, all_iters)
            #imgs = torch.squeeze(imgs)
            imgs = imgs.float()
            labels = labels.long()
            #imgs, labels = data
            if args['gpu']:
                imgs = imgs.cuda().float()
                labels = labels.cuda().long()

            optimizer.zero_grad()
            outputs, aux = net(imgs)  #
            assert outputs.shape[1] == RS.num_classes + 1
            loss_main = criterion(outputs, labels)
            loss_aux = criterion(aux, labels)
            loss = loss_main * 0.7 + loss_aux * 0.3
            loss.backward()
            optimizer.step()

            labels = labels.cpu().detach().numpy()
            outputs = outputs.cpu().detach()
            _, preds = torch.max(outputs, dim=1)
            preds = preds.numpy()
            # batch_valid_sum = 0
            acc_curr_meter = AverageMeter()
            for (pred, label) in zip(preds, labels):
                acc, _ = accuracy(pred, label)
                acc_curr_meter.update(acc)
            acc_meter.update(acc_curr_meter.avg)
            train_main_loss.update(loss.cpu().detach().numpy())
            train_aux_loss.update(loss_aux.cpu().detach().numpy())

            curr_time = time.time() - start

            if (i + 1) % args['print_freq'] == 0:
                print(
                    '[epoch %d] [iter %d / %d %.1fs] [lr %f] [train loss %.4f acc %.2f]'
                    % (curr_epoch, i + 1, len(train_loader), curr_time,
                       optimizer.param_groups[0]['lr'], train_main_loss.val,
                       acc_meter.val * 100))
                writer.add_scalar('train_main_loss', train_main_loss.val,
                                  running_iter)
                writer.add_scalar('train_accuracy', acc_meter.val,
                                  running_iter)
                writer.add_scalar('train_aux_loss', train_aux_loss.avg,
                                  running_iter)
                writer.add_scalar('lr', optimizer.param_groups[0]['lr'],
                                  running_iter)

        acc_v, fwiou_v, loss_v = validate(val_loader, net, criterion,
                                          curr_epoch)
        if acc_meter.avg > bestaccT: bestaccT = acc_meter.avg

        if fwiou_v > bestfwiou:
            bestfwiou = fwiou_v
            bestloss = loss_v
            bestaccV = acc_v
            torch.save(
                net.state_dict(),
                os.path.join(
                    args['chkpt_dir'], NET_NAME + '_%de_OA%.2f_fwIoU%.2f.pth' %
                    (curr_epoch, acc_v * 100, fwiou_v * 100)))
        print(
            'Total time: %.1fs Best rec: Train acc %.2f, Val acc %.2f fwiou %.2f, Val_loss %.4f'
            % (time.time() - begin_time, bestaccT * 100, bestaccV * 100,
               bestfwiou * 100, bestloss))
        curr_epoch += 1
        #scheduler.step()
        if curr_epoch >= args['epochs']:
            return
예제 #27
0
def train_moco(epoch, train_loader, model, model_ema, contrast, criterion,
               optimizer, opt, recorder):
    """
    one epoch training for instance discrimination
    """
    print("==> (MoCo) training...")
    model_ema.eval()
    model.train()
    loss_meter = AverageMeter()
    prob_meter = AverageMeter()
    batch_time = AverageMeter()
    data_time = AverageMeter()

    end = time.time()
    for idx, (inputs, _, index) in enumerate(train_loader):
        data_time.update(time.time() - end)
        # print(inputs[0].size())
        bsz = inputs[0].size(0)
        # fixed args.batch_size
        for i in range(len(inputs)):
            inputs[i] = inputs[i].float()
            inputs[i] = inputs[i].cuda()
        if bsz < opt.pt_batch_size:
            print("batch less than 16, continue")
            continue
        index = index.cuda(non_blocking=True)
        # ===================forward=====================
        anchor, positive, negative = inputs
        feat_n, _ = model(negative)
        feat_q, _ = model(anchor)
        feat_k, _ = model_ema(positive)
        if feat_k.size(0) > feat_n.size(0):
            print("wrong bsz")
        out = contrast(feat_q, feat_k, feat_n, index)
        prob = out[:, 0].mean()
        loss = criterion(out)
        # ===================backward=====================
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # ===================meters=====================
        loss_meter.update(loss.item(), bsz)
        prob_meter.update(prob.item(), bsz)
        moment_update(model, model_ema, opt.pt_alpha)
        torch.cuda.synchronize()
        batch_time.update(time.time() - end)
        message = ('MoCo Train: [{0}][{1}/{2}]\t'
                   'BT {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                   'DT {data_time.val:.3f} ({data_time.avg:.3f})\t'
                   'loss {loss.val:.3f} ({loss.avg:.3f})\t'
                   'prob {prob.val:.3f} ({prob.avg:.3f})'.format(
                       epoch,
                       idx + 1,
                       len(train_loader),
                       batch_time=batch_time,
                       data_time=data_time,
                       loss=loss_meter,
                       prob=prob_meter))
        # print info
        if (idx + 1) % opt.pt_print_freq == 0:
            print(message)
            recorder.record_message('a', message)
            # print(out.shape)
            sys.stdout.flush()
        end = time.time()
    return loss_meter.avg, prob_meter.avg
예제 #28
0
def test(epoch, cfg, data_loader, model, obj_vtx, obj_info, criterions):

    model.eval()
    Eval = Evaluation(cfg.dataset, obj_info, obj_vtx)
    if 'trans' in cfg.pytorch.task.lower():
        Eval_trans = Evaluation(cfg.dataset, obj_info, obj_vtx)

    if not cfg.test.ignore_cache_file:
        est_cache_file = cfg.test.cache_file
        # gt_cache_file = cfg.test.cache_file.replace('pose_est', 'pose_gt')
        gt_cache_file = cfg.test.cache_file.replace('_est', '_gt')
        if os.path.exists(est_cache_file) and os.path.exists(gt_cache_file):
            Eval.pose_est_all = np.load(est_cache_file,
                                        allow_pickle=True).tolist()
            Eval.pose_gt_all = np.load(gt_cache_file,
                                       allow_pickle=True).tolist()
            fig_save_path = os.path.join(cfg.pytorch.save_path, str(epoch))
            mkdir_p(fig_save_path)
            if 'all' in cfg.test.test_mode.lower():
                Eval.evaluate_pose()
                Eval.evaluate_pose_add(fig_save_path)
                Eval.evaluate_pose_arp_2d(fig_save_path)
            elif 'pose' in cfg.test.test_mode.lower():
                Eval.evaluate_pose()
            elif 'add' in cfg.test.test_mode.lower():
                Eval.evaluate_pose_add(fig_save_path)
            elif 'arp' in cfg.test.test_mode.lower():
                Eval.evaluate_pose_arp_2d(fig_save_path)
            else:
                raise Exception("Wrong test mode: {}".format(
                    cfg.test.test_mode))

            return None, None

        else:
            logger.info("test cache file {} and {} not exist!".format(
                est_cache_file, gt_cache_file))
            userAns = input("Generating cache file from model [Y(y)/N(n)]:")
            if userAns.lower() == 'n':
                sys.exit(0)
            else:
                logger.info("Generating test cache file!")

    preds = {}
    Loss = AverageMeter()
    Loss_rot = AverageMeter()
    Loss_trans = AverageMeter()
    num_iters = len(data_loader)
    bar = Bar('{}'.format(cfg.pytorch.exp_id[-60:]), max=num_iters)

    time_monitor = False
    vis_dir = os.path.join(cfg.pytorch.save_path, 'test_vis_{}'.format(epoch))
    if not os.path.exists(vis_dir):
        os.makedirs(vis_dir)
    for i, (obj, obj_id, inp, pose, c_box, s_box, box,
            trans_local) in enumerate(data_loader):
        if cfg.pytorch.gpu > -1:
            inp_var = inp.cuda(cfg.pytorch.gpu, async=True).float()
        else:
            inp_var = inp.float()

        bs = len(inp)
        # forward propagation
        T_begin = time.time()
        pred_rot, pred_trans = model(inp_var)
        T_end = time.time() - T_begin
        if time_monitor:
            logger.info(
                "time for a batch forward of resnet model is {}".format(T_end))

        if i % cfg.test.disp_interval == 0:
            # input image
            inp_rgb = (inp[0].cpu().numpy().copy() *
                       255)[[2, 1, 0], :, :].astype(np.uint8)
            cfg.writer.add_image('input_image', inp_rgb, i)
            cv2.imwrite(os.path.join(vis_dir, '{}_inp.png'.format(i)),
                        inp_rgb.transpose(1, 2, 0)[:, :, ::-1])
            if 'rot' in cfg.pytorch.task.lower():
                # coordinates map
                pred_coor = pred_rot[0, 0:3].data.cpu().numpy().copy()
                pred_coor[0] = im_norm_255(pred_coor[0])
                pred_coor[1] = im_norm_255(pred_coor[1])
                pred_coor[2] = im_norm_255(pred_coor[2])
                pred_coor = np.asarray(pred_coor, dtype=np.uint8)
                cfg.writer.add_image('test_coor_x_pred',
                                     np.expand_dims(pred_coor[0], axis=0), i)
                cfg.writer.add_image('test_coor_y_pred',
                                     np.expand_dims(pred_coor[1], axis=0), i)
                cfg.writer.add_image('test_coor_z_pred',
                                     np.expand_dims(pred_coor[2], axis=0), i)
                # gt_coor = target[0, 0:3].data.cpu().numpy().copy()
                # gt_coor[0] = im_norm_255(gt_coor[0])
                # gt_coor[1] = im_norm_255(gt_coor[1])
                # gt_coor[2] = im_norm_255(gt_coor[2])
                # gt_coor = np.asarray(gt_coor, dtype=np.uint8)
                # cfg.writer.add_image('test_coor_x_gt', np.expand_dims(gt_coor[0], axis=0), i)
                # cfg.writer.add_image('test_coor_y_gt', np.expand_dims(gt_coor[1], axis=0), i)
                # cfg.writer.add_image('test_coor_z_gt', np.expand_dims(gt_coor[2], axis=0), i)
                # confidence map
                pred_conf = pred_rot[0, 3].data.cpu().numpy().copy()
                pred_conf = (im_norm_255(pred_conf)).astype(np.uint8)
                cfg.writer.add_image('test_conf_pred',
                                     np.expand_dims(pred_conf, axis=0), i)
                # gt_conf = target[0, 3].data.cpu().numpy().copy()
                # cfg.writer.add_image('test_conf_gt', np.expand_dims(gt_conf, axis=0), i)
            if 'trans' in cfg.pytorch.task.lower():
                pred_trans_ = pred_trans[0].data.cpu().numpy().copy()
                gt_trans_ = trans_local[0].data.cpu().numpy().copy()
                cfg.writer.add_scalar('test_trans_x_gt', gt_trans_[0],
                                      i + (epoch - 1) * num_iters)
                cfg.writer.add_scalar('test_trans_y_gt', gt_trans_[1],
                                      i + (epoch - 1) * num_iters)
                cfg.writer.add_scalar('test_trans_z_gt', gt_trans_[2],
                                      i + (epoch - 1) * num_iters)
                cfg.writer.add_scalar('test_trans_x_pred', pred_trans_[0],
                                      i + (epoch - 1) * num_iters)
                cfg.writer.add_scalar('test_trans_y_pred', pred_trans_[1],
                                      i + (epoch - 1) * num_iters)
                cfg.writer.add_scalar('test_trans_z_pred', pred_trans_[2],
                                      i + (epoch - 1) * num_iters)
                cfg.writer.add_scalar('test_trans_x_err',
                                      np.abs(pred_trans_[0] - gt_trans_[0]),
                                      i + (epoch - 1) * num_iters)
                cfg.writer.add_scalar('test_trans_y_err',
                                      np.abs(pred_trans_[1] - gt_trans_[1]),
                                      i + (epoch - 1) * num_iters)
                cfg.writer.add_scalar('test_trans_z_err',
                                      np.abs(pred_trans_[2] - gt_trans_[2]),
                                      i + (epoch - 1) * num_iters)

        if 'rot' in cfg.pytorch.task.lower():
            pred_coor = pred_rot[:, 0:3].data.cpu().numpy().copy()
            pred_conf = pred_rot[:, 3].data.cpu().numpy().copy()
        else:
            pred_coor = np.zeros(bs)
            pred_conf = np.zeros(bs)

        if 'trans' in cfg.pytorch.task.lower():
            pred_trans = pred_trans.data.cpu().numpy().copy()
        else:
            pred_trans = np.zeros(bs)

        col = list(
            zip(obj, obj_id.numpy(), pred_coor, pred_conf, pred_trans,
                pose.numpy(), c_box.numpy(), s_box.numpy(), box.numpy()))
        for idx in range(len(col)):
            obj_, obj_id_, pred_coor_, pred_conf_, pred_trans_, pose_gt, c_box_, s_box_, box_ = col[
                idx]
            T_begin = time.time()
            if 'rot' in cfg.pytorch.task.lower():
                # building 2D-3D correspondences
                pred_coor_ = pred_coor_.transpose(1, 2, 0)
                pred_coor_[:, :, 0] = pred_coor_[:, :, 0] * abs(
                    obj_info[obj_id_]['min_x'])
                pred_coor_[:, :, 1] = pred_coor_[:, :, 1] * abs(
                    obj_info[obj_id_]['min_y'])
                pred_coor_[:, :, 2] = pred_coor_[:, :, 2] * abs(
                    obj_info[obj_id_]['min_z'])
                pred_coor_ = pred_coor_.tolist()
                eroMask = False
                if eroMask:
                    kernel = np.ones((3, 3), np.uint8)
                    pred_conf_ = cv2.erode(pred_conf_, kernel)
                pred_conf_ = (pred_conf_ - pred_conf_.min()) / (
                    pred_conf_.max() - pred_conf_.min())
                pred_conf_ = pred_conf_.tolist()

                select_pts_2d = []
                select_pts_3d = []
                c_w = int(c_box_[0])
                c_h = int(c_box_[1])
                s = int(s_box_)
                w_begin = c_w - s / 2.
                h_begin = c_h - s / 2.
                w_unit = s * 1.0 / cfg.dataiter.out_res
                h_unit = s * 1.0 / cfg.dataiter.out_res

                min_x = 0.001 * abs(obj_info[obj_id_]['min_x'])
                min_y = 0.001 * abs(obj_info[obj_id_]['min_y'])
                min_z = 0.001 * abs(obj_info[obj_id_]['min_z'])
                for x in range(cfg.dataiter.out_res):
                    for y in range(cfg.dataiter.out_res):
                        if pred_conf_[x][y] < cfg.test.mask_threshold:
                            continue
                        if abs(pred_coor_[x][y][0]) < min_x  and abs(pred_coor_[x][y][1]) < min_y  and \
                            abs(pred_coor_[x][y][2]) < min_z:
                            continue
                        select_pts_2d.append(
                            [w_begin + y * w_unit, h_begin + x * h_unit])
                        select_pts_3d.append(pred_coor_[x][y])

                model_points = np.asarray(select_pts_3d, dtype=np.float32)
                image_points = np.asarray(select_pts_2d, dtype=np.float32)

            if 'trans' in cfg.pytorch.task.lower():
                # compute T from translation head
                ratio_delta_c = pred_trans_[:2]
                ratio_depth = pred_trans_[2]
                pred_depth = ratio_depth * (cfg.dataiter.out_res / s_box_)
                pred_c = ratio_delta_c * box_[2:] + c_box_
                pred_x = (pred_c[0] - cfg.dataset.camera_matrix[0, 2]
                          ) * pred_depth / cfg.dataset.camera_matrix[0, 0]
                pred_y = (pred_c[1] - cfg.dataset.camera_matrix[1, 2]
                          ) * pred_depth / cfg.dataset.camera_matrix[1, 1]
                T_vector_trans = np.asarray([pred_x, pred_y, pred_depth])
                pose_est_trans = np.concatenate(
                    (np.eye(3), np.asarray((T_vector_trans).reshape(3, 1))),
                    axis=1)

            try:
                if 'rot' in cfg.pytorch.task.lower():
                    dist_coeffs = np.zeros(
                        (4, 1))  # Assuming no lens distortion
                    if cfg.test.pnp == 'iterPnP':  # iterative PnP algorithm
                        success, R_vector, T_vector = cv2.solvePnP(
                            model_points,
                            image_points,
                            cfg.dataset.camera_matrix,
                            dist_coeffs,
                            flags=cv2.SOLVEPNP_ITERATIVE)
                    elif cfg.test.pnp == 'ransac':  # ransac algorithm
                        _, R_vector, T_vector, inliers = cv2.solvePnPRansac(
                            model_points,
                            image_points,
                            cfg.dataset.camera_matrix,
                            dist_coeffs,
                            flags=cv2.SOLVEPNP_EPNP)
                    else:
                        raise NotImplementedError(
                            "Not support PnP algorithm: {}".format(
                                cfg.test.pnp))
                    R_matrix = cv2.Rodrigues(R_vector, jacobian=0)[0]
                    pose_est = np.concatenate(
                        (R_matrix, np.asarray(T_vector).reshape(3, 1)), axis=1)
                    if 'trans' in cfg.pytorch.task.lower():
                        pose_est_trans = np.concatenate(
                            (R_matrix,
                             np.asarray((T_vector_trans).reshape(3, 1))),
                            axis=1)
                    Eval.pose_est_all[obj_].append(pose_est)
                    Eval.pose_gt_all[obj_].append(pose_gt)
                    Eval.num[obj_] += 1
                    Eval.numAll += 1
                if 'trans' in cfg.pytorch.task.lower():
                    Eval_trans.pose_est_all[obj_].append(pose_est_trans)
                    Eval_trans.pose_gt_all[obj_].append(pose_gt)
                    Eval_trans.num[obj_] += 1
                    Eval_trans.numAll += 1
            except:
                Eval.num[obj_] += 1
                Eval.numAll += 1
                if 'trans' in cfg.pytorch.task.lower():
                    Eval_trans.num[obj_] += 1
                    Eval_trans.numAll += 1
                logger.info('error in solve PnP or Ransac')

            T_end = time.time() - T_begin
            if time_monitor:
                logger.info(
                    "time spend on PnP+RANSAC for one image is {}".format(
                        T_end))

        Bar.suffix = 'test Epoch: [{0}][{1}/{2}]| Total: {total:} | ETA: {eta:} | Loss {loss.avg:.4f} | Loss_rot {loss_rot.avg:.4f} | Loss_trans {loss_trans.avg:.4f}'.format(
            epoch,
            i,
            num_iters,
            total=bar.elapsed_td,
            eta=bar.eta_td,
            loss=Loss,
            loss_rot=Loss_rot,
            loss_trans=Loss_trans)
        bar.next()

    epoch_save_path = os.path.join(cfg.pytorch.save_path, str(epoch))
    if not os.path.exists(epoch_save_path):
        os.makedirs(epoch_save_path)
    if 'rot' in cfg.pytorch.task.lower():
        logger.info("{} Evaluate of Rotation Branch of Epoch {} {}".format(
            '-' * 40, epoch, '-' * 40))
        preds['poseGT'] = Eval.pose_gt_all
        preds['poseEst'] = Eval.pose_est_all
        if cfg.pytorch.test:
            np.save(os.path.join(epoch_save_path, 'pose_est_all_test.npy'),
                    Eval.pose_est_all)
            np.save(os.path.join(epoch_save_path, 'pose_gt_all_test.npy'),
                    Eval.pose_gt_all)
        else:
            np.save(
                os.path.join(epoch_save_path,
                             'pose_est_all_epoch{}.npy'.format(epoch)),
                Eval.pose_est_all)
            np.save(
                os.path.join(epoch_save_path,
                             'pose_gt_all_epoch{}.npy'.format(epoch)),
                Eval.pose_gt_all)
        # evaluation
        if 'all' in cfg.test.test_mode.lower():
            Eval.evaluate_pose()
            Eval.evaluate_pose_add(epoch_save_path)
            Eval.evaluate_pose_arp_2d(epoch_save_path)
        else:
            if 'pose' in cfg.test.test_mode.lower():
                Eval.evaluate_pose()
            if 'add' in cfg.test.test_mode.lower():
                Eval.evaluate_pose_add(epoch_save_path)
            if 'arp' in cfg.test.test_mode.lower():
                Eval.evaluate_pose_arp_2d(epoch_save_path)

    if 'trans' in cfg.pytorch.task.lower():
        logger.info("{} Evaluate of Translation Branch of Epoch {} {}".format(
            '-' * 40, epoch, '-' * 40))
        preds['poseGT'] = Eval_trans.pose_gt_all
        preds['poseEst'] = Eval_trans.pose_est_all
        if cfg.pytorch.test:
            np.save(
                os.path.join(epoch_save_path, 'pose_est_all_test_trans.npy'),
                Eval_trans.pose_est_all)
            np.save(
                os.path.join(epoch_save_path, 'pose_gt_all_test_trans.npy'),
                Eval_trans.pose_gt_all)
        else:
            np.save(
                os.path.join(epoch_save_path,
                             'pose_est_all_trans_epoch{}.npy'.format(epoch)),
                Eval_trans.pose_est_all)
            np.save(
                os.path.join(epoch_save_path,
                             'pose_gt_all_trans_epoch{}.npy'.format(epoch)),
                Eval_trans.pose_gt_all)
        # evaluation
        if 'all' in cfg.test.test_mode.lower():
            Eval_trans.evaluate_pose()
            Eval_trans.evaluate_pose_add(epoch_save_path)
            Eval_trans.evaluate_pose_arp_2d(epoch_save_path)
        else:
            if 'pose' in cfg.test.test_mode.lower():
                Eval_trans.evaluate_pose()
            if 'add' in cfg.test.test_mode.lower():
                Eval_trans.evaluate_pose_add(epoch_save_path)
            if 'arp' in cfg.test.test_mode.lower():
                Eval_trans.evaluate_pose_arp_2d(epoch_save_path)

    bar.finish()
    return {
        'Loss': Loss.avg,
        'Loss_rot': Loss_rot.avg,
        'Loss_trans': Loss_trans.avg
    }, preds
예제 #29
0
def train_dsm_triplet(epoch, train_loader, model, optimizer, opt, pos_aug,
                      neg_aug, recorder):
    """
    one epoch training for instance discrimination
    """
    print("==> (DSM triplet) training...")
    model.train()

    def set_bn_train(m):
        classname = m.__class__.__name__
        if classname.find('BatchNorm') != -1:
            m.train()

    batch_time = AverageMeter()
    data_time = AverageMeter()
    loss_meter = AverageMeter()
    triplet_loss = nn.TripletMarginLoss(margin=0.5, p=2)
    end = time.time()
    for idx, (inputs, _, index) in enumerate(train_loader):
        data_time.update(time.time() - end)

        bsz = inputs[0].size(0)
        # fixed args.batch_size
        if bsz < opt.pt_batch_size:
            print("batch less than 16, continue")
            continue
        for i in range(len(inputs)):
            inputs[i] = inputs[i].float()
            inputs[i] = inputs[i].cuda()
        # ===================forward=====================
        anchor_old, positive, negative = inputs

        # here a series of data augmentation
        # ====================================================postive operation=======================
        anchor = pos_aug(anchor_old)
        feat_k = model(positive)
        feat_n = model(negative)
        feat_q = model(anchor)
        if feat_k.size(0) > feat_q.size(0):
            print("wrong bsz")
        intra_loss = triplet_loss(feat_q, feat_k, feat_n)
        inter_loss = triplet_loss(feat_q, feat_k, flip(feat_n, 0))
        # for j in range(bsz-2):
        #     inter_loss += triplet_loss(feat_q, feat_k, shift(feat_n, 0))
        alpha_1 = 1
        alpha_2 = 1
        loss = alpha_1 * intra_loss + alpha_2 * inter_loss
        # print(loss)
        # ===================backward=====================
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # ===================meters=====================
        loss_meter.update(loss.item(), bsz)
        torch.cuda.synchronize()
        batch_time.update(time.time() - end)
        end = time.time()
        message = ('DSM triplet Train: [{0}][{1}/{2}]\t'
                   'BT {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                   'DT {data_time.val:.3f} ({data_time.avg:.3f})\t'
                   'loss {loss.val:.3f} ({loss.avg:.3f})'.format(
                       epoch,
                       idx + 1,
                       len(train_loader),
                       batch_time=batch_time,
                       data_time=data_time,
                       loss=loss_meter))
        # print info
        if (idx + 1) % opt.pt_print_freq == 0:
            print(message)
            recorder.record_message('a', message)
            # print(out.shape)
            sys.stdout.flush()
    return loss_meter.avg
예제 #30
0
def train_model(epoch, model, optimizer, lr_scheduler, loader, test_loader):
    global GLOBAL_STEP

    test_loader_it = iter(test_loader)

    loss_meter = AverageMeter()
    val_loss_meter = AverageMeter()
    acc_meter = AverageMeter()
    val_acc_meter = AverageMeter()

    model.train()
    model.to(device)

    print('=' * 20 + "Model Training" + '=' * 20)

    loss_func = nn.CrossEntropyLoss()

    for i, batch in tqdm(enumerate(loader)):
        start = time.time()
        model.train()
        optimizer.zero_grad()
        model.zero_grad()
        sentence1, sentence2, label = batch
        label = label.to(device)
        pred = model((sentence1, sentence2))
        loss = loss_func(pred, label)
        loss.backward()
        grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip_bound)
        optimizer.step()
        acc = torch.mean((torch.max(pred, 1)[1] == label).type(torch.float))
        loss_meter.update(loss.item())
        acc_meter.update(acc.item())

        end = time.time()
        used_time = end - start

        if (GLOBAL_STEP) % args.log_every == 0:
            try:
                batch = next(test_loader_it)
            except:
                test_loader_it = iter(test_loader)
                batch = next(test_loader_it)
            eval_loss, eval_acc = batch_eval(batch, model)
            val_loss_meter.update(eval_loss.item())
            val_acc_meter.update(eval_acc.item())
            lr = optimizer.param_groups[0]['lr']
            display = 'epoch=' + str(epoch) + \
                      '\tglobal_step=%d' % (GLOBAL_STEP) + \
                      '\tloss=%.4f' % (loss_meter.val) + \
                      '\tloss_avg=%.4f' % (loss_meter.avg) + \
                      '\tval_loss=%.4f' % (val_loss_meter.avg) + \
                      '\tacc=%.4f' % (acc_meter.avg) + \
                      '\tval_acc=%.4f' % (val_acc_meter.avg) + \
                      '\tlr=%.6f' % (lr) + \
                      '\t|g|=%.4f' % (grad_norm) + \
                      '\ttime=%.2fit/s' % (1. / used_time)
            tqdm.write(display)
            save_mode(epoch=epoch,
                      model=model,
                      optimizer=optimizer,
                      lr_scheduler=lr_scheduler)
        GLOBAL_STEP += 1
    return
예제 #31
0
def train_dsm(epoch, train_loader, model, model_ema, contrast, criterion,
              optimizer, opt, pos_aug, neg_aug, recorder):
    """
    one epoch training for instance discrimination
    """
    print("==> (DSM) training...")
    model.train()
    model_ema.eval()

    def set_bn_train(m):
        classname = m.__class__.__name__
        if classname.find('BatchNorm') != -1:
            m.train()

    model_ema.apply(set_bn_train)

    batch_time = AverageMeter()
    data_time = AverageMeter()
    loss_meter = AverageMeter()
    prob_meter = AverageMeter()

    end = time.time()
    for idx, (inputs, _, index) in enumerate(train_loader):
        data_time.update(time.time() - end)
        # print(inputs[0].size())
        bsz = inputs[0].size(0)
        # fixed args.batch_size
        if bsz < opt.pt_batch_size:
            print("batch less than 16, continue")
            continue
        for i in range(len(inputs)):
            inputs[i] = inputs[i].float()
            inputs[i] = inputs[i].cuda()
        index = index.cuda(non_blocking=True)

        # ===================forward=====================
        anchor_old, positive, negative = inputs
        # print(anchor_old.size())
        # here a series of data augmentation
        # ====================================================postive operation=======================
        anchor = pos_aug(anchor_old)
        # positive = flip(anchor, 2)
        # shuffle_ids, reverse_ids = get_shuffle_ids(bsz)
        feat_q, map_q = model(anchor)
        feat_k, map_k = model_ema(positive)
        # tcr_loss = tcr(map_q, map_k)
        # with torch.no_grad():
        #     positive = positive[shuffle_ids]
        #     feat_k = model_ema(positive)
        #     feat_k = feat_k[reverse_ids]
        feat_n, _ = model(negative)
        if feat_n.size(0) > feat_q.size(0):
            print("wrong bsz")
        out = contrast(feat_q, feat_k, feat_n, index)
        contrast_loss = criterion(out)
        loss = contrast_loss  # + tcr_loss  # + sample_loss # + contrast_loss2 # + cls_loss + mixup_loss
        # print(contrast_loss, tcr_loss)

        # ===================backward=====================
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # ===================meters=====================
        prob = out[:, 0].mean()
        prob_meter.update(prob.item(), bsz)
        loss_meter.update(loss.item(), bsz)
        moment_update(model, model_ema, opt.pt_alpha)
        torch.cuda.synchronize()
        batch_time.update(time.time() - end)
        end = time.time()
        message = ('DSM Train: [{0}][{1}/{2}]\t'
                   'BT {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                   'DT {data_time.val:.3f} ({data_time.avg:.3f})\t'
                   'loss {loss.val:.3f} ({loss.avg:.3f})\t'
                   'prob {prob.val:.3f} ({prob.avg:.3f})'.format(
                       epoch,
                       idx + 1,
                       len(train_loader),
                       batch_time=batch_time,
                       data_time=data_time,
                       loss=loss_meter,
                       prob=prob_meter))
        # print info
        if (idx + 1) % opt.pt_print_freq == 0:
            print(message)
            recorder.record_message('a', message)
            # print(out.shape)
            sys.stdout.flush()
    return loss_meter.avg, prob_meter.avg