给出一个简单的pytorch程序
时间: 2023-07-04 17:16:18 浏览: 45
这里给出一个简单的PyTorch程序,用于实现一个简单的全连接神经网络进行MNIST手写数字图像分类任务:
```python
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
# 定义数据预处理
data_transforms = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
# 加载MNIST数据集
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=data_transforms)
test_dataset = datasets.MNIST('./data', train=False, download=True, transform=data_transforms)
# 定义数据加载器
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 Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(784, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = x.view(-1, 784)
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# 定义训练参数
lr = 0.01
momentum = 0.5
epochs = 10
# 定义模型、损失函数和优化器
model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)
# 训练模型
for epoch in range(epochs):
for batch_idx, (data, target) in enumerate(train_loader):
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_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
output = model(data)
test_loss += criterion(output, target).item()
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
```
该程序首先定义了数据预处理、加载MNIST数据集和数据加载器,然后定义了一个简单的全连接神经网络模型,包含一个输入层、一个隐层和一个输出层。接着定义了训练的参数、损失函数和优化器,并进行了模型的训练和测试。
在训练过程中,程序使用torch.utils.data.DataLoader来进行数据批量加载,每个批次的数据量为64。然后将每个批次的数据输入模型进行训练,计算损失函数并进行反向传播优化模型参数。在每个epoch结束后,程序会输出该epoch的训练进度和损失值。在测试过程中,程序使用with torch.no_grad()来关闭梯度计算,从而提高计算效率。然后将测试集数据输入模型进行推理,并计算模型的准确率和平均损失值。