crossover-sgd/unzip_imagenet.py at master · Yeosangho/crossover-sgd · GitHub

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
33
34
35
36
37
import numpy as np
import torchvision
import torch
import torchvision.transforms as transforms

import random
def make_dataloader():
    ii64 = np.iinfo(np.int64)
    r = random.randint(0, ii64.max)
    torch.manual_seed(r)
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    train_dataset = torchvision.datasets.ImageNet(root='/scratch/x2026a02/', split='train',
                                        download=False, transform=transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]))


def make_validation_dataloader():
    ii64 = np.iinfo(np.int64)
    r = random.randint(0, ii64.max)
    torch.manual_seed(r)
    torch.cuda.manual_seed(r)
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
    val_set = torchvision.datasets.ImageNet(root='/scratch/x2026a02/', split='val', download=False, transform=transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            normalize,
        ]))
make_validation_dataloader()
make_dataloader()