Using this repo : GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite i am trying to train yolov5 on a custom dataset containing 3 classes.once the training is done,during validation i get this error :

Validating runs/train/Result4/weights/best.pt... Fusing layers... Model Summary: 476 layers, 87212152 parameters, 0 gradients, 217.1 GFLOPs Class Images Labels P R [email protected] [email protected]:.95: 100% 1/1 [00:04<00:00, 4.19s/it] Traceback (most recent call last): File "train.py", line 625, in <module> main(opt) File "train.py", line 522, in main train(opt.hyp, opt, device, callbacks) File "train.py", line 429, in train compute_loss=compute_loss) # val best model with plots File "/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py", line 28, in decorate_context return func(*args, **kwargs) File "/content/val.py", line 240, in run p, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names) File "/content/utils/metrics.py", line 78, in ap_per_class plot_mc_curve(px, f1, Path(save_dir) / 'F1_curve.png', names, ylabel='F1') File "/content/utils/metrics.py", line 323, in plot_mc_curve ax.plot(px, y, linewidth=1, label=f'{names[i]}') # plot(confidence, metric) KeyError: 1

even though i wrote custom yaml file like this :

%%writetemplate /content/DataFile/data.yaml
train: /content/convertor/fold0/images/train 
val: /content/convertor/fold0/images/val 
test: /content/convertor/fold0/images/val 
nc: 3 
names: ["name","age","mobile"] 
i still get the error above and while inferring, all the bounding boxes are given “name” as label ,instead of predicting these 3 classes names: [“name”,“age”,“mobile”]

my model always predicts “name” for all 3 fields(please check the last image attached in the colab code):

here is the code with outputs and errors : Google Colab
how can i solve this issue? thanks in advance
              
I don’t see the error in the linked notebook, but would guess that names[i] is failing for i=1.

Print names right before calling the plot function and make sure it’s containing the expected values.
Your second issue sounds as if the model is overfitting to a specific class. I’m not familiar with your use case but in case you are working with an imbalanced dataset, try to use a WeightedRandomSampler if possible or a weighted loss. It’s also a good idea to try to overfit a small dataset first (e.g. just 10 samples) by playing around with the hyperparameters to see if any obvious bugs are still in the code.
              
output of 17th cell of my linked notebook here : Google Colab giving this error after all epoch training is finish.
if you check the visualized image attached in the last code cell of that notebook then you will see that my model is able to detect all those 3 classes perfectly and with good performance but it’s saying all those 3 classes are “name”
i printed names in train.py file right before calling it and this dictionary was printed [“name”,“age”,“mobile”]