示例#1
0
def get_data(dataname,
             data_dir,
             height,
             width,
             batch_size,
             combine_all=False,
             min_size=0.,
             max_size=0.8,
             workers=8,
             test_batch=64):
    root = osp.join(data_dir, dataname)

    dataset = datasets.create(dataname, root, combine_all=combine_all)

    num_classes = dataset.num_train_ids

    train_transformer = T.Compose([
        T.RandomHorizontalFlip(),
        T.Resize((height, width), interpolation=3),
        T.RandomOcclusion(min_size, max_size),
        T.ToTensor(),
    ])

    test_transformer = T.Compose([
        T.Resize((height, width), interpolation=3),
        T.ToTensor(),
    ])

    train_loader = DataLoader(Preprocessor(dataset.train,
                                           root=osp.join(
                                               dataset.images_dir,
                                               dataset.train_path),
                                           transform=train_transformer),
                              batch_size=batch_size,
                              num_workers=workers,
                              shuffle=True,
                              pin_memory=True,
                              drop_last=True)

    query_loader = DataLoader(Preprocessor(dataset.query,
                                           root=osp.join(
                                               dataset.images_dir,
                                               dataset.query_path),
                                           transform=test_transformer),
                              batch_size=test_batch,
                              num_workers=32,
                              shuffle=False,
                              pin_memory=True)

    gallery_loader = DataLoader(Preprocessor(dataset.gallery,
                                             root=osp.join(
                                                 dataset.images_dir,
                                                 dataset.gallery_path),
                                             transform=test_transformer),
                                batch_size=test_batch,
                                num_workers=32,
                                shuffle=False,
                                pin_memory=True)

    return dataset, num_classes, train_loader, query_loader, gallery_loader
示例#2
0
def get_data(dataname, data_dir, model, matcher, save_path, args):
    root = osp.join(data_dir, dataname)

    dataset = datasets.create(dataname, root, combine_all=args.combine_all)

    num_classes = dataset.num_train_ids

    train_transformer = T.Compose([
        T.Resize((args.height, args.width),
                 interpolation=InterpolationMode.BICUBIC),
        T.Pad(10),
        T.RandomCrop((args.height, args.width)),
        T.RandomHorizontalFlip(0.5),
        T.RandomRotation(5),
        T.ColorJitter(brightness=(0.5, 2.0),
                      contrast=(0.5, 2.0),
                      saturation=(0.5, 2.0),
                      hue=(-0.1, 0.1)),
        T.RandomOcclusion(args.min_size, args.max_size),
        T.ToTensor(),
    ])

    test_transformer = T.Compose([
        T.Resize((args.height, args.width),
                 interpolation=InterpolationMode.BICUBIC),
        T.ToTensor(),
    ])

    train_path = osp.join(dataset.images_dir, dataset.train_path)
    train_loader = DataLoader(
        Preprocessor(dataset.train,
                     root=train_path,
                     transform=train_transformer),
        batch_size=args.batch_size,
        num_workers=args.workers,
        sampler=GraphSampler(dataset.train, train_path, test_transformer,
                             model, matcher, args.batch_size,
                             args.num_instance, args.test_gal_batch,
                             args.test_prob_batch, save_path, args.gs_verbose),
        pin_memory=True)

    query_loader = DataLoader(Preprocessor(dataset.query,
                                           root=osp.join(
                                               dataset.images_dir,
                                               dataset.query_path),
                                           transform=test_transformer),
                              batch_size=args.test_fea_batch,
                              num_workers=args.workers,
                              shuffle=False,
                              pin_memory=True)

    gallery_loader = DataLoader(Preprocessor(dataset.gallery,
                                             root=osp.join(
                                                 dataset.images_dir,
                                                 dataset.gallery_path),
                                             transform=test_transformer),
                                batch_size=args.test_fea_batch,
                                num_workers=args.workers,
                                shuffle=False,
                                pin_memory=True)

    return dataset, num_classes, train_loader, query_loader, gallery_loader
