Edits to make the script actually compile and achieve 99% on MNIST
This commit is contained in:
parent
60e12b8eca
commit
57eb6ecea7
23
mnist.py
23
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))
|
||||
|
|
Loading…
Reference in a new issue