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

图中是本人存放数据集的目录

--datasets
	--COCO
		--test#当做验证集
			--annotations#coco格式的标签json
			--images#图片
		--train-aug#训练集（数据增强后）
			--annotations
			--images

yolox GitHub 搭建自己的数据集教程

yolox_coco_l.py（exp/example/yolox_coco/ 目录下）

设置训练集、验证集和评估方法
GitHub 里面提到这部分需要参考 exps/example/yolox_voc/yolox_voc_s.py

#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# Copyright (c) Megvii, Inc. and its affiliates.
import os
import torch.distributed as dist
from yolox.exp import Exp as MyExp
import torch


class Exp(MyExp):
    def __init__(self):
        super(Exp, self).__init__()
        self.depth = 1
        self.width = 1
        self.exp_name = os.path.split(os.path.realpath(__file__))[1].split(".")[0]
        # self.input_size = (960, 540)###

        # Define yourself dataset path
        '''修改
        self.data_dir\
        self.val_ann\
        self.train_ann\'''
        self.data_dir = "/home/server/xcg/YOLOX/datasets/COCO"
        self.val_ann = "pig_annotation_test.json"
        self.train_ann = "pig_annotation_train_aug.json"

        '''修改self.num_classes'''
        self.num_classes = 3

        self.max_epoch = 300
        self.data_num_workers = 16
        self.eval_interval = 1

    '''get_data_loader设置训练集'''
    '''修改
    VOCDetection 为 COCODataset
    COCODataset 的 data_dir\json_file\name
    name是训练集图片保存的文件夹名称'''
    def get_data_loader(self, batch_size, is_distributed, no_aug=False):
            from yolox.data import (
                COCODataset,#VOCDetection
                TrainTransform,
                YoloBatchSampler,
                DataLoader,
                InfiniteSampler,
                MosaicDetection,
            )

            dataset = COCODataset(
                data_dir="/home/server/xcg/YOLOX/datasets/COCO/train-aug",
                json_file=self.train_ann,
                name="images",
                img_size=self.input_size,
                preproc=TrainTransform(
                    rgb_means=(0.485, 0.456, 0.406),
                    std=(0.229, 0.224, 0.225),
                    max_labels=50,
                ),
            )

            dataset = MosaicDetection(
                dataset,
                mosaic=not no_aug,
                img_size=self.input_size,
                preproc=TrainTransform(
                    rgb_means=(0.485, 0.456, 0.406),
                    std=(0.229, 0.224, 0.225),
                    max_labels=120,
                ),
                degrees=self.degrees,
                translate=self.translate,
                scale=self.scale,
                shear=self.shear,
                perspective=self.perspective,
                enable_mixup=self.enable_mixup,
            )

            self.dataset = dataset

            if is_distributed:
                batch_size = batch_size // dist.get_world_size()

            sampler = InfiniteSampler(
                len(self.dataset), seed=self.seed if self.seed else 0
            )

            batch_sampler = YoloBatchSampler(
                sampler=sampler,
                batch_size=batch_size,
                drop_last=False,
                input_dimension=self.input_size,
                mosaic=not no_aug,
            )

            dataloader_kwargs = {"num_workers": self.data_num_workers, "pin_memory": True}
            dataloader_kwargs["batch_sampler"] = batch_sampler
            train_loader = DataLoader(self.dataset, **dataloader_kwargs)

            return train_loader

    '''get_eval_loader设置验证集'''
    '''修改
    VOCDetection 为 COCODataset
    COCODataset 的 data_dir\json_file\name
    name是验证集图片保存的文件夹名称'''
    def get_eval_loader(self, batch_size, is_distributed, testdev=False):
            from yolox.data import COCODataset, ValTransform

            valdataset = COCODataset(
                data_dir="/home/server/xcg/YOLOX/datasets/COCO/test",
                json_file=self.val_ann,##
                name="images",###
                img_size=self.test_size,
                preproc=ValTransform(
                    rgb_means=(0.485, 0.456, 0.406),
                    std=(0.229, 0.224, 0.225),
                ),
            )

            if is_distributed:
                batch_size = batch_size // dist.get_world_size()
                sampler = torch.utils.data.distributed.DistributedSampler(
                    valdataset, shuffle=False
                )
            else:
                sampler = torch.utils.data.SequentialSampler(valdataset)

            dataloader_kwargs = {
                "num_workers": self.data_num_workers,
                "pin_memory": True,
                "sampler": sampler,
            }
            dataloader_kwargs["batch_size"] = batch_size
            val_loader = torch.utils.data.DataLoader(valdataset, **dataloader_kwargs)

            return val_loader

    '''get_evaluator 设置评估方法'''
    '''替换 VOCEvaluator 为 COCOEvaluator '''
    def get_evaluator(self, batch_size, is_distributed, testdev=False):
            from yolox.evaluators import COCOEvaluator #VOCEvaluator

            val_loader = self.get_eval_loader(batch_size, is_distributed, testdev=testdev)
            evaluator = COCOEvaluator(
                dataloader=val_loader,
                img_size=self.test_size,
                confthre=self.test_conf,
                nmsthre=self.nmsthre,
                num_classes=self.num_classes,
            )
            return evaluator

参考图片
标注的 json 文件放在 annotations 文件夹里面，代码默认的操作（yolox/data/datasets/coco.py）
这里的 self.data_dir 好像没怎么用到
self.train_ann 和 self.val_ann 只要填json文件名就行

训练命令

python tools/train.py 
-f /home/server/xcg/YOLOX/exps/example/yolox_coco/yolox_coco_l.py #网络配置文件和数据加载
-d 0 ##cuda
-b 8 #batch_size
--fp16 #混合精度训练
-o #抢占gpu资源
-c /home/server/xcg/YOLOX/yolox_l.pth #加载预训练权重