示例#3
0
def main(args):
    cudnn.deterministic = False
    cudnn.benchmark = True

    exp_database_dir = osp.join(args.exp_dir, string.capwords(args.dataset))
    output_dir = osp.join(exp_database_dir, args.method, args.sub_method)
    log_file = osp.join(output_dir, 'log.txt')
    # Redirect print to both console and log file
    sys.stdout = Logger(log_file)

    # Create data loaders
    dataset, num_classes, train_loader, _, _ = \
        get_data(args.dataset, args.data_dir, args.height, args.width, args.batch_size, args.combine_all,
                 args.workers, args.test_fea_batch)

    # Create model
    model = resmap.create(args.arch,
                          ibn_type=args.ibn,
                          final_layer=args.final_layer,
                          neck=args.neck).cuda()
    num_features = model.num_features
    # print(model)
    # print('\n')

    for arg in sys.argv:
        print('%s ' % arg, end='')
    print('\n')

    # Criterion

    feamap_factor = {'layer2': 8, 'layer3': 16, 'layer4': 32}
    hei = args.height // feamap_factor[args.final_layer]
    wid = args.width // feamap_factor[args.final_layer]
    criterion = QAConvLoss(num_classes, num_features, hei, wid,
                           args.mem_batch_size).cuda()

    # Optimizer
    base_param_ids = set(map(id, model.base.parameters()))
    new_params = [p for p in model.parameters() if id(p) not in base_param_ids]
    param_groups = [{
        'params': model.base.parameters(),
        'lr': 0.1 * args.lr
    }, {
        'params': new_params,
        'lr': args.lr
    }, {
        'params': criterion.parameters(),
        'lr': args.lr
    }]

    optimizer = torch.optim.SGD(param_groups,
                                lr=args.lr,
                                momentum=0.9,
                                weight_decay=5e-4,
                                nesterov=True)

    # Decay LR by a factor of 0.1 every step_size epochs
    lr_scheduler = StepLR(optimizer, step_size=args.step_size, gamma=0.1)

    # Load from checkpoint
    start_epoch = 0

    if args.resume or args.evaluate:
        print('Loading checkpoint...')
        if args.resume and (args.resume != 'ori'):
            checkpoint = load_checkpoint(args.resume)
        else:
            checkpoint = load_checkpoint(
                osp.join(output_dir, 'checkpoint.pth.tar'))
        model.load_state_dict(checkpoint['model'])
        criterion.load_state_dict(checkpoint['criterion'])
        optimizer.load_state_dict(checkpoint['optim'])
        start_epoch = checkpoint['epoch']
        print("=> Start epoch {} ".format(start_epoch))
    else:
        pre_tr = PreTrainer(model, criterion, optimizer, train_loader,
                            args.pre_epochs, args.max_steps, args.num_trials)
        result_file = osp.join(exp_database_dir, args.method,
                               'pretrain_metric.txt')
        model, criterion, optimizer = pre_tr.train(result_file, args.method,
                                                   args.sub_method)

    model = nn.DataParallel(model).cuda()
    criterion = nn.DataParallel(criterion).cuda()

    enhance_data_aug = False

    if not args.evaluate:
        # Trainer
        trainer = Trainer(model, criterion)

        t0 = time.time()
        # Start training
        for epoch in range(start_epoch, args.epochs):
            loss, acc = trainer.train(epoch, train_loader, optimizer)

            lr = list(map(lambda group: group['lr'], optimizer.param_groups))
            lr_scheduler.step()
            train_time = time.time() - t0

            print(
                '* Finished epoch %d at lr=[%g, %g, %g]. Loss: %.3f. Acc: %.2f%%. Training time: %.0f seconds.                  \n'
                %
                (epoch + 1, lr[0], lr[1], lr[2], loss, acc * 100, train_time))

            save_checkpoint(
                {
                    'model': model.module.state_dict(),
                    'criterion': criterion.module.state_dict(),
                    'optim': optimizer.state_dict(),
                    'epoch': epoch + 1,
                },
                fpath=osp.join(output_dir, 'checkpoint.pth.tar'))

            if not enhance_data_aug and epoch < args.epochs - 1 and acc > args.acc_thr:
                enhance_data_aug = True
                print('\nAcc = %.2f%% > %.2f%%. Start to Flip and Block.\n' %
                      (acc * 100, args.acc_thr * 100))

                train_transformer = T.Compose([
                    T.RandomHorizontalFlip(),
                    T.Resize((args.height, args.width), interpolation=3),
                    T.RandomOcclusion(args.min_size, args.max_size),
                    T.ToTensor(),
                ])

                train_loader = DataLoader(Preprocessor(
                    dataset.train,
                    root=osp.join(dataset.images_dir, dataset.train_path),
                    transform=train_transformer),
                                          batch_size=args.batch_size,
                                          num_workers=args.workers,
                                          shuffle=True,
                                          pin_memory=True,
                                          drop_last=True)

    # Final test
    print('Evaluate the learned model:')
    t0 = time.time()

    # Evaluator
    evaluator = Evaluator(model)

    avg_rank1 = 0
    avg_mAP = 0
    num_testsets = 0

    test_names = args.testset.strip().split(',')
    for test_name in test_names:
        if test_name not in datasets.names():
            print('Unknown dataset: %s.' % test_name)
            continue

        testset, test_query_loader, test_gallery_loader = \
            get_test_data(test_name, args.data_dir, args.height, args.width, args.workers, args.test_fea_batch)

        test_rank1, test_mAP, test_rank1_rerank, test_mAP_rerank, test_rank1_tlift, test_mAP_tlift, test_dist, \
        test_dist_rerank, test_dist_tlift, pre_tlift_dict = \
            evaluator.evaluate(test_query_loader, test_gallery_loader, testset, criterion.module,
                               args.test_gal_batch, args.test_prob_batch,
                               args.tau, args.sigma, args.K, args.alpha)

        if test_name != args.dataset:
            avg_rank1 += test_rank1
            avg_mAP += test_mAP
            num_testsets += 1

        if testset.has_time_info:
            print(
                '  %s: rank1=%.1f, mAP=%.1f, rank1_rerank=%.1f, mAP_rerank=%.1f,'
                ' rank1_rerank_tlift=%.1f, mAP_rerank_tlift=%.1f.\n' %
                (test_name, test_rank1 * 100, test_mAP * 100,
                 test_rank1_rerank * 100, test_mAP_rerank * 100,
                 test_rank1_tlift * 100, test_mAP_tlift * 100))
        else:
            print('  %s: rank1=%.1f, mAP=%.1f.\n' %
                  (test_name, test_rank1 * 100, test_mAP * 100))

        result_file = osp.join(exp_database_dir, args.method,
                               test_name + '_results.txt')
        with open(result_file, 'a') as f:
            f.write('%s/%s:\n' % (args.method, args.sub_method))
            if testset.has_time_info:
                f.write(
                    '\t%s: rank1=%.1f, mAP=%.1f, rank1_rerank=%.1f, mAP_rerank=%.1f, rank1_rerank_tlift=%.1f, '
                    'mAP_rerank_tlift=%.1f.\n\n' %
                    (test_name, test_rank1 * 100, test_mAP * 100,
                     test_rank1_rerank * 100, test_mAP_rerank * 100,
                     test_rank1_tlift * 100, test_mAP_tlift * 100))
            else:
                f.write('\t%s: rank1=%.1f, mAP=%.1f.\n\n' %
                        (test_name, test_rank1 * 100, test_mAP * 100))

        if args.save_score:
            test_gal_list = np.array(
                [fname for fname, _, _, _ in testset.gallery], dtype=np.object)
            test_prob_list = np.array(
                [fname for fname, _, _, _ in testset.query], dtype=np.object)
            test_gal_ids = [pid for _, pid, _, _ in testset.gallery]
            test_prob_ids = [pid for _, pid, _, _ in testset.query]
            test_gal_cams = [c for _, _, c, _ in testset.gallery]
            test_prob_cams = [c for _, _, c, _ in testset.query]
            test_score_file = osp.join(exp_database_dir, args.method,
                                       args.sub_method,
                                       '%s_score.mat' % test_name)
            sio.savemat(test_score_file, {
                'score': 1. - test_dist,
                'score_rerank': 1. - test_dist_rerank,
                'score_tlift': 1. - test_dist_tlift,
                'gal_time': pre_tlift_dict['gal_time'],
                'prob_time': pre_tlift_dict['prob_time'],
                'gal_list': test_gal_list,
                'prob_list': test_prob_list,
                'gal_ids': test_gal_ids,
                'prob_ids': test_prob_ids,
                'gal_cams': test_gal_cams,
                'prob_cams': test_prob_cams
            },
                        oned_as='column',
                        do_compression=True)

    avg_rank1 /= num_testsets
    avg_mAP /= num_testsets
    result_file = osp.join(exp_database_dir, args.method, 'avg_results.txt')
    with open(result_file, 'a') as f:
        f.write('%s/%s:\n' % (args.method, args.sub_method))
        f.write('\t rank1=%.2f, mAP=%.2f.\n\n' %
                (avg_rank1 * 100, avg_mAP * 100))

    test_time = time.time() - t0
    if not args.evaluate:
        print('Finished training at epoch %d, loss %.3f, acc %.2f%%.\n' %
              (epoch + 1, loss, acc * 100))
        print(
            "Total training time: %.3f sec. Average training time per epoch: %.3f sec."
            % (train_time, train_time / (args.epochs - start_epoch + 1)))
    print("Total testing time: %.3f sec.\n" % test_time)

    for arg in sys.argv:
        print('%s ' % arg, end='')
    print('\n')