最近有个项目需要在Jetson Nano 上部署NanoSAM模型，在环境搭建上面遇到了不少问题，搞了好几天，终于搞完了，现在浅浅的记录一下。
这里需要注意，因为训练需要，我在Jetson上和本地Linux都进行了安装及部署，所以可能会在架构上搞错了，jetson一般是aarch64架构，而Linux是X86_64架构
创建并激活新环境（可选）
conda create -n myenv python=3.8
conda activate myenv
myenv 是你搭建的环境名字，可以更改，如果觉得下载太慢，可以配置清华大学镜像源
配置清华大学镜像源
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/conda config --dd channeconda config --set show_channel_urls yes
1. 安装环境配置
● 安装Pytorch、torchvision
安装Pytorch
安装依赖库
sudo apt-get -y update
sudo apt-get -y install autoconf bc build-essential g++-8 gcc-8 clang-8 lld-8 gettext-base gfortran-8 iputils-ping libbz2-dev libc++-dev libcgal-dev libffi-dev libfreetype6-dev libhdf5-dev libjpeg-dev liblzma-dev libncurses5-dev libncursesw5-dev libpng-dev libreadline-dev libssl-dev libsqlite3-dev libxml2-dev libxslt-dev locales moreutils openssl python-openssl rsync scons python3-pip libopenblas-dev
进入网址PyTorch for Jetson，根据Jetpack系统版本选择合适的whl文件，whl文件名称开头代表Pytorch版本，cp38代表Python版本是3.8
具体可参考Jeston Orin Nnao 安装pytorch与torchvision环境
安装torchvision 0.15.1
根据Pytorch版本选择适配的torchvision版本，我选择的Pytorch版本是2.0.0，对应的torchvision是0.15.1，进入GitHub - pytorch/vision at v0.15.1 ，在Tags选择合适的版本，下载zip文件
考虑到Github网络不稳定，我是先下到自己的电脑上，然后再将文件拷到Linux系统的
# 解压
unzip vision-v0.15.1.zip
# 进入解压文件
cd vision-0.15.1
# 安装
python3 setup.py install --user
● 安装Tensor RT
环境要求
这是TensorRT的GitHub仓库GitHub - NVIDIA/TensorRT
这里说明Tensor RT安装需要的环境要求，其中比较重要的就是cuda和cuDNN，这里推荐cuda11.8 + cuDNN 8.9.0
查看是否安装cuda 以及 cuDNN
查看cuda版本
# 查看cuda版本
nvcc --version
# 查看cuda安装路径
which nvcc
whereis nvcc
这里的cuda版本是11.8
查看cuDNN版本
cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2
如果执行上面的指令后没有显示，说明没有安装或环境未配置正确，这里以没安装为例
查看当前Ubuntu版本
lsb_release -a
安装cuda（可选）
安装依赖包
sudo apt-get update
sudo apt-get install gcc g++ freeglut3-dev build-essential libx11-dev libxmu-dev libxi-dev libglu1-mesa libglu1-mesa-dev
进入官方下载网址选择合适的cuda版本
jetson一般是aarch64架构，而Linux是X86_64架构
这里有三种方式可以安装，一般选择deb网上下载或本地下载，不要选runfile安装！！！（如果后面安装Tensor RT使用的也是deb包安装的话，否则会出现依赖问题无法成功安装，这个问题搞了我整整一天！！！）
下面会有一个安装指引，可以直接复制上面的指令，这里使用的deb的本地安装
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-ubuntu2004-11-8-local_11.8.0-520.61.05-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2004-11-8-local_11.8.0-520.61.05-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2004-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda
配置环境变量
echo 'export PATH=/usr/local/cuda-11.8/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc
再次执行 查看cuda版本，如果没有显示，查看环境变量是否配置正确
nvcc --version
安装cuDNN（可选）
进入网址 cuDNN Archive | NVIDIA Developer选择合适的cuDNN版本进行安装
jetson一般是aarch64架构，而Linux是X86_64架构
下载好后，进入下载文件夹，执行安装deb、导入GPG密钥、更新包列表
sudo dpkg -i cudnn-local-repo-ubuntu2004-8.9.0.26_1.0-1_amd64.deb
# 导入GPG密钥 这个安装完后系统会自动提醒的，直接复制建议的指令就行
sudo cp /var/cudnn-local-repo-ubuntu2004-8.9.0.26/cudnn-local-repo-ubuntu2004-8.9.0.26-keyring.gpg /usr/share/keyrings/
# 更新包列表
sudo apt-get update
安装运行库、开发人员库以及代码示例
sudo apt-get install libcudnn8
sudo apt-get install libcudnn8-dev
sudo apt-get install libcudnn8-samples
验证安装
cat /usr/local/cuda/include/cudnn.h | grep CUDNN_MAJOR -A 2
安装Tensor RT
进入网址下载对应的TensorRT版本，注意Tensor RT、cuda、cuDNN的版本必须适配，否则无法正常工作，以及需要确定自己的系统架构-我现在的是X86_64架构
jetson一般是aarch64架构，而Linux是X86_64架构
下载完成后，进入下载文件夹，执行以下指令，将 -i后面的文件名替换成下载的文件名即可
sudo dpkg -i nv-tensorrt-local-repo-ubuntu2004-8.x.x-cuda-11.8_1.0-1_amd64.deb
sudo cp /var/nv-tensorrt-local-repo-ubuntu2004-8.x.x-cuda-11.8/*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
安装Tensor RT包
sudo apt-get install tensorrt
配置环境变量，将bin、lib加入环境变量
echo 'export LD_LIBRARY_PATH=/usr/src/tensorrt/bin:$LD_LIBRARY_PATH' >> ~/.bashrc
echo 'export PATH=$PATH:/usr/src/tensorrt/bin' >> ~/.bashrc
source ~/.bashrc
验证安装
dpkg -l | grep tensorrt
安装Python包（需要与系统安装的TensorRT版本一致）
pip install tensorrt=8.5.3
# 安装依赖
sudo apt-get install python3-libnvinfer-dev
验证安装
python3 -c "import tensorrt; print(tensorrt.__version__)"
● 安装tensor2trt
去Github仓库下载压缩包或者用git clone指令下载，下载完成后进入文件夹，执行setup.py
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install
2. 安装NanoSAM Python包
下载NanoSAM Python包，执行以下指令，下载压缩包或者用git clone指令下载
# 下载NanaSAM Python 至本地
git clone https://github.com/NVIDIA-AI-IOT/nanosam
cd nanosam
python3 setup.py develop --user
3. 为掩码解码器构建TensorRT引擎
进入解压文件夹nanosam-main
● 导出MobileSAM掩码解码器ONNX文件
python3 -m nanosam.tools.export_sam_mask_decoder_onnx \ --model-type=vit_t \ --checkpoint=assets/mobile_sam.pt \ --output=data/mobile_sam_mask_decoder.onnx
执行过程中报错了，经查询是urllib版本太高了，降级即可
降级 urllib3（可选）
# 卸载当前版本
pip uninstall urllib3
# 安装兼容版本
pip install urllib3==1.26.15
● 构建TensorRT引擎
trtexec \ --onnx=data/mobile_sam_mask_decoder.onnx \ --saveEngine=data/mobile_sam_mask_decoder.engine \ --minShapes=point_coords:1x1x2,point_labels:1x1 \ --optShapes=point_coords:1x1x2,point_labels:1x1 \ --maxShapes=point_coords:1x10x2,point_labels:1x10
4. 为NanoSAM图像编码器构建TensorRT引擎
i. 下载图像编码器：resnet18_image_encoder.onnx
我在目录中发现有这个文件，所以就没下载了；不清楚是不是无意中下到中，没有的话下载就行了
ii. 构建TensorRT引擎
trtexec \ --onnx=data/resnet18_image_encoder.onnx \ --saveEngine=data/resnet18_image_encoder.engine \ --fp16
5. 运行基本用法示例
python3 examples/basic_usage.py \ --image_encoder=data/resnet18_image_encoder.engine \ --mask_decoder=data/mobile_sam_mask_decoder.engine
运行完后会在data文件夹输出一个分割图片
训练
● 下载并提取COCO 2017列车图像
# mkdir -p data/coco # uncomment if it doesn't existmkdir -p data/cococd data/cocowget http://images.cocodataset.org/zips/train2017.zipunzip train2017.zipcd ../..
● 构建MobileSAM图像编码器（用作教师模型）
1.导出到ONNX
不同的batch_size
batch_size =16
python3 -m nanosam.tools.export_sam_image_encoder_onnx \ --checkpoint="assets/mobile_sam.pt" \ --output="data/mobile_sam_image_encoder_bs16.onnx" \ --model_type=vit_t \ --batch_size=16
batch_size =8
python3 -m nanosam.tools.export_sam_image_encoder_onnx \ --checkpoint="assets/mobile_sam.pt" \ --output="data/mobile_sam_image_encoder_bs8.onnx" \ --model_type=vit_t \ --batch_size=8
batch_size =4
python3 -m nanosam.tools.export_sam_image_encoder_onnx \ --checkpoint="assets/mobile_sam.pt" \ --output="data/mobile_sam_image_encoder_bs4.onnx" \ --model_type=vit_t \ --batch_size=4
batch_size =1
python3 -m nanosam.tools.export_sam_image_encoder_onnx \ --checkpoint="assets/mobile_sam.pt" \ --output="data/mobile_sam_image_encoder_bs1.onnx" \ --model_type=vit_t \ --batch_size=1
2.构建批量大小为16的TensorRT引擎
不同的batch_size
trtexec \ --onnx=data/mobile_sam_image_encoder_bs16.onnx \ --shapes=image:16x3x1024x1024 \ --saveEngine=data/mobile_sam_image_encoder_bs16.engine
trtexec \ --onnx=data/mobile_sam_image_encoder_bs8.onnx \ --shapes=image:8x3x1024x1024 \ --saveEngine=data/mobile_sam_image_encoder_bs8.engine
trtexec \ --onnx=data/mobile_sam_image_encoder_bs4.onnx \ --shapes=image:4x3x1024x1024 \ --saveEngine=data/mobile_sam_image_encoder_bs4.engine
trtexec \ --onnx=data/mobile_sam_image_encoder_bs1.onnx \ --shapes=image:1x3x1024x1024 \ --saveEngine=data/mobile_sam_image_encoder_bs1.engine
3.通过提取MobileSAM训练NanoSAM图像编码器
不同的batch_size
python3 -m nanosam.tools.train \ --images=data/coco/train2017 \ --output_dir=data/models/resnet18 \ --model_name=resnet18 \ --teacher_image_encoder_engine=data/mobile_sam_image_encoder_bs16.engine \ --batch_size=16
python3 -m nanosam.tools.train \ --images=data/coco/train2017 \ --output_dir=data/models/resnet18 \ --model_name=resnet18 \ --teacher_image_encoder_engine=data/mobile_sam_image_encoder_bs8.engine \ --batch_size=8
python3 -m nanosam.tools.train \ --images=data/coco/train2017 \ --output_dir=data/models/resnet18 \ --model_name=resnet18 \ --teacher_image_encoder_engine=data/mobile_sam_image_encoder_bs4.engine \ --batch_size=4
python3 -m nanosam.tools.train \ --images=data/coco/train2017 \ --output_dir=data/models/resnet18 \ --model_name=resnet18 \ --teacher_image_encoder_engine=data/mobile_sam_image_encoder_bs1.engine \ --batch_size=1
4.将训练好的NanoSAM图像编码器导出到ONNX
python3 -m nanosam.tools.export_image_encoder_onnx \ --model_name=resnet18 \ --checkpoint="data/models/resnet18/checkpoint.pth" \ --output="data/resnet18_image_encoder.onnx"
python3 -m nanosam.tools.export_image_encoder_onnx \ --model_name=resnet18 \ --checkpoint="data/models/resnet18/checkpoint.pth" \ --output="data/resnet18_image_encoder_new.onnx"
过程中因为网络问题无法获取到训练模型，这是因为默认网站需要魔法上网，只需换成镜像网站就行
临时导入环境变量
export HF_ENDPOINT=https://hf-mirror.com
一些有用的指令
环境变量
查看当前环境变量
echo $LD_LIBRARY_PATH
删除重复的路径
如果 LD_LIBRARY_PATH 中有重复的路径，可以使用以下命令删除重复的路径并只保留一个：
export LD_LIBRARY_PATH=$(echo $LD_LIBRARY_PATH | tr ':' '\n' | sort -u | tr '\n' ':' | sed 's/:$//')
删除指定环境变量
如果要删除 /path/to/TensorRT-8.6.1/lib，可以使用以下命令：
export LD_LIBRARY_PATH=$(echo $LD_LIBRARY_PATH | tr ':' '\n' | grep -v '/path/to/TensorRT-8.6.1/lib' | tr '\n' ':' | sed 's/:$//')
添加环境变量
为了使添加的环境变量永久生效，可以将上述命令添加到 ~/.bashrc 文件中：
echo 'export LD_LIBRARY_PATH="/home/gdc434/anaconda3/envs/NanoSAM/lib/python3.8/site-packages/tensorrt:$LD_LIBRARY_PATH"' >> ~/.bashrcecho 'export LD_LIBRARY_PATH="/usr/src/tensorrt/bin:$LD_LIBRARY_PATH"' >> ~/.bashrcsource ~/.bashrc
使用 gedit（图形界面编辑器）查看环境变量配置文件
gedit ~/.bashrc
加载配置环境
source ~/.bashrc
查找文件
相当于搜索文件，而且效率会高很多
sudo find / -name "cudnn_version.h" 2>/dev/null
打开文件夹（图形界面）
这个可以结合查找文件一起用，效果杠杠滴
xdg-open 文件路径
扫一扫
2752
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
