Example #1
0
def get_data_loader(opt, batch_size, num_workers, logger, kvstore=None):
    data_dir = opt.data_dir
    val_data_dir = opt.val_data_dir
    scale_ratios = [float(i) for i in opt.scale_ratios.split(',')]
    input_size = opt.input_size
    default_mean = [0.485, 0.456, 0.406]
    default_std = [0.229, 0.224, 0.225]

    def batch_fn(batch, ctx):
        data = split_and_load(batch[0],
                              ctx_list=ctx,
                              batch_axis=0,
                              even_split=False)
        label = split_and_load(batch[1],
                               ctx_list=ctx,
                               batch_axis=0,
                               even_split=False)
        return data, label

    if opt.data_aug == 'v1':
        # GluonCV style, not keeping aspect ratio, multi-scale crop
        transform_train = video.VideoGroupTrainTransform(
            size=(input_size, input_size),
            scale_ratios=scale_ratios,
            mean=default_mean,
            std=default_std)
        transform_test = video.VideoGroupValTransform(size=input_size,
                                                      mean=default_mean,
                                                      std=default_std)
    elif opt.data_aug == 'v2':
        # GluonCV style, keeping aspect ratio, multi-scale crop, same as mmaction style
        transform_train = video.VideoGroupTrainTransformV2(
            size=(input_size, input_size),
            short_side=opt.new_height,
            scale_ratios=scale_ratios,
            mean=default_mean,
            std=default_std)
        transform_test = video.VideoGroupValTransformV2(
            crop_size=(input_size, input_size),
            short_side=opt.new_height,
            mean=default_mean,
            std=default_std)
    elif opt.data_aug == 'v3':
        # PySlowFast style, keeping aspect ratio, random short side scale jittering
        transform_train = video.VideoGroupTrainTransformV3(
            crop_size=(input_size, input_size),
            min_size=opt.new_height,
            max_size=opt.new_width,
            mean=default_mean,
            std=default_std)
        transform_test = video.VideoGroupValTransformV2(
            crop_size=(input_size, input_size),
            short_side=opt.new_height,
            mean=default_mean,
            std=default_std)
    elif opt.data_aug == 'v4':
        # mmaction style, keeping aspect ratio, random crop and resize, only for SlowFast family models, similar to 'v3'
        transform_train = video.VideoGroupTrainTransformV4(size=(input_size,
                                                                 input_size),
                                                           mean=default_mean,
                                                           std=default_std)
        transform_test = video.VideoGroupValTransformV2(
            crop_size=(input_size, input_size),
            short_side=opt.new_height,
            mean=default_mean,
            std=default_std)
    else:
        logger.info('Data augmentation %s is not supported yet.' %
                    (opt.data_aug))

    if opt.dataset == 'kinetics400':
        train_dataset = Kinetics400(
            setting=opt.train_list,
            root=data_dir,
            train=True,
            new_width=opt.new_width,
            new_height=opt.new_height,
            new_length=opt.new_length,
            new_step=opt.new_step,
            target_width=input_size,
            target_height=input_size,
            video_loader=opt.video_loader,
            use_decord=opt.use_decord,
            slowfast=opt.slowfast,
            slow_temporal_stride=opt.slow_temporal_stride,
            fast_temporal_stride=opt.fast_temporal_stride,
            data_aug=opt.data_aug,
            num_segments=opt.num_segments,
            transform=transform_train)
        val_dataset = Kinetics400(
            setting=opt.val_list,
            root=val_data_dir,
            train=False,
            new_width=opt.new_width,
            new_height=opt.new_height,
            new_length=opt.new_length,
            new_step=opt.new_step,
            target_width=input_size,
            target_height=input_size,
            video_loader=opt.video_loader,
            use_decord=opt.use_decord,
            slowfast=opt.slowfast,
            slow_temporal_stride=opt.slow_temporal_stride,
            fast_temporal_stride=opt.fast_temporal_stride,
            data_aug=opt.data_aug,
            num_segments=opt.num_segments,
            transform=transform_test)
    elif opt.dataset == 'ucf101':
        train_dataset = UCF101(setting=opt.train_list,
                               root=data_dir,
                               train=True,
                               new_width=opt.new_width,
                               new_height=opt.new_height,
                               new_length=opt.new_length,
                               target_width=input_size,
                               target_height=input_size,
                               data_aug=opt.data_aug,
                               num_segments=opt.num_segments,
                               transform=transform_train)
        val_dataset = UCF101(setting=opt.val_list,
                             root=data_dir,
                             train=False,
                             new_width=opt.new_width,
                             new_height=opt.new_height,
                             new_length=opt.new_length,
                             target_width=input_size,
                             target_height=input_size,
                             data_aug=opt.data_aug,
                             num_segments=opt.num_segments,
                             transform=transform_test)
    elif opt.dataset == 'somethingsomethingv2':
        train_dataset = SomethingSomethingV2(setting=opt.train_list,
                                             root=data_dir,
                                             train=True,
                                             new_width=opt.new_width,
                                             new_height=opt.new_height,
                                             new_length=opt.new_length,
                                             new_step=opt.new_step,
                                             target_width=input_size,
                                             target_height=input_size,
                                             video_loader=opt.video_loader,
                                             use_decord=opt.use_decord,
                                             data_aug=opt.data_aug,
                                             num_segments=opt.num_segments,
                                             transform=transform_train)
        val_dataset = SomethingSomethingV2(setting=opt.val_list,
                                           root=data_dir,
                                           train=False,
                                           new_width=opt.new_width,
                                           new_height=opt.new_height,
                                           new_length=opt.new_length,
                                           new_step=opt.new_step,
                                           target_width=input_size,
                                           target_height=input_size,
                                           video_loader=opt.video_loader,
                                           use_decord=opt.use_decord,
                                           data_aug=opt.data_aug,
                                           num_segments=opt.num_segments,
                                           transform=transform_test)
    elif opt.dataset == 'hmdb51':
        train_dataset = HMDB51(setting=opt.train_list,
                               root=data_dir,
                               train=True,
                               new_width=opt.new_width,
                               new_height=opt.new_height,
                               new_length=opt.new_length,
                               new_step=opt.new_step,
                               target_width=input_size,
                               target_height=input_size,
                               video_loader=opt.video_loader,
                               use_decord=opt.use_decord,
                               data_aug=opt.data_aug,
                               num_segments=opt.num_segments,
                               transform=transform_train)
        val_dataset = HMDB51(setting=opt.val_list,
                             root=data_dir,
                             train=False,
                             new_width=opt.new_width,
                             new_height=opt.new_height,
                             new_length=opt.new_length,
                             new_step=opt.new_step,
                             target_width=input_size,
                             target_height=input_size,
                             video_loader=opt.video_loader,
                             use_decord=opt.use_decord,
                             data_aug=opt.data_aug,
                             num_segments=opt.num_segments,
                             transform=transform_test)
    elif opt.dataset == 'custom':
        train_dataset = VideoClsCustom(
            setting=opt.train_list,
            root=data_dir,
            train=True,
            new_width=opt.new_width,
            new_height=opt.new_height,
            new_length=opt.new_length,
            new_step=opt.new_step,
            target_width=input_size,
            target_height=input_size,
            video_loader=opt.video_loader,
            use_decord=opt.use_decord,
            slowfast=opt.slowfast,
            slow_temporal_stride=opt.slow_temporal_stride,
            fast_temporal_stride=opt.fast_temporal_stride,
            data_aug=opt.data_aug,
            num_segments=opt.num_segments,
            transform=transform_train)
        val_dataset = VideoClsCustom(
            setting=opt.val_list,
            root=val_data_dir,
            train=False,
            new_width=opt.new_width,
            new_height=opt.new_height,
            new_length=opt.new_length,
            new_step=opt.new_step,
            target_width=input_size,
            target_height=input_size,
            video_loader=opt.video_loader,
            use_decord=opt.use_decord,
            slowfast=opt.slowfast,
            slow_temporal_stride=opt.slow_temporal_stride,
            fast_temporal_stride=opt.fast_temporal_stride,
            data_aug=opt.data_aug,
            num_segments=opt.num_segments,
            transform=transform_test)
    else:
        logger.info('Dataset %s is not supported yet.' % (opt.dataset))

    logger.info('Load %d training samples and %d validation samples.' %
                (len(train_dataset), len(val_dataset)))

    if kvstore is not None:
        train_data = gluon.data.DataLoader(
            train_dataset,
            batch_size=batch_size,
            num_workers=num_workers,
            sampler=ShuffleSplitSampler(len(train_dataset),
                                        num_parts=kvstore.num_workers,
                                        part_index=kvstore.rank),
            prefetch=int(opt.prefetch_ratio * num_workers),
            last_batch='rollover')
        val_data = gluon.data.DataLoader(
            val_dataset,
            batch_size=batch_size,
            num_workers=num_workers,
            sampler=ShuffleSplitSampler(len(val_dataset),
                                        num_parts=kvstore.num_workers,
                                        part_index=kvstore.rank),
            prefetch=int(opt.prefetch_ratio * num_workers),
            last_batch='discard')
    else:
        train_data = gluon.data.DataLoader(train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True,
                                           num_workers=num_workers,
                                           prefetch=int(opt.prefetch_ratio *
                                                        num_workers),
                                           last_batch='rollover')
        val_data = gluon.data.DataLoader(val_dataset,
                                         batch_size=batch_size,
                                         shuffle=False,
                                         num_workers=num_workers,
                                         prefetch=int(opt.prefetch_ratio *
                                                      num_workers),
                                         last_batch='discard')

    return train_data, val_data, batch_fn
