def get_train_data(rec_train, batch_size, data_nthreads, input_size,
crop_ratio, args):
def train_batch_fn(batch, ctx):
data = batch[0].as_in_context(ctx)
label = batch[1].as_in_context(ctx)
return data, label
jitter_param = 0.4
lighting_param = 0.1
resize = int(math.ceil(input_size / crop_ratio))
train_transforms = []
if args.auto_aug:
print('Using AutoAugment')
from autogluon.utils.augment import AugmentationBlock, autoaug_imagenet_policies
train_transforms.append(AugmentationBlock(autoaug_imagenet_policies()))
from gluoncv.utils.transforms import EfficientNetRandomCrop
from autogluon.utils import pil_transforms
if input_size >= 320:
train_transforms.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:
train_transforms.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(train_transforms)
train_set = mx.gluon.data.vision.ImageRecordDataset(
rec_train).transform_first(transform_train)
train_sampler = SplitSampler(len(train_set),
num_parts=num_workers,
part_index=rank)
train_data = gluon.data.DataLoader(
train_set,
batch_size=batch_size, # shuffle=True,
last_batch='discard',
num_workers=data_nthreads,
sampler=train_sampler)
return train_data, train_batch_fn
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