-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathloader.py
127 lines (103 loc) · 4.38 KB
/
loader.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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# Define the path to your dataset
celeb_data_dir = r"/teamspace/studios/this_studio/datasets/celeba_hq_root"
bridge_data_dir = r"/teamspace/studios/this_studio/datasets/lsun/bridge_root"
def get_loaders(name: str, batch_sz: int, img_sz: int):
"""
Given the name of a dataset, the batch size, and the image size, returns the corresponding data loaders for training and validation.
Parameters:
name (str): The name of the dataset. Must be one of 'celeb', 'bridge', 'cifar10', or 'fashion'.
batch_sz (int): The size of each batch.
img_sz (int): The size of each image.
Returns:
tuple: A tuple containing the data loaders for training and validation. If the dataset is 'celeb' or 'bridge', the tuple contains two data loaders. If the dataset is 'cifar10' or 'fashion', the tuple contains two data loaders for training and testing.
Raises:
ValueError: If the dataset name is not one of 'celeb', 'bridge', 'cifar10', or 'fashion'.
"""
transform = transforms.Compose(
[
transforms.Resize((img_sz, img_sz)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Lambda(lambda t: (t * 2) - 1),
]
)
assert name in [
"celeb",
"celeb_gray",
"bridge",
"cifar10",
"fashion",
], "Dataset name must be one of 'celeb', 'bridge', 'cifar10', or 'fashion'"
if name == "celeb":
dataset = datasets.ImageFolder(celeb_data_dir, transform=transform)
train, val = torch.utils.data.random_split(dataset, [0.8, 0.2])
trainloader = DataLoader(train, batch_size=batch_sz, shuffle=True)
valloader = DataLoader(val, batch_size=batch_sz, shuffle=False)
return trainloader, valloader
elif name == "celeb_gray":
custom_transform = transforms.Compose(
[
transforms.Resize((img_sz, img_sz)),
transforms.RandomHorizontalFlip(),
transforms.Grayscale(),
transforms.ToTensor(),
transforms.Lambda(lambda t: (t * 2) - 1),
]
)
dataset = datasets.ImageFolder(celeb_data_dir, transform=custom_transform)
train, val = torch.utils.data.random_split(dataset, [0.8, 0.2])
trainloader = DataLoader(train, batch_size=batch_sz, shuffle=True)
valloader = DataLoader(val, batch_size=batch_sz, shuffle=False)
return trainloader, valloader
elif name == "bridge":
dataset = datasets.ImageFolder(bridge_data_dir, transform=transform)
train, val = torch.utils.data.random_split(dataset, [0.8, 0.2])
trainloader = DataLoader(train, batch_size=batch_sz, shuffle=True)
valloader = DataLoader(val, batch_size=batch_sz, shuffle=False)
return trainloader, valloader
elif name == "cifar10":
# Download and load the training dataset
trainset = datasets.CIFAR10(
root="/teamspace/studios/this_studio/datasets",
train=True,
download=True,
transform=transform,
)
trainloader = DataLoader(
trainset, batch_size=batch_sz, shuffle=True, num_workers=2
)
# Download and load the test dataset
testset = datasets.CIFAR10(
root="/teamspace/studios/this_studio/datasets",
train=False,
download=True,
transform=transform,
)
testloader = DataLoader(
testset, batch_size=batch_sz, shuffle=False, num_workers=2
)
return trainloader, testloader
elif name == "fashion":
# Download and load the training dataset
trainset = datasets.FashionMNIST(
root="/teamspace/studios/this_studio/datasets",
train=True,
download=True,
transform=transform,
)
trainloader = DataLoader(
trainset, batch_size=batch_sz, shuffle=True, num_workers=2
)
testset = datasets.FashionMNIST(
root="/teamspace/studios/this_studio/datasets",
train=False,
download=True,
transform=transform,
)
testloader = DataLoader(
testset, batch_size=batch_sz, shuffle=False, num_workers=2
)
return trainloader, testloader