-
Notifications
You must be signed in to change notification settings - Fork 1
/
data_loader.py
27 lines (22 loc) · 1.5 KB
/
data_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
import torch
from torchvision import datasets, transforms
class DataLoader:
train_loader, test_loader = None, None
def __init__(self, args):
self.use_cuda = not args.no_cuda and torch.cuda.is_available()
self.kwargs = {'num_workers': 1, 'pin_memory': True} if self.use_cuda else {}
if args.data_name == 'MNIST':
self.loadMNIST(args)
def loadMNIST(self, args):
self.train_loader = torch.utils.data.DataLoader(datasets.MNIST(args.data_dir,
train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])), batch_size=args.train_batch_size, shuffle=True, **self.kwargs)
self.test_loader = torch.utils.data.DataLoader(datasets.MNIST(args.data_dir,
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])), batch_size=args.test_batch_size, shuffle=True, **self.kwargs)