File size: 1,087 Bytes
3589275 | 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 | import os
import torch
import torchvision.datasets as datasets
class STL10:
def __init__(self,
preprocess,
location=os.path.expanduser('~/data'),
batch_size=128,
num_workers=16):
location = os.path.join(location, 'stl10')
self.train_dataset = datasets.STL10(
root=location,
download=True,
split='train',
transform=preprocess
)
self.train_loader = torch.utils.data.DataLoader(
self.train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers
)
self.test_dataset = datasets.STL10(
root=location,
download=True,
split='test',
transform=preprocess
)
self.test_loader = torch.utils.data.DataLoader(
self.test_dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers
)
self.classnames = self.train_dataset.classes |