1、数据集接口增加
打开datasets/wrappers.py,在def _dataset中增加新的数据集项,例如:
打开_submodules/neuron/data/datasets/_init_.py,添加from .uavtir import *
在_submodules/neuron/data/datasets路径下新建文件uavtir.py,该文件根据其它数据集文件进行仿写,存在了服务器里
2、训练anti-uava
在tools下新建train_uavtir.py,与train_qg_rcnn.py相比,变动如下
(1)加入sys的两行代码来避免找不到其它文件,如果不加,也可以在运行前现在命令行输入
export PYTHONPATH=$PYTHONPATH:/home1/users/GlobalTrack-master/
(2)加了限制GPUs可见的命令,这里使用两卡,省的之后每次运行都需要在命令行设定
os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
(3)屏蔽warnings
import warnings
warnings.filterwarnings("ignore", category=UserWarning)
(4)更改args里的参数
base_dataset的默认改为uavtir_train
sampling_prob的默认改为1.0
gpus的默认改为2
3、运行
先测试单卡训练是否通过
python tools/train_uavtir.py --config configs/qg_rcnn_r50_fpn.py --load_from initial_weights/qg_rcnn_r50_fpn_2x_20181010-443129e1.pth --gpus 1
单卡(3090)预计训练1小时,内存占用为21/24G