yolox 训练自己的数据集 （COCO格式）

准备工作

1. 准备自己的数据集

Animals_Coco
├─annotations
├─train2017
└─val2017

在annotations 文件夹下包含两个重要的文件instances_train2017.json，instances_val2017.json
在train2017和val2017 包含的是训练和验证的图片数据。

2. YOLOX环境搭建

下载 YOLOX，克隆不下来的话，也可以自己下载zip文件然后解压 。

cd ~/code
git clone https://github.com/Megvii-BaseDetection/YOLOX.git

搭建环境

conda create -n yolox  python=3.7 # 创建名称为yolox的新环境
conda activate yolox # 进入环境
cd YOLOX # 进入YOLOX文件夹
pip3 install -U pip && pip3 install -r requirements.txt   # 安装代码依赖的库文件
python3 setup.py develop  # 通过setup.py安装一些库文件

安装 apex

git clone https://github.com/NVIDIA/apex.git  # 将apex 下载到~/code/YOLOX/， 克隆不下来的话，也可以自己下载zip文件然后解压到YOLOX中
 cd apex
 sudo pip3 install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

bug:

 pip.exceptions.InstallationError: Command "python setup.py egg_info" failed with error code 1 in /tmp/pip-76g90y3m-build/

解决方法：python3 setup.py install

安装 pycocotools

pip3 install cython
pip3 install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'
# 如果上一步下载失败，也可以尝试先将cocoapi下载并解压到YOLOX文件夹，然后执行本地安装命令
pip3 install ~/code/YOLOX/cocoapi/PythonAPI

3. 测试

下载yolox_s.pth 放到 ~/code/YOLOX/preModels/ 文件夹

测试

# 单图cpu测试
python tools/demo.py image -n yolox-s -c preModels/yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result 

# 单图GPU测试
python tools/demo.py image -n yolox-s -c preModels/yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device gpu
# 或者 用指定 文件的方式 specify your detector's config
# 区别： -n 指定 yolo 模型名字 | -f 指定  yolo  config 文件 ，这两种方式用其中一种即可
python tools/demo.py image -f exps/default/yolox_s.py -c preModels/yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result  --device gpu

输出如下，表示运行成功

python tools/demo.py image -n yolox-s -c preModels/yolox_s.pth --path assets/dog.jpg --conf 0.25 --nms 0.45 --tsize 640 --save_result --device 0

2021-08-25 01:59:29.086 | INFO     | __main__:main:249 - Args: Namespace(camid=0, ckpt='preModels/yolox_s.pth', conf=0.25, demo='image', device='0', exp_file=None, experiment_name='yolox_s', fp16=False, fuse=False, legacy=False, name='yolox-s', nms=0.45, path='assets/dog.jpg', save_result=True, trt=False, tsize=640)

2021-08-25 01:59:29.285 | INFO     | __main__:main:259 - Model Summary: Params: 8.97M, Gflops: 26.81
2021-08-25 01:59:29.286 | INFO     | __main__:main:270 - loading checkpoint
2021-08-25 01:59:29.427 | INFO     | __main__:main:274 - loaded checkpoint done.

2021-08-25 01:59:29.619 | INFO     | __main__:inference:159 - Infer time: 0.1775s
2021-08-25 01:59:29.621 | INFO     | __main__:image_demo:196 - Saving detection result in ./YOLOX_outputs/yolox_s/vis_res/2021_08_25_01_59_29/dog.jpg

4. 训练自己的coco格式数据集

1.将 yolox/data/datasets/coco_classes.py 中的 COCO_CLASSES 修改为自己数据集的类别。

2.修改exps/example/custom/yolox_s.py。 这个文件是实验的配置文件，包含数据集地址，类别数，max_epoch等。更多参数及其默认值见yolox/exp/yolox_base.py。

yolox_s.py 中的Exp 类继承自yolox_base.py中的Exp类，因此可以将yolox_base.py中需要修改的配置参数添加至yolox_s.py进行修改。

tools/train.py 中的"–exp_file"参数值即yolox_s.py的文件路径。

self.data_dir = "D:/Z_Data/Animals_Coco" # 修改数据集地址 self.data_dir
self.train_ann = "instances_train2017.json"
self.val_ann = "instances_val2017.json"
self.num_classes = 2  # 修改类别 self.num_classes
# 剩下的 self.max_epoch，self.data_num_workers，self.eval_interval 可自行选择修改

3.修改 tools/train.py 中的参数配置

# 设置 default="Animals_Coco"， 训练后结果就会保存在 tools/YOLOX_outputs/Animals_Coco下
parser.add_argument("-expn", "--experiment-name", type=str, default=None)

# 设置 model_name，如果--exp_file参数为None，则通过此参数加载exps/default/中的默认的实验配置
parser.add_argument("-n", "--name", type=str, default="yolox-s", help="model name")

# 设置 batch_size
parser.add_argument("-b", "--batch-size", type=int, default=64, help="batch size")

# 设置gpu数量，因为我只有一张卡，所以设 default=0(如果只有一张卡的话，设置0或1都会启用gpu)
parser.add_argument(
    "-d", "--devices", default=0, type=int, help="device for training"
)

# 设置你的数据配置的路径，default="../exps/example/custom/yolox_s.py"
parser.add_argument(
    "-f",
    "--exp_file",
    default="../exps/example/custom/yolox_s.py",  # 如果出现报错 doesn't contains class named 'Exp', 将此处改为绝对路径即可。
    type=str,
    help="plz input your expriment description file",
)

# 设置预训练权重路径， default="../weights/yolox_s.pth"
parser.add_argument("-c", "--ckpt", default="../weights/yolox_s.pth", type=str, help="checkpoint file")

4.终端运行

python tools/train.py  # 也可以在后面添加配置参数  -b 64 -d 4 等。

如果要指定使用哪几颗GPU，可在 python 前添加环境变量 CUDA_VISIBLE_DEVICES。注意与 --devices 参数数量相匹配。示例

CUDA_VISIBLE_DEVICES=1，2  python tools/train.py  -b 32 -d 2 

5. 训练结果测试

在训练完成后，可以对训练的YOLOX算法模型进行测试，测试文件保存在tools/文件夹下，测试可以运行demo.py和eval.py这两个文件。
测试时需要修改配置参数--exp_file，保持与训练时的配置文件一致.

版权声明：本文为CSDN博主「zxxRobot」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/zxxxiazai/article/details/120906366