Skip to content

Commit d2acdb9

Browse files
authored
Merge pull request #18 from pytorch/mnistcleanup
cleanup mnist
2 parents 4d7c05c + 031392f commit d2acdb9

File tree

5 files changed

+72
-65
lines changed

5 files changed

+72
-65
lines changed

mnist/README.md

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,5 @@
22

33
```bash
44
pip install -r requirements.txt
5-
python data.py
65
python main.py
76
```

mnist/main.py

Lines changed: 64 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -1,52 +1,64 @@
11
from __future__ import print_function
2-
import os
2+
import os, argparse
33
import torch
44
import torch.nn as nn
55
import torch.optim as optim
66
from torch.autograd import Variable
77

88
cuda = torch.cuda.is_available()
99

10-
def print_header(msg):
11-
print('===>', msg)
10+
# Training settings
11+
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
12+
parser.add_argument('--batchSize', type=int, default=64, metavar='input batch size')
13+
parser.add_argument('--testBatchSize', type=int, default=1000, metavar='input batch size for testing')
14+
parser.add_argument('--trainSize', type=int, default=1000, metavar='Train dataset size (max=60000). Default: 1000')
15+
parser.add_argument('--nEpochs', type=int, default=2, metavar='number of epochs to train')
16+
parser.add_argument('--lr', type=float, default=0.01, metavar='Learning Rate. Default=0.01')
17+
parser.add_argument('--momentum', type=float, default=0.5, metavar='Default=0.5')
18+
parser.add_argument('--seed', type=int, default=123, metavar='Random Seed to use. Default=123')
19+
opt = parser.parse_args()
20+
print(opt)
21+
22+
torch.manual_seed(opt.seed)
23+
if cuda == True:
24+
torch.cuda.manual_seed(opt.seed)
1225

1326
if not os.path.exists('data/processed/training.pt'):
1427
import data
1528

1629
# Data
17-
print_header('Loading data')
30+
print('===> Loading data')
1831
with open('data/processed/training.pt', 'rb') as f:
1932
training_set = torch.load(f)
2033
with open('data/processed/test.pt', 'rb') as f:
2134
test_set = torch.load(f)
2235

2336
training_data = training_set[0].view(-1, 1, 28, 28).div(255)
37+
training_data = training_data[:opt.trainSize]
2438
training_labels = training_set[1]
2539
test_data = test_set[0].view(-1, 1, 28, 28).div(255)
2640
test_labels = test_set[1]
2741

2842
del training_set
2943
del test_set
3044

31-
# Model
32-
print_header('Building model')
45+
print('===> Building model')
3346
class Net(nn.Container):
3447
def __init__(self):
35-
super(Net, self).__init__(
36-
conv1 = nn.Conv2d(1, 20, 5),
37-
pool1 = nn.MaxPool2d(2, 2),
38-
conv2 = nn.Conv2d(20, 50, 5),
39-
pool2 = nn.MaxPool2d(2, 2),
40-
fc1 = nn.Linear(800, 500),
41-
fc2 = nn.Linear(500, 10),
42-
relu = nn.ReLU(),
43-
softmax = nn.LogSoftmax(),
44-
)
48+
super(Net, self).__init__()
49+
self.conv1 = nn.Conv2d(1, 10, 5)
50+
self.pool1 = nn.MaxPool2d(2,2)
51+
self.conv2 = nn.Conv2d(10, 20, 5)
52+
self.pool2 = nn.MaxPool2d(2, 2)
53+
self.fc1 = nn.Linear(320, 50)
54+
self.fc2 = nn.Linear(50, 10)
55+
self.relu = nn.ReLU()
56+
self.softmax = nn.LogSoftmax()
4557

4658
def forward(self, x):
4759
x = self.relu(self.pool1(self.conv1(x)))
4860
x = self.relu(self.pool2(self.conv2(x)))
49-
x = x.view(-1, 800)
61+
x = x.view(-1, 320)
5062
x = self.relu(self.fc1(x))
5163
x = self.relu(self.fc2(x))
5264
return self.softmax(x)
@@ -56,60 +68,59 @@ def forward(self, x):
5668
model.cuda()
5769

