一、模型选择介绍
近期需要做一个针对分类图像的模型，并部署到RK3588公版的开发板上，可选择的有YOLOv8、ResNet、MobileNet…
而瑞芯微的部署demo：rknn_model_zoo中可以看到，并没有yolov8的图像分类demo，只有ResNet和MobileNet的demo。因此为了省事，就按照瑞芯微的方法来，看了下其官网ONNX模型示例的大小，ResNet比MobileNet大了一个数量级，根据经验，肯定模型越小运行速度越快，因此选择MobileNet进行分类任务。
注：rknn_model_zoo中只有ResNet50v2和MobileNetv2这两个版本，其余的分类模型版本并未公布，所以大家如果为了省事，最好就直接选择MobileNetv2版本吧，这样在部署时，和官方demo一致，非常方便。
另外，本篇文章参考了博主@nice-wyh，原文链接。
下面进入正题，先开始模型训练的必要准备。
二、文件及环境配置
由于MobileNetv2结构较为简单，且图像分类也没有目标检测的复杂后处理流程，因此训练MobileNetv2所需要的文件很少，不同于以前训练YOLO系列模型的繁琐流程。
1.所需文件
在这里，我们所需的文件都是py类型的，只有四个。
其中model_v2.py是MobileNetv2的定义文件
train.py和predict.py是模型训练和预测用文件
pt2onnx.py是训练得到的pt模型转onnx的文件
各文件详细内容如下，可直接复制粘贴
1.model_v2.py：
from torch import nn
import torch
# 将channel调整为离8最近的整数倍，这样的处理对硬件更加的友好，也有一定训练速度的提升
def _make_divisible(ch, divisor=8, min_ch=None):
"""
This function is taken from the original tf repo.
It ensures that all layers have a channel number that is divisible by 8
It can be seen here:
https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
"""
if min_ch is None:
min_ch = divisor
new_ch = max(min_ch, int(ch + divisor / 2) // divisor * divisor)
# Make sure that round down does not go down by more than 10%.
if new_ch < 0.9 * ch:
new_ch += divisor
return new_ch
class ConvBNReLU(nn.Sequential):
def __init__(self, in_channel, out_channel, kernel_size=3, stride=1, groups=1):
padding = (kernel_size - 1) // 2
super(ConvBNReLU, self).__init__(
# 这里如果group=1，则为普通卷积；group=输入特征矩阵的深度时，则为DW卷积
nn.Conv2d(in_channel, out_channel, kernel_size, stride, padding, groups=groups, bias=False),
nn.BatchNorm2d(out_channel),
nn.ReLU6(inplace=True)
)
class InvertedResidual(nn.Module):
def __init__(self, in_channel, out_channel, stride, expand_ratio):
super(InvertedResidual, self).__init__()
hidden_channel = in_channel * expand_ratio
# 当步长为1，且输入输出维度相同时，使用捷径分支
self.use_shortcut = stride == 1 and in_channel == out_channel
layers = []
if expand_ratio != 1:
# 1x1 pointwise conv
layers.append(ConvBNReLU(in_channel, hidden_channel, kernel_size=1))
layers.extend([
# 3x3 depthwise conv
ConvBNReLU(hidden_channel, hidden_channel, stride=stride, groups=hidden_channel),
# 1x1 pointwise conv(linear)
nn.Conv2d(hidden_channel, out_channel, kernel_size=1, bias=False),
nn.BatchNorm2d(out_channel),
])
self.conv = nn.Sequential(*layers)
def forward(self, x):
if self.use_shortcut:
return x + self.conv(x)
else:
return self.conv(x)
class MobileNetV2(nn.Module):
def __init__(self, num_classes=1000, alpha=1.0, round_nearest=8):
super(MobileNetV2, self).__init__()
block = InvertedResidual
input_channel = _make_divisible(32 * alpha, round_nearest)
last_channel = _make_divisible(1280 * alpha, round_nearest)
inverted_residual_setting = [
# t, c, n, s
# t:将输入特征矩阵深度调整t倍
# c:输入channel
# n:bottle(倒残差结构重复的次数)
# s:每个block中，第一个bottleneck的步长
[1, 16, 1, 1],
[6, 24, 2, 2],
[6, 32, 3, 2],
[6, 64, 4, 2],
[6, 96, 3, 1],
[6, 160, 3, 2],
[6, 320, 1, 1],
]
features = []
# conv1 layer
features.append(ConvBNReLU(3, input_channel, stride=2))
# building inverted residual residual blockes
for t, c, n, s in inverted_residual_setting:
output_channel = _make_divisible(c * alpha, round_nearest)
for i in range(n):
stride = s if i == 0 else 1
features.append(block(input_channel, output_channel, stride, expand_ratio=t))
input_channel = output_channel
# building last several layers
features.append(ConvBNReLU(input_channel, last_channel, 1))
# combine feature layers
self.features = nn.Sequential(*features)
# building classifier
self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
self.classifier = nn.Sequential(
nn.Dropout(0.2),
nn.Linear(last_channel, num_classes)
)
# weight initialization
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out')
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.BatchNorm2d):
nn.init.ones_(m.weight)
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Linear):
nn.init.normal_(m.weight, 0, 0.01)
nn.init.zeros_(m.bias)
def forward(self, x):
x = self.features(x)
x = self.avgpool(x)
x = torch.flatten(x, 1)
x = self.classifier(x)
return x
if __name__ == "__main__":
net = MobileNetV2(num_classes=10)
in_data = torch.randn(1, 3, 224, 224)
out = net(in_data)
print(out)
2. train.py：
import os
import sys
import json
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, datasets
from tqdm import tqdm
from model_v2 import MobileNetV2
def main():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"using {device} device.")
data_transform = {
"train": transforms.Compose([transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])]),
"val": transforms.Compose([transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])}
# 获取数据集路径
image_path = os.path.join(os.getcwd(), "data_set", "raw_data") # flower data set path
assert os.path.exists(image_path), "{} path does not exist.".format(image_path)
# 加载数据集，准备读取
train_dataset = datasets.ImageFolder(root=os.path.join(image_path, "train"), transform=data_transform["train"])
validate_dataset = datasets.ImageFolder(root=os.path.join(image_path, "val"), transform=data_transform["val"])
nw = min([os.cpu_count(), 16 if 16 > 1 else 0, 8]) # number of workers
print('Using {} dataloader workers every process'.format(nw))
# 加载数据集
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=nw)
validate_loader = torch.utils.data.DataLoader(validate_dataset, batch_size=16, shuffle=False, num_workers=nw)
train_num = len(train_dataset)
val_num = len(validate_dataset)
print(f"using {train_num} images for training, {val_num} images for validation.")
# {'cane':0, 'carpet':1, 'cavallo':2, 'elefante':3, 'farfalla':4, 'gallina':5, 'gatto':6, 'mucca':7, 'pecora':8, 'ragno':9, 'scoiattolo':10}
flower_list = train_dataset.class_to_idx
cla_dict = dict((val, key) for key, val in flower_list.items())
# write dict into json file
json_str = json.dumps(cla_dict, indent=4)
with open('class_indices.json', 'w') as json_file:
json_file.write(json_str)
# load pretrain weights
# download url: https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
net = MobileNetV2(num_classes=11)
model_weight_path = "./mobilenet_v2.pth"
assert os.path.exists(model_weight_path), f"file {model_weight_path} dose not exist."
pre_weights = torch.load(model_weight_path, map_location='cpu')
# delete classifier weights,因为预训练参数是基于ImageNet数据集训练的，类别为1000，所以需要删掉最后一层参数，只保留其他部分
pre_dict = {k: v for k, v in pre_weights.items() if net.state_dict()[k].numel() == v.numel()}
missing_keys, unexpected_keys = net.load_state_dict(pre_dict, strict=False)
# freeze features weights
# 预训练模型中，我们只希望微调最后几层，因此冻结前面的权重和偏置参数
for param in net.features.parameters():
param.requires_grad = False
net.to(device)
# define loss function
loss_function = nn.CrossEntropyLoss()
# construct an optimizer
params = [p for p in net.parameters() if p.requires_grad]
optimizer = optim.Adam(params, lr=0.0001)
epochs = 5
best_acc = 0.0
train_steps = len(train_loader)
for epoch in range(epochs):
# train
net.train()
running_loss = 0.0
train_bar = tqdm(train_loader, file=sys.stdout)
for step, data in enumerate(train_bar):
images, labels = data
optimizer.zero_grad()
logits = net(images.to(device))
loss = loss_function(logits, labels.to(device))
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
train_bar.desc = f"train epoch[{epoch + 1}/{epochs}] loss:{loss:.3f}"
# validate
net.eval()
acc = 0.0 # accumulate accurate number / epoch
with torch.no_grad():
val_bar = tqdm(validate_loader, file=sys.stdout)
for val_data in val_bar:
val_images, val_labels = val_data
outputs = net(val_images.to(device))
# loss = loss_function(outputs, test_labels)
predict_y = torch.max(outputs, dim=1)[1]
acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
val_bar.desc = f"valid epoch[{epoch + 1}/{epochs}]"
val_accurate = acc / val_num
print('[epoch %d] train_loss: %.3f val_accuracy: %.3f' %
(epoch + 1, running_loss / train_steps, val_accurate))
if val_accurate > best_acc:
best_acc = val_accurate
torch.save(net, "./MobileNetV2.pth")
print('Finished Training')
if __name__ == '__main__':
main()
3. predict.py：
import os
import json
import torch
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt
from model_v2 import MobileNetV2
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
data_transform = transforms.Compose(
[transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
# load image
img_path = "./ea34b0062bf4063ed1584d05fb1d4e9fe777ead218ac104497f5c978a6ebb3bf_640.jpg"
img = Image.open(img_path)
plt.imshow(img)
# [N, C, H, W]
img = data_transform(img)
# expand batch dimension
img = torch.unsqueeze(img, dim=0)
# read class_indict
json_path = './class_indices.json'
assert os.path.exists(json_path), "file: '{}' dose not exist.".format(json_path)
with open(json_path, "r") as f:
class_indict = json.load(f)
# create model
model = MobileNetV2(num_classes=10).to(device)
# load model weights
model = torch.load("./MobileNetV2.pth")
model.eval()
with torch.no_grad():
# predict class
output = torch.squeeze(model(img.to(device))).cpu()
predict = torch.softmax(output, dim=0)
predict_cla = torch.argmax(predict).numpy()
print_res = f"class: {class_indict[str(predict_cla)]} prob: {predict[predict_cla].numpy():.3}"
plt.title(print_res)
for i in range(len(predict)):
print(f"class: {class_indict[str(i)]:10} prob: {predict[i].numpy():.3}")
plt.show()
if __name__ == '__main__':
main()
4. pt2onnx.py:
import torch
import torchvision
from model_v2 import MobileNetV2
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = MobileNetV2(num_classes=10).to(device)
model=torch.load("/home/lm/MobileNet/MobileNetV2.pth")
model.eval()
example = torch.ones(1, 3, 244, 244)
example = example.to(device)
torch.onnx.export(model, example, "MobileNetV2.onnx", verbose=True, opset_version=12)
注意这里的“opset_version=12”，最好选择12版本的op指令集，因为瑞芯微示例的onnx模型的操作集版本为12。
5. 预训练模型：
mobilenetv2.pth预训练模型可以去官网下载:
https://download.pytorch.org/models/mobilenet_v2-b0353104.pth，下不了的科学上网试试，另外CSDN应该也有下载链接。
6. 所需数据集：
我这边共有11个类别，需要注意，图像分类和目标检测的数据集格式不同，目标检测数据集需要txt、xml、json等标注数据格式…
但是图像分类是不用标注数据集的，而你的类别名，就是通过所在的子文件夹名来定义的。具体格式如下：
还是分为train和val两个文件夹，两个文件夹下各有11个类别的子文件夹，子文件夹名即其中所有图片代表的类别名，打开其中子文件夹内容如下：
就是正常图片格式，无图片命名要求。
2. 环境配置
这里我用的环境就是Yolov8训练时的环境，因为根据train.py所需导入的文件，几乎没啥要求，直接安装torch、torchvision和tqdm…就行了。
三、模型训练
将train.py中的模型路径、数据集路径、epoch次数设置好后，执行train.py，终端结果如下：
四、模型转换
1. PTH转换ONNX
将训练得到的pth模型，修改pt2onnx.py中的pth模型路径，执行pt2onnx.py，终端结果如下：
终端显示如上后即可得到ONNX模型，用netron打开后看到模型结构如下：
2. ONNX转换RKNN
打开PC虚拟机，激活自己的转化RKNN的环境：rknn210，至于这个环境是啥以及怎么来的，可以参考我的这篇文章，其中有详细的该转换环境的介绍：【YOLOv8n部署至RK3588】模型训练→转换rknn→部署全流程
一、mobilenet.py文件修改
打开rknn_model_zoo/examples/mobilenet目录如下：
修改其中的python/mobilenet.py文件，修改内容如下：
注意，这里一定要改成input.1，因为netron打开onnx模型可以看到模型的输入名就是input.1（可看上述onnx模型图），这样确保与模型输入名一致，执行后不会报错。
二、激活环境，执行mobilenet.py：
三、终端显示如下，得到RKNN模型：
四、RKNN模型结构：
打开netron，查看rknn模型结构如下，输出的 “11” 是指有11个类别：
五、板端部署
打开rknn_model_zoo/examples/mobilenet/cpp目录
修改官方的demo中源文件的参数目录（这里的路径是我后期自己改的，原官方demo中的源文件是没有的）：
将头文件配置好后，文件夹中有官方的CMakeLists.txt。
我的CMakeLists.txt内容如下：
make_minimum_required(VERSION 3.10)
project(rknn_mobilenet_demo)
include_directories(include)
if (ENABLE_ASAN)
message(STATUS "BUILD WITH ADDRESS SANITIZER")
set (CMAKE_C_FLAGS_DEBUG "${CMAKE_C_FLAGS_DEBUG} -fno-omit-frame-pointer -fsanitize=address")
set (CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -fno-omit-frame-pointer -fsanitize=address")
set (CMAKE_LINKER_FLAGS_DEBUG "${CMAKE_LINKER_FLAGS_DEBUG} -fno-omit-frame-pointer -fsanitize=address")
endif ()
set(mobilenet_file src/mobilenet.cc)
if (TARGET_SOC STREQUAL "rv1106" OR TARGET_SOC STREQUAL "rv1103")
set(mobilenet_file rknpu2/mobilenet_rv1106_1103.cc)
add_definitions(-DRV1106_1103)
#dma
include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../../../3rdparty/allocator/dma)
endif()
if (TARGET_SOC STREQUAL "rk1808" OR TARGET_SOC STREQUAL "rv1109" OR TARGET_SOC STREQUAL "rv1126")
set(mobilenet_file rknpu1/mobilenet.cc)
endif()
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/ 3rdparty.out)
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/utils/ utils.out)
set(CMAKE_INSTALL_RPATH "$ORIGIN/../lib")
file(GLOB SRCS ${CMAKE_CURRENT_SOURCE_DIR}/*.cc)
add_executable(${PROJECT_NAME}
src/main.cc
${mobilenet_file}
)
# rknn api
set(RKNN_API_PATH ${CMAKE_SOURCE_DIR}/rknn_lib/librknn_api)
set(LIB_ARCH aarch64)
set(LIBRKNNRT ${RKNN_API_PATH}/${LIB_ARCH}/librknnrt.so)
include_directories(${RKNN_API_PATH}/include)
include_directories(${CMAKE_SOURCE_DIR}/3rdparty)
#opencv
find_package(OpenCV REQUIRED)
include_directories(${OpenCV_INCLUDE_DIRS})
target_link_libraries(${PROJECT_NAME}
fileutils
imageutils
${LIBRKNNRT}
${OpenCV_LIBS}
dl
)
if (CMAKE_SYSTEM_NAME STREQUAL "Android")
target_link_libraries(${PROJECT_NAME} log )
endif()
if (CMAKE_SYSTEM_NAME STREQUAL "Linux")
set(THREADS_PREFER_PTHREAD_FLAG ON)
find_package(Threads REQUIRED)
target_link_libraries(${PROJECT_NAME} Threads::Threads)
endif()
target_include_directories(${PROJECT_NAME} PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}
${LIBRKNNRT_INCLUDES}
)
install(TARGETS ${PROJECT_NAME} DESTINATION .)
install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/../model/bell.jpg DESTINATION model)
install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/../model/synset.txt DESTINATION model)
set(file_path ${CMAKE_CURRENT_SOURCE_DIR}/../model/mobilenet_v2.rknn)
if (EXISTS ${file_path})
install(FILES ${file_path} DESTINATION model)
endif()
终端输入cmake . 生成Makefile，再make编译，如下所示：
可执行文件推理如下：
共有335张1920*1080图片，推理耗时17562ms，平均耗时52.4ms
六、推理结果
取三张原图进行推理，结果如下：
原图1：
结果1：
原图2：
结果2：
原图3：
结果3：
总体来说，推理得分还是很高的，MobileNetv2在各项任务上足以胜任，效果不比ResNet50v2差，而且推理速度理论上要快很多。
以上即为MobileNetv2图像分类任务部署至RK3588的全流程。
扫一扫
889
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
