| import torch
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)
# 设定模型和优化器
model = nn.Sequential(nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10), nn.LogSoftmax(dim=1))
optimizer = optim.SGD(model.parameters(), lr=0.01)
# 设定实验参数
batch_sizes = [32, 64, 128, 256, 512]
num_epochs = 10
validate_freq = 100
# 进行实验
results = []
for batch_size in batch_sizes:
train_accuracy = 0.0
test_accuracy = 0.0
for epoch in range(num_epochs):
for images, labels in train_loader:
images = images.view(images.shape[0], -1)
logps = model(images)
loss = -torch.sum(logps) / images.shape[0]
optimizer.zero_grad()
loss.backward()
optimizer.step()
for images, labels in test_loader:
images = images.view(images.shape[0], -1)
logps = model(images)
ps = torch.exp(logps)
topk = torch.topk(ps, k=1)
labels = labels.view(-1)
corrects = (topk[1] == labels).sum().item() / labels.shape[0]
accuracy = corrects / labels.shape[0] * 100.0
train_accuracy += accuracy / num_epochs * batch_size * validate_freq / len(train_loader) * 1000000 / (1 + epoch) / 1e6 # 计算训练精度（百万次）
test_accuracy += accuracy / num_epochs * batch_size * validate_freq / len(test_loader) * 1000000 / (1 + epoch) / 1e6 # 计算测试精度（百万次）
results.append((batch_size, train_accuracy, test_accuracy)) # 将结果添加到列表中，每个元素包含batch_size、训练精度和测试精度
|