def torch_loader(data_path, size):
if not os.path.exists(data_path + '/train'): download_cifar10(data_path)
# Data loading code
traindir = os.path.join(data_path, 'train')
valdir = os.path.join(data_path, 'test')
normalize = transforms.Normalize(mean=[0.4914, 0.48216, 0.44653],
std=[0.24703, 0.24349, 0.26159])
tfms = [transforms.ToTensor(), normalize]
scale_size = 40
padding = int((scale_size - size) / 2)
train_tfms = transforms.Compose([
pad, # TODO: use `padding` rather than assuming 4
transforms.RandomCrop(size),
transforms.ColorJitter(.25, .25, .25),
transforms.RandomRotation(2),
transforms.RandomHorizontalFlip(),
] + tfms)
train_dataset = datasets.ImageFolder(traindir, train_tfms)
train_sampler = (
torch.utils.data.distributed.DistributedSampler(train_dataset)
if args.distributed else None)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
val_tfms = transforms.Compose(tfms)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir, val_tfms),
batch_size=args.batch_size * 2,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
aug_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir, train_tfms),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
train_loader = DataPrefetcher(train_loader)
val_loader = DataPrefetcher(val_loader)
aug_loader = DataPrefetcher(aug_loader)
if args.prof:
train_loader.stop_after = 200
val_loader.stop_after = 0
data = ModelData(data_path, train_loader, val_loader)
data.sz = args.sz
data.aug_dl = aug_loader
if train_sampler is not None: data.trn_sampler = train_sampler
return data
def torch_loader(data_path, size):
# Data loading code
traindir = os.path.join(data_path, 'train')
valdir = os.path.join(data_path, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_tfms = transforms.Compose([
transforms.RandomResizedCrop(size),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(.4, .4, .4),
transforms.ToTensor(),
Lighting(0.1, __imagenet_pca['eigval'], __imagenet_pca['eigvec']),
normalize,
])
train_dataset = datasets.ImageFolder(traindir, train_tfms)
train_sampler = (
torch.utils.data.distributed.DistributedSampler(train_dataset)
if args.distributed else None)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
val_tfms = transforms.Compose([
transforms.Resize(int(size * 1.14)),
transforms.CenterCrop(size),
transforms.ToTensor(),
normalize,
])
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir, val_tfms),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
aug_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir, train_tfms),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
train_loader = DataPrefetcher(train_loader)
val_loader = DataPrefetcher(val_loader)
aug_loader = DataPrefetcher(aug_loader)
if args.prof:
train_loader.stop_after = 200
val_loader.stop_after = 0
data = TorchModelData(data_path, train_loader, val_loader, aug_loader)
if train_sampler is not None: data.trn_sampler = train_sampler
return data, train_sampler