基于EfficientDet训练自己的数据集

github：Zylo

使用的torch和torchvision版本，torch和torchvision的版本要对应，cuda版本一致

数据集

将数据集做成coco格式

# for example, coco2017
datasets/
    -coco2017/
        -train2017/
            -000000000001.jpg
            -000000000002.jpg
            -000000000003.jpg
        -val2017/
            -000000000004.jpg
            -000000000005.jpg
            -000000000006.jpg
        -annotations
            -instances_train2017.json
            -instances_val2017.json
           

其中，annotations中的类别需要从1开始，不从0开始

修改project文件下的yml文件

project_name: logo  # also the folder name of the dataset that under data_path folder
train_set: train
val_set: val
num_gpus: 1

obj_list: [ 'adidas0', 'chanel','gucci','hh','lacoste','mk','nike','prada','puma','supreme' ]
           

训练

• from scratch

python train.py -c 0 --batch_size 64 --optim sgd --lr 8e-2
           

• Train a custom dataset with pretrained weights

python train.py -c 2 -p your_project_name --batch_size 8 --lr 1e-3 --num_epochs 10 --load_weights /path/to/your/weights/efficientdet-d2.pth --head_only True
           

评估

python coco_eval.py -p your_project_name -c 5 -w /path/to/your/weights
           

inference

python efficientdet_test.py