IndexError: list index out of range - YOLOR Custom Training

YOLOR github: https://github.com/WongKinYiu/yolor

🚨 YOLOR train.py 실행시 index error

6클래스로 custom training을 시도했으나, train이 난 후 test 시 plot 그리는 부분에서 에러 발생

Traceback (most recent call last):
  File "train.py", line 537, in <module>
    train(hyp, opt, device, tb_writer, wandb)
  File "train.py", line 344, in train
    log_imgs=opt.log_imgs if wandb else 0)
  File "/yolo/yolor/test.py", line 226, in test
    plot_images(img, output_to_target(output, width, height), paths, f, names)  # predictions
  File "/yolo/yolor/utils/plots.py", line 163, in plot_images
    cls = names[cls] if names else cls
IndexError: list index out of range
Traceback (most recent call last):
  File "/opt/conda/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/opt/conda/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/opt/conda/lib/python3.6/site-packages/torch/distributed/launch.py", line 260, in <module>
    main()
  File "/opt/conda/lib/python3.6/site-packages/torch/distributed/launch.py", line 256, in main
    cmd=cmd)

💡 class predict를 COCO Dataset의 class 개수인 80개 기준으로 했기 때문!

utils/plots.py

163 line의 cls는 결국 targets값 (모델이 예측한 라벨)에서 오는데, 여기서 80 classes 기준으로 결과를 뽑아내기 때문에 인덱스 에러 발생

사용하는 config 파일의 head layer 및 implicit 부분을 클래스 개수에 맞게 바꿔줘야 한다.
cfg/yolor_p6.py

default : 3 * (5 + 80) = 255
custom : anchor * (5 + classes)

yolor_p6.cfg 기준 수정해야 할 line은 다음과 같다.
#L1569, #L1573, #L1577, #L1581, #L1605, #L1649, #L1693, #L1737

implicit_mul의 filter의 역할은 정확히 모르겠지만 implicit knowledge와 관련된 것으로 추정
각 head의 convolutional의 filter는 최종 classification 결과 뱉기 전에 class에 따라 shape를 맞춰주는 부분

내가 사용하는 dataset은 6 classes 이기 때문에 3 * (5 + 6) 으로 고쳐 주었다.