dataset_mnist.py 文件源码

def __init__(self, batchsize, train=True):
        Dataset.__init__(self)
        data_root = join(dirname(realpath(__file__)), 'MNIST_data')
        self.name = "mnist"
        self.range = [0.0, 1.0]
        self.data_dims = [1, 28, 28]
        self.batchsize = batchsize
        self.data = dsets.MNIST(root=data_root,
                           download=True,
                           train=train,
                           transform=transforms.Compose([
                                transforms.ToTensor()]))
        self.dataloder = tdata.DataLoader(self.data, self.batchsize, shuffle=True)
        self.iter = iter(self.dataloder)
        self._index = 0