File size: 945 Bytes
e857f97 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | import torch
from data.build_datasets import DataInfo
from open_clip import image_transform, OPENAI_DATASET_STD, OPENAI_DATASET_MEAN, get_tokenizer
from torchvision import datasets
def get_imagenet(args, split):
assert split in ["val", "v2"]
preprocess_val = image_transform(
args.image_size,
is_train=False,
mean=OPENAI_DATASET_MEAN,
std=OPENAI_DATASET_STD,
)
if split == "v2":
from imagenetv2_pytorch import ImageNetV2Dataset
dataset = ImageNetV2Dataset(location=args.imagenet_v2, transform=preprocess_val)
else:
data_path = args.imagenet_val
assert data_path
dataset = datasets.ImageFolder(data_path, transform=preprocess_val)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
num_workers=args.workers,
sampler=None,
)
return DataInfo(dataloader=dataloader, sampler=None)
|