Example #2
0
def main(logger):
    opt = parse_args()
    print(opt)

    # Garbage collection, default threshold is (700, 10, 10).
    # Set threshold lower to collect garbage more frequently and release more CPU memory for heavy data loading.
    gc.set_threshold(100, 5, 5)

    # set env
    num_gpus = opt.num_gpus
    batch_size = opt.batch_size
    context = [mx.cpu()]
    if num_gpus > 0:
        batch_size *= max(1, num_gpus)
        context = [mx.gpu(i) for i in range(num_gpus)]

    num_workers = opt.num_workers
    print('Total batch size is set to %d on %d GPUs' % (batch_size, num_gpus))

    # get data
    default_mean = [0.485, 0.456, 0.406]
    default_std = [0.229, 0.224, 0.225]
    if opt.ten_crop:
        if opt.data_aug == 'v1':
            transform_test = transforms.Compose([
                video.VideoTenCrop(opt.input_size),
                video.VideoToTensor(),
                video.VideoNormalize(default_mean, default_std)
            ])
        else:
            transform_test = transforms.Compose([
                video.ShortSideRescale(opt.input_size),
                video.VideoTenCrop(opt.input_size),
                video.VideoToTensor(),
                video.VideoNormalize(default_mean, default_std)
            ])
        opt.num_crop = 10
    elif opt.three_crop:
        if opt.data_aug == 'v1':
            transform_test = transforms.Compose([
                video.VideoThreeCrop(opt.input_size),
                video.VideoToTensor(),
                video.VideoNormalize(default_mean, default_std)
            ])
        else:
            transform_test = transforms.Compose([
                video.ShortSideRescale(opt.input_size),
                video.VideoThreeCrop(opt.input_size),
                video.VideoToTensor(),
                video.VideoNormalize(default_mean, default_std)
            ])
        opt.num_crop = 3
    else:
        if opt.data_aug == 'v1':
            transform_test = video.VideoGroupValTransform(size=opt.input_size,
                                                          mean=default_mean,
                                                          std=default_std)
        else:
            transform_test = video.VideoGroupValTransformV2(
                crop_size=(input_size, input_size),
                short_side=input_size,
                mean=default_mean,
                std=default_std)
        opt.num_crop = 1

    if not opt.deploy:
        # get model
        if opt.use_pretrained and len(opt.hashtag) > 0:
            opt.use_pretrained = opt.hashtag
        classes = opt.num_classes
        model_name = opt.model
        # Currently, these is no hashtag for int8 models.
        if opt.quantized:
            model_name += '_int8'
            opt.use_pretrained = True

        net = get_model(name=model_name,
                        nclass=classes,
                        pretrained=opt.use_pretrained,
                        num_segments=opt.num_segments,
                        num_crop=opt.num_crop)
        net.cast(opt.dtype)
        net.collect_params().reset_ctx(context)
        if opt.mode == 'hybrid':
            net.hybridize(static_alloc=True, static_shape=True)
        if opt.resume_params is not '' and not opt.use_pretrained:
            net.load_parameters(opt.resume_params, ctx=context)
            print('Pre-trained model %s is successfully loaded.' %
                  (opt.resume_params))
        else:
            print(
                'Pre-trained model is successfully loaded from the model zoo.')
    else:
        model_name = 'deploy'
        net = mx.gluon.SymbolBlock.imports(
            '{}-symbol.json'.format(opt.model_prefix), ['data'],
            '{}-0000.params'.format(opt.model_prefix))
        net.hybridize(static_alloc=True, static_shape=True)

    print("Successfully loaded model {}".format(model_name))
    # dummy data for benchmarking performance
    if opt.benchmark:
        benchmarking(opt, net, context)
        sys.exit()

    if opt.dataset == 'ucf101':
        val_dataset = UCF101(setting=opt.val_list,
                             root=opt.data_dir,
                             train=False,
                             new_width=opt.new_width,
                             new_height=opt.new_height,
                             new_length=opt.new_length,
                             target_width=opt.input_size,
                             target_height=opt.input_size,
                             test_mode=True,
                             data_aug=opt.data_aug,
                             num_segments=opt.num_segments,
                             transform=transform_test)
    elif opt.dataset == 'kinetics400':
        val_dataset = Kinetics400(
            setting=opt.val_list,
            root=opt.data_dir,
            train=False,
            new_width=opt.new_width,
            new_height=opt.new_height,
            new_length=opt.new_length,
            new_step=opt.new_step,
            target_width=opt.input_size,
            target_height=opt.input_size,
            video_loader=opt.video_loader,
            use_decord=opt.use_decord,
            slowfast=opt.slowfast,
            slow_temporal_stride=opt.slow_temporal_stride,
            fast_temporal_stride=opt.fast_temporal_stride,
            test_mode=True,
            data_aug=opt.data_aug,
            num_segments=opt.num_segments,
            num_crop=opt.num_crop,
            transform=transform_test)
    elif opt.dataset == 'somethingsomethingv2':
        val_dataset = SomethingSomethingV2(setting=opt.val_list,
                                           root=opt.data_dir,
                                           train=False,
                                           new_width=opt.new_width,
                                           new_height=opt.new_height,
                                           new_length=opt.new_length,
                                           new_step=opt.new_step,
                                           target_width=opt.input_size,
                                           target_height=opt.input_size,
                                           video_loader=opt.video_loader,
                                           use_decord=opt.use_decord,
                                           data_aug=opt.data_aug,
                                           num_segments=opt.num_segments,
                                           transform=transform_test)
    elif opt.dataset == 'hmdb51':
        val_dataset = HMDB51(setting=opt.val_list,
                             root=opt.data_dir,
                             train=False,
                             new_width=opt.new_width,
                             new_height=opt.new_height,
                             new_length=opt.new_length,
                             new_step=opt.new_step,
                             target_width=opt.input_size,
                             target_height=opt.input_size,
                             video_loader=opt.video_loader,
                             use_decord=opt.use_decord,
                             data_aug=opt.data_aug,
                             num_segments=opt.num_segments,
                             transform=transform_test)
    else:
        logger.info('Dataset %s is not supported yet.' % (opt.dataset))

    val_data = gluon.data.DataLoader(val_dataset,
                                     batch_size=batch_size,
                                     shuffle=False,
                                     num_workers=num_workers,
                                     prefetch=int(opt.prefetch_ratio *
                                                  num_workers),
                                     last_batch='discard')
    print('Load %d test samples in %d iterations.' %
          (len(val_dataset), len(val_data)))

    # calibrate FP32 model into INT8 model
    if opt.calibration:
        calibration(net, val_data, opt, context, logger)
        sys.exit()

    start_time = time.time()
    acc_top1_val, acc_top5_val = test(context, val_data, opt, net)
    end_time = time.time()

    print('Test accuracy: acc-top1=%f acc-top5=%f' %
          (acc_top1_val * 100, acc_top5_val * 100))
    print('Total evaluation time is %4.2f minutes' %
          ((end_time - start_time) / 60))