在Ubuntu 20.04上从零搭建Faster R-CNN PyTorch环境（避坑CUDA 11.1 + PyTorch 1.9）

在Ubuntu 20.04上从零搭建Faster R-CNN PyTorch环境避坑CUDA 11.1 PyTorch 1.9当深度学习遇上目标检测Faster R-CNN无疑是这个领域的重要里程碑。而PyTorch作为当下最受欢迎的深度学习框架之一其灵活性和易用性让研究者趋之若鹜。本文将带你从零开始在Ubuntu 20.04系统上搭建Faster R-CNN的PyTorch实现环境特别针对CUDA 11.1和PyTorch 1.9版本进行优化配置避开那些让无数开发者头疼的版本兼容性问题。1. 环境准备与基础配置1.1 系统要求与初始检查在开始之前确保你的系统满足以下基本要求操作系统Ubuntu 20.04 LTS推荐使用纯净安装GPUNVIDIA显卡至少4GB显存推荐8GB以上驱动NVIDIA驱动版本≥450.80.02存储空间至少50GB可用空间用于数据集和模型首先更新系统并安装基础工具sudo apt update sudo apt upgrade -y sudo apt install -y build-essential cmake git wget unzip验证NVIDIA驱动是否正确安装nvidia-smi你应该看到类似如下的输出其中包含你的GPU型号和驱动版本信息----------------------------------------------------------------------------- | NVIDIA-SMI 450.80.02 Driver Version: 450.80.02 CUDA Version: 11.0 | |--------------------------------------------------------------------------- | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | | | | MIG M. | || | 0 GeForce RTX 208... On | 00000000:01:00.0 On | N/A | | 30% 45C P8 18W / 250W | 687MiB / 11019MiB | 0% Default | | | | N/A | ---------------------------------------------------------------------------1.2 CUDA 11.1与cuDNN安装CUDA 11.1是一个相对较新的版本与PyTorch 1.9有良好的兼容性。以下是安装步骤下载CUDA 11.1安装包wget https://developer.download.nvidia.com/compute/cuda/11.1.0/local_installers/cuda_11.1.0_455.23.05_linux.run运行安装程序sudo sh cuda_11.1.0_455.23.05_linux.run注意安装时不要选择安装驱动如果已经安装了更新的驱动只选择CUDA Toolkit。安装完成后将CUDA路径添加到环境变量echo export PATH/usr/local/cuda-11.1/bin:$PATH ~/.bashrc echo export LD_LIBRARY_PATH/usr/local/cuda-11.1/lib64:$LD_LIBRARY_PATH ~/.bashrc source ~/.bashrc验证CUDA安装nvcc --version接下来安装cuDNN这是NVIDIA提供的深度神经网络加速库从NVIDIA官网下载cuDNN for CUDA 11.1需要注册账号解压并安装tar -xzvf cudnn-11.1-linux-x64-v8.0.5.39.tgz sudo cp cuda/include/cudnn*.h /usr/local/cuda-11.1/include sudo cp cuda/lib64/libcudnn* /usr/local/cuda-11.1/lib64 sudo chmod ar /usr/local/cuda-11.1/include/cudnn*.h /usr/local/cuda-11.1/lib64/libcudnn*2. Python环境与PyTorch安装2.1 创建隔离的Python环境为了避免与系统Python环境冲突我们使用conda创建一个独立环境wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh bash Miniconda3-latest-Linux-x86_64.sh # 按照提示完成安装后 conda create -n fasterrcnn python3.8 -y conda activate fasterrcnn2.2 安装PyTorch 1.9及相关依赖PyTorch 1.9对CUDA 11.1有官方支持使用以下命令安装conda install pytorch1.9.0 torchvision0.10.0 torchaudio0.9.0 cudatoolkit11.1 -c pytorch -c conda-forge验证PyTorch是否正确识别了CUDAimport torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))你应该看到类似如下的输出1.9.0 True NVIDIA GeForce RTX 2080 Ti2.3 安装其他必要依赖Faster R-CNN需要一些额外的Python包pip install numpy scipy matplotlib opencv-python pillow scikit-image pyyaml tqdm特别要注意的是scipy的版本1.3.0以上版本移除了imread函数我们需要使用imageio替代pip install imageio3. Faster R-CNN源码获取与编译3.1 获取源代码推荐使用PyTorch 1.0兼容的分支git clone -b pytorch-1.0 --single-branch https://github.com/jwyang/faster-rcnn.pytorch.git cd faster-rcnn.pytorch3.2 编译CUDA扩展Faster R-CNN需要编译一些CUDA扩展cd lib python setup.py build develop这个步骤可能会遇到几个常见问题错误invalid command develop解决方案确保你已经安装了最新版本的setuptoolspip install --upgrade setuptoolsImportError: cannot import name _mask解决方案这通常是因为编译不完整尝试完全删除build目录后重新编译rm -rf build python setup.py build developscipy.misc.imread缺失解决方案修改代码中使用imageio.imread替代# 在需要的地方替换 from imageio import imread3.3 准备数据集我们将使用PASCAL VOC 2007数据集作为示例创建数据目录结构mkdir -p data/VOCdevkit2007 cd data/VOCdevkit2007下载数据集wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCdevkit_08-Jun-2007.tar解压数据集tar xvf VOCtrainval_06-Nov-2007.tar tar xvf VOCtest_06-Nov-2007.tar tar xvf VOCdevkit_08-Jun-2007.tar创建软链接cd ../.. ln -s data/VOCdevkit2007/VOCdevkit data/VOCdevkit20074. 模型训练与验证4.1 下载预训练模型Faster R-CNN通常使用在ImageNet上预训练的基础网络如VGG16或ResNet101mkdir -p data/pretrained_model cd data/pretrained_model下载VGG16预训练模型wget https://s3-us-west-2.amazonaws.com/jwyang/faster-rcnn-vgg16-1.0.0.pth或者ResNet101预训练模型wget https://s3-us-west-2.amazonaws.com/jwyang/faster-rcnn-res101-1.0.0.pth4.2 训练配置训练脚本的主要参数如下表所示参数说明推荐值--dataset使用的数据集pascal_voc--net基础网络架构vgg16或res101--bs批处理大小根据GPU内存调整通常1-4--nw数据加载线程数根据CPU核心数调整通常4-8--lr初始学习率0.001-0.01--lr_decay_step学习率衰减步长5-10--cuda使用GPU加速总是启用使用VGG16训练示例CUDA_VISIBLE_DEVICES0 python trainval_net.py \ --dataset pascal_voc \ --net vgg16 \ --bs 4 \ --nw 4 \ --lr 0.001 \ --lr_decay_step 5 \ --cuda4.3 常见训练问题解决GPU内存不足降低批处理大小--bs参数使用更小的基础网络如VGG16代替ResNet101启用梯度检查点需要修改代码Loss不收敛调整学习率尝试0.01, 0.001, 0.0001检查数据预处理是否正确确保预训练模型正确加载训练速度慢增加数据加载线程数--nw参数使用更快的存储如SSD代替HDD考虑使用混合精度训练4.4 模型测试与评估训练完成后可以使用以下命令评估模型在测试集上的表现python test_net.py \ --dataset pascal_voc \ --net vgg16 \ --checksession 1 \ --checkepoch 20 \ --checkpoint 10021 \ --cuda参数说明checksession: 训练会话ID通常为1checkepoch: 训练的epoch数checkpoint: 模型检查点编号5. 自定义数据集训练5.1 准备自定义数据集要训练自己的数据集需要按照PASCAL VOC的格式组织数据VOCdevkit └── VOC2007 ├── Annotations # XML标注文件 ├── ImageSets │ └── Main # 训练/验证/测试划分文件 └── JPEGImages # 原始图像5.2 修改类别定义在lib/datasets/pascal_voc.py中修改类别列表self._classes ( __background__, # always index 0 class1, class2, # 添加你的类别... )5.3 清除缓存在开始新训练前删除旧的缓存文件rm data/cache/voc_2007_trainval_gt_roidb.pkl5.4 启动训练使用与前面类似的命令开始训练确保指定正确的数据集路径和网络架构。6. 模型部署与推理6.1 单张图像测试使用demo.py脚本可以对单张图像进行测试python demo.py \ --net vgg16 \ --checksession 1 \ --checkepoch 20 \ --checkpoint 10021 \ --cuda \ --load_dir models \ --image_dir images6.2 批量推理要批量处理一个目录下的所有图像可以修改demo.py脚本或使用如下命令for img in images/*.jpg; do python demo.py \ --net vgg16 \ --checksession 1 \ --checkepoch 20 \ --checkpoint 10021 \ --cuda \ --load_dir models \ --image $img done6.3 性能优化技巧启用CuDNN基准测试在代码开始处添加torch.backends.cudnn.benchmark True使用混合精度训练安装apex库git clone https://github.com/NVIDIA/apex cd apex pip install -v --no-cache-dir --global-option--cpp_ext --global-option--cuda_ext ./模型量化训练后可以对模型进行量化以减少内存占用和提高推理速度model torch.quantization.quantize_dynamic( model, {torch.nn.Linear}, dtypetorch.qint8 )7. 高级配置与调优7.1 修改锚点(Anchor)设置Faster R-CNN的性能很大程度上依赖于锚点的设置。要修改锚点参数编辑lib/model/config.py__C.ANCHOR_SCALES [8, 16, 32] # 默认锚点尺度 __C.ANCHOR_RATIOS [0.5, 1, 2] # 默认宽高比7.2 调整NMS阈值非极大值抑制(NMS)是目标检测中的关键步骤可以调整其阈值__C.TEST.NMS 0.3 # NMS阈值默认0.37.3 学习率调度策略默认使用步长衰减可以尝试其他策略如余弦退火scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max10)7.4 多GPU训练要使用多个GPU进行训练可以使用PyTorch的DataParallelmodel torch.nn.DataParallel(model)然后在训练命令中指定多个GPUCUDA_VISIBLE_DEVICES0,1 python trainval_net.py ...8. 可视化与调试8.1 训练过程可视化使用TensorBoard可以方便地监控训练过程安装TensorBoardpip install tensorboard在代码中添加记录from torch.utils.tensorboard import SummaryWriter writer SummaryWriter() writer.add_scalar(Loss/total, loss.item(), iteration)启动TensorBoardtensorboard --logdirruns8.2 检测结果可视化修改demo.py脚本可以保存带有检测框的图像cv2.imwrite(output.jpg, im2show)8.3 常见错误调试CUDA out of memory降低批处理大小使用更小的输入图像尺寸清理不必要的GPU缓存NaN loss检查数据中是否有无效值降低学习率添加梯度裁剪预测框位置异常检查锚点设置是否合理验证数据标注是否正确调整回归损失权重