当前位置： 首页 > news >正文

【Unet++ MobileNetv2语义分割部署至RK3588】模型训练→转换RKNN→开发板部署

news 2025/11/5 10:47:57

请添加图片描述

已在GitHub开源与本博客同步的Unetplusplus_MobileNetv2-Semantic-Segmentation项目，
地址：Unetplusplus_MobileNetv2-Semantic-Segmentation
详细使用教程，可参考README.md或参考本博客第七章 模型部署

注：本文是以瑞芯微RK3588 SoC进行示例，同时也可支持瑞芯微其他系列SoC：RK3562、RK3566、RK3568、RK3576、RK3399PRO、RK1808、RV1126、RV1126B、RV1103、RV1106、RV1109等，部署流程也基本一致，如需帮助，可通过Github仓库的 README.md 沟通。

文章目录

  • 一、项目回顾
  • 二、文件梳理
  • 三、语义分割数据集的准备
  • 四、SMP框架的模型训练
  • 五、PT转ONNX
  • 六、ONNX转RKNN
  • 七、模型部署

一、项目回顾

博主之前有写过YOLO11、YOLOv8目标检测&图像分割、YOLOv10、v5目标检测、MoblieNetv2、ResNet50图像分类的模型训练、转换、部署文章，感兴趣的小伙伴可以了解下：
【YOLO11-obb部署至RK3588】模型训练→转换RKNN→开发板部署
【YOLO11部署至RK3588】模型训练→转换RKNN→开发板部署
【YOLOv10部署RK3588】模型训练→转换rknn→部署流程
【YOLOv8-obb部署至RK3588】模型训练→转换RKNN→开发板部署
【YOLOv8-pose部署至RK3588】模型训练→转换RKNN→开发板部署
【YOLOv8seg部署RK3588】模型训练→转换rknn→部署全流程
【YOLOv8部署至RK3588】模型训练→转换rknn→部署全流程
【YOLOv7部署至RK3588】模型训练→转换RKNN→开发板部署
【YOLOv6部署至RK3588】模型训练→转换RKNN→开发板部署
【YOLOv5部署至RK3588】模型训练→转换RKNN→开发板部署
【MobileNetv2图像分类部署至RK3588】模型训练→转换rknn→部署流程
【ResNet50图像分类部署至RK3588】模型训练→转换RKNN→开发板部署
YOLOv8n部署RK3588开发板全流程（pt→onnx→rknn模型转换、板端后处理检测）

最近在搞语义分割的项目，涉及到模型选型（比如Unet++配MobileNetV2）、训练调试、以及后续转ONNX和RKNN，最终部署到RK3588开发板上。在原型验证阶段，发现如果手动去实现各种分割模型（如Unet, FPN, DeepLabV3+）并适配不同的Backbone（如ResNet, MobileNet），非常繁琐且容易出错。

查了下GitHub，发现 qubvel-org / segmentation_models.pytorch (SMP) 这个项目非常优秀。它把主流的分割架构和海量的预训练Backbone都集成好了，API调用极其简洁，几行代码就能构建和训练一个强大的分割模型。用这个库的人非常多，但专门结合嵌入式部署到瑞芯微的RK3588上的文章不多，遂开此文，重点介绍下这个利器，相互学习。

二、文件梳理

