yolov5转换成.bmodel的教程(使用官方教程、使用yolov5原始模型)
1 环境配置
Linux开发环境,Ubuntu18.04,运行内存8G
1.1 下载SophonSDK 3.0.0开发包
sudo apt update
sudo apt install unzip
sudo wget https://sophon-file.sophon.cn/sophon-prod-s3/drive/22/07/18/11/sophonsdk_v3.0.0_20220716.zip
sudo unzip sophonsdk_v3.0.0_20220716.zip
#操作可能加锁,需要解锁
sudo chmod -R 777 *
1.2 docker安装
# 安装docker
sudo apt-get install docker.io
# docker命令免root权限执行
# 创建docker用户组,若已有docker组会报错,没关系可忽略
sudo groupadd docker
# 将当前用户加入docker组
sudo gpasswd -a ${USER} docker
# 重启docker服务
sudo service docker restart
# 切换当前会话到新group或重新登录重启X会话
newgrp docker
1.3 下载docker镜像
cd sophonsdk_v3.0.0_20220716
sudo wget https://sophon-file.sophon.cn/sophon-prod-s3/drive/22/07/19/10/x86_sophonsdk3_ubuntu18.04_py37_dev_22.06_docker.zip
bash unzip x86_sophonsdk3_ubuntu18.04_py37_dev_22.06_docker
1.4 挂载docker镜像
cd x86_sophonsdk3_ubuntu18.04_py37_dev_22.06_docker
sudo docker load -i x86_sophonsdk3_ubuntu18.04_py37_dev_22.06.docker
1.5 创建docker容器并进入docker
cd ..
sudo tar -zxvf sophonsdk_v3.0.0.tar.gz
cd sophonsdk_v3.0.0
sudo ./docker_run_sophonsdk.sh
1.6 在docker容器内安装依赖库和设计环境变量
cd scripts
./install_lib.sh nntc # 设置环境变量,注意此命令只对当前终端有效,重新进入需要重新执行
source envsetup_cmodel.sh # for CMODEL MODE
#source envsetup_pcie.sh # for PCIE MODE
2 准备模型
#参考文章链接
#https://github.com/sophon-ai-algo/examples/blob/3.0.0/simple/yolov5/README.md
cd ..
wget https://github.com/sophon-ai-algo/examples/archive/refs/heads/3.0.0.zip
unzip 3.0.0.zip
chmod -R 777 *
2.1 下载yolov5源码
cd yolov5
# 下载yolov5源码
git clone https://github.com/ultralytics/yolov5.git yolov5_github
# 切换到yolov5工程目录
cd yolov5_github
# 使用tag从远程创建本地v6.1 分支
git branch v6.1 v6.1
# 下载yolov5s v6.1版本
wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s.pt
2.2 修改models/yolo.py(使用原始模型这一步不修改)
#修改Detect类的forward函数的最后return语句,实现不同的输出
# 以三个输出模型为例
....
def forward(self, x):
z = [] # inference output
for i in range(self.nl):
x[i] = self.m[i](x[i]) # conv
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
if not self.training: # inference
if self.onnx_dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
y = x[i].sigmoid()
if self.inplace:
y[..., 0:2] = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i] # xy
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
else: # for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953
xy = (y[..., 0:2] * 2 - 0.5 + self.grid[i]) * self.stride[i] # xy
wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh
y = torch.cat((xy, wh, y[..., 4:]), -1)
z.append(y.view(bs, -1, self.no))
return x if self.training else x # 3个输出
# return x if self.training else (torch.cat(z, 1)) # 1个输出
# return x if self.training else (torch.cat(z, 1), x) # 4个输出
....
2.3 导出JIT模型
SophonSDK中的PyTorch模型编译工具BMNETP只接受PyTorch的JIT模型(TorchScript模型)。
JIT(Just-In-Time)是一组编译工具,用于弥合PyTorch研究与生产之间的差距。它允许创建可以在不依赖Python解释器的情况下运行的模型,并且可以更积极地进行优化。在已有PyTorch的Python模型(基类为torch.nn.Module)的情况下,通过torch.jit.trace就可以得到JIT模型,如torch.jit.trace(python_model, torch.rand(input_shape)).save(‘jit_model’)。BMNETP暂时不支持带有控制流操作(如if语句或循环)的JIT模型,因此不能使用torch.jit.script,而要使用torch.jit.trace,它仅跟踪和记录张量上的操作,不会记录任何控制流操作。这部分操作yolov5已经为我们写好,只需运行如下命令即可导出符合要求的JIT模型:
# 创建python虚拟环境virtualenv
pip3 install virtualenv
# 切换到虚拟环境
virtualenv -p python3 --system-site-packages env_yolov5
source env_yolov5/bin/activate
修改requirements.txt里的内容
# YOLOv5 requirements
# Usage: pip install -r requirements.txt
# Base ------------------------------------------------------------------------
gitpython>=3.1.30
matplotlib>=3.3
numpy>=1.21.5
opencv-python>=4.1.1
Pillow>=9.4.0
psutil # system resources
PyYAML>=5.3.1
requests>=2.23.0
scipy>=1.4.1
thop>=0.1.1 # FLOPs computation
torch>=1.8.0 # see https://pytorch.org/get-started/locally (recommended)
torchvision>=0.9.0
tqdm>=4.64.0
ultralytics>=8.0.120
# protobuf<=3.20.1 # https://github.com/ultralytics/yolov5/issues/8012
# Logging ---------------------------------------------------------------------
# tensorboard>=2.4.1
# clearml>=1.2.0
# comet
# Plotting --------------------------------------------------------------------
pandas>=1.1.4
seaborn>=0.11.0
# Export ----------------------------------------------------------------------
# coremltools>=6.0 # CoreML export
# onnx>=1.10.0 # ONNX export
# onnx-simplifier>=0.4.1 # ONNX simplifier
# nvidia-pyindex # TensorRT export
# nvidia-tensorrt # TensorRT export
# scikit-learn<=1.1.2 # CoreML quantization
# tensorflow>=2.4.0,<=2.13.1 # TF exports (-cpu, -aarch64, -macos)
# tensorflowjs>=3.9.0 # TF.js export
# openvino-dev>=2023.0 # OpenVINO export
# Deploy ----------------------------------------------------------------------
setuptools>=65.5.1 # Snyk vulnerability fix
# tritonclient[all]~=2.24.0
# Extras ----------------------------------------------------------------------
# ipython # interactive notebook
# mss # screenshots
# albumentations>=1.0.3
# pycocotools>=2.0.6 # COCO mAP
# 安装依赖
pip3 install -r requirements.txt
# 导出jit模型
python3 export.py --weights yolov5s.pt --include torchscript
chmod -R 777 *
# 退出虚拟环境
deactivate
# 将生成好的jit模型yolov5s.torchscript拷贝到${YOLOv5}/build文件夹下
mkdir ../build
cp yolov5s.torchscript ../build/yolov5s_coco_v6.1_3output.trace.pt
# 拷贝一份到${YOLOv5}/data/models文件夹下
mkdir ../data
mkdir ../data/models
cp yolov5s.torchscript ../data/models/yolov5s_coco_v6.1_3output.trace.pt
2.4 准备量化数据集
cd ..
#把sample/yolov5/scripts文件夹复制到和data build平级的文件夹里
cd scripts
./00_prepare.sh
3 模型转换
3.1 生成FP32 BModel
./1_gen_fp32bmodel.sh
上述脚本会在${YOLOv5}/data/models下生成yolov5s_640_coco_v6.1_3output_fp32_1b.bmodel文件,即转换好的FP32 BModel,使用bm_model.bin --info查看模型具体信息
3.2 生成INT8 BModel
INT8 BModel的生成请注意修改model_info.sh中的模型名称、生成模型目录和输入大小shapes、使用的量化LMDB文件目录、batch_size、img_size等参数。
执行以下命令,使用一键量化工具cali_model,生成INT8 BModel:
./2_gen_int8bmodel.sh
上述脚本会在${YOLOv5}/data/models下生成yolov5s_640_coco_v6.1_3output_int8_1b.bmodel,即转换好的INT8 BModel,使用bm_model.bin --info查看模型具体信息