File size: 998 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 | import os
import torch
import torchvision.datasets as datasets
class KMNIST:
def __init__(
self,
preprocess,
location=os.path.expanduser("~/data"),
batch_size=128,
num_workers=6,
):
location = os.path.join(location, "KMNIST")
self.train_dataset = datasets.KMNIST(
root=location, download=True, train=True, 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.KMNIST(
root=location, download=True, train=False, 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
|