Skip to content

Commit 8e7dff3

Browse files
authored
Merge pull request #212 from astonzhang/nin
Add transform=None in utils.DataLoader
2 parents f111a6c + f120361 commit 8e7dff3

1 file changed

Lines changed: 6 additions & 5 deletions

File tree

utils.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ class DataLoader(object):
1616
time. But the limits are 1) all examples in dataset have the same shape, 2)
1717
data transfomer needs to process multiple examples at each time
1818
"""
19-
def __init__(self, dataset, batch_size, shuffle, transform):
19+
def __init__(self, dataset, batch_size, shuffle, transform=None):
2020
self.dataset = dataset
2121
self.batch_size = batch_size
2222
self.shuffle = shuffle
@@ -47,7 +47,7 @@ def __len__(self):
4747
def load_data_fashion_mnist(batch_size, resize=None, root="~/.mxnet/datasets/fashion-mnist"):
4848
"""download the fashion mnist dataest and then load into memory"""
4949
def transform_mnist(data, label):
50-
# transform a batch of examples
50+
# Transform a batch of examples.
5151
if resize:
5252
n = data.shape[0]
5353
new_data = nd.zeros((n, resize, resize, data.shape[3]))
@@ -56,11 +56,12 @@ def transform_mnist(data, label):
5656
data = new_data
5757
# change data from batch x height x width x channel to batch x channel x height x width
5858
return nd.transpose(data.astype('float32'), (0,3,1,2))/255, label.astype('float32')
59-
59+
6060
mnist_train = gluon.data.vision.FashionMNIST(root=root, train=True, transform=None)
6161
mnist_test = gluon.data.vision.FashionMNIST(root=root, train=False, transform=None)
62-
train_data = DataLoader(mnist_train, batch_size, shuffle=True, transform = transform_mnist)
63-
test_data = DataLoader(mnist_test, batch_size, shuffle=False, transform = transform_mnist)
62+
# Transform later to avoid memory explosion.
63+
train_data = DataLoader(mnist_train, batch_size, shuffle=True, transform=transform_mnist)
64+
test_data = DataLoader(mnist_test, batch_size, shuffle=False, transform=transform_mnist)
6465
return (train_data, test_data)
6566

6667
def try_gpu():

0 commit comments

Comments
 (0)