已在GitHub开源与本博客同步的ResNet50v2_RK3588_Classificationt项目,地址:https://github.com/A7bert777/ResNet50v2_RK3588_Classification
详细使用教程,可参考README.md或参考本博客第八章 模型部署
一、项目回顾
博主之前有写过YOLO11、YOLOv8目标检测&图像分割、YOLOv10目标检测、MoblieNetv2图像分类的模型训练、转换、部署文章,感兴趣的小伙伴可以了解下:
【YOLO11部署至RK3588】模型训练→转换RKNN→开发板部署
【YOLOv10部署RK3588】模型训练→转换rknn→部署流程
【YOLOv8-obb部署至RK3588】模型训练→转换RKNN→开发板部署
【YOLOv8-pose部署至RK3588】模型训练→转换RKNN→开发板部署
【YOLOv8部署至RK3588】模型训练→转换rknn→部署全流程
【YOLOv8seg部署RK3588】模型训练→转换rknn→部署全流程
【MobileNetv2图像分类部署至RK3588】模型训练→转换rknn→部署流程
YOLOv8n部署RK3588开发板全流程(pt→onnx→rknn模型转换、板端后处理检测)
二、模型选择介绍
近期需要做一个针对分类图像的模型,并部署到RK3588公版的开发板上,可选择的有YOLOv8、ResNet50、MobileNet…
而瑞芯微的部署demo:rknn_model_zoo中可以看到,并没有yolov8的图像分类demo,只有ResNet和MobileNet的demo。因此为了省事,就按照瑞芯微的方法来,看了下其官网ONNX模型示例的大小,ResNet比MobileNet大了一个数量级,根据经验,肯定模型越小运行速度越快,而MobileNet的部署流程已经在之前的博客中展示过了,大家可以去看之前博主的相关文章。最近没事做了个指示灯亮灭的ResNet50的图像分类项目,涉及模型训练、转ONNX、转RKNN量化以及RK3588开发板调试部署,查了下CSDN上暂未有关于ResNet在RK系列开发板的免费详细教程,遂开此文,相互学习。
三、文件梳理
Resnet50的训练、转换、部署所需四个项目文件:
第一个:train050.py和pt2onnx.py(下文中会给详细的脚本内容);
第二个:用于在虚拟机中进行onnx转rknn的虚拟环境配置项目文件(链接在此);
第三个:在开发板上做模型部署的项目文件(链接在此)。
注:
1.第三个项目文件中的内容很多,里面涉及到rknn模型转换以及模型部署的所有内容,所以该文件在模型转换中也要与第三个文件配合使用。
版本如下:
第二个文件rknn-toolkit2为v2.1.0
第三个文件rknn_model_zoo也用v2.1.0(rknn-toolkit2尽量和rknn_model_zoo版本一致)
如图所示:
下面进入正题,先开始模型训练的必要准备。
四、文件及环境配置
由于ResNet50结构较为简单,且图像分类也没有目标检测的复杂后处理流程,因此训练ResNet50所需要的文件很少,不同于以前训练YOLO系列模型的繁琐流程。
以下是训练所用脚本:
其中train_050.py的内容如下所示:
import torch # 导入PyTorch库,用于深度学习
from torchvision import datasets, models, transforms # 导入torchvision中的数据集、模型和图像转换模块
import torch.nn as nn # 导入PyTorch的神经网络模块
import torch.optim as optim # 导入PyTorch的优化器模块
from torch.utils.data import DataLoader # 导入DataLoader,用于批量加载数据
import time # 导入时间模块,用于计算训练时间
import numpy as np # 导入NumPy库,用于数值计算
import matplotlib.pyplot as plt # 导入Matplotlib库,用于绘图
import os # 导入操作系统模块,用于处理文件路径
from tqdm import tqdm # 导入tqdm模块,用于显示进度条
image_transforms = {
'train': transforms.Compose([
transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)), # 随机裁剪并调整大小为256x256
transforms.RandomRotation(degrees=15), # 随机旋转15度
transforms.RandomHorizontalFlip(), # 随机水平翻转
#transforms.CenterCrop(size=224), # 中心裁剪为224x224
transforms.CenterCrop(size=112),
transforms.ToTensor(), # 将图像转换为Tensor
transforms.Normalize([0.485, 0.456, 0.406], # 归一化,使用ImageNet的均值和标准差
[0.229, 0.224, 0.225])
]),
'valid': transforms.Compose([
transforms.Resize(size=256), # 调整大小为256x256
#transforms.CenterCrop(size=224), # 中心裁剪为224x224
transforms.CenterCrop(size=112),
transforms.ToTensor(), # 将图像转换为Tensor
transforms.Normalize([0.485, 0.456, 0.406], # 归一化,使用ImageNet的均值和标准差
[0.229, 0.224, 0.225])
])
}
# 三、加载数据
# torchvision.transforms包DataLoader是 Pytorch 重要的特性,它们使得数据增加和加载数据变得非常简单。
# 使用 DataLoader 加载数据的时候就会将之前定义的数据 transform 就会应用的数据上了。
dataset = 'xxx' # 数据集名称
train_directory = '/xxx/train' # 训练集路径
valid_directory = '/xxx/val' # 验证集路径
batch_size = 128 # 批量大小
num_classes = 2 #分类种类数
print(train_directory) # 打印训练集路径
print(valid_directory) # 打印验证集路径
data = {
'train': datasets.ImageFolder(root=train_directory, transform=image_transforms['train']), # 加载训练集
'valid': datasets.ImageFolder(root=valid_directory, transform=image_transforms['valid']) # 加载验证集
}
# 打印训练集类别及其对应编号
print("训练集图片类别及其对应编号(种类名:编号):",data['train'].class_to_idx)
print("测试集图片类别及其对应编号:",data['valid'].class_to_idx)
train_data_size = len(data['train']) # 训练集大小
valid_data_size = len(data['valid']) # 验证集大小
train_data = DataLoader(data['train'], batch_size=batch_size, shuffle=True, num_workers=8) # 训练集DataLoader
valid_data = DataLoader(data['valid'], batch_size=batch_size, shuffle=True, num_workers=8) # 验证集DataLoader
# 打印训练集和验证集的大小
print("训练集图片数量:",train_data_size, "测试集图片数量:",valid_data_size)
# 四、迁移学习
# 这里使用ResNet-50的预训练模型。
#resnet50 = models.resnet50(pretrained=True) # 旧版加载预训练模型的方式
resnet50 = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1) # 加载ResNet-50预训练模型
# 在PyTorch中加载模型时,所有参数的‘requires_grad’字段默认设置为true。这意味着对参数值的每一次更改都将被存储,以便在用于训练的反向传播图中使用。
# 这增加了内存需求。由于预训练的模型中的大多数参数已经训练好了,因此将requires_grad字段重置为false。
for param in resnet50.parameters():
param.requires_grad = False # 冻结模型参数,不更新梯度
# 为了适应自己的数据集,将ResNet-50的最后一层替换为,将原来最后一个全连接层的输入喂给一个有256个输出单元的线性层,接着再连接ReLU层和Dropout层,然后是256 x 6的线性层,输出为6通道的softmax层。
fc_inputs = resnet50.fc.in_features # 获取全连接层的输入特征数
resnet50.fc = nn.Sequential(
nn.Linear(fc_inputs, 256), # 全连接层,输入为fc_inputs,输出为256
nn.ReLU(), # ReLU激活函数
nn.Dropout(0.4), # Dropout层,丢弃率为0.4
nn.Linear(256, num_classes), # 全连接层,输入为256,输出为num_classes
nn.LogSoftmax(dim=1) # LogSoftmax层,用于多分类任务
)
# 用GPU进行训练。
resnet50 = resnet50.to('cuda:0') # 将模型移动到GPU
# 定义损失函数和优化器。
loss_func = nn.NLLLoss() # 负对数似然损失函数
optimizer = optim.Adam(resnet50.parameters()) # Adam优化器
# 五、训练
def train_and_valid(model, loss_function, optimizer, epochs=25):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 设置设备
history = [] # 用于存储训练和验证的损失和准确率
best_acc = 0.0 # 最佳验证准确率
best_epoch = 0 # 最佳验证准确率对应的epoch
for epoch in range(epochs):
epoch_start = time.time() # 记录epoch开始时间
print("Epoch: {}/{}".format(epoch+1, epochs)) # 打印当前epoch
model.train() # 设置模型为训练模式
train_loss = 0.0 # 训练损失
train_acc = 0.0 # 训练准确率
valid_loss = 0.0 # 验证损失
valid_acc = 0.0 # 验证准确率
for i, (inputs, labels) in enumerate(tqdm(train_data)): # 遍历训练集
inputs = inputs.to(device) # 将输入数据移动到GPU
labels = labels.to(device) # 将标签数据移动到GPU
#因为这里梯度是累加的,所以每次记得清零
optimizer.zero_grad() # 清空梯度
outputs = model(inputs) # 前向传播,计算输出
loss = loss_function(outputs, labels) # 计算损失
#print("标签值:",labels) # 打印标签值
#print("输出值:",outputs) # 打印输出值
loss.backward() # 反向传播,计算梯度
optimizer.step() # 更新参数
train_loss += loss.item() * inputs.size(0) # 累计训练损失
ret, predictions = torch.max(outputs.data, 1) # 获取预测结果
correct_counts = predictions.eq(labels.data.view_as(predictions)) # 计算正确预测的数量
acc = torch.mean(correct_counts.type(torch.FloatTensor)) # 计算准确率
train_acc += acc.item() * inputs.size(0) # 累计训练准确率
with torch.no_grad(): # 不计算梯度
model.eval() # 设置模型为评估模式
for j, (inputs, labels) in enumerate(tqdm(valid_data)): # 遍历验证集
inputs = inputs.to(device) # 将输入数据移动到GPU
labels = labels.to(device) # 将标签数据移动到GPU
outputs = model(inputs) # 前向传播,计算输出
loss = loss_function(outputs, labels) # 计算损失
valid_loss += loss.item() * inputs.size(0) # 累计验证损失
ret, predictions = torch.max(outputs.data, 1) # 获取预测结果
correct_counts = predictions.eq(labels.data.view_as(predictions)) # 计算正确预测的数量
acc = torch.mean(correct_counts.type(torch.FloatTensor)) # 计算准确率
valid_acc += acc.item() * inputs.size(0) # 累计验证准确率
avg_train_loss = train_loss/train_data_size # 计算平均训练损失
avg_train_acc = train_acc/train_data_size # 计算平均训练准确率
avg_valid_loss = valid_loss/valid_data_size # 计算平均验证损失
avg_valid_acc = valid_acc/valid_data_size # 计算平均验证准确率
# 将结果存入history
history.append([avg_train_loss, avg_valid_loss, avg_train_acc, avg_valid_acc])
# 如果当前验证准确率大于最佳验证准确率
if best_acc < avg_valid_acc:
best_acc = avg_valid_acc # 更新最佳验证准确率
best_epoch = epoch + 1 # 更新最佳验证准确率对应的epoch
# 记录epoch结束时间
epoch_end = time.time()
# 打印当前epoch的训练和验证结果
print("Epoch: {:03d}, Training: Loss: {:.4f}, Accuracy: {:.4f}%, \n\t\tValidation: Loss: {:.4f}, Accuracy: {:.4f}%, Time: {:.4f}s".format(
epoch+1, avg_valid_loss, avg_train_acc*100, avg_valid_loss, avg_valid_acc*100, epoch_end-epoch_start
))
# 打印最佳验证准确率及其对应的epoch
print("Best Accuracy for validation : {:.4f} at epoch {:03d}".format(best_acc, best_epoch))
# 保存模型
torch.save(model, 'models/'+dataset+'_model_'+str(epoch+1)+'.pt')
# 返回训练好的模型和history
return model, history
num_epochs = 100 #训练周期数
trained_model, history = train_and_valid(resnet50, loss_func, optimizer, num_epochs) # 训练模型
torch.save(history, 'models/'+dataset+'_history.pt') # 保存训练历史
history = np.array(history) # 将history转换为NumPy数组
plt.plot(history[:, 0:2]) # 绘制训练和验证损失曲线
plt.legend(['Tr Loss', 'Val Loss']) # 设置图例
plt.xlabel('Epoch Number') # 设置x轴标签
plt.ylabel('Loss') # 设置y轴标签
plt.ylim(0, 1) # 设置y轴范围
plt.savefig(dataset+'_loss_curve.png') # 保存损失曲线图
plt.show() # 显示损失曲线图
plt.plot(history[:, 2:4]) # 绘制训练和验证准确率曲线
plt.legend(['Tr Accuracy', 'Val Accuracy']) # 设置图例
plt.xlabel('Epoch Number') # 设置x轴标签
plt.ylabel('Accuracy') # 设置y轴标签
plt.ylim(0, 1) # 设置y轴范围
plt.savefig(dataset+'_accuracy_curve.png') # 保存准确率曲线图
plt.show() # 显示准确率曲线图
pt2onnx.py的完整内容如下所示:
'''
该脚本功能为将训练好的ResNet50图像分类模型(PyTorch .pt格式)转换为ONNX格式(opset_version=12)
在脚本末尾的main函数中直接配置输入模型路径、输出路径和参数
'''
import torch
import torch.onnx
def convert_pt_to_onnx(pt_model_path, onnx_model_path, input_size=(112, 112), opset_version=12):
"""
将PyTorch模型转换为ONNX格式
参数:
pt_model_path: PyTorch模型文件路径(.pt)
onnx_model_path: 输出ONNX模型文件路径(.onnx)
input_size: 模型输入尺寸 (H, W)
opset_version: ONNX opset版本
"""
# 1. 加载PyTorch模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.load(pt_model_path, map_location=device)
model.eval() # 设置为评估模式
# 2. 准备虚拟输入 (batch_size固定为1)
batch_size = 1
channels = 3
dummy_input = torch.randn(batch_size, channels, *input_size).to(device)
# 3. 导出ONNX模型(使用静态batch_size)
torch.onnx.export(
model, # 要导出的模型
dummy_input, # 模型输入
onnx_model_path, # 输出文件路径
export_params=True, # 导出模型参数
opset_version=opset_version, # 指定opset版本
do_constant_folding=True, # 优化常量折叠
input_names=['input'], # 输入节点名称
output_names=['output'], # 输出节点名称
# 移除dynamic_axes参数以固定batch_size为1
)
print(f"[成功] 模型已导出到: {onnx_model_path}")
print(f"输入尺寸: {batch_size}x{channels}x{input_size[0]}x{input_size[1]} (静态batch_size)")
print(f"输出尺寸: {batch_size}x2 (静态batch_size)")
print(f"ONNX opset版本: {opset_version}")
if __name__ == "__main__":
# 直接配置参数(无需命令行参数)
pt_model_path = "./xxx.pt"
onnx_model_path = "./xxx.onnx"
input_height = 112
input_width = 112
opset_version = 12
# 执行转换
convert_pt_to_onnx(
pt_model_path=pt_model_path,
onnx_model_path=onnx_model_path,
input_size=(input_height, input_width),
opset_version=opset_version
)
然后是训练环境:
大家可以按照图中版本自行pip或conda安装。
五、模型训练
激活并配置好环境后,直接执行脚本:
python train_050.py
如下所示:
训练完成后,就能在当前目录下的model文件夹下看到我们各个epoch得到的模型了:
六、PT转ONNX
然后根据终端显示的最优epoch的模型拿出来,使用pt2onnx.py对其进行转换,注意,要在pt2onnx.py脚本中修改你的pt模型路径和转出的onnx模型的保存路径,如下所示:
然后执行脚本:
python pt2onnx.py
执行后,在终端得到转换后的onnx模型,onnx模型可用netron工具打开,如下所示:
因为我只有两个类别:on、off,所以输出就是1,2(1是输出的batchsize)
七、ONNX转RKNN
在进行这一步的时候,如果你是在云服务器上运行,请先确保你租的卡能支持RKNN的转换运行。博主是在自己的虚拟机中进行转换。
先安装转换环境
这里我们先创建环境:
conda create -n rknn210 python=3.8
创建完成如下所示:
现在需要用到rknn-toolkit2-2.1.0文件。
进入rknn-toolkit2-2.1.0\rknn-toolkit2-2.1.0\rknn-toolkit2\packages文件夹下,看到如下内容:
在终端激活环境,在终端输入
pip install -r requirements_cp38-2.1.0.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
然后再输入
pip install rknn_toolkit2-2.1.0+708089d1-cp38-cp38-linux_x86_64.whl
然后,我们的转rknn环境就配置完成了。
现在要进行模型转换,其实大家可以参考rknn_model_zoo-2.1.0\examples\ResNet下的README指导进行转换,先打开rknn_model_zoo下的resnet.py:
进行如下修改:
↑注意要把输入从默认的224改成你自己的size,另外可以在model文件夹下放一张自己的数据集图片,用来测量量化后的推理效果。
然后在终端执行rknn转换命令:
python resnet.py ../model/indicator_light_on_off_v3_classifier_model_82_p0.9945.onnx rk3588
然后在model文件夹下就能看到我们转换得到的rknn模型了:
然后将rknn模型复制到win下,用netron打开,如下所示:
可以看到,模型很简洁,输入输出和onnx是一样的。
另外可以看到,fp32的onnx模型在转换成int8的rknn模型后,对比如下所示,模型体积只有单精的四分之一:
八、模型部署
如果前面流程都已实现,模型的结构也没问题的话,则可以进行最后一步:模型端侧部署。
我已经帮大家做好了所有的环境适配工作,科学上网后访问博主GitHub仓库:ResNet50v2_RK3588_Classification ,进行简单的路径修改就即可编译运行。
统一声明:
1、这个仓库的项目只能做图片检测,不支持视频流检测,没时间做这个,有需要的自己修改代码。
2、从GitHub的README.md中加QQ后直接说问题和小星星截图,对于常见的相同问题,很多都已在CSDN博客中提到了(RKNN转换流程是统一的,可去博主所有的RKNN相关博客下去翻评论),已在评论中详细解释过的问题,不予回复。
重点:请大家举手之劳,帮我的仓库点个小星星
点了小星星的同学可以免费帮你解决转模型与部署过程中遇到的问题。
注:保密原因,博主此博客的PT、ONNX、RKNN模型均保密不予提供,请大家拿自己的模型进行测试
git clone后把项目复制到开发板上,按如下流程操作:
①:cd build,删除所有build文件夹下的内容
②:cd src 修改main.cc,修改main函数中的如下四个内容:
把标签名的txt的路径改成你自己的路径:
其中txt文档的内容如下所示,其实就是你在训练ResNet50v2分类模型时的那几个类别名(即文件夹名,注意要按个按照文件夹名的先后顺序来,不然即便推理得分很高,但类名是错的)
③:把你之前训练好并已转成RKNN格式的模型放到ResNetv2_RK3588_Classification/model文件夹下,然后把你要检测的所有图片都放到ResNetv2_RK3588_Classification/inputimage下。
④:进入build文件夹进行编译
cd build
cmake ..
make
在build下生成可执行文件文件:rknn_resnet50_demo
完整流程如下所示,执行完后即可在build文件夹下看到生成的可执行文件:
拿出一个图片进行测试:
然后在build下打开终端输入如下命令(可执行文件 模型路径 输入图片路径):
./rknn_resnet50_demo ../model/indicator_light_on_off_v3_classifier_model_82_p0.9945.rknn ../inputimage/022684.png
终端结果如下所示:
总体来说,推理得分还是很高的,ResNetv2在各项任务上足以胜任,以上即为MobileNetv2图像分类任务部署至RK3588的全流程。