/
mnist_class_28.py
121 lines (93 loc) · 4.14 KB
/
mnist_class_28.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
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchbearer
import torchvision
from torchbearer import Trial, callbacks
from torchvision import transforms
from memory import Memory
class Block(nn.Module):
def __init__(self, in_planes, out_planes, stride=1, padding=0):
super(Block, self).__init__()
self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=3, padding=padding, stride=stride, bias=False)
self.bn = nn.BatchNorm2d(out_planes)
torch.nn.init.xavier_uniform_(self.conv.weight)
def forward(self, x):
out = F.relu(self.bn(self.conv(x)))
return out
class ContextNet(nn.Module):
def __init__(self):
super(ContextNet, self).__init__()
self.conv1 = Block(1, 64, stride=2)
self.conv2 = Block(64, 128, stride=2)
self.conv3 = Block(128, 256, stride=2)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = x.view(x.size(0), -1)
return x
class GlimpseNet(nn.Module):
def __init__(self):
super(GlimpseNet, self).__init__()
self.conv1 = Block(1, 64, stride=2)
self.conv2 = Block(64, 128, stride=2)
self.conv3 = Block(128, 256, stride=2)
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = self.conv3(x)
x = x.view(x.size(0), -1)
return x
class MnistClassifier(nn.Module):
def __init__(self, count, memory_size):
super(MnistClassifier, self).__init__()
self.memory = Memory(
hidden_size=memory_size * 2,
memory_size=memory_size,
glimpse_size=28,
g_down=1024,
c_down=1024,
context_net=ContextNet(),
glimpse_net=GlimpseNet()
)
self.count = count
self.drop = nn.Dropout(0.5)
self.qdown = nn.Linear(1024, memory_size)
self.classifier = nn.Linear(memory_size, 10)
self.soft = nn.LogSoftmax(dim=1)
def forward(self, x, state=None):
image = x
x, context = self.memory.init(image)
query = F.relu6(self.drop(self.qdown(context.detach())))
for i in range(self.count):
x = self.memory.glimpse(x, image)
myp = self.memory(query)
return self.soft(self.classifier(myp))
def run(count, memory_size, device='cuda'):
traintransform = transforms.Compose([transforms.RandomRotation(20), transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
trainset = torchvision.datasets.MNIST(root='./data/mnist', train=True,
download=True, transform=traintransform)
trainloader = torch.utils.data.DataLoader(trainset, pin_memory=True, batch_size=128,
shuffle=True, num_workers=10)
testtransform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
testset = torchvision.datasets.MNIST(root='./data/mnist', train=False,
download=True, transform=testtransform)
testloader = torch.utils.data.DataLoader(testset, pin_memory=True, batch_size=128,
shuffle=False, num_workers=10)
base_dir = os.path.join('mnist_' + str(memory_size), str(count))
model = MnistClassifier(count, memory_size)
optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001)
trial = Trial(model, optimizer, nn.NLLLoss(), ['acc', 'loss'], callbacks=[
callbacks.MostRecent(os.path.join(base_dir, '{epoch:02d}.pt')),
callbacks.GradientClipping(5),
callbacks.MultiStepLR(milestones=[50, 100, 150, 190, 195]),
callbacks.ExponentialLR(0.99),
callbacks.TensorBoard(write_graph=False, comment=base_dir)
]).with_train_generator(trainloader).to(device)
trial.run(200)
trial.with_test_generator(testloader).evaluate(data_key=torchbearer.TEST_DATA)
if __name__ == "__main__":
run(8, 512)