CSDN上已有很多的Unet++的训练文章，但原有的Unet++训练是比较封闭的，所以查看了这个SMP项目，可以自己自定义编码器的解码器，当然，该项目只适用于语义分割任务。
目前需要的文件有如下几个：
1：SMP项目（尽量选0.5.0版本的）(链接在此）
2：在虚拟机中进行onnx转rknn的虚拟环境配置项目文件(链接在此)；
3：onnx转rknn的文件(链接在此)；（rknn-toolkit2尽量和rknn_model_zoo版本一致）

如下所示：
在这里插入图片描述

在这里插入图片描述

在这里插入图片描述

三、语义分割数据集的准备

博主主要用的是Labelme，用多边形把每个目标勾出来，分配好类别（比如背景=0，人=1，车=2）。但Labelme导出的JSON还不能直接拿去训练，关键一步是需要写脚本，把JSON批量转成模型能认的单通道PNG Mask图，图里每个像素的值就是它的类别索引。（至于脚本如何写，建议直接让DeepSeek生成）
此处示例如下：
原图：
请添加图片描述
标签图：
请添加图片描述

标签图基本都是黑的，以像素值进行类别区分。
数据集的标准格式如下所示：
在这里插入图片描述
train和val是原图，trainannot和valannot是黑色的标签图

四、SMP框架的模型训练

先进性环境配置，参考requirements下的required.txt、docs.txt、minimum.old：
在这里插入图片描述
在这里插入图片描述
在这里插入图片描述

完成环境配置。博主的环境如下所示：

# packages in environment at /root/anaconda3/envs/smp050:
#
# Name                    Version                   Build  Channel
_libgcc_mutex             0.1                        main  
albucore                  0.0.24                    <pip>
albumentations            2.0.8                     <pip>
annotated-types           0.7.0                     <pip>
ca-certificates           2025.2.25            h06a4308_0  
certifi                   2025.4.26                 <pip>
charset-normalizer        3.4.2                     <pip>
contourpy                 1.3.0                     <pip>
cycler                    0.12.1                    <pip>
eval_type_backport        0.2.2                     <pip>
filelock                  3.18.0                    <pip>
fonttools                 4.58.2                    <pip>
fsspec                    2025.5.1                  <pip>
hf-xet                    1.1.3                     <pip>
huggingface-hub           0.33.0                    <pip>
idna                      3.10                      <pip>
imageio                   2.37.0                    <pip>
importlib_resources       6.5.2                     <pip>
kiwisolver                1.4.7                     <pip>
ld_impl_linux-64          2.40                 h12ee557_0  
libffi                    3.3                  he6710b0_2  
libgcc-ng                 9.1.0                hdf63c60_0  
libstdcxx-ng              9.1.0                hdf63c60_0  
matplotlib                3.9.4                     <pip>
ncurses                   6.3                  h7f8727e_2  
numpy                     1.24.4                    <pip>
opencv-python             4.11.0.86                 <pip>
opencv-python-headless    4.11.0.86                 <pip>
openssl                   1.1.1w               h7f8727e_0  
packaging                 25.0                      <pip>
pillow                    11.2.1                    <pip>
pip                       25.1               pyhc872135_2  
pydantic                  2.11.5                    <pip>
pydantic_core             2.33.2                    <pip>
pyparsing                 3.2.3                     <pip>
python                    3.9.12               h12debd9_1  
python-dateutil           2.9.0.post0               <pip>
PyYAML                    6.0.2                     <pip>
readline                  8.1.2                h7f8727e_1  
requests                  2.32.4                    <pip>
safetensors               0.5.3                     <pip>
scipy                     1.13.1                    <pip>
setuptools                78.1.1           py39h06a4308_0  
simsimd                   6.4.9                     <pip>
six                       1.17.0                    <pip>
sqlite                    3.38.5               hc218d9a_0  
stringzilla               3.12.5                    <pip>
timm                      1.0.15                    <pip>
tk                        8.6.12               h1ccaba5_0  
torch                     1.10.0+cu111              <pip>
torchvision               0.11.0+cu111              <pip>
tqdm                      4.67.1                    <pip>
typing-inspection         0.4.1                     <pip>
typing_extensions         4.14.0                    <pip>
tzdata                    2025b                h04d1e81_0  
urllib3                   2.4.0                     <pip>
wheel                     0.45.1           py39h06a4308_0  
xz                        5.2.5                h7f8727e_1  
zipp                      3.23.0                    <pip>
zlib                      1.2.12               h7f8727e_2

环境配置完成后，需要训练脚本，如果直接下载SMP，是没有合适的训练脚本的，博主这里提供自己的训练脚本train050.py，如下所示：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'import numpy as np
import torch
import torch.optim as optim
import segmentation_models_pytorch as smp
from segmentation_models_pytorch.utils.metrics import IoU
from torch.utils.data import DataLoader
from torch.utils.data import Dataset as BaseDataset
import albumentations as albu
import time
import cv2
from tqdm import tqdm  # 添加进度条库
import warnings# 忽略权重加载警告
warnings.filterwarnings("ignore", message="Error loading mobilenet_v2 `imagenet` weights.*")# 数据集目录
DATA_DIR = 'xxxxx'
x_train_dir = os.path.join(DATA_DIR, 'train')
y_train_dir = os.path.join(DATA_DIR, 'trainannot')
x_valid_dir = os.path.join(DATA_DIR, 'val')
y_valid_dir = os.path.join(DATA_DIR, 'valannot')# 模型参数
ENCODER = 'mobilenet_v2'
ENCODER_WEIGHTS = 'imagenet'
CLASSES = ['xxx']
ACTIVATION = 'sigmoid'
DEVICE = 'cuda'
BATCH_SIZE = 16
LR = 0.0001
EPOCHS = 800# 定义目标图像尺寸 (方便管理) 
TARGET_SIZE = 640# 数据加载类
class SegmentationDataset(BaseDataset):CLASSES = ['carpet']def __init__(self, images_dir, masks_dir, augmentation=None, preprocessing=None):self.ids = os.listdir(images_dir)self.images_fps = [os.path.join(images_dir, img_id) for img_id in self.ids]self.masks_fps = [os.path.join(masks_dir, img_id) for img_id in self.ids]self.augmentation = augmentationself.preprocessing = preprocessingself.class_value = 1  # 二分类任务def __getitem__(self, i):image = cv2.imread(self.images_fps[i])image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)mask = cv2.imread(self.masks_fps[i], 0)mask = (mask == self.class_value).astype('float')if self.augmentation:sample = self.augmentation(image=image, mask=mask)image, mask = sample['image'], sample['mask']if self.preprocessing:sample = self.preprocessing(image=image, mask=mask)image, mask = sample['image'], sample['mask']return image, maskdef __len__(self):return len(self.ids)def get_training_augmentation():return albu.Compose([albu.HorizontalFlip(p=0.5),albu.Affine(scale=(0.8, 1.2),  # 调整缩放范围以适应letterboxtranslate_percent=(-0.1, 0.1),rotate=10, # 可以加入轻微旋转p=0.7),# 核心修改：先缩放（保持比例），再填充albu.LongestMaxSize(max_size=TARGET_SIZE, interpolation=cv2.INTER_LINEAR, p=1), albu.PadIfNeeded(                                                               min_height=TARGET_SIZE, min_width=TARGET_SIZE, p=1, border_mode=cv2.BORDER_CONSTANT, value=[128, 128, 128],       # 图像填充值 (灰色)mask_value=0                 # mask填充值 (黑色)),albu.GaussNoise(p=0.2),albu.RandomBrightnessContrast(p=0.2),])def get_validation_augmentation():return albu.Compose([albu.LongestMaxSize(max_size=TARGET_SIZE, interpolation=cv2.INTER_LINEAR, p=1), albu.PadIfNeeded(                                                               min_height=TARGET_SIZE, min_width=TARGET_SIZE, p=1, border_mode=cv2.BORDER_CONSTANT, value=[128, 128, 128],       # 图像填充值 (灰色)mask_value=0                 # mask填充值 (黑色))])# 处理单通道Mask和多通道Image
def to_tensor(x, **kwargs):if x.ndim == 2:  # 处理单通道Mask (H, W)x = np.expand_dims(x, axis=0)  # 增加通道维 -> (1, H, W)else:  # 处理RGB图像 (H, W, C)x = x.transpose(2, 0, 1)      # 转置为 (C, H, W)return x.astype('float32')# 预处理
preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)def get_preprocessing(preprocessing_fn):return albu.Compose([albu.Lambda(image=preprocessing_fn),albu.Lambda(image=to_tensor, mask=to_tensor),])# 创建数据加载器
train_dataset = SegmentationDataset(x_train_dir, y_train_dir,augmentation=get_training_augmentation(),preprocessing=get_preprocessing(preprocessing_fn),
)valid_dataset = SegmentationDataset(x_valid_dir, y_valid_dir,augmentation=get_validation_augmentation(),preprocessing=get_preprocessing(preprocessing_fn),
)train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=4)
valid_loader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=4)# 创建模型
model = smp.create_model('unetplusplus',encoder_name=ENCODER, encoder_weights=ENCODER_WEIGHTS,in_channels=3,classes=1,activation=ACTIVATION,
)
model.to(DEVICE)# 损失函数
loss_fn = smp.losses.DiceLoss(mode='binary')# 优化器
optimizer = optim.Adam(model.parameters(), lr=LR)# 创建 IoU 计算器实例
iou_metric = IoU(threshold=0.5).to(DEVICE)# 在训练循环之前创建模型目录
MODEL_DIR = 'xxx'
os.makedirs(MODEL_DIR, exist_ok=True)  # 确保目录存在# 训练循环
max_iou = 0
best_model_path = None  # 用于记录当前最佳模型的路径for epoch in range(EPOCHS):print(f'\nEpoch {epoch+1}/{EPOCHS}:')# ===== 训练阶段 =====model.train()train_loss = 0.0train_bar = tqdm(train_loader, desc='Training', unit='batch', dynamic_ncols=True)for images, masks in train_bar:images, masks = images.to(DEVICE), masks.to(DEVICE)optimizer.zero_grad()outputs = model(images)loss = loss_fn(outputs, masks)loss.backward()optimizer.step()train_loss += loss.item()train_bar.set_postfix(loss=f'{loss.item():.4f}', lr=f"{optimizer.param_groups[0]['lr']:.6f}")# ===== 验证阶段 =====model.eval()valid_loss = 0.0iou_scores = []valid_bar = tqdm(valid_loader, desc='Validating', unit='batch', dynamic_ncols=True)with torch.no_grad():for images, masks in valid_bar:images, masks = images.to(DEVICE), masks.to(DEVICE)outputs = model(images)loss_val = loss_fn(outputs, masks)valid_loss += loss_val.item()batch_iou = iou_metric(outputs, masks).item()iou_scores.append(batch_iou)valid_bar.set_postfix(val_loss=f'{loss_val.item():.4f}', iou=f'{batch_iou:.4f}')# 计算平均指标train_loss /= len(train_loader)valid_loss /= len(valid_loader)mean_iou = np.mean(iou_scores)print(f'[Summary] Train Loss: {train_loss:.4f} | Valid Loss: {valid_loss:.4f} | IoU: {mean_iou:.4f}')# 保存最佳模型if mean_iou > max_iou:max_iou = mean_iouif best_model_path and os.path.exists(best_model_path):os.remove(best_model_path)print(f"删除旧的最佳模型: {best_model_path}")model_name = f'best_model_iou_{mean_iou:.4f}.pth'model_path = os.path.join(MODEL_DIR, model_name)torch.save(model.state_dict(), model_path)best_model_path = model_pathprint(f'[Model Saved] 新的最佳模型保存至: {model_path}')# 学习率调整if epoch == 50:optimizer.param_groups[0]['lr'] = LR / 10print('Learning rate decreased to 1e-5')

