神经网络分类模型代码

以下是一个使用 PyTorch 实现的简单神经网络分类模型代码示例：

import torch
import torch.nn as nn
import torch.optim as optim

# 定义神经网络模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 64)
        self.fc2 = nn.Linear(64, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = x.view(x.size(0), -1)  # 将输入数据展开成一维向量
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 加载数据
train_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('/data', train=True, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize((0.1307,), (0.3081,))
                               ])),
    batch_size=32, shuffle=True)

# 实例化模型和优化器
model = Net()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型
for epoch in range(10):
    for data, target in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = nn.CrossEntropyLoss()(output, target)
        loss.backward()
        optimizer.step()

# 测试模型
test_loader = torch.utils.data.DataLoader(
    torchvision.datasets.MNIST('/data', train=False, download=True,
                               transform=torchvision.transforms.Compose([
                                   torchvision.transforms.ToTensor(),
                                   torchvision.transforms.Normalize((0.1307,), (0.3081,))
                               ])),
    batch_size=1000, shuffle=True)

model.eval()
correct = 0
total = 0
with torch.no_grad():
    for data, target in test_loader:
        output = model(data)
        _, predicted = torch.max(output.data, 1)
        total += target.size(0)
        correct += (predicted == target).sum().item()

print('Accuracy on test set: %d %%' % (100 * correct / total))

这个代码示例定义了一个两层的全连接神经网络模型，用于将 MNIST 手写数字数据集中的图片分类成 0 到 9 的数字。在训练过程中使用了随机梯度下降优化器，并在测试集上计算了模型的准确率。