Problem with adding class/label to pretrained model

Hi, I think that I’ve tried everything I could, so I’d be glad if someone could help me. Also I didn’t find the answer in the Internet or the forums, it seems like a bug for me, because it should just work.

I’m trying just to add new labels to pretrained dataset, I’ve made many tries and got to the place, where simplest change in yaml is getting me nowhere, so:

My train1.py file:

from ultralytics import YOLO

model = YOLO(‘yolov5n.pt’)
results = model.train(data=‘coco2017_custom.yaml’, epochs=1, pretrained=True)

results

I’ve downloaded yolov5/data/coco128.yaml at master · ultralytics/yolov5 · GitHub (which is often refered in tutorials), saved to my workspace directory as coco2017_custom.yaml, ran training and everything works well:

YOLOv5n summary (fused): 193 layers, 2649200 parameters, 0 gradients, 7.7 GFLOPs
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 4/4 [00:18<00
                   all        128        929      0.658       0.49      0.578      0.419
                person        128        254      0.787      0.673      0.759      0.505
               bicycle        128          6      0.525      0.192       0.48      0.255

(and so on)

BUT when I add just one class to the dowloaded yaml file:


# YOLOv5 🚀 by Ultralytics, AGPL-3.0 license
# COCO128 dataset https://www.kaggle.com/ultralytics/coco128 (first 128 images from COCO train2017) by Ultralytics
# Example usage: python train.py --data coco128.yaml
# parent
# ├── yolov5
# └── datasets
#     └── coco128  ← downloads here (7 MB)


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../datasets/coco128  # dataset root dir
train: images/train2017  # train images (relative to 'path') 128 images
val: images/train2017  # val images (relative to 'path') 128 images
test:  # test images (optional)

# Classes
names:
  0: person
  1: bicycle
  2: car
  3: motorcycle
  4: airplane
  5: bus
  6: train
  7: truck
  8: boat
  9: traffic light
  10: fire hydrant
  11: stop sign
  12: parking meter
  13: bench
  14: bird
  15: cat
  16: dog
  17: horse
  18: sheep
  19: cow
  20: elephant
  21: bear
  22: zebra
  23: giraffe
  24: backpack
  25: umbrella
  26: handbag
  27: tie
  28: suitcase
  29: frisbee
  30: skis
  31: snowboard
  32: sports ball
  33: kite
  34: baseball bat
  35: baseball glove
  36: skateboard
  37: surfboard
  38: tennis racket
  39: bottle
  40: wine glass
  41: cup
  42: fork
  43: knife
  44: spoon
  45: bowl
  46: banana
  47: apple
  48: sandwich
  49: orange
  50: broccoli
  51: carrot
  52: hot dog
  53: pizza
  54: donut
  55: cake
  56: chair
  57: couch
  58: potted plant
  59: bed
  60: dining table
  61: toilet
  62: tv
  63: laptop
  64: mouse
  65: remote
  66: keyboard
  67: cell phone
  68: microwave
  69: oven
  70: toaster
  71: sink
  72: refrigerator
  73: book
  74: clock
  75: vase
  76: scissors
  77: teddy bear
  78: hair drier
  79: toothbrush
  80: xxx


# Download script/URL (optional)
download: https://ultralytics.com/assets/coco128.zip

(please note last entry in names field), run test, then NONE of labels are detected by the training. (I use the same train1.py as before):

Validating runs/detect/train/weights/best.pt...
Ultralytics YOLOv8.0.83 🚀 Python-3.10.8 torch-1.13.1 CPU
YOLOv5n summary (fused): 193 layers, 2658071 parameters, 0 gradients, 7.8 GFLOPs
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 4/4 [00:18<00
                   all        128        929          0          0          0          0

(and thats all the labels in the summary, no other which were present in the previous example)

Also using this model in inference produces NO DETECTIONS, while original, trained on unchanged yaml file (but loaded from my workspace local directory) is working well.

I didn’t change anything in train data at to this point, but later I’ve edited one of label files just to include new label in datasets, and it produces the same results. With added label to the yaml file, the training script is detecting change in config:

Overriding model.yaml nc=80 with nc=81

I’ve tried this using yolov5n.pt and yolov8n.pt pretrained models, the results are the same.

Edit: some typos and clarifications

Ok, I think that I resolved problem.

I tried with 10 epochs, and success rate (mAP50-95 metric) started to grow above 0 value. I’ve retrained with 50 epochs and got ~0.5 of mAP50-95 after 30. epoch.

Not sure though why adding a label would destroy model’s inference rate after retraining the model with low number of epoch, it’s probably the internal mechanics of the algorithm.

I hope that my experience could help someone struggling with try and error approach.