注意：需要将脚本中的各个数据集路径都适配好。
这里要着重说一下，因为Unet++作为解码器的效果是比较好的，而MobileNetv2作为轻量级网络的编码能力也很优秀，所以在当前脚本中，是将Unet++和MobileNetv2结合起来，实现语义分割模型的整体架构，既保证性能，又实现轻量化。

现在可以开始训练，如下所示：
在这里插入图片描述
训练完成后，会在model文件夹下保存最优模型：
在这里插入图片描述

五、PT转ONNX

同样博主在此处提供专用的pt转onnx的脚本：pt2onnx.py：

import os
import torch
import segmentation_models_pytorch as smp# 配置参数
MODEL_PATH = 'xxx.pth'
ONNX_PATH = 'xx.onnx'
ENCODER = 'mobilenet_v2'
ENCODER_WEIGHTS = 'imagenet'
CLASSES = ['xxx']
ACTIVATION = 'sigmoid'
INPUT_SHAPE = (1, 3, 640, 640)  # 固定尺寸
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
OP_SET_VERSION = 12def main():print(f"使用设备: {DEVICE}")print("正在创建模型结构...")model = smp.create_model('unetplusplus',encoder_name=ENCODER,encoder_weights=ENCODER_WEIGHTS,in_channels=3,classes=1,activation=ACTIVATION,)print(f"加载训练权重: {MODEL_PATH}")state_dict = torch.load(MODEL_PATH, map_location=DEVICE)model.load_state_dict(state_dict)model.to(DEVICE)model.eval()# 准备固定尺寸的虚拟输入print("准备固定尺寸的虚拟输入数据...")dummy_input = torch.randn(INPUT_SHAPE, device=DEVICE)# 导出模型到ONNX格式（固定尺寸）print(f"导出固定尺寸ONNX模型...")torch.onnx.export(model,dummy_input,ONNX_PATH,export_params=True,opset_version=OP_SET_VERSION,do_constant_folding=True,input_names=['input'],output_names=['output'],dynamic_axes=None  # 关键修改：移除动态维度)print(f"✓ 固定尺寸ONNX模型已保存至: {ONNX_PATH}")# 验证输出print("\n验证输出:")with torch.no_grad():torch_output = model(dummy_input)print(f"输入形状: {dummy_input.shape}")print(f"输出形状: {torch_output.shape}")print(f"输出应为固定尺寸: (1, 1, 640, 640)")if __name__ == '__main__':main()