5870
criterion = nn.NLLLoss()
59-
60-
# Training settings
61-
BATCH_SIZE = 150
62-
TEST_BATCH_SIZE = 1000
63-
NUM_EPOCHS = 2
64-
65-
optimizer = optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)
71+
optimizer = optim.SGD(model.parameters(), lr=opt.lr, momentum=opt.momentum)
6672

6773
def train(epoch):
68-
batch_data_t = torch.FloatTensor(BATCH_SIZE, 1, 28, 28)
69-
batch_targets_t = torch.LongTensor(BATCH_SIZE)
74+
# create buffers for mini-batch
75+
batch_data = torch.FloatTensor(opt.batchSize, 1, 28, 28)
76+
batch_targets = torch.LongTensor(opt.batchSize)
7077
if cuda:
71-
batch_data_t = batch_data_t.cuda()
72-
batch_targets_t = batch_targets_t.cuda()
73-
batch_data = Variable(batch_data_t, requires_grad=False)
74-
batch_targets = Variable(batch_targets_t, requires_grad=False)
75-
for i in range(0, training_data.size(0), BATCH_SIZE):
78+
batch_data, batch_targets = batch_data.cuda(), batch_targets.cuda()
79+
80+
# create autograd Variables over these buffers
81+
batch_data, batch_targets = Variable(batch_data), Variable(batch_targets)
82+
83+
for i in range(0, training_data.size(0)-opt.batchSize+1, opt.batchSize):
84+
start, end = i, i+opt.batchSize
7685
optimizer.zero_grad()
77-
batch_data.data[:] = training_data[i:i+BATCH_SIZE]
78-
batch_targets.data[:] = training_labels[i:i+BATCH_SIZE]
79-
loss = criterion(model(batch_data), batch_targets)
86+
batch_data.data[:] = training_data[start:end]
87+
batch_targets.data[:] = training_labels[start:end]
88+
output = model(batch_data)
89+
loss = criterion(output, batch_targets)
8090
loss.backward()
8191
loss = loss.data[0]
8292
optimizer.step()
83-
print('Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.4f}'.format(epoch,
84-
i+BATCH_SIZE, training_data.size(0),
85-
float(i+BATCH_SIZE)/training_data.size(0)*100, loss))
93+
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.4f}'
94+
.format(epoch, end, opt.trainSize, float(end)/opt.trainSize*100, loss))
8695

8796
def test(epoch):
88-
test_loss = 0
89-
batch_data_t = torch.FloatTensor(TEST_BATCH_SIZE, 1, 28, 28)
90-
batch_targets_t = torch.LongTensor(TEST_BATCH_SIZE)
97+
# create buffers for mini-batch
98+
batch_data = torch.FloatTensor(opt.testBatchSize, 1, 28, 28)
99+
batch_targets = torch.LongTensor(opt.testBatchSize)
91100
if cuda:
92-
batch_data_t = batch_data_t.cuda()
93-
batch_targets_t = batch_targets_t.cuda()
94-
batch_data = Variable(batch_data_t, volatile=True)
95-
batch_targets = Variable(batch_targets_t, volatile=True)
101+
batch_data, batch_targets = batch_data.cuda(), batch_targets.cuda()
102+
103+
# create autograd Variables over these buffers
104+
batch_data = Variable(batch_data, volatile=True)
105+
batch_targets = Variable(batch_targets, volatile=True)
106+
107+
test_loss = 0
96108
correct = 0
97-
for i in range(0, test_data.size(0), TEST_BATCH_SIZE):
98-
print('Testing model: {}/{}'.format(i, test_data.size(0)), end='\r')
99-
batch_data.data[:] = test_data[i:i+TEST_BATCH_SIZE]
100-
batch_targets.data[:] = test_labels[i:i+TEST_BATCH_SIZE]
109+
110+
for i in range(0, test_data.size(0), opt.testBatchSize):
111+
batch_data.data[:] = test_data[i:i+opt.testBatchSize]
112+
batch_targets.data[:] = test_labels[i:i+opt.testBatchSize]
101113
output = model(batch_data)
102114
test_loss += criterion(output, batch_targets)
103-
pred = output.data.max(1)[1]
115+
pred = output.data.max(1)[1] # get the index of the max log-probability
104116
correct += pred.long().eq(batch_targets.data.long()).cpu().sum()
105117

