98 lines
3.0 KiB
Python
98 lines
3.0 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F # Add this line
|
|
import torchvision
|
|
import torchvision.transforms as transforms
|
|
|
|
|
|
# Device configuration
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
# Hyper-parameters
|
|
num_epochs = 5
|
|
batch_size = 4
|
|
learning_rate = 0.001
|
|
|
|
# Dataset has PILImage images of range [0, 1].
|
|
# We transform them to Tensors of normalized range [-1, 1]
|
|
transform = transforms.Compose([transforms.ToTensor()])
|
|
|
|
# CIFAR10: 60000 32x32 color images in 10 classes, with 6000 images per class
|
|
# train_dataset = torchvision.datasets.CIFAR10(root="./data", train=True, download=True, transform=transform)
|
|
train_dataset = torchvision.datasets.MNIST(root="./data", train=True, download=True, transform=transform)
|
|
|
|
|
|
# test_dataset = torchvision.datasets.CIFAR10(root="./data", train=False, download=True, transform=transform)
|
|
test_dataset = torchvision.datasets.MNIST(root="./data", train=False, download=True, transform=transform)
|
|
|
|
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
|
|
|
|
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
|
|
|
|
|
|
class ConvNet(nn.Module):
|
|
def __init__(self):
|
|
super(ConvNet, self).__init__()
|
|
self.conv1 = nn.Conv2d(1, 6, 3, 1, 1)
|
|
self.pool = nn.MaxPool2d(2, 2)
|
|
self.conv2 = nn.Conv2d(6, 16, 5)
|
|
self.fc1 = nn.Linear(16 * 5 * 5, 10)
|
|
# self.fc2 = nn.Linear(120, 84)
|
|
# self.fc3 = nn.Linear(84, 10)
|
|
|
|
def forward(self, x):
|
|
x = self.pool(F.relu(self.conv1(x)))
|
|
x = self.pool(F.relu(self.conv2(x)))
|
|
x = x.view(-1, 16 * 5 * 5)
|
|
# x = F.relu(self.fc1(x))
|
|
# x = F.relu(self.fc2(x))
|
|
# x = self.fc3(x)
|
|
|
|
x = self.fc1(x)
|
|
return x
|
|
|
|
|
|
model = ConvNet().to(device)
|
|
|
|
criterion = nn.CrossEntropyLoss()
|
|
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
|
|
|
|
|
|
# Train the model
|
|
n_total_steps = len(train_loader)
|
|
for epoch in range(num_epochs):
|
|
for i, (images, labels) in enumerate(train_loader):
|
|
images = images.to(device)
|
|
labels = labels.to(device)
|
|
|
|
# Forward pass
|
|
outputs = model(images)
|
|
loss = criterion(outputs, labels)
|
|
|
|
# Backward and optimize
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
if (i + 1) % 2000 == 0:
|
|
print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.4f}")
|
|
|
|
print("Finished Training")
|
|
|
|
# Test the model
|
|
with torch.no_grad():
|
|
n_correct = 0
|
|
n_samples = 0
|
|
for images, labels in test_loader:
|
|
images = images.to(device)
|
|
labels = labels.to(device)
|
|
outputs = model(images)
|
|
|
|
# max returns (value ,index)
|
|
_, predicted = torch.max(outputs.data, 1)
|
|
n_samples += labels.size(0)
|
|
n_correct += (predicted == labels).sum().item()
|
|
|
|
acc = 100.0 * n_correct / n_samples
|
|
print(f"Accuracy of the network on the 10000 test images: {acc} %")
|