Pytorch 实践手写数字识别深度学习网络 LeNet-5
文章目录
- [Pytorch 实践手写数字识别深度学习网络 LeNet-5](#Pytorch 实践手写数字识别深度学习网络 LeNet-5)
训练手写体识别任务是一个非常简单的学习任务,理论简单是简单,但是我相信很多人都和我一样,想体验一下一个学习任务的全流程,有原始数据,处理数据,编写网络,训练模型,测试模型,使用模型这个过程。今天我们就由我来带大家体验一下。
认识 LeNet-5
LeNet-5出自论文《Gradient-Based Learning Applied to Document Recognition》, 原本是一种用于手写体字符识别的非常高效的卷积神经网络,包含了深度学习的基本模块:卷 积层,池化层,全连接层。
- INPUT(输入层) :输入28∗28的图片。
- C1(卷积层):选取6个5∗5卷积核(不包含偏置),得到6个特征图,每个特征
- 图的一个边为28−5+1=24。
- S2(池化层):池化层是一个下采样层,输出12∗12∗6的特征图。
- C3(卷积层):选取16个大小为5∗5卷积核,得到特征图大小为8∗8∗16。
- S4(池化层):窗口大小为2∗2,输出4∗4∗16的特征图。
- F5(全连接层):120个神经元。
- F6(全连接层):84个神经元。
- OUTPUT(输出层):10个神经元,10分类问题。
认识数据集
MNIST数据集来自美国国家标准与技术研究所,National Institute of Standards and Technology(NIST),数据集由来自250个不同人手写的数字构 成,其中50%是高中学生,50%来自人口普查局(the Census Bureau)的工 作人员。
训练集:60000,测试集:10000
MNIST数据集可在 http://yann.lecun.com/exdb/mnist/ 获取
大家如果想要的话可以联系我邮箱2837468248@qq.com,也可以直接发给你。
处理数据集
一般我们进行这个实践的时候,因为这个太经典了,很多的框架就直接集成了这个数据集,不用我们自己处理了,直接拿来用就好了。如下:
python
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)
其中的 torchvision
的包的 MNIST 数据集已经替我们解析好了一切。下载、解压、加载都做好了,效果如下:
虽然这样我们也能用,但是我们今天是来体验全流程的,我们要自己处理。
下载数据集
可以通过上面的数据集下载网址进行下载
下载好后是 gz 格式,你可以选择用程序对其进行解压,或者直接用解压软件解压。
解压后我是这样存放数据的
读取数据
一开始看到这个 ubyte 形式的数据我直接震惊,这啥数据呀,哪里有图片呢,真不清楚,后面了解到这是把数据进行压缩了。 具体解释请看文章
读取图片数据:
python
# 读取图像文件
def load_mnist_images(file_path):
with open(file_path, 'rb') as f:
magic, num_images, rows, cols = struct.unpack('>IIII', f.read(16))
if magic != 2051:
raise ValueError(f'Invalid magic number {magic} in file: {file_path}')
images = np.frombuffer(f.read(), dtype=np.uint8).reshape(num_images, 1, rows, cols)
return images
读取标签数据
python
# 读取标签文件
def load_mnist_labels(file_path):
with open(file_path, 'rb') as f:
magic, num_labels = struct.unpack('>II', f.read(8))
if magic != 2049:
raise ValueError(f'Invalid magic number {magic} in file: {file_path}')
labels = np.frombuffer(f.read(), dtype=np.uint8)
return labels
定义Dataset的继承类
原来直接用 torch vision 的数据的话,他会直接包装好,但是我们这里是要全部体验全流程,所以我们自己包装。
为什么要有 Dataset 类呢?
因为在进行深度学习训练的时候都是一批一批进行训练的,需要把数据载入到 Pytorch 提供的 dataloader 中去,方便 pytorch 后面对我们的数据方便进行操作。
我们自己定义一个 Dataset 类的话,一定要实现三个函数
python
__init__
__len__
__getitem__
这里我们的实现如下:
python
# 自定义 Dataset 类
class MNISTDataset(Dataset):
def __init__(self, images, labels, transform=None):
self.images = images
self.labels = labels
self.transform = transform
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
image = self.images[idx]
label = self.labels[idx]
if self.transform:
image = self.transform(image)
return image, label
把数据进行载入
这里载入的时候我们有个小地方需要注意一下,对数据进行一个简单的处理,就是把数据的维度进行放小,原来图片的维度如下:
我们需要的数据维度是 (60000,28,28)
的数据,所以要进行降维
python
train_images = train_images.squeeze(axis=1) # 把第一维度进行减掉
后面测试集的数据的话就是一样的处理。
载入dataloader
python
transform = transforms.Compose([
transforms.ToTensor(), #把numpy数据转化为tensor
transforms.Normalize((0.5,), (0.5,)) #对数据进行归一化处理
])
from torch.utils.data import DataLoader
train_dataset = MNISTDataset(train_images, train_labels,transform=transform)
test_dataset = MNISTDataset(test_images,test_labels,transform=transform)
train_loader = DataLoader(dataset=train_dataset, batch_size=64,shuffle=True)
test_loader = DataLoader(dataset=test_dataset,batch_size=64,shuffle=False)
编写网络
这个网络是非常简单的,直接定义一个类继承 torch.nn.Module
进行实现就好了,代码如下:
python
# 网络定义
import torch
import torch.nn as nn
from torch.utils.data import Dataset
import torch.optim as optim
import torchvision.transforms as transforms
class LeNet5(nn.Module):
def __init__(self):
super(LeNet5, self).__init__()
self.conv1 = nn.Conv2d(1, 6, kernel_size=5, stride=1, padding=2)
self.conv2 = nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = torch.relu(self.conv1(x)) #卷积以后进行激活
x = torch.max_pool2d(x, kernel_size=2, stride=2) #最大池化,提取特征
x = torch.relu(self.conv2(x))
x = torch.max_pool2d(x, kernel_size=2, stride=2)
x = x.view(-1, 16 * 5 * 5) #把数据进行展平方便全连接层的输入
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
x = self.fc3(x)
return x
编写训练与测试代码
训练的话就是一般深度学习的流程,直接调用 pytorch 的API进行解决了。
python
# 训练和测试函数
def train(model, device, train_loader, optimizer, criterion, epoch):
model.train()
best_model = model
min_loss = 1
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()
if min_loss > loss.item():
best_model, best_loss = model, loss.item()
print("update")
if batch_idx % 100 == 0:
print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')
print("模型训练结束")
print("保存最好 loss 模型,loss:",min_loss)
torch.save(best_model.state_dict(),'best-lenet5.pth')
def test(model, device, test_loader, criterion):
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 += 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(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset):.0f}%)\n')
# 训练和测试模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = LeNet5().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
num_epochs = 10
for epoch in range(1, num_epochs + 1):
train(model, device, train_loader, optimizer, criterion, epoch)
test(model, device, test_loader, criterion)
实践结果展示
载入保存的模型 pth 文件,然后手写一个 28 * 28 的数字图片进行处理识别效果如下
代码如下:
python
model = LeNet5()
model.load_state_dict(torch.load('best-lenet5.pth',map_location=torch.device('cpu')))
# 假设图像路径
image_path = '5.png' # 替换为你的图像路径
from PIL import Image
# 使用 PIL 库打开图像
image = Image.open(image_path)
# 使用 torchvision.transforms 进行数据转换和归一化
transform = transforms.Compose([
transforms.Resize((28, 28)), # 调整大小到 28x28
transforms.Grayscale(),
transforms.ToTensor(), # 转为 Tensor
transforms.Normalize((0.5,), (0.5,)) # 归一化
])
# 应用转换
image_tensor = transform(image).unsqueeze(0)
predict_output = model(image_tensor)
pred_num = predict_output.argmax(dim=1,keepdim=True)
print(pred_num) # 数据要写满28*28的格子才能预测)
如果需要 pth 文件也可以联系我,不过这个训练很快,可以自己训练玩!
完整代码
python
# 网络定义
import torch
import torch.nn as nn
from torch.utils.data import Dataset
import torch.optim as optim
import torchvision.transforms as transforms
class LeNet5(nn.Module):
def __init__(self):
super(LeNet5, self).__init__()
self.conv1 = nn.Conv2d(1, 6, kernel_size=5, stride=1, padding=2)
self.conv2 = nn.Conv2d(6, 16, kernel_size=5, stride=1, padding=0)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = torch.relu(self.conv1(x))
x = torch.max_pool2d(x, kernel_size=2, stride=2)
x = torch.relu(self.conv2(x))
x = torch.max_pool2d(x, kernel_size=2, stride=2)
x = x.view(-1, 16 * 5 * 5)
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
x = self.fc3(x)
return x
# 自定义 Dataset 类
class MNISTDataset(Dataset):
def __init__(self, images, labels, transform=None):
self.images = images
self.labels = labels
self.transform = transform
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
image = self.images[idx]
label = self.labels[idx]
if self.transform:
image = self.transform(image)
return image, label
import numpy as np
import struct
# 读取标签文件
def load_mnist_labels(file_path):
with open(file_path, 'rb') as f:
magic, num_labels = struct.unpack('>II', f.read(8))
if magic != 2049:
raise ValueError(f'Invalid magic number {magic} in file: {file_path}')
labels = np.frombuffer(f.read(), dtype=np.uint8)
return labels
# 读取图像文件
def load_mnist_images(file_path):
with open(file_path, 'rb') as f:
magic, num_images, rows, cols = struct.unpack('>IIII', f.read(16))
if magic != 2051:
raise ValueError(f'Invalid magic number {magic} in file: {file_path}')
images = np.frombuffer(f.read(), dtype=np.uint8).reshape(num_images, 1, rows, cols)
return images
# 获取到标签,图像数据
train_images = load_mnist_images('./data/train/train-images-idx3-ubyte')
train_labels = load_mnist_labels('./data/train/train-labels-idx1-ubyte')
print('train_images.shape', train_images.shape)
print('label.shape', train_labels.shape)
train_images = train_images.squeeze(axis=1)
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
test_images = load_mnist_images('./data/test/t10k-images-idx3-ubyte')
test_labels = load_mnist_labels('./data/test/t10k-labels-idx1-ubyte')
test_images = test_images.squeeze(axis = 1)
print(test_images.shape)
print(test_labels.shape)
from torch.utils.data import DataLoader
train_dataset = MNISTDataset(train_images, train_labels,transform=transform)
test_dataset = MNISTDataset(test_images,test_labels,transform=transform)
train_loader = DataLoader(dataset=train_dataset, batch_size=64,shuffle=True)
test_loader = DataLoader(dataset=test_dataset,batch_size=64,shuffle=False)
# 训练和测试函数
def train(model, device, train_loader, optimizer, criterion, epoch):
model.train()
best_model = model
min_loss = 100000.0
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()
if min_loss > loss.item():
best_model, best_loss = model, loss.item()
print("update")
if batch_idx % 100 == 0:
print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}')
print("模型训练结束")
print("保存最好 loss 模型,loss:",min_loss)
torch.save(best_model.state_dict(),'best-lenet5.pth')
def test(model, device, test_loader, criterion):
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 += 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(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset):.0f}%)\n')
# 训练和测试模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = LeNet5().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
num_epochs = 10
for epoch in range(1, num_epochs + 1):
train(model, device, train_loader, optimizer, criterion, epoch)
test(model, device, test_loader, criterion)