mnist dataset download setting

This commit is contained in:
Kawa Yg
2017-11-22 08:13:22 +08:00
parent 2ae0c0fa6a
commit 895ad9a26c

View File

@ -7,6 +7,11 @@ torch: 0.1.11
torchvision
matplotlib
"""
# library
# standard library
import os
# third-party library
import torch
import torch.nn as nn
from torch.autograd import Variable
@ -20,16 +25,20 @@ torch.manual_seed(1) # reproducible
EPOCH = 1 # train the training data n times, to save time, we just train 1 epoch
BATCH_SIZE = 50
LR = 0.001 # learning rate
DOWNLOAD_MNIST = True # set to False if you have downloaded
DOWNLOAD_MNIST = False
# Mnist digits dataset
if not(os.path.exists('./mnist/')) or not os.listdir('./mnist/'):
# not mnist dir or mnist is empyt dir
DOWNLOAD_MNIST = True
train_data = torchvision.datasets.MNIST(
root='./mnist/',
train=True, # this is training data
transform=torchvision.transforms.ToTensor(), # Converts a PIL.Image or numpy.ndarray to
# torch.FloatTensor of shape (C x H x W) and normalize in the range [0.0, 1.0]
download=DOWNLOAD_MNIST, # download it if you don't have it
download=DOWNLOAD_MNIST,
)
# plot one example