summaryrefslogtreecommitdiff
path: root/mnist.py
diff options
context:
space:
mode:
authorAnthony Wang2021-08-25 21:01:46 -0500
committerAnthony Wang2021-08-25 21:01:46 -0500
commit57eb6ecea7695dbc9e5abbbc5695201982b00da5 (patch)
tree65735151c3475bd2acbd76f8a525b4fd2c029f7c /mnist.py
parent60e12b8eca838f9aba7f632c43ccc73a47c8ed99 (diff)
Edits to make the script actually compile and achieve 99% on MNIST
Diffstat (limited to 'mnist.py')
-rw-r--r--mnist.py23
1 files changed, 11 insertions, 12 deletions
diff --git a/mnist.py b/mnist.py
index 2b4085f..f6aba1d 100644
--- a/mnist.py
+++ b/mnist.py
@@ -6,6 +6,7 @@ from torchvision import datasets
from torchvision.transforms import ToTensor, Lambda, Compose
import matplotlib.pyplot as plt
+
training_data = datasets.MNIST(
root="data",
train=True,
@@ -20,7 +21,7 @@ test_data = datasets.MNIST(
transform=ToTensor(),
)
-batch_size = 64
+batch_size = 100
train_loader = DataLoader(training_data, batch_size=batch_size)
test_loader = DataLoader(test_data, batch_size=batch_size)
@@ -47,16 +48,15 @@ class CNN(nn.Module):
self.fc2 = nn.Linear(in_features=600, out_features=120)
self.fc3 = nn.Linear(in_features=120, out_features=10)
-
-def forward(self, x):
- out = self.layer1(x)
- out = self.layer2(out)
- out = out.view(out.size(0), -1)
- out = self.fc1(out)
- out = self.drop(out)
- out = self.fc2(out)
- out = self.fc3(out)
- return out
+ def forward(self, x):
+ out = self.layer1(x)
+ out = self.layer2(out)
+ out = out.view(out.size(0), -1)
+ out = self.fc1(out)
+ out = self.drop(out)
+ out = self.fc2(out)
+ out = self.fc3(out)
+ return out
model = CNN()
@@ -93,7 +93,6 @@ for epoch in range(num_epochs):
total = 0
correct = 0
for images, labels in test_loader:
- images, labels = images.to(device), labels.to(device)
labels_list.append(labels)
test = Variable(images.view(batch_size, 1, 28, 28))