个人环境
操作系统:Ubuntu 16.04.1 x86_64
CUDA驱动:11.1
Python:3.8
torch=1.8.1+cu111
torchvision=0.9.1+cu111
torchaudio=0.8.1
设置环境变量
- vim ~/.bashrc
- 在.bashrc中最后添加以下内容:
export CUDA_HOME=/usr/local/cuda-11.1
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH - source ~/.bashrc
spconv版本
pip install spconv-cu111
训练
指定单GPU训练
CUDA_VISIBLE_DEVICES=0 python train.py --cfg_file cfgs/kitti_models/pv_rcnn.yaml
指定多GPU训练
CUDA_VISIBLE_DEVICES=4,5,6,7 python -m torch.distributed.launch --nproc_per_node=4 train.py --cfg_file cfgs/kitti_models/pv_rcnn.yaml --launcher pytorch