search

ultralytics / yolov5

YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite
https://docs.ultralytics.com
GNU Affero General Public License v3.0
50.09k stars 16.18k forks source link

Is it possible to use COCO pretrained weights with my 3 class dataset? #5449

Closed dariogonle closed 2 years ago

dariogonle commented 2 years ago

Search before asking

Question

I have a 3 classes datasets, being one of it the class Person. I would like to use the COCO pretrained weights to improve the results of my Person class.

The problem I see is that COCO has 80 classes and I only have 3. I'm having problems when I try to use the COCO pretrained weights with my dataset, and I was wondering if it is as simple as it cannot be done. If you have a 3 classes dataset you cannot use COCO pretrained weights. Is this correct?

Thank you in advance.

Additional

No response

Zengyf-CVer commented 2 years ago

@dariogonle It has nothing to do with the number of categories in your custom data set. From your description, I guess that there is a problem with your custom data set. Carefully check the category assignment and instance assignment issues of the training set and validation set.

Le0v1n commented 2 years ago

Actually yes, you need modify some configurations and match the label requirements of YOLO v5.

dariogonle commented 2 years ago

@Zengyf-CVer, @Le0v1n or @glenn-jocher Is it possible to train from COCO pretraiend weights but using only the class Person?

I mean, do a fine tuning of COCO weights with the class Person of my dataset.

glenn-jocher commented 2 years ago

@dariogonle YOLOv5 can train on any dataset using any YOLOv5 model as pretrained weights, class counts are irrelevant. Nothing special is needed, just run a training command:

i.e. train VOC.yaml (10 classes) starting from YOLOv5s.pt (80 classes):

python train.py --data VOC.yaml --weights yolov5s.pt
dariogonle commented 2 years ago

Thank you for your answer @glenn-jocher, but I am a bit confused.

If I use my own dataset.yaml

nc: 3
names: ['person', 'VE', 'HE']

I get the following results: image

but if I change my dataset.yaml (the paths to the dataset are the same) to

nc: 80  # number of classes
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']  # class names

I get the following: image

In both cases I have used the following:

python train.py --epochs 1 --batch 8 --data dataset.yaml --hyp data/hyps/hyp.finetune.yaml --weights 'yolov5m.pt' --img-size 704

So something weird must be happening, because with my dataset.yaml it is not working.

glenn-jocher commented 2 years ago

@dariogonle if your dataset is COCO then you should use the nc and names field you showed, if you have a custom dataset you supply your custom dataset nc and names.

Start from Train Custom Dataset tutorial:

YOLOv5 Tutorials

dariogonle commented 2 years ago

@dariogonle if your dataset is COCO then you should use the nc and names field you showed, if you have a custom dataset you supply your custom dataset nc and names.

Start from Train Custom Dataset tutorial:

YOLOv5 Tutorials

@glenn-jocher it is true that if I do it as you said, I can train a model with my dataset starting from COCO's pretrained weights, but I don't want to train a new model, what I want to do is a fine tunning, I mean get COCO's weights of the Person class and fine tuning it during few epochs with my Person images.

My dataset is a custom one, but I want to use the pretrained weights of COCO, as they'd improve the Person class metrics . So my questions are the following: Shoud I use my own .yaml file?

nc: 3
names: ['person', 'VE', 'HE']

Should I use the following command to train?

python train.py --epochs 15 --batch 8 --data customDataset.yaml --hyp data/hyps/hyp.finetune.yaml --weights yolov5m.pt --img-size 704

Because if I do it like this I get poor results, but in case I use coco.yaml file I get good ones.

Thank you in advance.

dariogonle commented 2 years ago

Ok, I think I got it.

Thank you all.

Amulya2110 commented 2 years ago

Hi @dariogonle I am also working on similar problem, where-in i want to improve the performance of a "dog" class , by taking pre-trained coco weights as initial weights and train on top my custom dog data.

By labelling my custom images with label 16 (dog class from pre-trained model) and by pointing to coco.yaml file and re-training with my custom images will help? Orelse could you please let me know if there is any other way that you have figured out.

Thankyou

dariogonle commented 2 years ago

Hi @Amulya2110,

I can tell you what have I done.

My original customDataset.yaml was:

train: ../../Datasets/CHV/train/images
val: ../../Datasets/CHV/test/images

nc: 3
names: ['Person', 'VE', 'HE']

and I changed it to this:

train: ../../Datasets/CHV/train/images
val: ../../Datasets/CHV/test/images

nc: 80  # number of classes
names: ['Person', 'VE', 'HE', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush']  # class names

And I just fine tunning it with the following command:

python train.py --epochs 15 --batch 8 --data customDataset.yaml --hyp data/hyps/hyp.finetune.yaml --weights yolov5m.pt --img-size 704

The result weights file will only have 3 classes as in my dataset there are only labels for classes 0, 1 and 2.

Amulya2110 commented 2 years ago

@glenn-jocher @dariogonle Thankyou for the reply @dariogonle . I understood what you have done , but i am unable to understand this point "The result weights file will only have 3 classes as in my dataset there are only labels for classes 0, 1 and 2.".

As the config file you pointed to while training says you have nc:80 , 80classes , so ideally shouldn't the weights file be of 80 class output even though in your custom dataset you have given labels for only 3 classes?

Does the weight file have no.of classes with respect to no.of classes labelled in our custom data images or with respect to the nc defined in config file(data.yaml)?

dariogonle commented 2 years ago

At the beginning it is true that you're setting 80 class, but as in this case I only have 3 clases in my dataset (0, 1 and 2) the resulting weights will only have 3 classes.

github-actions[bot] commented 2 years ago

👋 Hello, this issue has been automatically marked as stale because it has not had recent activity. Please note it will be closed if no further activity occurs.

Access additional YOLOv5 🚀 resources:

Access additional Ultralytics ⚡ resources:

Feel free to inform us of any other issues you discover or feature requests that come to mind in the future. Pull Requests (PRs) are also always welcomed!

Thank you for your contributions to YOLOv5 🚀 and Vision AI ⭐!