torch 示例程序

下面是一个简单的 PyTorch 示例程序，用于训练一个简单的神经网络来对手写数字进行分类：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor, Compose

# 定义神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 加载数据集
train_data = MNIST(root='data', train=True, download=True, transform=Compose([ToTensor()]))
test_data = MNIST(root='data', train=False, download=True, transform=Compose([ToTensor()]))

# 定义数据加载器
train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
test_loader = DataLoader(test_data, batch_size=128, shuffle=False)

# 定义损失函数和优化器
net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)

# 训练网络
for epoch in range(10):
    for i, (inputs, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        if i % 100 == 0:
            print(f'Epoch {epoch}, Iteration {i}, Loss: {loss.item():.4f}')

# 在测试集上评估模型性能
correct = 0
total = 0

with torch.no_grad():
    for inputs, labels in test_loader:
        outputs = net(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Test Accuracy: {100 * correct / total:.2f}%')

该程序首先定义了一个简单的神经网络，其中包含一个具有128个隐藏单元的全连接层和一个输出层。然后，它加载了 MNIST 数据集，并使用数据加载器将训练和测试数据分批加载到神经网络中。接下来，它定义了损失函数和优化器，并使用训练数据训练了神经网络。最后，它在测试集上评估了模型的性能