pytorch实现LeNet5手写数字识别+各层特征图可视化
- LeNet5网络结构
LeNet-5共有7层，不包含输入，每层都包含可训练参数；每个层有多个Feature Map，每个FeatureMap通过一种卷积滤波器提取输入的一种特征，然后每个FeatureMap有多个神经元。
在论文上的LeNet5的结构如下，由于论文的数据集是32x32的，mnist数据集是28x28的，所有只有INPUT变了，其余地方会严格按照LeNet5的结构编写程序：
- 网络模型
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Sequential( # input_size=(1*28*28)
nn.Conv2d(1, 6, 5, 1, 2), # padding=2保证输入输出尺寸相同
nn.ReLU(), # input_size=(6*28*28)
nn.MaxPool2d(kernel_size=2, stride=2), # output_size=(6*14*14)
)
self.conv2 = nn.Sequential(
nn.Conv2d(6, 16, 5),
nn.ReLU(), # input_size=(16*10*10)
nn.MaxPool2d(2, 2) # output_size=(16*5*5)
)
self.fc1 = nn.Sequential(
nn.Linear(16 * 5 * 5, 120),
nn.ReLU()
)
self.fc2 = nn.Sequential(
nn.Linear(120, 84),
nn.ReLU()
)
self.fc3 = nn.Linear(84, 10)
# 定义前向传播过程，输入为x
def forward(self, x):
x = self.conv1(x)
fp1 = x.detach() # 核心代码
x = self.conv2(x)
# nn.Linear()的输入输出都是维度为一的值，所以要把多维度的tensor展平成一维
x = x.view(x.size()[0], -1)
x = self.fc1(x)
x = self.fc2(x)
x = self.fc3(x)
return x # F.softmax(x, dim=1)
- 训练代码
def train(epoch): # 定义每个epoch的训练细节
model.train() # 设置为trainning模式
for batch_idx, (data, target) in enumerate(train_loader):
data = data.to(device)
target = target.to(device)
data, target = Variable(data), Variable(target) # 把数据转换成Variable
optimizer.zero_grad() # 优化器梯度初始化为零
output = model(data) # 把数据输入网络并得到输出，即进行前向传播
loss = F.cross_entropy(output, target) # 交叉熵损失函数
loss.backward() # 反向传播梯度
optimizer.step() # 结束一次前传+反传之后，更新参数
if batch_idx % log_interval == 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()))
- 测试代码
def test():
model.eval() # 设置为test模式
test_loss = 0 # 初始化测试损失值为0
correct = 0 # 初始化预测正确的数据个数为0
for data, target in test_loader:
data = data.to(device)
target = target.to(device)
data, target = Variable(data), Variable(target) # 计算前要把变量变成Variable形式，因为这样子才有梯度
output = model(data)
test_loss += F.cross_entropy(output, target, size_average=False).item() # sum up batch loss 把所有loss值进行累加
pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
correct += pred.eq(target.data.view_as(pred)).cpu().sum() # 对预测正确的数据个数进行累加
test_loss /= len(test_loader.dataset) # 因为把所有loss值进行过累加，所以最后要除以总得数据长度才得平均loss
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
- LeNet.py完整代码
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable
lr = 0.01 # 学习率
momentum = 0.5
log_interval = 10 # 跑多少次batch进行一次日志记录
epochs = 10
batch_size = 64
test_batch_size = 1000
class LeNet(nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = nn.Sequential( # input_size=(1*28*28)
nn.Conv2d(1, 6, 5, 1, 2), # padding=2保证输入输出尺寸相同
nn.ReLU(), # input_size=(6*28*28)
nn.MaxPool2d(kernel_size=2, stride=2), # output_size=(6*14*14)
)
self.conv2 = nn.Sequential(
nn.Conv2d(6, 16, 5),
nn.ReLU(), # input_size=(16*10*10)
nn.MaxPool2d(2, 2) # output_size=(16*5*5)
)
self.fc1 = nn.Sequential(
nn.Linear(16 * 5 * 5, 120),
nn.ReLU()
)
self.fc2 = nn.Sequential(
nn.Linear(120, 84),
nn.ReLU()
)
self.fc3 = nn.Linear(84, 10)
# 定义前向传播过程，输入为x
def forward(self, x):
x = self.conv1(x)
fp1 = x.detach() # 核心代码
x = self.conv2(x)
# nn.Linear()的输入输出都是维度为一的值，所以要把多维度的tensor展平成一维
x = x.view(x.size()[0], -1)
x = self.fc1(x)
x = self.fc2(x)
x = self.fc3(x)
return x # F.softmax(x, dim=1)
def train(epoch): # 定义每个epoch的训练细节
model.train() # 设置为trainning模式
for batch_idx, (data, target) in enumerate(train_loader):
data = data.to(device)
target = target.to(device)
data, target = Variable(data), Variable(target) # 把数据转换成Variable
optimizer.zero_grad() # 优化器梯度初始化为零
output = model(data) # 把数据输入网络并得到输出，即进行前向传播
loss = F.cross_entropy(output, target) # 交叉熵损失函数
loss.backward() # 反向传播梯度
optimizer.step() # 结束一次前传+反传之后，更新参数
if batch_idx % log_interval == 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()))
def test():
model.eval() # 设置为test模式
test_loss = 0 # 初始化测试损失值为0
correct = 0 # 初始化预测正确的数据个数为0
for data, target in test_loader:
data = data.to(device)
target = target.to(device)
data, target = Variable(data), Variable(target) # 计算前要把变量变成Variable形式，因为这样子才有梯度
output = model(data)
test_loss += F.cross_entropy(output, target, size_average=False).item() # sum up batch loss 把所有loss值进行累加
pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability
correct += pred.eq(target.data.view_as(pred)).cpu().sum() # 对预测正确的数据个数进行累加
test_loss /= len(test_loader.dataset) # 因为把所有loss值进行过累加，所以最后要除以总得数据长度才得平均loss
print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
if __name__ == '__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 启用GPU
train_loader = torch.utils.data.DataLoader( # 加载训练数据
datasets.MNIST('../data', train=True, download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)) # 数据集给出的均值和标准差系数，每个数据集都不同的，都数据集提供方给出的
])),
batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader( # 加载训练数据，详细用法参考我的Pytorch打怪路（一）系列-（1）
datasets.MNIST('../data', train=False, transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,)) # 数据集给出的均值和标准差系数，每个数据集都不同的，都数据集提供方给出的
])),
batch_size=test_batch_size, shuffle=True)
model = LeNet() # 实例化一个网络对象
model = model.to(device)
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum) # 初始化优化器
for epoch in range(1, epochs + 1): # 以epoch为单位进行循环
train(epoch)
test()
torch.save(model, 'model.pth') # 保存模型
- 中间层特征提取代码
class FeatureExtractor(nn.Module):
def __init__(self, submodule, extracted_layers):
super(FeatureExtractor, self).__init__()
self.submodule = submodule
self.extracted_layers = extracted_layers
def forward(self, x):
outputs = []
print(self.submodule._modules.items())
for name, module in self.submodule._modules.items():
if "fc" in name:
print(name)
x = x.view(x.size(0), -1)
print(module)
x = module(x)
print(name)
if name in self.extracted_layers:
outputs.append(x)
return outputs
- LeNet5Pre.py完整代码
import torch
import torch.nn as nn
import cv2
import torch.nn.functional as F
from LeNet5 import LeNet # 重要，虽然显示灰色(即在次代码中没用到)，但若没有引入这个模型代码，加载模型时会找不到模型
from torch.autograd import Variable
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt
# 中间特征提取
class FeatureExtractor(nn.Module):
def __init__(self, submodule, extracted_layers):
super(FeatureExtractor, self).__init__()
self.submodule = submodule
self.extracted_layers = extracted_layers
def forward(self, x):
outputs = []
print(self.submodule._modules.items())
for name, module in self.submodule._modules.items():
if "fc" in name:
print(name)
x = x.view(x.size(0), -1)
print(module)
x = module(x)
print(name)
if name in self.extracted_layers:
outputs.append(x)
return outputs
if __name__ == '__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = torch.load('model.pth') # 加载模型
model = model.to(device)
model.eval() # 把模型转为test模式
img = cv2.imread("8.jpg") # 读取要预测的图片
trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # 图片转为灰度图，因为mnist数据集都是灰度图
img = trans(img)
img = img.to(device)
img = img.unsqueeze(0) # 图片扩展多一维,因为输入到保存的模型中是4维的[batch_size,通道,长，宽]，而普通图片只有三维，[通道,长，宽]
# 扩展后，为[1，1，28，28]
output = model(img)
prob = F.softmax(output, dim=1)
prob = Variable(prob)
prob = prob.cpu().numpy() # 用GPU的数据训练的模型保存的参数都是gpu形式的，要显示则先要转回cpu，再转回numpy模式
print(prob) # prob是10个分类的概率
pred = np.argmax(prob) # 选出概率最大的一个
print(pred.item())
# 特征输出
net = LeNet().to(device)
exact_list = ["conv1", "conv2"]
myexactor = FeatureExtractor(net, exact_list)
x = myexactor(img)
# 特征输出可视化
for i in range(6):
ax = plt.subplot(1, 6, i + 1)
ax.set_title('Feature {}'.format(i))
ax.axis('off')
plt.imshow(x[0].data.cpu()[0, i, :, :], cmap='jet')
plt.show()
-
LeNet5Pre.py输入图像
输入图像为28*28像素的黑底白字手写数字图像。
-
运行结果
-
参考文献：
pytorch用LeNet5识别Mnist手写体数据集(训练+预测单张输入图片代码)
基于Pytorch的特征图提取