Convolutional neural network

5_convolutional_net.py

@@ -0,0 +1,93 @@
+import numpy as np
+
+import torch
+from torch.autograd import Variable
+from torch import optim
+
+from data_util import load_mnist
+
+
+# Separately create two sequential here since PyTorch doesn't have nn.View()
+class ConvNet(torch.nn.Module):
+    def __init__(self, output_dim):
+        super(ConvNet, self).__init__()
+
+        self.conv = torch.nn.Sequential()
+        self.conv.add_module("conv_1", torch.nn.Conv2d(1, 10, kernel_size=5))
+        self.conv.add_module("maxpool_1", torch.nn.MaxPool2d(kernel_size=2))
+        self.conv.add_module("relu_1", torch.nn.ReLU())
+        self.conv.add_module("conv_2", torch.nn.Conv2d(10, 20, kernel_size=5))
+        self.conv.add_module("dropout_2", torch.nn.Dropout())
+        self.conv.add_module("maxpool_2", torch.nn.MaxPool2d(kernel_size=2))
+        self.conv.add_module("relu_2", torch.nn.ReLU())
+
+        self.fc = torch.nn.Sequential()
+        self.fc.add_module("fc1", torch.nn.Linear(320, 50))
+        self.fc.add_module("relu_3", torch.nn.ReLU())
+        self.fc.add_module("dropout_3", torch.nn.Dropout())
+        self.fc.add_module("fc2", torch.nn.Linear(50, output_dim))
+        self.fc.add_module("relu_4", torch.nn.ReLU())
+        self.fc.add_module("softmax", torch.nn.Softmax())
+
+    def forward(self, x):
+        x = self.conv.forward(x)
+        x = x.view(-1, 320)
+        return self.fc.forward(x)
+
+
+def train(model, loss, optimizer, x, y):
+    x = Variable(x, requires_grad=False)
+    y = Variable(y, requires_grad=False)
+
+    # Reset gradient
+    optimizer.zero_grad()
+
+    # Forward
+    fx = model.forward(x)
+    output = loss.forward(fx, y)
+
+    # Backward
+    output.backward()
+
+    # Update parameters
+    optimizer.step()
+
+    return output.data[0]
+
+
+def predict(model, x):
+    var_x = Variable(x, requires_grad=False)
+    output = model.forward(var_x)
+    return output.data.numpy().argmax(axis=1)
+
+
+def main():
+    torch.manual_seed(42)
+    trX, teX, trY, teY = load_mnist(onehot=False)
+    trX = trX.reshape(-1, 1, 28, 28)
+    teX = teX.reshape(-1, 1, 28, 28)
+
+    trX = torch.from_numpy(trX).float()
+    teX = torch.from_numpy(teX).float()
+    trY = torch.from_numpy(trY).long()
+
+    n_examples = len(trX)
+    n_classes = 10
+    model = ConvNet(output_dim=n_classes)
+    loss = torch.nn.CrossEntropyLoss(size_average=True)
+    optimizer = optim.RMSprop(model.parameters(), lr=0.001)
+    batch_size = 100
+
+    for i in range(100):
+        cost = 0.
+        num_batches = n_examples / batch_size
+        for k in range(num_batches):
+            start, end = k * batch_size, (k + 1) * batch_size
+            cost += train(model, loss, optimizer, trX[start:end], trY[start:end])
+        predY = predict(model, teX)
+        print("Epoch %d, cost = %f, acc = %.2f%%"
+              % (i + 1, cost / num_batches, 100. * np.mean(predY == teY)))
+
+
+if __name__ == "__main__":
+    main()