How to solve KeyError: 1 while validating best custom torch yolov5 model with plots?

username: mobassir94
cooked: <p>Using this repo : <a href="https://github.com/ultralytics/yolov5" class="inline-onebox" rel="noopener nofollow ugc">GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch &gt; ONNX &gt; CoreML &gt; TFLite</a> i am trying to train yolov5 on a custom dataset containing 3 classes.once the training is done,during validation i get this error :</p>
<pre><code class="lang-auto">
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     mAP@.5 mAP@.5:.95: 100% 1/1 [00:04&lt;00:00,  4.19s/it]
Traceback (most recent call last):
  File "train.py", line 625, in &lt;module&gt;
    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
</code></pre>
<p>even though i wrote custom yaml file like this :</p>
<pre><code class="lang-auto">%%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"] 

</code></pre>
<p>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”]<br>
my model always predicts “name” for all 3 fields(please check the last image attached in the colab code):<br>
here is the code with outputs and errors : <a href="https://colab.research.google.com/drive/1Pyz5hCi0fpmQ4lFV9rvAcdMp2yg9KN1P#scrollTo=ogZdrD4OY2Wg" class="inline-onebox" rel="noopener nofollow ugc">Google Colab</a></p>
<p>how can i solve this issue? thanks in advance</p>

username: ptrblck
cooked: <p>I don’t see the error in the linked notebook, but would guess that <code>names[i]</code> is failing for <code>i=1</code>.<br>
Print <code>names</code> right before calling the plot function and make sure it’s containing the expected values.</p>
<p>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 <code>WeightedRandomSampler</code> 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.</p>

username: mobassir94
cooked: <p>output of 17th cell of my linked notebook here : <a href="https://colab.research.google.com/drive/1Pyz5hCi0fpmQ4lFV9rvAcdMp2yg9KN1P#scrollTo=uVnDpfE0ZwB_" class="inline-onebox" rel="noopener nofollow ugc">Google Colab</a> giving this error after all epoch training is finish.</p>
<p>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”</p>
<p>i printed names in train.py file right before calling it and this dictionary was printed [“name”,“age”,“mobile”]</p>