-
Notifications
You must be signed in to change notification settings - Fork 0
/
mnist.py
151 lines (112 loc) · 4.15 KB
/
mnist.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)
kwargs = {} # if gpu is disabled
# kwargs = {'num_workers': 1, 'pin_memory: True} # if gpu is enabled
train_data = torch.utils.data.DataLoader(
datasets.MNIST('data', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(), # first, convert image to PyTorch tensor
transforms.Normalize((0.1307,), (0.3081,)) # normalize inputs
])
),
batch_size=64,
shuffle=True,
**kwargs
)
test_data = torch.utils.data.DataLoader(
datasets.MNIST('data', train=False,
transform=transforms.Compose([
transforms.ToTensor(), # first, convert image to PyTorch tensor
transforms.Normalize((0.1307,), (0.3081,)) # normalize inputs
])
),
batch_size=64,
shuffle=True,
**kwargs
)
class Netz(nn.Module):
def __init__(self):
super(Netz, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv_dropout = nn.Dropout2d()
# 320: 20 pictures with 4x4 pixels
self.fc1 = nn.Linear(320, 60)
self.fc2 = nn.Linear(60, 10)
def forward(self, x):
x = self.conv1(x)
x = F.max_pool2d(x, 2)
x = F.relu(x)
x = self.conv2(x)
x = self.conv_dropout(x)
x = F.max_pool2d(x, 2)
x = F.relu(x)
#print(x.size())
# batch_size; each picture splitted to 20 pictures with 4x4 pixels
# torch.Size([64, 20, 4, 4])
# convert 20 pictures with 4x4 pixels to a linear vector with 320 neurons (20 x 4 x 4)
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
x = self.fc2(x)
# softmax: only on of each 10 outputs (labels) should be 1; the rest should be 0
return F.log_softmax(x, dim=1)
model = Netz()
# it's always a good idea to start with lr=0.1 and momentum=0.8
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.8)
def train(epoch):
model.train()
counter = 0
# batch_id from 0 to 63, (64 x data, 64 x target -> 0 to 9)
for batch_id, (data, target) in enumerate(train_data):
# if we do have a gpu
# data = data.cuda()
# target = target.cuda()
# convert data and target from tensors to variable
data = Variable(data)
target = Variable(target)
optimizer.zero_grad()
out = model(data)
criterion = F.nll_loss
loss = criterion(out, target)
loss.backward()
optimizer.step()
if counter % (2 * 64) == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch,
batch_id * len(data),
len(train_data.dataset),
100. * batch_id / len(train_data),
loss.item()
))
counter += 1
def test(epoch):
model.eval()
loss = 0
correct = 0
for data, target in test_data:
# if we do have a gpu
# data = data.cuda()
# target = target.cuda()
# convert data and target from tensors to variable
data = Variable(data)
target = Variable(target)
out = model(data)
loss += F.nll_loss(out, target, reduction='sum').item()
# 64 x 10 labels (the label with max number -> class)
#print(out.data)
prediction = out.data.max(1, keepdim=True)[1]
correct += prediction.eq(target.data.view_as(prediction)).sum().item()
loss = loss / len(test_data.dataset)
print('Test epoch: {} \tLoss: {:.6f}\tCorrect: {:.2f} %'.format(epoch, loss, 100. * correct / len(test_data.dataset)))
#print('Correct: ', correct)
#print('Records all: ', len(test_data.dataset))
#print('Percent Correct: {:.2f} %'.format(100. * correct / len(test_data.dataset)))
for epoch in range(1, 10):
train(epoch)
test(epoch)