HandReader3/modern_net.py at master · shotaroikeda/HandReader3 · GitHub

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
## Preprocessing
import tflearn.datasets.mnist as mnist
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable

CUDA = False

## Constants
NUM_EPOCH = 1000
BATCH_SIZE = 1024

## Data
train_x, train_y, test_x, test_y = mnist.load_data(one_hot=True)
train_x = train_x.reshape([-1, 28, 28, 1])
test_x = test_x.reshape([-1, 28, 28, 1])

if CUDA:
    test_x = Variable(torch.Tensor(test_x).cuda().view(test_x.shape[0], 1, 28, 28))
    test_y = Variable(torch.Tensor(np.argmax(test_y, axis=1)).cuda().view(test_y.shape[0])).long()
else:
    test_x = Variable(torch.Tensor(test_x).view(test_x.shape[0], 1, 28, 28))
    test_y = Variable(torch.Tensor(np.argmax(test_y, axis=1)).view(test_y.shape[0])).long()

## Network
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16*4*4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, 16*4*4)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)


    def accuracy(self):
        result = self.forward(test_x)
        _, preds = result.max(1)
        return (preds == test_y).float().mean()


net = LeNet()

if CUDA:
    net = net.cuda()

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), lr=1e-3, momentum=0.9)

## Training
for epoch in xrange(10000):
    batch = np.random.choice(train_x.shape[0], BATCH_SIZE)
    if CUDA:
        batch_x = Variable(torch.Tensor(train_x[batch]).cuda().view(BATCH_SIZE, 1, 28, 28))
        batch_y = Variable(torch.Tensor(np.argmax(train_y[batch], axis=1)).cuda().view(BATCH_SIZE)).long()
    else:
        batch_x = Variable(torch.Tensor(train_x[batch]).view(BATCH_SIZE, 1, 28, 28))
        batch_y = Variable(torch.Tensor(np.argmax(train_y[batch], axis=1)).view(BATCH_SIZE)).long()

    optimizer.zero_grad()

    result = net.forward(batch_x)
    loss = criterion(result, batch_y)

    loss.backward()
    optimizer.step()

    if epoch % 100 == 0:
        print "Loss:"
        print loss
        print "Accuracy:"
        print net.accuracy()