forked from kuc2477/pytorch-splitnet
-
Notifications
You must be signed in to change notification settings - Fork 0
/
data.py
67 lines (57 loc) · 1.76 KB
/
data.py
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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
from torchvision import datasets, transforms
_MNIST_TRAIN_TRANSFORMS = _MNIST_TEST_TRANSFORMS = [
transforms.ToTensor(),
transforms.ToPILImage(),
transforms.Pad(2),
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
]
_CIFAR_TRAIN_TRANSFORMS = [
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(
(0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)
),
]
_CIFAR_TEST_TRANSFORMS = [
transforms.ToTensor(),
transforms.Normalize(
(0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)
),
]
TRAIN_DATASETS = {
'mnist': datasets.MNIST(
'./datasets/mnist', train=True, download=True,
transform=transforms.Compose(_MNIST_TRAIN_TRANSFORMS)
),
'cifar10': datasets.CIFAR10(
'./datasets/cifar10', train=True, download=True,
transform=transforms.Compose(_CIFAR_TRAIN_TRANSFORMS)
),
'cifar100': datasets.CIFAR100(
'./datasets/cifar100', train=True, download=True,
transform=transforms.Compose(_CIFAR_TRAIN_TRANSFORMS)
)
}
TEST_DATASETS = {
'mnist': datasets.MNIST(
'./datasets/mnist', train=False,
transform=transforms.Compose(_MNIST_TEST_TRANSFORMS)
),
'cifar10': datasets.CIFAR10(
'./datasets/cifar10', train=False,
transform=transforms.Compose(_CIFAR_TEST_TRANSFORMS)
),
'cifar100': datasets.CIFAR100(
'./datasets/cifar100', train=False,
transform=transforms.Compose(_CIFAR_TEST_TRANSFORMS)
)
}
DATASET_CONFIGS = {
'mnist': {'size': 32, 'channels': 1, 'classes': 10},
'cifar10': {'size': 32, 'channels': 3, 'classes': 10},
'cifar100': {'size': 32, 'channels': 3, 'classes': 100},
}