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
54
55
56
57
58
59
60
61
| import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
# 数据预处理
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
# 加载数据集
train_dataset = datasets.MNIST('data/', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('data/', train=False, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)
# 定义神经网络模型
class NeuralNetwork(nn.Module):
def __init__(self):
super(NeuralNetwork, self).__init__()
self.layer1 = nn.Linear(784, 128)
self.layer2 = nn.Linear(128, 64)
self.layer3 = nn.Linear(64, 10)
def forward(self, x):
x = torch.relu(self.layer1(x))
x = torch.relu(self.layer2(x))
x = self.layer3(x)
return x
# 实例化模型
model = NeuralNetwork()
# 定义损失函数和优化器
loss_func = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
# 训练模型
epochs = 5
for epoch in range(epochs):
for batch, (images, labels) in enumerate(train_loader):
optimizer.zero_grad()
outputs = model(images.view(-1, 784))
loss = loss_func(outputs, labels)
loss.backward()
optimizer.step()
# 在测试集上评估模型
with torch.no_grad():
correct = 0
total = 0
for images, labels in test_loader:
outputs = model(images.view(-1, 784))
_, predicted = torch.max(outputs, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
accuracy = correct / total
print(f"Accuracy: {accuracy}")
|