修改pt2onnx.py中的pt和onnx模型的路径，然后执行脚本，结果如下所示：
在这里插入图片描述
执行完脚本后，在脚本中的指定路径下生成onnx模型

六、ONNX转RKNN

在进行这一步的时候，如果你是在云服务器上运行，请先确保你租的卡能支持RKNN的转换运行。博主是在自己的虚拟机中进行转换。
先安装转换环境
这里我们先conda create -n rknn232 python=3.8创建环境，创建完成后激活该环境，打开第三个文件：rknn_toolkit2-2.3.2，在图示路径下找到这两个文件：
在这里插入图片描述

在终端激活环境，在终端输入

pip install -r requirements_cp38-2.3.2.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

然后再输入

pip install rknn_toolkit2-2.3.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

然后，我们的转rknn环境就配置完成了。

现在开始转换RKNN模型：
先进入rknn_model_zoo-2.3.2/examples/ppseg/model文件夹，将刚才得到的ONNX模型复制进来（我演示的位置是2.1版本的rknn_model_zoo，但操作流程是一致的）：
在这里插入图片描述

打开ppseg/python下的convert.py，修改参数，如下所示：

在这里插入图片描述
激活环境，执行该脚本开始量化，如下所示：在这里插入图片描述

然后在model文件夹下生成我们所要的rknn模型：在这里插入图片描述

