好的,首先我们需要导入一些必要的库:
```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
```
然后我们需要定义我们的神经网络模型,这里我们使用一个简单的多层感知机(Multilayer Perceptron, MLP)模型:
```python
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(784, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, 10)
def forward(self, x):
x = x.view(-1, 784)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return F.log_softmax(x, dim=1)
```
在上面的代码中,我们定义了一个包含三个全连接层的神经网络模型,它的输入是一个大小为784的向量(28x28的图片展开成一维向量),输出是一个大小为10的向量,表示图片对应的数字的概率分布。
接下来我们需要定义一些超参数,并且加载和预处理 MNIST 数据集:
```python
batch_size = 64
test_batch_size = 1000
epochs = 10
lr = 0.01
momentum = 0.5
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dataset = datasets.MNIST('data', train=False, transform=transform)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=test_batch_size, shuffle=True)
```
在上面的代码中,我们定义了一些超参数,包括每次训练和测试的 batch size、训练轮数、学习率和动量,以及使用的设备(如果有 CUDA 设备则使用 CUDA)。我们还定义了一个数据预处理的 Pipeline,将图片转换为张量并进行归一化。然后我们使用 `datasets.MNIST` 加载 MNIST 数据集,并使用 `torch.utils.data.DataLoader` 将数据集转换成一个可迭代的数据加载器,方便我们进行批量的训练和测试。
接下来我们定义损失函数和优化器:
```python
model = Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)
criterion = nn.CrossEntropyLoss()
```
在上面的代码中,我们使用 `Net()` 初始化我们的神经网络模型,并将模型移动到指定的设备上。我们使用 `optim.SGD` 定义了一个随机梯度下降优化器,使用 `nn.CrossEntropyLoss` 定义了交叉熵损失函数。
最后,我们使用训练集进行训练,并使用测试集进行测试:
```python
for epoch in range(1, epochs + 1):
# Train
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
if batch_idx % 100 == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
# Test
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += criterion(output, target).item() * data.size(0)
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
```
在上面的代码中,我们首先进入训练模式,遍历训练集中的所有批次,并计算损失和梯度,然后使用优化器更新模型参数。我们还在每个训练轮次的前100个批次打印一次训练损失。
接着我们进入测试模式,遍历测试集中的所有批次,并计算损失和准确率。我们还将测试损失和准确率打印出来。
最后,我们训练模型并测试模型的准确率。如果你想要尝试更复杂的模型,可以自己修改模型的结构。