配置GlobalTrack到Anti-UAV(三)

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

你可能感兴趣的:(python,开发语言)