Merge pull request #12 from YJH-666/master
mnist dataset download setting
This commit is contained in:
@ -7,6 +7,11 @@ torch: 0.1.11
|
|||||||
torchvision
|
torchvision
|
||||||
matplotlib
|
matplotlib
|
||||||
"""
|
"""
|
||||||
|
# library
|
||||||
|
# standard library
|
||||||
|
import os
|
||||||
|
|
||||||
|
# third-party library
|
||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
from torch.autograd import Variable
|
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
|
EPOCH = 1 # train the training data n times, to save time, we just train 1 epoch
|
||||||
BATCH_SIZE = 50
|
BATCH_SIZE = 50
|
||||||
LR = 0.001 # learning rate
|
LR = 0.001 # learning rate
|
||||||
DOWNLOAD_MNIST = True # set to False if you have downloaded
|
DOWNLOAD_MNIST = False
|
||||||
|
|
||||||
|
|
||||||
# Mnist digits dataset
|
# 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(
|
train_data = torchvision.datasets.MNIST(
|
||||||
root='./mnist/',
|
root='./mnist/',
|
||||||
train=True, # this is training data
|
train=True, # this is training data
|
||||||
transform=torchvision.transforms.ToTensor(), # Converts a PIL.Image or numpy.ndarray to
|
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]
|
# 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
|
# plot one example
|
||||||
|
|||||||
Reference in New Issue
Block a user