使用Python编写LeNet()模型进行手写数字识别
发布时间:2023-12-24 16:26:34
LeNet是一个经典的卷积神经网络模型,最早由Yann LeCun在1998年设计用于手写数字识别。该模型在MNIST数据集上取得了很好的效果,并且奠定了卷积神经网络的基本架构。
下面我将介绍如何使用Python编写LeNet()模型进行手写数字识别,并给出一个简单的使用例子。
首先,我们需要导入必要的库:
import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F
然后,我们可以定义LeNet()模型的结构:
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Conv2d(1, 6, kernel_size=5)
self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
self.fc1 = nn.Linear(16*4*4, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2)
x = x.view(-1, 16*4*4)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
LeNet()模型的结构包括两个卷积层和三个全连接层。其中,卷积层(conv)用于提取图像的特征,全连接层(fc)用于分类。
下一步是定义训练和测试的函数:
def train(model, device, train_loader, optimizer, criterion):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
def test(model, device, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.cross_entropy(output, target, reduction='sum').item()
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
accuracy = 100. * correct / len(test_loader.dataset)
return test_loss, accuracy
训练函数(train)通过循环遍历训练数据集进行模型的训练,并更新模型的参数。测试函数(test)则通过遍历测试数据集计算模型在测试集上的损失值和准确率。
最后,我们可以创建一个训练和测试的主函数来调用上述的函数:
def main():
# 设置随机种子
torch.manual_seed(0)
# 定义设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 加载数据集
train_loader = torch.utils.data.DataLoader(
torchvision.datasets.MNIST('data', train=True, download=True,
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(
torchvision.datasets.MNIST('data', train=False, transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize((0.1307,), (0.3081,))
])),
batch_size=1000, shuffle=True)
# 初始化模型
model = LeNet().to(device)
# 定义优化器和损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()
# 训练模型
for epoch in range(1, 11):
train(model, device, train_loader, optimizer, criterion)
test_loss, accuracy = test(model, device, test_loader)
print('Epoch {}: Test Loss = {:.4f}, Accuracy = {:.2f}%'.format(
epoch, test_loss, accuracy))
在主函数中,我们首先设置随机种子,然后加载MNIST数据集。接着,初始化LeNet()模型,并定义Adam优化器和交叉熵损失函数。最后,通过循环训练模型,并输出每个epoch的测试损失值及准确率。
现在我们可以调用main()函数开始训练和测试LeNet()模型了:
if __name__ == '__main__':
main()
这样,我们就完成了使用Python编写LeNet()模型进行手写数字识别的过程,并给出了一个简单的使用例子。该例子可以帮助您更好地理解LeNet模型的结构和训练流程。希望对您有所帮助!
