def create_test_dataloaders(config, dataset_dir):
input_info = create_input_infos(config)[0]
image_size = input_info.shape[-1]
batch_size = input_info.shape[0]
normalize = transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
train_transforms = transforms.Compose([
transforms.CenterCrop(image_size),
transforms.ToTensor(),
normalize,
])
dummy_config = type('dummy', (object, ), {'dataset_dir': dataset_dir})()
train_dataset = create_cifar(dummy_config,
dataset_config='cifar10',
is_train=True,
transform=train_transforms)
# Do not set num_workers > 0 here - random hangs occur during pytest runs of this files
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
shuffle=False,
pin_memory=True,
drop_last=True)
return train_loader, train_dataset
def create_test_dataloaders(model_size, dataset_dir, batch_size):
normalize = transforms.Normalize(mean=(0.5, 0.5, 0.5),
std=(0.5, 0.5, 0.5))
train_transforms = transforms.Compose([
transforms.RandomResizedCrop(model_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
dummy_config = type('dummy', (object,), {'dataset_dir': dataset_dir})()
train_dataset = create_cifar(dummy_config, dataset_config='cifar10', is_train=True, transform=train_transforms)
pin_memory = True
workers = 1
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=False, num_workers=workers,
pin_memory=pin_memory)
return train_loader, train_dataset