106118
test_loss = test_loss.data[0]
107-
test_loss /= (test_data.size(0) / TEST_BATCH_SIZE) # criterion averages over batch size
108-
print('TEST SET RESULTS:' + ' ' * 20)
109-
print('Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)'.format(
119+
test_loss /= (test_data.size(0) / opt.testBatchSize) # criterion averages over batch size
120+
print('\nTest Set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
110121
test_loss, correct, test_data.size(0),
111122
float(correct)/test_data.size(0)*100))
112123

113-
for epoch in range(1, NUM_EPOCHS+1):
124+
for epoch in range(1, opt.nEpochs+1):
114125
train(epoch)
115126
test(epoch)

mnist/requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
torch
22
six
3+
tqdm

word_language_model/main.py

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,18 +5,15 @@
55
# test set.
66
###############################################################################
77

8-
import argparse
9-
import time
10-
import math
11-
8+
import argparse, time, math
129
import torch
1310
import torch.nn as nn
1411
from torch.autograd import Variable
1512

1613
import data
1714
import model
1815

19-
parser = argparse.ArgumentParser(description='PyTorch PTB Language Model')
16+
parser = argparse.ArgumentParser(description='PyTorch PennTreeBank RNN/LSTM Language Model')
2017

2118
# Data parameters
2219
parser.add_argument('-data' , type=str, default='./data/penn', help='Location of the data corpus' )
@@ -41,8 +38,7 @@
4138

4239
# Set the random seed manually for reproducibility.
4340
torch.manual_seed(args.seed)
44-
# If the GPU is enabled, do some plumbing.
45-
41+
# If the GPU is enabled, warn the user to use it
4642
if torch.cuda.is_available() and not args.cuda:
4743
print("WARNING: You have a CUDA device, so you should probably run with -cuda")
4844

word_language_model/model.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,10 @@ class RNNModel(nn.Container):
77
and a decoder. Runs one RNN step at a time.
88
"""
99

10-
def __init__(self, rnnType, ntoken, ninp, nhid, nlayers):
10+
def __init__(self, rnn_type, ntoken, ninp, nhid, nlayers):
1111
super(RNNModel, self).__init__(
1212
encoder = nn.sparse.Embedding(ntoken, ninp),
13-
rnn = nn.RNNBase(rnnType, ninp, nhid, nlayers, bias=False),
13+
rnn = nn.RNNBase(rnn_type, ninp, nhid, nlayers, bias=False),
1414
decoder = nn.Linear(nhid, ntoken),
1515
)
1616

@@ -21,7 +21,7 @@ def __init__(self, rnnType, ntoken, ninp, nhid, nlayers):
2121
self.decoder.bias.data.fill_(0)
2222
self.decoder.weight.data.uniform_(-initrange, initrange)
2323

24-
self.rnnType = rnnType
24+
self.rnn_type = rnn_type
2525
self.nhid = nhid
2626
self.nlayers = nlayers
2727

@@ -33,7 +33,7 @@ def forward(self, input, hidden):
3333

3434
def initHidden(self, bsz):
3535
weight = next(self.parameters()).data
36-
if self.rnnType == 'LSTM':
36+
if self.rnn_type == 'LSTM':
3737
return (Variable(weight.new(self.nlayers, bsz, self.nhid).zero_()),
3838
Variable(weight.new(self.nlayers, bsz, self.nhid).zero_()))
3939
else:

0 commit comments

Comments
 (0)