用netron打开，如下所示，我只有一个类别输出通道：
在这里插入图片描述

七、模型部署

如果前面流程都已实现，模型的结构也没问题的话，则可以进行最后一步：模型端侧部署。
我已经帮大家做好了所有的环境适配工作，科学上网后访问博主GitHub仓库:Unetplusplus_MobileNetv2-Semantic-Segmentation
在这里插入图片描述

git clone后把项目复制到开发板上，按如下流程操作：
①：cd build，删除所有build文件夹下的内容
②：cd src 修改main.cc，修改main函数中的如下三个内容，将这三个参数改成自己的绝对路径：
在这里插入图片描述
③：把你之前训练好并已转成RKNN格式的模型放到Unetplusplus_MobileNetv2-Semantic-Segmentation/model文件夹下，然后把你要检测的所有图片都放到Unetplusplus_MobileNetv2-Semantic-Segmentation/inputimage下。
在运行程序后，生成的结果图片在Unetplusplus_MobileNetv2-Semantic-Segmentation/outputimage下

④：进入build文件夹进行编译

cd build

cmake ..

make

在build下生成可执行文件文件：rknn_yolo11obb_demo
在build路径下输入

./rknn_Unet++_MobileNetv2_demo

注：我已经在GIthub的仓库中按照路径放置好了本博客所使用到的所有文件，包括rknn模型、输入图片、输出图片以及txt文档，只需修改cpp中的文件路径即可，或者可参考main.cc中的main函数，按照如下命令亦可：

原inputimage文件夹下的图片：
在这里插入图片描述

在修改完main.cc中main函数的路径参数后执行完.rknn_Unet++_MobileNetv2_demo后在outputimage下的输出结果图片：
在这里插入图片描述

根据输出结果，可以看到语义分割效果还是可以的。

上述即博主此次更新的Unetplusplus_MobileNetv2-Semantic-Segmentation部署RK3588，包含PT转ONNX转RKNN的全流程步骤，欢迎交流！

http://www.dtcms.com/a/569748.html

相关文章：

  • 深圳做网站做app我想自己做网站可以赚钱
  • 新颖网站页面设计wordpress弹窗登录注册插件
  • 服务端开发的基本概念
  • 快递鸟电子面单打印接口技术对接文档
  • FreeRTOS事件组全解析：多任务同步核心技巧
  • 网站的头尾和导航的公用文件wordpress主题 问答
  • 【GlobalMapper精品教程】096：连接PostGIS数据库（Postgresql）
  • 网站备案服务商查询网站域名使用期
  • 机器学习日报12
  • 【文档】Stomp 协议
  • 自己的网站怎么接广告联盟设计库
  • 网站怎么解析域名解析网站制作公司官网南京
  • liunx文件及目录管理和vim编辑
  • [leetcode]对顶堆，对数时间添加元素，常数时间取中位数(或者第K大的数)
  • 公司软件网站建设免费dw网页模板
  • 力扣146LRU缓存
  • 网站怎么做充值系统下载网站需要写哪些内容
  • 网站有没有做网站地图怎么看vi设计案例ppt
  • 网站页面设计培训班长沙人才招聘网最新招聘2024
  • 虚幻引擎5 GAS开发俯视角RPG游戏 P07-02 授予能力
  • 真实的大模型中，embedding映射的高维矩阵维度和 attention矩阵运算的规模尺寸？
  • 中山网站设计与建设北京网上注册公司
  • ctf show-misc
  • 电子商务网站开发与管理实验报告青岛seo网络优化公司
  • 广告投放网站网站flash代码
  • 【Solidity 从入门到精通】第2章 Solidity 语言概览与环境搭建
  • 前端缓存战争：回车与刷新按钮的终极对决！
  • 做会计题目的网站手机网站推荐大全
  • 【论文精读】AVID：基于扩散模型的任意长度视频修复
  • 电子学会青少年软件编程（C/C++）1级等级考试真题试卷（2025年9月）