File size: 1,619 Bytes
0d94b00 6afcd6e 0d94b00 6afcd6e 0d94b00 |
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 |
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from models import Net
# Load the MNIST dataset
train_set = torchvision.datasets.MNIST(root='./data', train=True,
download=True, transform=transforms.ToTensor())
test_set = torchvision.datasets.MNIST(root='./data', train=False,
download=True, transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(train_set, batch_size=32,
shuffle=True)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=32,
shuffle=False)
net = Net()
# Use CrossEntropyLoss for multi-class classification
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)
# Train the model
for epoch in range(50): # Loop over the dataset multiple times
for i, data in enumerate(train_loader, 0):
inputs, labels = data
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
print('Finished Training')
# Test the model
correct = 0
total = 0
with torch.no_grad():
for data in test_loader:
images, labels = data
outputs = net(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
print(f'Accuracy of the network on test images: {100 * correct / total}%')
torch.save(net,'mnist.pth') |