前言
在优化之前我就简单的实现了这两个模型对图像的预测,起初我考虑的是通过opencv提取标志牌中的主体内容交给模型训练和预测。 标志牌分别是:
以上四种。
我在通过A4纸彩印之后通过无人车的摄像头录制视频并且通过opencv实时处理提取其中白色的部分生成二值图像。
opencv的代码:
import cv2
import numpy as np
import os
import torch
from net import *
from torchvision import transforms
output_folder = 'data/R'
if not os.path.exists(output_folder):
os.makedirs(output_folder)
frame_count = 1
def process_frame(frame):
global frame_count
frame = frame[int(frame.shape[0]//2):,:]
# 转换为HSV颜色空间
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
# 定义蓝色范围
lower_blue = np.array([100, 40, 0])
upper_blue = np.array([140, 255, 255])
# 创建蓝色掩码
mask_blue = cv2.inRange(hsv, lower_blue, upper_blue)
# cv2.imshow("mask",mask_blue)
# 在蓝色区域内查找白色区域
blue_only = cv2.bitwise_and(frame, frame, mask=mask_blue)
gray_blue = cv2.cvtColor(blue_only, cv2.COLOR_BGR2GRAY)
# cv2.imshow("gray_blue",gray_blue)
# 查找轮廓
contours, _ = cv2.findContours(gray_blue, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cv2.drawContours(blue_only, contours, -1, (0, 255, 0), 1)
# 寻找面积最大的轮廓
max_contour = None
max_area = 0
for contour in contours:
area = cv2.contourArea(contour)
if area > max_area:
max_contour = contour
max_area = area
cv2.imshow("blue_only",blue_only)
x,y,w,h = cv2.boundingRect(max_contour)
# 裁剪图像
res = frame[y:y+h, x:x+w]
res = cv2.cvtColor(res,cv2.COLOR_BGR2GRAY)
# cv2.imshow("rect",res)
_, mask_white = cv2.threshold(res, 160,255, cv2.THRESH_BINARY)
mask_white = cv2.resize(mask_white,(60,40))
text_area = mask_white[5:35,5:55]
cv2.imshow("text_area",text_area)
# text_area = cv2.resize(text_area,(240,160))
# frame_filename = os.path.join(output_folder, f'frame_{frame_count:04d}.png') # while os.path.exists(frame_filename): # frame_count += 1 # frame_filename = os.path.join(output_folder, f'frame_{frame_count:04d}.png') # 保存当前帧为图片
# cv2.imwrite(frame_filename, text_area)
# text_area = cv2.cvtColor(text_area,cv2.COLOR_GRAY2RGB)
conf_a = cv2.matchTemplate(text_area,template_A,cv2.TM_CCOEFF_NORMED)
conf_b = cv2.matchTemplate(text_area,template_B,cv2.TM_CCOEFF_NORMED)
conf = max(conf_a,conf_b)
print(conf)
if conf > 0.65 and conf_a > conf_b:
print("识别到标志A")
elif conf > 0.65 and conf_a < conf_b:
print("识别到标志B")
cv2.imshow("text_area",text_area)
# cv2.imshow("template_a",template_A)
return frame
# 初始化摄像头
cap = cv2.VideoCapture(2)
template_A = cv2.imread('A.png', 0)
# template_A = cv2.resize(template_A, (240,160))
template_B = cv2.imread('B.png', 0)
# template_B = cv2.resize(template_B, (240,160))
while True:
# 捕获一帧图像
ret, frame = cap.read()
# frame = cv2.imread("../tools/image/A/frame_0012.png")
# if not ret: # break
# 处理并显示文本
result_frame = process_frame(frame)
# 显示结果
# cv2.imshow('Text Detection', result_frame)
# 按'q'键退出循环
if cv2.waitKey(1) & 0xFF == ord('q'):
break
# 清理资源
cap.release()
cv2.destroyAllWindows()
经过处理后的图像类似这种:
这是图像提取的过程:
按照我预想的,通过Lenet或者Alexnet这类简单的卷积神经网络,并且仅仅识别这些简单且只有四类的图像实现较为准确的预测应该表现会不错,并且由于他们较小的体积,也能够部署到树莓派上支持实时运行检测。
现实情况是,我遇到了一系列的问题。
例如:
或者是严重的过拟合:
后来经过我反思,我决定重新实现这个功能。
现在我打算从数据集优化和网络结构优化来入手。
目前的我大致的思路是:
-
优化数据集
-
改变激活函数和引入正则化
-
调整网络结构。## 准备工作
由于接下来需要经过较多的调整和测试,为了方便比较我首先编写在训练过程中需要用到的一些工具函数和类。 他们包括:
-
构建网络模型
-
数据集加载类
-
划分训练集和测试集的工具函数
-
记录训练数据的记录函数
-
用于训练模型的训练函数
-
用于验证模型以及对模型评估的函数。
在对模型训练和评估完成之后,编写适用此模型的预测函数,可以通过实例化模型和输入数据预测图像类别。
构建网络模型
我首先使用原始的Lenet来进行训练和预测,因此构建和最初的Lenet一致架构的网络类: (这里我微调了一下卷积层的输入输出的尺寸)
import torch
import torch.nn as nn
from torch import optim
import torch.nn.functional as F
class Lenet(nn.Module):
def __init__(self, input_size=32, num_classes=4):
super(Lenet, self).__init__()
self.conv1 = nn.Conv2d(in_channels=1,out_channels=6,kernel_size=5,stride=1,padding=0)
self.pool1 = nn.MaxPool2d(kernel_size=2,stride=2)
self.conv2 = nn.Conv2d(in_channels=6,out_channels=16,kernel_size=5)
self.pool2 = nn.MaxPool2d(kernel_size=2,stride=2)
self.fc1 = nn.Linear(in_features=16*5*5,out_features=120)
self.fc2 = nn.Linear(in_features=120,out_features=84)
self.fc3 = nn.Linear(in_features=84,out_features=num_classes)
def forward(self, x):
x = F.relu(self.conv1(x))
x = self.pool1(x)
x = F.relu(self.conv2(x))
x = self.pool2(x)
x = x.view(-1, 16*5*5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
划分训练集和测试集
依照上面的数据集目录结构,将图片路径存放到train.txt和test.txt中来划分训练集和测试集。
import os
import random
def split_dataset(dataset_path, train_ratio, output_dir):
"""
将数据集划分为训练集和测试集,并将文件路径写入 train.txt 和 test.txt。
参数:
- dataset_path: str, 数据集根路径,每个类别为一个子文件夹。
- train_ratio: float, 训练集比例(0.0 到 1.0)。
- output_dir: str, 输出目录,存放 train.txt 和 test.txt 文件。
""" # 检查输出目录是否存在,不存在则创建
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# 定义输出文件路径
train_file = os.path.join(output_dir, "train.txt")
test_file = os.path.join(output_dir, "test.txt")
# 清空输出文件
with open(train_file, "w") as f:
pass
with open(test_file, "w") as f:
pass
# 遍历每个类别文件夹
for category in os.listdir(dataset_path):
category_path = os.path.join(dataset_path, category)
if not os.path.isdir(category_path):
continue # 跳过非文件夹项
print(f"Processing category: {category}")
# 获取该类别的所有图片文件名
image_files = [f for f in os.listdir(category_path) if os.path.isfile(os.path.join(category_path, f))]
# 打乱文件顺序
random.shuffle(image_files)
# 按比例划分训练集和测试集
split_idx = int(len(image_files) * train_ratio)
train_files = image_files[:split_idx]
test_files = image_files[split_idx:]
# 将文件路径写入 train.txt 和 test.txt with open(train_file, "a") as train_f:
for file in train_files:
train_f.write(f"{category}/{file}\n")
with open(test_file, "a") as test_f:
for file in test_files:
test_f.write(f"{category}/{file}\n")
print(f"Dataset split completed. Train list saved to {train_file}, Test list saved to {test_file}")
数据集加载类
这一部分可以之前引用之前的代码,我使用Dataset类来构建自定义的数据集类。 我的目录结构大致如下:
/data
- /img
- /A
- /B
- /Left
- /Right
label.txt
train.txt
val.txt
其中在label.txt是四个类别的名称,每个类别一行。 在train.txt是经过划分后的训练集 在val.txt中是经过划分后的验证集 他们的内容类似:类别(目录)/文件名
根据这些定义数据集加载类:
import cv2
import os
from torch.utils.data import Dataset
class ProkingDataset(Dataset):
def __init__(self, root, mode, transform=None):
"""
:param root: 数据集根路径
:param mode: train or val :param transform: 自定义的变换对象
""" self.root = root
self.transform = transform
# 加载图片路径
with open(os.path.join(self.root, f'{mode}.txt'),"r") as f:
self.files = [file.strip() for file in f]
print(self.files)
# 加载标签
with open(os.path.join(root, "label.txt"), "r") as labels:
self.labels = [label.strip() for label in labels]
def __len__(self):
return len(self.files)
def __getitem__(self, idx):
filename = self.files[idx]
# 获取图片内容
path = os.path.join(self.root, filename)
img = cv2.imread(path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # 由于Lenet输入通道为一,使用灰度图输入
if self.transform is not None:
img = self.transform(img)
# 拆分标签
label = filename.split("/")[0] # 类别名称
label_index = self.labels.index(label)
return img, label_index
记录函数
由于后续我打算对模型进行多次调整来优化表现,因此考虑定义多个记录函数,包括:
-
初始化时用于记录网络信息、初始超参等信息的初始化记录函数,这个函数应该创建一个日志文件,并返回文件路径
-
训练过程中每个epoch调用一次的过程函数,用于记录每次epoch的数据
-
训练结束后对结果进行评估保存
最终我构造了一个记录类:
import time
import datetime
import os
import logging
class Recorder():
def __init__(self,model,train_loader,optimizer,criterion,num_epoches,output_path):
self.model=model
self.train_loader=train_loader
self.optimizer=optimizer
self.criterion=criterion
self.num_epoches=num_epoches
self.lr=optimizer.defaults['lr']
self.start_time = time.time()
self.total_loss = 0
self.total_acc = 0
logging.basicConfig(
level=logging.INFO,
filename=output_path,
format="%(asctime)s : %(message)s",
)
def start_record(self):
logging.info("start record train logs")
logging.info(f"model structure:\r\n{self.model}\r\n")
logging.info(f"optimizer structure:\r\n{self.optimizer}\r\n")
logging.info(f"criterion structure:\r\n{self.criterion}\r\n")
logging.info(f"epoch nums:{self.num_epoches}")
logging.info(f"learning rate : {self.lr}")
logging.info(f"train data nums:{len(self.train_loader.dataset)}")
logging.info(f"train data batch_size:{len(self.train_loader.batch_size)}")
def record_epoch(self,epoch,loss,acc):
self.total_loss += loss
self.total_acc += acc
logging.info(f"epoch : {epoch}")
logging.info(f"epoch loss : {loss} , total loss : {self.total_loss}")
logging.info(f"train acc : {acc} , total acc : {self.total_acc}")
logging.info(f"spend time: {time.time() - self.start_time}")
def end_record(self):
end_time = time.time() # 获取结束时间
elapsed_time = end_time - self.start_time
elapsed_hours = elapsed_time // 3600
elapsed_minutes = (elapsed_time % 3600) // 60
elapsed_seconds = elapsed_time % 60
logging.info("Training ended")
logging.info(f"Total time taken: {int(elapsed_hours)}h {int(elapsed_minutes)}m {elapsed_seconds:.2f}s")
logging.info(f"avg acc : {self.total_acc / self.num_epoches}")
logging.info(f"avg loss : {self.total_loss / self.num_epoches}")
训练函数
训练函数只需要编写一个用于单次训练调用的函数来通过循环调用即可。
def train_epoch(net, train_loader, optimizer, criterion, epoch, recorder=None):
"""
用于训练中迭代的函数
:param net: 网络实例
:param train_loader: 训练数据加载器
:param optimizer: 优化函数
:param criterion: 损失函数
:param epoch: 当前训练轮数
:param recorder: 可选,用于记录日志的对象
""" net.train() # 切换模型到训练模式
total_loss = 0
correct = 0
total = 0
for images, labels in train_loader:
optimizer.zero_grad() # 梯度清零
output = net(images) # 计算输出
loss = criterion(output, labels) # 计算损失
loss.backward() # 反向传播
optimizer.step() # 更新参数
total_loss += loss.item() # 累加损失
# 计算准确率
_, predicted = torch.max(output, 1) # 获取预测类别
correct += (predicted == labels).sum().item() # 累加正确预测数
total += labels.size(0) # 累加样本数
# 计算本轮平均损失和准确率
avg_loss = total_loss / len(train_loader)
acc = correct / total
# 打印训练日志
print(f"Epoch {epoch} - Loss: {avg_loss:.4f}, Accuracy: {acc:.2%}")
# 记录日志到 recorder(可选)
if recorder:
recorder.record_epoch(epoch, avg_loss, acc)
模型评估
针对模型评估,编写与训练函数类似预测程序,加载测试集输入到训练好的模型(权重)中预测类别,统计预测中的损失和准确率。
与训练不同的是,评估需要禁用梯度计算,并且将torch定义的模型实例设置为评估模式(eval
)
在模型评估中,只需要前向传播,不需要反向传播。
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from ProkingDataset import ProkingDataset
from functions import *
from torch.utils.data import DataLoader
from Lenet import Lenet
import torch.optim as optim
from Recorder import Recorder
from sklearn.metrics import confusion_matrix, classification_report
# 全局参数
data_dir = 'clean_data'
batch_size = 32
model_path = 'mark_classify_model.pth' # 模型权重路径
# 实例化模型
model = Lenet()
model = model.float() # 确保模型参数是 float 类型
print("模型初始化完成")
# 加载模型权重
try:
model.load_state_dict(torch.load(model_path))
print(f"模型权重加载完成:{model_path}")
except FileNotFoundError:
print(f"模型权重文件未找到:{model_path}")
exit()
# 设置为评估模式
model.eval()
# 定义数据变化器
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5), (0.5)), # 标准化
])
# 加载测试集
test_dataset = ProkingDataset(data_dir, "test", transform)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# 定义损失函数
criterion = nn.CrossEntropyLoss()
# 评估模型
def evaluate_model(model, test_loader, criterion):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
total_loss = 0
correct = 0
total = 0
all_labels = []
all_predictions = []
with torch.no_grad(): # 禁用梯度计算
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
# 前向传播
outputs = model(images)
loss = criterion(outputs, labels)
total_loss += loss.item()
# 获取预测类别
_, predicted = torch.max(outputs, 1)
correct += (predicted == labels).sum().item()
total += labels.size(0)
# 收集所有标签和预测结果
all_labels.extend(labels.cpu().numpy())
all_predictions.extend(predicted.cpu().numpy())
# 计算准确率和平均损失
accuracy = correct / total
avg_loss = total_loss / len(test_loader)
print(f"测试集平均损失: {avg_loss:.4f}, 准确率: {accuracy:.2%}")
# 混淆矩阵和分类报告
conf_matrix = confusion_matrix(all_labels, all_predictions)
print("混淆矩阵:")
print(conf_matrix)
class_report = classification_report(all_labels, all_predictions, target_names=test_dataset.labels)
print("分类报告:")
print(class_report)
return accuracy, avg_loss, conf_matrix, class_report
# 调用评估函数
accuracy, avg_loss, conf_matrix, class_report = evaluate_model(model, test_loader, criterion)
预测函数
考虑到不同模型的预测,在编写预测函数的时候需要接受如下参数:
-
模型实例
-
预训练权重路径
-
图像数据
-
图像变换器
-
类别映射
预测函数需要做的事情是将输入的图像输入放到模型中进行前向传播,获取类别。 为了确保预测的准确性,需要对图像数据进行处理,使用与训练时一致的transform。 由于模型输出的类别是整数,如果需要的话,指定类别映射,即在训练时的类别索引对应的类别名称。
代码:
def predict(model,pth,image,transform=None,label_dict=None):
"""
使用预训练权重进行模型预测
:param model: 模型实例
:param pth: 预训练权重
:param image: 图像数据(应预先处理适应模型输入)
:param transform: 图像变换器
:param label_dict: 标签映射(可选)
:return: 预测的类别索引(以及名称)
"""
# 加载模型以及权重
model.eval()
model.load_state_dict(torch.load(pth))
# 对图像进行处理
if transform is not None:
image = transform(image)
# 进行预测
with torch.no_grad():
output = model(image)
label_idx = torch.max(output, 1).indices[0].item()
if label_dict is not None:
return label_idx,label_dict[label_idx]
else:
return label_idx,"unknown"
优化数据集
我的数据集存在的问题很显而易见:
在我之前训练模型的时候我并没有对其进行清洗,其中存在一些纯黑的图像和非常残缺的图像极其容易导致模型训练产生错误的结果,并且由于其他的类别也存在类似的数据,这很可能直接导致预测类别的错误。
这里我考虑两种方式:
-
对数据进行清洗,删除异常的数据(例如全黑的图像)
-
使用原始的RGB图像,即保留蓝色区域而非只提取白色区域作为数据集,这应该有助于训练过程中学习更多的特征。
这里思路是首先读取指定路径下所有的图像,这些图像是彩色的图像,然后将彩色图像转为灰度图,这时候得到的是32\*32的灰度图像,然后计算每个类别的文件夹中所有图像中不为0的像素点在图像中的均值和方差,并且剔除离群值,将距离均值和方差较远的图片删除。
import cv2
import os
import numpy as np
from tqdm import tqdm
def clean_dataset(dataset_path, output_path, threshold=1):
"""
清洗图像数据集。
参数:
- dataset_path: str, 数据集路径,每个类别存放在不同的文件夹中。
- output_path: str, 清洗后数据集保存的路径。
- threshold: float, 离群值的阈值(单位:标准差,默认值为3)。
""" # 创建输出目录
if not os.path.exists(output_path):
os.makedirs(output_path)
# 遍历每个类别文件夹
for category in os.listdir(dataset_path):
category_path = os.path.join(dataset_path, category)
if not os.path.isdir(category_path):
continue
print(f"Processing category: {category}")
# 获取类别的输出路径
category_output_path = os.path.join(output_path, category)
if not os.path.exists(category_output_path):
os.makedirs(category_output_path)
# 存储每张图片的非零像素点统计
nonzero_stats = []
images = []
# 读取所有图像并计算非零像素点均值和方差
for filename in tqdm(os.listdir(category_path)):
image_path = os.path.join(category_path, filename)
try:
# 读取彩色图像
image = cv2.imread(image_path)
if image is None:
continue
# 转为灰度图并调整大小为 32x32 gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
resized_image = cv2.resize(gray_image, (32, 32))
# 计算非零像素点的均值
nonzero_count = np.count_nonzero(resized_image)
nonzero_stats.append(nonzero_count)
images.append((filename, resized_image, nonzero_count))
except Exception as e:
print(f"Error processing {image_path}: {e}")
continue
# 计算均值和方差
nonzero_mean = np.mean(nonzero_stats)
nonzero_std = np.std(nonzero_stats)
print(f"Category {category} - Nonzero Mean: {nonzero_mean}, Std: {nonzero_std}")
# 剔除离群值
for filename, resized_image, nonzero_count in images:
# 判断是否为离群值
if abs(nonzero_count - nonzero_mean) > threshold * nonzero_std:
print(f"Removing outlier: {filename}")
continue
# 保存清洗后的图像
output_image_path = os.path.join(category_output_path, filename)
cv2.imwrite(output_image_path, resized_image)
print("Dataset cleaning completed.")
clean_dataset("data","clean_data",0.8)
经过清洗由于录制车辆移动过程中产生的黑色图像已经全部被删除,并且剔除了部分扭曲的图像。
清洗数据之后,编写用于训练模型的train.py,首先划分训练集和测试集:
# 划分训练集和测试集
data_dir = 'clean_data'
split_dataset(data_dir,0.85,data_dir)
编写完成的训练程序:
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from ProkingDataset import ProkingDataset
from functions import *
from torch.utils.data import DataLoader
from Lenet import Lenet
import torch.optim as optim
from Recorder import Recorder
# 定义超参数
num_epochs = 20
lr = 0.002
batch_size = 32
data_dir = 'clean_data'
log_path = "log"
# 划分训练集和测试集
split_dataset(data_dir,0.85,data_dir)
print("数据集划分完成")
# 定义数据变化器
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5), (0.5)), # 标准化
transforms.RandomHorizontalFlip(), # 随机水平裁剪
])
# 加载数据
train_dataset = ProkingDataset(data_dir,"train", transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
print("数据加载完成")
# 创建模型
model = Lenet()
model = model.float() # 确保模型参数是 float 类型
print("模型初始化完成")
# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr) # 记得调学习率
print("损失函数与优化器初始化完成")
# 记录器
recorder = Recorder(model,train_loader,optimizer,criterion,num_epochs,log_path)
# 训练模型
recorder.start_record()
for epoch in range(num_epochs):
train_epoch(model,train_loader,optimizer, criterion, epoch, recorder)
print('训练完成.')
recorder.end_record()
# 保存模型参数
torch.save(model.state_dict(), 'mark_classify_model.pth')
print('Model parameters saved to lenet_model.pth')
查看日志:
Lenet比较小,并且我的数据集也很少,这里训练只花费15秒就完成了。
接下来使用之前编写好的评估程序进行评估:
可以看到模型在测试集的表现非常好。
然后使用之前编写的预测函数进行预测,在使用预测函数之前需要进行简单的加载,这里编写predict.py文件:
import torch
from functions import *
import numpy as np
from Lenet import Lenet
from torchvision import transforms
import cv2
# 加载图像
img = cv2.imread("clean_data/B/frame_0051.png")
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# 定义变化器
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5), (0.5)), # 标准化
transforms.RandomHorizontalFlip(), # 随机水平裁剪
])
# 标签映射
label_dict = {
0:"A",
1:"B",
2:"Left",
3:"Right",
}
# 预测
model = Lenet()
pth = "mark_classify_model.pth"
idx,name = predict(model,pth,img,transform,label_dict)
print(f"预测索引:{idx},类别:{name}")
运行结果:
使用实际数据测试: 原图像:
截取其中B的部分,再使用opencv进行简单的处理:
对图像灰度化之后重设尺寸再提取阈值:
img = cv2.imread("img.png")
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img = cv2.resize(img,(32,32))
_, img = cv2.threshold(img, 180,255, cv2.THRESH_BINARY)
最终的图像:
预测输出:
经过测试这里仍然存在一些问题,在对图像预处理的时候需要确保提取出主要部分,否则会导致预测错误。
为了解决这个问题,我尝试在预测中对每个类别输出的置信度进行判断,当置信度最大的值和第二大的值相同的时候,应该输出unknown,因为这很可能是错误的预测,这符合我的业务场景,我需要实时进行检测并且需要严格过滤错误的预测。
另一个办法是,我通过在网络最后输出的线性层中加入一个softmax使其输出为概率分布,检查最大的概率与1的差值来决定本次预测的可靠程度。
我通过加入softmax调整输出为概率分布来解决这个问题,部署到树莓派测试:
表现非常好!
使用AlexNet
为了更加优化我打算改用AlexNet直接对彩色的RGB图像进行优化,由于AlexNet的网络更深并且接受RGB通道的图像,可以学习到的特征更多,或许对于抑制错误的判断会有帮助。
我的数据集如下:
首先搭建AlexNet网络结构:
import os
import cv2
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torch import nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
class AlexNet(nn.Module):
def __init__(self, num_classes=4):
super(AlexNet, self).__init__()
self.features = nn.Sequential(
# 第一个卷积层
nn.Conv2d(3, 96, kernel_size=11, stride=4),
nn.ReLU(),
nn.MaxPool2d(kernel_size=3, stride=2),
# 第二个卷积层
nn.Conv2d(96, 256, kernel_size=5, padding=2),
nn.ReLU(),
nn.MaxPool2d(kernel_size=3, stride=2),
# 第三个卷积层
nn.Conv2d(256, 384, kernel_size=3, padding=1),
nn.ReLU(),
# 第四个卷积层
nn.Conv2d(384, 384, kernel_size=3, padding=1),
nn.ReLU(),
# 第五个卷积层
nn.Conv2d(384, 256, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(kernel_size=3, stride=2)
)
self.classifier = nn.Sequential(
nn.Dropout(p=0.5),
nn.Linear(6400, 4096),
nn.ReLU(),
nn.Dropout(p=0.5),
nn.Linear(4096, 4096),
nn.ReLU(),
nn.Linear(4096, 1000),
nn.ReLU(),
nn.Linear(1000, num_classes),
)
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), -1) # Flatten the tensor
x = self.classifier(x)
return x
由于彩色图像参数庞大,修改训练函数使用GPU训练:
def train_epoch(net, train_loader, optimizer, criterion, epoch, recorder=None):
"""
用于训练中迭代的函数
:param net: 网络实例
:param train_loader: 训练数据加载器
:param optimizer: 优化函数
:param criterion: 损失函数
:param epoch: 当前训练轮数
:param recorder: 可选,用于记录日志的对象
""" net.train() # 切换模型到训练模式
total_loss = 0
correct = 0
total = 0
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
net.to(device)
for images, labels in train_loader:
images = images.to(device)
labels = labels.to(device)
optimizer.zero_grad() # 梯度清零
output = net(images) # 计算输出
loss = criterion(output, labels) # 计算损失
loss.backward() # 反向传播
optimizer.step() # 更新参数
total_loss += loss.item() # 累加损失
# 计算准确率
_, predicted = torch.max(output, 1) # 获取预测类别
correct += (predicted == labels).sum().item() # 累加正确预测数
total += labels.size(0) # 累加样本数
# 计算本轮平均损失和准确率
avg_loss = total_loss / len(train_loader)
acc = correct / total
# 打印训练日志
print(f"Epoch {epoch} - Loss: {avg_loss:.4f}, Accuracy: {acc:.2%}")
# 记录日志到 recorder(可选)
if recorder:
recorder.record_epoch(epoch, avg_loss, acc)
编写用户训练ALexNet的程序,与Lenet类似,主要修改点是修改标准化的参数适应三通道图像:
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from AlexDataset import AlexDataset
from functions import *
from torch.utils.data import DataLoader
from alexnet import AlexNet
import torch.optim as optim
from Recorder import Recorder
# 定义超参数
num_epochs = 50
lr = 0.01
batch_size = 32
data_dir = 'alex_data'
log_path = "alex_log"
# 划分训练集和测试集
split_dataset(data_dir,0.85,data_dir)
print("数据集划分完成")
# 定义数据变化器
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5)), # 标准化
transforms.RandomHorizontalFlip(), # 随机水平裁剪
])
# 加载数据
train_dataset = AlexDataset(data_dir,"train", transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
print("数据加载完成")
# 创建模型
model = AlexNet()
model = model.float() # 确保模型参数是 float 类型
print("模型初始化完成")
# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.RAdam(model.parameters(), lr=lr) # 记得调学习率
print("损失函数与优化器初始化完成")
# 记录器
recorder = Recorder(model,train_loader,optimizer,criterion,num_epochs,log_path)
# 训练模型
recorder.start_record()
for epoch in range(num_epochs):
train_epoch(model,train_loader,optimizer, criterion, epoch, recorder)
print('训练完成.')
recorder.end_record()
# 保存模型参数
torch.save(model.state_dict(), 'alex_model.pth')
print('Model parameters saved to alex_model.pth')
到这里已经可以开始训练了,但是在我训练的时候出现了这个情况:
Epoch 8 - Loss: 0.4467, Accuracy: 84.94%
Epoch 9 - Loss: 0.3610, Accuracy: 88.96%
Epoch 10 - Loss: 0.2635, Accuracy: 91.66%
Epoch 11 - Loss: 0.6465, Accuracy: 86.95%
Epoch 12 - Loss: 3.2391, Accuracy: 32.94%
Epoch 13 - Loss: 1.4066, Accuracy: 29.61%
Epoch 14 - Loss: 1.3726, Accuracy: 29.61%
Epoch 15 - Loss: 1.3723, Accuracy: 29.61%
Epoch 16 - Loss: 1.3727, Accuracy: 29.61%
Epoch 17 - Loss: 1.3722, Accuracy: 29.61%
Epoch 18 - Loss: 1.3722, Accuracy: 29.61%
Epoch 19 - Loss: 1.3721, Accuracy: 29.61%
Epoch 20 - Loss: 1.3722, Accuracy: 29.61%
Epoch 21 - Loss: 1.3723, Accuracy: 29.61%
Epoch 22 - Loss: 1.3722, Accuracy: 29.61%
Epoch 23 - Loss: 1.3729, Accuracy: 29.61%
Epoch 24 - Loss: 1.3722, Accuracy: 29.61%
Epoch 25 - Loss: 1.3723, Accuracy: 29.61%
在训练过程中出现了loss和准确率震荡和停止的问题。
从epoch11到epoch12的loss显著上升,这可能是由于梯度爆炸导致的,随后停滞不前,loss始终在一个值上下波动,acc没有变化,这可能是模型陷入了局部最优的问题。
针对这两个问题进行调整,首先是梯度爆炸: 优化梯度爆炸的问题,最首先考虑的是调整学习率,使用较小的学习率多次训练。 另一个方法是在训练过程中加入梯度裁剪限制梯度的最大值。
这里将学习率调整为0.002,然后在训练中加入:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
以及还有一个原因是epoch过多,根据日志信息,其实在epoch10的时候模型精确到就已经达到90了,这时候再接着训练更新权重,就可能导致参数的反向调整,然后出现一系列问题。
所以在调小学习率之后,我又将epoch调整为15.
训练日志:
Epoch 1 - Loss: 0.3901, Accuracy: 87.26%
Epoch 2 - Loss: 0.0931, Accuracy: 96.74%
Epoch 3 - Loss: 0.0631, Accuracy: 98.06%
Epoch 4 - Loss: 0.0970, Accuracy: 97.62%
Epoch 5 - Loss: 0.0630, Accuracy: 98.12%
Epoch 6 - Loss: 0.0317, Accuracy: 99.25%
Epoch 7 - Loss: 0.1070, Accuracy: 97.18%
Epoch 8 - Loss: 0.0751, Accuracy: 98.06%
Epoch 9 - Loss: 0.1887, Accuracy: 96.36%
Epoch 10 - Loss: 0.2801, Accuracy: 93.29%
Epoch 11 - Loss: 0.1598, Accuracy: 96.68%
Epoch 12 - Loss: 0.1161, Accuracy: 97.74%
Epoch 13 - Loss: 0.0937, Accuracy: 97.93%
Epoch 14 - Loss: 0.0682, Accuracy: 98.37%
训练完成.
Model parameters saved to alex_model.pth
最后收敛的loss为0.06,为了检查时候存在过拟合,使用测试集进行测试:
import torch
import torch.nn as nn
import torchvision.transforms as transforms
from AlexDataset import AlexDataset
from functions import *
from torch.utils.data import DataLoader
from alexnet import AlexNet
from sklearn.metrics import confusion_matrix, classification_report
# 全局参数
data_dir = 'alex_data'
batch_size = 32
model_path = 'alex_model.pth' # 模型权重路径
# 实例化模型
model = AlexNet()
model = model.float() # 确保模型参数是 float 类型
print("模型初始化完成")
# 加载模型权重
try:
model.load_state_dict(torch.load(model_path))
print(f"模型权重加载完成:{model_path}")
except FileNotFoundError:
print(f"模型权重文件未找到:{model_path}")
exit()
# 设置为评估模式
model.eval()
# 定义数据变化器
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5)), # 标准化
])
# 加载测试集
test_dataset = AlexDataset(data_dir, "test", transform)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# 定义损失函数
criterion = nn.CrossEntropyLoss()
# 评估模型
def evaluate_model(model, test_loader, criterion):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
total_loss = 0
correct = 0
total = 0
all_labels = []
all_predictions = []
with torch.no_grad(): # 禁用梯度计算
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
# 前向传播
outputs = model(images)
loss = criterion(outputs, labels)
total_loss += loss.item()
# 获取预测类别
_, predicted = torch.max(outputs, 1)
correct += (predicted == labels).sum().item()
total += labels.size(0)
# 收集所有标签和预测结果
all_labels.extend(labels.cpu().numpy())
all_predictions.extend(predicted.cpu().numpy())
# 计算准确率和平均损失
accuracy = correct / total
avg_loss = total_loss / len(test_loader)
print(f"测试集平均损失: {avg_loss:.4f}, 准确率: {accuracy:.2%}")
# 混淆矩阵和分类报告
conf_matrix = confusion_matrix(all_labels, all_predictions)
print("混淆矩阵:")
print(conf_matrix)
class_report = classification_report(all_labels, all_predictions, target_names=test_dataset.labels)
print("分类报告:")
print(class_report)
return accuracy, avg_loss, conf_matrix, class_report
# 调用评估函数
accuracy, avg_loss, conf_matrix, class_report = evaluate_model(model, test_loader, criterion)
使用之前编写的验证代码稍加改动即可,运行输出:
模型权重加载完成:alex_model.pth
['A/0425.png', 'A/0374.png', 'A/0108.png', 'A/0282.png', 'A/0424.png', 'A/0118.png', 'A/0500.png', 'A/0322.png', 'A/0179.png', 'A/0190.png', 'A/0375.png', 'A/0045.png', 'A/0307.png', 'A/0297.png', 'A/0228.png', 'A/0273.png', 'A/0267.png', 'A/0539.png', 'A/0268.png', 'A/0238.png', 'A/0256.png', 'A/0408.png', 'A/0127.png', 'A/0453.png', 'A/0025.png', 'A/0212.png', 'A/0354.png', 'A/0495.png', 'A/0352.png', 'A/0398.png', 'A/0364.png', 'A/0295.png', 'A/0415.png', 'A/0330.png', 'A/0523.png', 'A/0249.png', 'A/0235.png', 'A/0465.png', 'A/0086.png', 'A/0436.png', 'A/0117.png', 'A/0323.png', 'A/0515.png', 'A/0365.png', 'A/0522.png', 'A/0302.png', 'A/0260.png', 'A/0030.png', 'A/0444.png', 'A/0358.png', 'A/0247.png', 'A/0027.png', 'A/0391.png', 'A/0377.png', 'A/0366.png', 'A/0517.png', 'A/0196.png', 'A/0081.png', 'A/0499.png', 'A/0396.png', 'A/0064.png', 'A/0379.png', 'A/0548.png', 'A/0036.png', 'A/0263.png', 'A/0265.png', 'A/0305.png', 'A/0537.png', 'A/0199.png', 'A/0395.png', 'A/0012.png', 'A/0094.png', 'A/0224.png', 'A/0384.png', 'A/0217.png', 'A/0418.png', 'A/0308.png', 'A/0327.png', 'A/0331.png', 'B/0512.png', 'B/0042.png', 'B/0059.png', 'B/0109.png', 'B/0129.png', 'B/0322.png', 'B/0074.png', 'B/0195.png', 'B/0355.png', 'B/0482.png', 'B/0061.png', 'B/0262.png', 'B/0086.png', 'B/0362.png', 'B/0249.png', 'B/0236.png', 'B/0186.png', 'B/0039.png', 'B/0203.png', 'B/0491.png', 'B/0201.png', 'B/0281.png', 'B/0420.png', 'B/0072.png', 'B/0464.png', 'B/0323.png', 'B/0053.png', 'B/0455.png', 'B/0202.png', 'B/0376.png', 'B/0550.png', 'B/0250.png', 'B/0213.png', 'B/0088.png', 'B/0251.png', 'B/0018.png', 'B/0122.png', 'B/0045.png', 'B/0271.png', 'B/0505.png', 'B/0316.png', 'B/0359.png', 'B/0499.png', 'B/0465.png', 'B/0137.png', 'B/0536.png', 'B/0168.png', 'B/0155.png', 'B/0220.png', 'B/0237.png', 'B/0145.png', 'B/0545.png', 'B/0265.png', 'B/0258.png', 'B/0205.png', 'B/0181.png', 'B/0218.png', 'B/0546.png', 'B/0458.png', 'B/0509.png', 'B/0390.png', 'B/0206.png', 'B/0192.png', 'B/0522.png', 'B/0229.png', 'B/0138.png', 'B/0211.png', 'B/0068.png', 'B/0243.png', 'B/0529.png', 'B/0438.png', 'B/0038.png', 'B/0190.png', 'B/0309.png', 'B/0365.png', 'B/0381.png', 'B/0478.png', 'B/0006.png', 'B/0462.png', 'B/0493.png', 'B/0526.png', 'B/0240.png', 'B/0210.png', 'B/0440.png', 'Left/0122.png', 'Left/0398.png', 'Left/0393.png', 'Left/0149.png', 'Left/0311.png', 'Left/0008.png', 'Left/0189.png', 'Left/0134.png', 'Left/0054.png', 'Left/0110.png', 'Left/0364.png', 'Left/0389.png', 'Left/0029.png', 'Left/0159.png', 'Left/0336.png', 'Left/0356.png', 'Left/0084.png', 'Left/0332.png', 'Left/0329.png', 'Left/0091.png', 'Left/0306.png', 'Left/0068.png', 'Left/0154.png', 'Left/0246.png', 'Left/0204.png', 'Left/0339.png', 'Left/0042.png', 'Left/0426.png', 'Left/0077.png', 'Left/0236.png', 'Left/0196.png', 'Left/0205.png', 'Left/0210.png', 'Left/0419.png', 'Left/0407.png', 'Left/0304.png', 'Left/0043.png', 'Left/0105.png', 'Left/0070.png', 'Left/0380.png', 'Left/0061.png', 'Left/0168.png', 'Left/0103.png', 'Left/0421.png', 'Left/0326.png', 'Left/0141.png', 'Left/0434.png', 'Left/0056.png', 'Left/0268.png', 'Left/0185.png', 'Left/0242.png', 'Left/0368.png', 'Left/0027.png', 'Left/0038.png', 'Left/0264.png', 'Left/0163.png', 'Left/0048.png', 'Left/0330.png', 'Left/0431.png', 'Left/0247.png', 'Left/0272.png', 'Left/0172.png', 'Left/0039.png', 'Left/0136.png', 'Left/0386.png', 'Left/0087.png', 'Left/0107.png', 'Left/0269.png', 'Right/0041.png', 'Right/0234.png', 'Right/0174.png', 'Right/0330.png', 'Right/0204.png', 'Right/0215.png', 'Right/0332.png', 'Right/0280.png', 'Right/0216.png', 'Right/0338.png', 'Right/0070.png', 'Right/0333.png', 'Right/0172.png', 'Right/0159.png', 'Right/0345.png', 'Right/0073.png', 'Right/0091.png', 'Right/0209.png', 'Right/0166.png', 'Right/0103.png', 'Right/0193.png', 'Right/0077.png', 'Right/0317.png', 'Right/0116.png', 'Right/0235.png', 'Right/0301.png', 'Right/0071.png', 'Right/0194.png', 'Right/0327.png', 'Right/0191.png', 'Right/0308.png', 'Right/0211.png', 'Right/0170.png', 'Right/0262.png', 'Right/0053.png', 'Right/0232.png', 'Right/0340.png', 'Right/0095.png', 'Right/0188.png', 'Right/0264.png', 'Right/0293.png', 'Right/0115.png', 'Right/0126.png', 'Right/0276.png', 'Right/0371.png', 'Right/0224.png', 'Right/0155.png', 'Right/0341.png', 'Right/0270.png', 'Right/0282.png', 'Right/0378.png', 'Right/0007.png', 'Right/0068.png']
测试集平均损失: 0.0242, 准确率: 98.59%
混淆矩阵:
[[79 0 0 0]
[ 1 83 0 0]
[ 0 0 67 1]
[ 0 0 2 51]]
分类报告:
precision recall f1-score support
A 0.99 1.00 0.99 79
B 1.00 0.99 0.99 84
Left 0.97 0.99 0.98 68
Right 0.98 0.96 0.97 53
accuracy 0.99 284
macro avg 0.98 0.98 0.98 284
weighted avg 0.99 0.99 0.99 284
从混淆矩阵中可以看出,只有B和Left出现一个错误。
到这里认为模型训练基本完成,然后部署到树莓派平台进行实际测试。
首先在本地编写一个用于预测的程序:
import torch
from functions import *
import numpy as np
from alexnet import AlexNet
from torchvision import transforms
import cv2
# 加载图像
img = cv2.imread("alex_data/A/0001.png")
# cv2.imshow("img",img)
# cv2.waitKey(3000)
# 定义变化器
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5)), # 标准化
transforms.RandomHorizontalFlip(), # 随机水平裁剪
])
# 标签映射
label_dict = {
0:"A",
1:"B",
2:"Left",
3:"Right",
}
print(img.shape)
# 预测
model = AlexNet()
pth = "alex_model.pth"
idx,name,prob = predict(model,pth,img,transform,label_dict)
print(f"预测索引:{idx},类别:{name},概率:{prob}")
注意,上面的代码会报出这个错误:
RuntimeError: mat1 and mat2 shapes cannot be multiplied (256x25 and 6400x4096)
这是因为预测时输入的图像没有使用加载器直接输入到网络,他的形状经过transform之后是[3,224,224]
而在训练和验证的时候,模型接受的形状实际上应该是一个四维的:[batch_size,3,224,224]
因此需要在第0维添加一个batch_size的维度。 最后修改的代码:
import torch
from functions import *
import numpy as np
from alexnet import AlexNet
from torchvision import transforms
import cv2
# 加载图像
img = cv2.imread("alex_data/A/0001.png")
img = torch.from_numpy(img).float()
img = img.unsqueeze(0)
img = torch.transpose(img,1,3)
print(img.shape)
# cv2.imshow("img",img)
# cv2.waitKey(3000)
# 定义变化器
transform = transforms.Compose([
transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5)), # 标准化
transforms.RandomHorizontalFlip(), # 随机水平裁剪
])
# 标签映射
label_dict = {
0:"A",
1:"B",
2:"Left",
3:"Right",
}
# 预测
model = AlexNet()
pth = "alex_model.pth"
idx,name,prob = predict(model,pth,img,transform,label_dict)
print(f"预测索引:{idx},类别:{name},概率:{prob}")
部署到树莓派之后:
差的离谱。。。 并且在调试的时候,由于alexnet训练后保存的模型比较大(接近200MB),树莓派加载需要一定的时间,而Lenet模型的加载就很快,在预测的时候Alexnet的速度相比Lenet也要慢上几倍。
对于这种场景,或许AlexNet不太适用,也可能是我在训练的时候没有做好优化。
目前还是采用训练好的Lenet进行预测,这样预测的准确度和速度都比较好。
AlexNet目前已有优化思路: 首先重新录制数据,先前的数据集是强制resize的,有些变形,这可能导致实际情况和训练时的误差较大,接下来我打算只保留掩膜,也就是类似这样的图像:
来进行训练,同时为了匹配输入图像,简单调整一下ALexNet的网络结构。
最终比赛中用的模型还是决定采用Lenet,AlexNet接下来就当做研究优化。
采用Lenet的好处是他很快,而且树莓派的性能羸弱,赛项对速度的要求又很高,并且在国赛的时候使用的是红底的标志牌,如果使用AlexNet实现任务还需要对红色的标志牌重新训练模型,而使用Lenet只需要对图像预处理的时候更改一下HSV阈值即可。综合对比下来最后还是使用LeNet比较好。