手写数字识别代码 pytorch
时间: 2023-09-14 21:06:27 浏览: 55
以下是一个基于PyTorch的手写数字识别代码。这个代码使用的是一个基于MNIST数据集的卷积神经网络模型。
```python
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
# 定义卷积神经网络模型
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=5, padding=2),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.MaxPool2d(2))
self.layer2 = nn.Sequential(
nn.Conv2d(16, 32, kernel_size=5, padding=2),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(2))
self.fc = nn.Linear(7 * 7 * 32, 10)
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = out.view(out.size(0), -1)
out = self.fc(out)
return out
# 加载数据集
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
train_data = MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = DataLoader(train_data, batch_size=100, shuffle=True)
# 定义模型、损失函数和优化器
model = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 训练模型
for epoch in range(10):
running_loss = 0.0
for i, (inputs, labels) in enumerate(train_loader):
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if (i+1) % 100 == 0:
print('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f' % (epoch+1, 10, i+1, len(train_data)//100, running_loss/100))
running_loss = 0.0
# 保存模型
torch.save(model.state_dict(), 'model.pth')
```
在训练完模型后,我们可以使用以下代码来测试我们的模型:
```python
import torch
import torchvision.transforms as transforms
from PIL import Image
# 加载模型
model = CNN()
model.load_state_dict(torch.load('model.pth'))
# 图像预处理
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
# 加载测试图像
image_path = 'test_image.png'
image = Image.open(image_path).convert('L')
image = transform(image)
image = image.unsqueeze(0)
# 使用模型进行预测
output = model(image)
_, predicted = torch.max(output.data, 1)
# 输出预测结果
print('Predicted Digit:', predicted.item())
```
在这个代码中,我们将一张手写数字的图像作为输入,使用我们训练好的模型来进行预测,并输出预测结果。