Пример #1
0
def get_val_data(rec_val, batch_size, data_nthreads, input_size, crop_ratio):
    def val_batch_fn(batch, ctx):
        data = batch[0].as_in_context(ctx)
        label = batch[1].as_in_context(ctx)
        return data, label

    normalize = transforms.Normalize([0.485, 0.456, 0.406],
                                     [0.229, 0.224, 0.225])
    crop_ratio = crop_ratio if crop_ratio > 0 else 0.875
    resize = int(math.ceil(input_size / crop_ratio))

    from gluoncv.utils.transforms import EfficientNetCenterCrop
    from autogluon.utils import pil_transforms

    if input_size >= 320:
        transform_test = transforms.Compose([
            pil_transforms.ToPIL(),
            EfficientNetCenterCrop(input_size),
            pil_transforms.Resize((input_size, input_size),
                                  interpolation=Image.BICUBIC),
            pil_transforms.ToNDArray(),
            transforms.ToTensor(), normalize
        ])
    else:
        transform_test = transforms.Compose([
            transforms.Resize(resize, keep_ratio=True),
            transforms.CenterCrop(input_size),
            transforms.ToTensor(), normalize
        ])

    val_set = mx.gluon.data.vision.ImageRecordDataset(rec_val).transform_first(
        transform_test)

    val_sampler = SplitSampler(len(val_set),
                               num_parts=num_workers,
                               part_index=rank)
    val_data = gluon.data.DataLoader(val_set,
                                     batch_size=batch_size,
                                     num_workers=data_nthreads,
                                     sampler=val_sampler)

    return val_data, val_batch_fn
Пример #2
0
    def get_data_loader(data_dir, batch_size, num_workers):
        normalize = transforms.Normalize([0.485, 0.456, 0.406],
                                         [0.229, 0.224, 0.225])
        jitter_param = 0.4
        lighting_param = 0.1
        input_size = opt.input_size
        crop_ratio = opt.crop_ratio if opt.crop_ratio > 0 else 0.875
        resize = int(math.ceil(input_size / crop_ratio))

        def batch_fn(batch, ctx):
            data = gluon.utils.split_and_load(batch[0],
                                              ctx_list=ctx,
                                              batch_axis=0)
            label = gluon.utils.split_and_load(batch[1],
                                               ctx_list=ctx,
                                               batch_axis=0)
            return data, label

        transform_train = []
        if opt.auto_aug:
            print('Using AutoAugment')
            from autogluon.utils.augment import AugmentationBlock, autoaug_imagenet_policies
            transform_train.append(
                AugmentationBlock(autoaug_imagenet_policies()))

        from gluoncv.utils.transforms import EfficientNetRandomCrop, EfficientNetCenterCrop
        from autogluon.utils import pil_transforms

        if input_size >= 320:
            transform_train.extend([
                EfficientNetRandomCrop(input_size),
                pil_transforms.Resize((input_size, input_size),
                                      interpolation=Image.BICUBIC),
                pil_transforms.RandomHorizontalFlip(),
                pil_transforms.ColorJitter(brightness=0.4,
                                           contrast=0.4,
                                           saturation=0.4),
                transforms.RandomLighting(lighting_param),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406],
                                     [0.229, 0.224, 0.225])
            ])
        else:
            transform_train.extend([
                transforms.RandomResizedCrop(input_size),
                transforms.RandomFlipLeftRight(),
                transforms.RandomColorJitter(brightness=jitter_param,
                                             contrast=jitter_param,
                                             saturation=jitter_param),
                transforms.RandomLighting(lighting_param),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406],
                                     [0.229, 0.224, 0.225])
            ])

        transform_train = transforms.Compose(transform_train)

        train_data = gluon.data.DataLoader(imagenet.classification.ImageNet(
            data_dir, train=True).transform_first(transform_train),
                                           batch_size=batch_size,
                                           shuffle=True,
                                           last_batch='discard',
                                           num_workers=num_workers)

        if input_size >= 320:
            transform_test = transforms.Compose([
                pil_transforms.ToPIL(),
                EfficientNetCenterCrop(input_size),
                pil_transforms.Resize((input_size, input_size),
                                      interpolation=Image.BICUBIC),
                pil_transforms.ToNDArray(),
                transforms.ToTensor(), normalize
            ])
        else:
            transform_test = transforms.Compose([
                transforms.Resize(resize, keep_ratio=True),
                transforms.CenterCrop(input_size),
                transforms.ToTensor(), normalize
            ])

        val_data = gluon.data.DataLoader(imagenet.classification.ImageNet(
            data_dir, train=False).transform_first(transform_test),
                                         batch_size=batch_size,
                                         shuffle=False,
                                         num_workers=num_workers)

        return train_data, val_data, batch_fn