from torch import nn, optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(dataset=train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=32, shuffle=False)
class SimpleNN(nn.Module):
super(SimpleNN, self).__init__()
self.fc1 = nn.Linear(28*28, 128)
self.fc2 = nn.Linear(128, 64)
self.fc3 = nn.Linear(64, 10)
x = x.view(-1, 28*28) # 将图像展平成一维向量
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
def train(model, train_loader, criterion, optimizer, epochs=10):
for epoch in range(epochs):
for images, labels in train_loader:
loss = criterion(outputs, labels)
running_loss += loss.item()
print(f'Epoch [{epoch+1}/{epochs}], Loss: {running_loss/len(train_loader):.4f}')
def evaluate(model, test_loader):
for images, labels in test_loader:
_, predicted = torch.max(outputs.data, 1)
correct += (predicted == labels).sum().item()
accuracy = 100 * correct / total
print(f'Test Accuracy: {accuracy:.2f}%')
train(model, train_loader, criterion, optimizer, epochs=10)
evaluate(model, test_loader)