卷积神经网络项目实现文档
关于项目实现的文档说明书,三个要素:数据、模型、训练
1、项目简介
关于项目的基本介绍。
大颗粒度分类:以物种作为分类,比如飞机、青蛙、狗、猫、马、鹿等。
小颗粒度分类:就是一个物种,比如汽车,同一个品牌的汽车又分很多车型,产线。比如奥迪里面就有很多分类。
实体颗粒度分类:具体到具体的人,比如指纹识别、人脸识别等具体的个体,具体的实体
1.1 项目名称
基于CNN实现奥迪车型的小颗粒度分类
1.2 项目简介
该项目旨在通过卷积神经网络(CNN)实现奥迪不同车型的小颗粒度分类,主要针对同一个品牌的汽车的不同车型来进行细粒度的视觉识别与分类。同一个品牌的汽车,整体外观风格大体一致,差别非常细微,因此传统的图像分类方法难以满足精确识别的需求。通过引入深度学习中的CNN模型,本项目将构建一个专门用于奥迪车型细粒度分类的网络架构,充分利用卷积层的局部特征提取能力以及深层网络的高阶特征融合能力。项目将包括数据预处理、特征提取、模型训练与评估等步骤,使用奥迪车型图像数据集进行模型训练,并通过精细的特征学习提升分类精度。此外,项目还将探索数据增强、迁移学习等技术手段,以提高模型的泛化能力和鲁棒性,最终实现高精度的小颗粒度车型分类。
2、数据
kaggle上的开源汽车数据集,截取了上面的10种奥迪车型作为数据集
2.3 数据增强
提升模型的泛化能力和鲁棒性。
# --- 2. 数据预处理 ---
transform_train = transforms.Compose([
transforms.Resize(256),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
transform_test = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
2.4 数据分割
训练集约占80%
测试集约占%20
3. 神经网络
使用了ResNet网络进行训练,然后在其基础上进行优化
4. 模型训练
def train_model(model, train_loader, test_loader, criterion, optimizer, device, num_epochs, writer):
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
running_corrects = 0
# 训练阶段
for inputs, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}"):
inputs = inputs.to(device)
labels = labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
_, preds = torch.max(outputs, 1)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / len(train_loader.dataset)
epoch_acc = running_corrects.double() / len(train_loader.dataset)
# 验证阶段
model.eval()
val_loss = 0.0
val_corrects = 0
with torch.no_grad():
for inputs, labels in test_loader:
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
_, preds = torch.max(outputs, 1)
val_loss += loss.item() * inputs.size(0)
val_corrects += torch.sum(preds == labels.data)
val_epoch_loss = val_loss / len(test_loader.dataset)
val_epoch_acc = val_corrects.double() / len(test_loader.dataset)
# 打印训练和验证结果
print(f"Epoch {epoch+1}/{num_epochs} - Train Loss: {epoch_loss:.4f}, Train Acc: {epoch_acc:.4f}, "
f"Val Loss: {val_epoch_loss:.4f}, Val Acc: {val_epoch_acc:.4f}")
# 记录训练和验证指标到 TensorBoard
writer.add_scalar("Loss/train", epoch_loss, epoch)
writer.add_scalar("Accuracy/train", epoch_acc, epoch)
writer.add_scalar("Loss/val", val_epoch_loss, epoch)
writer.add_scalar("Accuracy/val", val_epoch_acc, epoch)
return model
4.1 训练参数
轮次:ecpochs = 70
批次:batch_size=32
学习率:lr=1e-4
4.2 损失函数
交叉熵损失函数
4.3 优化器
使用动量优化器
optim.Adam()
4.4 训练过程可视化
使用tensorBoard 和 wandb
网络结构:
5. 模型验证
验证我们的模型的鲁棒性和泛化能力
5.1 验证过程数据化
生成Excel:
5.2 指标报表
准确度:0.7692307
精确度:0.7941640
召回率:0.7689926
5.3 混淆矩阵
可视化
|
6. 模型优化
在模型中加入CBMA模块,并且添加了批归一化,增加训练轮次
6.1 注意力机制
import torch
import torch.nn as nn
import torch.nn.functional as F
class ChannelAttention(nn.Module):
def __init__(self, in_planes, ratio=16):
super(ChannelAttention, self).__init__()
self.max_pool = nn.AdaptiveMaxPool2d(1)
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.fc1 = nn.Sequential(
nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False),
nn.ReLU(),
nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False),
)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
max_out = self.fc1(self.max_pool(x))
avg_out = self.fc1(self.avg_pool(x))
out = max_out + avg_out
return self.sigmoid(out)
class SpatialAttention(nn.Module):
def __init__(self, kernel_size=7):
super(SpatialAttention, self).__init__()
assert kernel_size in (3, 7), "kernel size must be 3 or 7"
self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
avg_out = torch.mean(x, dim=1, keepdim=True)
max_out, _ = torch.max(x, dim=1, keepdim=True)
x = torch.cat([avg_out, max_out], dim=1)
x = self.conv(x)
return self.sigmoid(x)
6.2 增加训练轮次和网络深度
num_epochs = 70 # 增加训练轮数
# 3. 添加一个更深的注意力机制模型
class ResNet50CBAMEnhanced(nn.Module):
def __init__(self, num_classes=10, pretrained=True):
super(ResNet50CBAMEnhanced, self).__init__()
# 使用我们自定义的带CBAM的ResNet50
self.resnet_cbam = resnet50CBAM(weights="IMAGENET1K_V1" if pretrained else None, num_classes=1000)
# 在最终分类层之前添加额外的注意力机制
self.final_ca = ChannelAttention(self.resnet_cbam.fc.in_features)
self.final_sa = SpatialAttention()
# 添加额外的全连接层(替换原来的fc层)
self.dropout = nn.Dropout(0.5)
self.fc_extra = nn.Linear(self.resnet_cbam.fc.in_features, 512)
self.bn_extra = nn.BatchNorm1d(512)
self.fc_final = nn.Linear(512, num_classes)
# 移除原来的fc层,因为我们用自己的分类器
del self.resnet_cbam.fc
def forward(self, x):
# 获取resnet的特征
x = self.resnet_cbam.conv1(x)
x = self.resnet_cbam.bn1(x)
x = self.resnet_cbam.relu(x)
x = self.resnet_cbam.maxpool(x)
x = self.resnet_cbam.layer1(x)
x = self.resnet_cbam.layer2(x)
x = self.resnet_cbam.layer3(x)
x = self.resnet_cbam.layer4(x)
# 应用最终的注意力机制
x = self.final_ca(x) * x
x = self.final_sa(x) * x
# 全局平均池化
x = self.resnet_cbam.avgpool(x)
x = torch.flatten(x, 1)
# 使用我们自己的分类器
x = self.dropout(x)
x = self.fc_extra(x)
x = self.bn_extra(x)
x = F.relu(x)
x = self.fc_final(x)
return x
6.3 预训练和更复杂的图片预处理
# 指定预训练权重,加载主干网络的权重(忽略CBAM层)
if weights is not None:
# 获取官方预训练权重
pretrained_state_dict = weights.get_state_dict(progress=progress)
# 获取当前模型的 state_dict
model_state_dict = model.state_dict()
# 筛选并加载匹配的权重
matched_keys = []
for key in pretrained_state_dict:
# 只加载那些在 model_state_dict 中存在且形状匹配的键
if key in model_state_dict and pretrained_state_dict[key].shape == model_state_dict[key].shape:
model_state_dict[key] = pretrained_state_dict[key]
matched_keys.append(key)
print(f"Successfully loaded {len(matched_keys)} pretrained parameters.")
# 将更新后的 state_dict 加载回模型
model.load_state_dict(model_state_dict)
7. 模型应用
7.1 图片预处理
# 图像预处理
img_tensor = imgread(image)
img_tensor = img_tensor.to(device)
7.2 模型推理
# 执行推理
with torch.no_grad():
outputs = model(img_tensor)
# 应用softmax获取概率
probabilities = torch.softmax(outputs, dim=1)
# 获取预测类别和置信度
confidence, predicted_class_id = torch.max(probabilities, 1)
# 转换为CPU numpy格式
confidence = confidence.cpu().item()
predicted_class_id = predicted_class_id.cpu().item()
7.3 类别显示
使用gradio实现交互式应用
8. 模型移植
使用ONNX
# ONNX.py
import os
import torch
import torch.nn as nn
from model import ResNet50CBAMEnhanced # 导入你的模型
if __name__ == "__main__":
# 获取脚本所在目录
dir = os.path.dirname(__file__)
# 权重文件路径 (确保这个 .pth 文件是你训练后保存的)
weightpath = os.path.join(dir, "resnet50CABM_enhanced_car_model_trained.pth")
# ONNX 输出路径
onnxpath = os.path.join(dir, "resnet50_cbam_enhanced.onnx")
# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# --- 关键修改开始 ---
# 1. 创建 ResNet50CBAMEnhanced 实例
num_classes = 10 # 替换成你训练时的实际类别数
model = ResNet50CBAMEnhanced(num_classes=num_classes, pretrained=False)
# 2. 加载训练好的权重
model.load_state_dict(torch.load(weightpath, map_location=device))
# --- 关键修改结束 ---
# 将模型移动到指定设备
model.to(device)
# 设置模型为评估模式 (非常重要!)
model.eval()
# 创建一个示例输入张量
# 尺寸应与训练时一致 (1, 3, 224, 224)
x = torch.randn(1, 3, 224, 224, device=device)
# 导出为 ONNX 格式
torch.onnx.export(
model,
x,
onnxpath,
verbose=True, # 输出转换过程的详细信息
input_names=["input"], # 指定输入节点名称
output_names=["output"], # 指定输出节点名称
opset_version=11, # 推荐指定一个 opset 版本 (如 11, 12, 13, 14, 15)
dynamic_axes={
'input': {0: 'batch_size'},
'output': {0: 'batch_size'}
} # 支持动态batch size
)
print("ONNX 导出成功!")
8.1 导出ONNX
8.2 使用ONNX推理
import cv2
import numpy as np
from PIL import Image
from torchvision import transforms
import onnxruntime as ort
# 定义数据预处理流程
# 注意:尺寸必须与训练/导出模型时一致!这里是 224x224
transformdata = transforms.Compose(
[
transforms.Resize((224, 224)), # 改为 224x224
transforms.ToTensor(),
# VGG 训练时使用的 ImageNet 的均值和标准差
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
]
)
def imgread(img_path):
# 使用OpenCV读取图像
imgdata = cv2.imread(img_path)
if imgdata is None:
raise FileNotFoundError(f"无法读取图像: {img_path}")
# 转换成RGB
imgdata = cv2.cvtColor(imgdata, cv2.COLOR_BGR2RGB)
# 将numpy数组转换成PIL Image
pil_img = Image.fromarray(imgdata)
# 应用预处理变换
imgdata = transformdata(pil_img)
# tensor ---> CHW ---> NCHW (增加 batch 维度)
imgdata = imgdata.unsqueeze(0).numpy() # 转为 numpy 数组,形状 (1, 3, 224, 224)
return imgdata
def inference():
# ONNX 模型路径 (确保路径正确)
onnx_model_path = "./resnet50_cbam_enhanced.onnx" # 与导出的文件名一致
# 加载ONNX模型
try:
model = ort.InferenceSession(onnx_model_path, providers=["CPUExecutionProvider"])
except Exception as e:
print(f"加载 ONNX 模型失败: {e}")
return
# 获取模型的输入名称
input_name = model.get_inputs()[0].name
print(f"模型输入名称: {input_name}")
# 图像路径 (替换为你要预测的图片路径)
img_path = "./images/Audi R8 Coupe 2012 001.jpg" # <-- 修改这里!
try:
# 图像预处理
imgdata = imgread(img_path)
print(f"输入数据形状: {imgdata.shape}")
except Exception as e:
print(f"图像预处理失败: {e}")
return
# 执行推理
try:
out = model.run(None, {input_name: imgdata})
# out 是一个列表,out[0] 是模型的输出 (numpy array)
print(f"模型输出形状: {out[0].shape}")
except Exception as e:
print(f"推理执行失败: {e}")
return
# 获取预测类别
# 假设输出是 (1, 10) 的 logits
predicted_class_id = np.argmax(out[0][0]) # out[0] 是 (1, 10), out[0][0] 是 (10,)
confidence = np.max(out[0][0]) # 获取最高置信度
# 类别标签 (根据你的训练数据顺序定义)
classlabels = ["Audi 100 Sedan 1994", "Audi 100 Wagon 1994", "Audi A5 Coupe 2012", "Audi R8 Coupe 2012", "Audi RS 4 Convertible 2008",
"Audi S4 Sedan 2007", "Audi S4 Sedan 2012", "Audi S5 Convertible 2012", "Audi S5 Convertible 2012", "Audi S6 Sedan 2011"] # <-- 替换为你的实际类别名!
predicted_label = classlabels[predicted_class_id]
print(f"预测类别: {predicted_label}")
print(f"类别 ID: {predicted_class_id}")
print(f"置信度: {confidence:.4f}")
if __name__ == "__main__":
inference()