请写一个基于pytorch的一个简单的分类神经网络代码
时间: 2023-03-28 20:01:52 浏览: 55
好的,以下是一个基于 PyTorch 的简单分类神经网络代码:
```python
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, 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
# 加载数据集
train_loader = torch.utils.data.DataLoader(
torchvision.datasets.MNIST('/data', train=True, download=True,
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(
(.1307,), (.3081,))
])),
batch_size=64, shuffle=True)
# 初始化模型和优化器
model = Net()
optimizer = optim.SGD(model.parameters(), lr=.01, momentum=.5)
# 训练模型
for epoch in range(10):
for batch_idx, (data, target) in enumerate(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(
(.1307,), (.3081,))
])),
batch_size=100, shuffle=True)
model.eval()
test_loss =
correct =
with torch.no_grad():
for data, target in test_loader:
output = model(data)
test_loss += nn.CrossEntropyLoss()(output, target).item()
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
print('Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.f}%)'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
```
希望对你有帮助!