Training YOLOv8 on a custom dataset

[Zunon]

Prerequisite

• [X] I have searched the existing and past issues but cannot get the expected help.
• [X] I have read the FAQ documentation but cannot get the expected help.
• [X] The bug has not been fixed in the latest version.

🐞 Describe the bug

I'm training a YOLOv8 model on a custom COCO-formatted dataset (mav-vid dataset). I came up with this config:

yolov8_n_syncbn_fast_8xb16-500e_mav-vid.py:

_base_ = './yolov8_s_syncbn_fast_8xb16-500e_coco.py'

deepen_factor = 0.33
num_classes = 1
widen_factor = 0.25

model = dict(
    backbone=dict(deepen_factor=deepen_factor, widen_factor=widen_factor),
    neck=dict(deepen_factor=deepen_factor, widen_factor=widen_factor),
    bbox_head=dict(head_module=dict(widen_factor=widen_factor, num_classes=num_classes)),
    train_cfg=dict(num_classes=num_classes))

data_root = 'data/mav_vid_dataset/'
metainfo = {
    'classes': ('drone', ),
    'palette': [(220, 20, 60)],
}
train_dataloader = dict(
    batch_size=1,
    dataset=dict(
        data_root=data_root,
        metainfo=metainfo,
        ann_file='annotations/train.json',
        data_prefix=dict(img='train/img/')))
val_dataloader = dict(
    dataset=dict(
        data_root=data_root,
        metainfo=metainfo,
        ann_file='annotations/val.json',
        data_prefix=dict(img='val/img/')))
test_dataloader = val_dataloader
val_evaluator = dict(ann_file=data_root + 'annotations/val.json')
test_evaluator = val_evaluator

However when I try to train it using tools/train.py I get the following error:

Traceback (most recent call last):
  File "lib/mmyolo/tools/train.py", line 115, in <module>
    main()
  File "lib/mmyolo/tools/train.py", line 111, in main
    runner.train()
  File "/home/user/.local/lib/python3.8/site-packages/mmengine/runner/runner.py", line 1721, in train
    model = self.train_loop.run()  # type: ignore
  File "/home/user/.local/lib/python3.8/site-packages/mmengine/runner/loops.py", line 96, in run
    self.run_epoch()
  File "/home/user/.local/lib/python3.8/site-packages/mmengine/runner/loops.py", line 112, in run_epoch
    self.run_iter(idx, data_batch)
  File "/home/user/.local/lib/python3.8/site-packages/mmengine/runner/loops.py", line 128, in run_iter
    outputs = self.runner.model.train_step(
  File "/home/user/.local/lib/python3.8/site-packages/mmengine/model/base_model/base_model.py", line 114, in train_step
    losses = self._run_forward(data, mode='loss')  # type: ignore
  File "/home/user/.local/lib/python3.8/site-packages/mmengine/model/base_model/base_model.py", line 340, in _run_forward
    results = self(**data, mode=mode)
  File "/home/user/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/Desktop/kalmheir-intern/OpTrainer/lib/mmdetection/mmdet/models/detectors/base.py", line 92, in forward
    return self.loss(inputs, data_samples)
  File "/home/user/Desktop/kalmheir-intern/OpTrainer/lib/mmdetection/mmdet/models/detectors/single_stage.py", line 78, in loss
    losses = self.bbox_head.loss(x, batch_data_samples)
  File "/home/user/Desktop/kalmheir-intern/OpTrainer/lib/mmyolo/mmyolo/models/dense_heads/yolov5_head.py", line 465, in loss
    losses = self.loss_by_feat(*loss_inputs)
  File "/home/user/Desktop/kalmheir-intern/OpTrainer/lib/mmyolo/mmyolo/models/dense_heads/yolov8_head.py", line 351, in loss_by_feat
    loss_cls = self.loss_cls(flatten_cls_preds, assigned_scores).sum()
  File "/home/user/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/user/Desktop/kalmheir-intern/OpTrainer/lib/mmdetection/mmdet/models/losses/cross_entropy_loss.py", line 291, in forward
    loss_cls = self.loss_weight * self.cls_criterion(
  File "/home/user/Desktop/kalmheir-intern/OpTrainer/lib/mmdetection/mmdet/models/losses/cross_entropy_loss.py", line 139, in binary_cross_entropy
    loss = F.binary_cross_entropy_with_logits(
  File "/home/user/.local/lib/python3.8/site-packages/torch/nn/functional.py", line 3163, in binary_cross_entropy_with_logits
    raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
ValueError: Target size (torch.Size([2, 8400, 80])) must be the same as input size (torch.Size([2, 8400, 1]))

I'm not sure how to resolve it and I've read #536 and #770 however they don't seem to have the same issue as me.

Environment

sys.platform: linux
Python: 3.8.16 (default, Jun 12 2023, 18:09:05) [GCC 11.2.0]
CUDA available: True
numpy_random_seed: 2147483648
GPU 0: NVIDIA RTX A5000
CUDA_HOME: /usr/local/cuda
NVCC: Cuda compilation tools, release 11.8, V11.8.89
GCC: gcc (Ubuntu 11.3.0-1ubuntu1~22.04.1) 11.3.0
PyTorch: 2.0.1+cu118
PyTorch compiling details: PyTorch built with:
  - GCC 9.3
  - C++ Version: 201703
  - Intel(R) oneAPI Math Kernel Library Version 2022.2-Product Build 20220804 for Intel(R) 64 architecture applications
  - Intel(R) MKL-DNN v2.7.3 (Git Hash 6dbeffbae1f23cbbeae17adb7b5b13f1f37c080e)
  - OpenMP 201511 (a.k.a. OpenMP 4.5)
  - LAPACK is enabled (usually provided by MKL)
  - NNPACK is enabled
  - CPU capability usage: AVX2
  - CUDA Runtime 11.8
  - NVCC architecture flags: -gencode;arch=compute_37,code=sm_37;-gencode;arch=compute_50,code=sm_50;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_70,code=sm_70;-gencode;arch=compute_75,code=sm_75;-gencode;arch=compute_80,code=sm_80;-gencode;arch=compute_86,code=sm_86;-gencode;arch=compute_90,code=sm_90
  - CuDNN 8.7
  - Magma 2.6.1
  - Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=11.8, CUDNN_VERSION=8.7.0, CXX_COMPILER=/opt/rh/devtoolset-9/root/usr/bin/c++, CXX_FLAGS= -D_GLIBCXX_USE_CXX11_ABI=0 -fabi-version=11 -Wno-deprecated -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOROCTRACER -DUSE_FBGEMM -DUSE_QNNPACK -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wall -Wextra -Werror=return-type -Werror=non-virtual-dtor -Werror=bool-operation -Wnarrowing -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wunused-local-typedefs -Wno-unused-parameter -Wno-unused-function -Wno-unused-result -Wno-strict-overflow -Wno-strict-aliasing -Wno-error=deprecated-declarations -Wno-stringop-overflow -Wno-psabi -Wno-error=pedantic -Wno-error=redundant-decls -Wno-error=old-style-cast -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Werror=cast-function-type -Wno-stringop-overflow, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_DISABLE_GPU_ASSERTS=ON, TORCH_VERSION=2.0.1, USE_CUDA=ON, USE_CUDNN=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=1, USE_NNPACK=ON, USE_OPENMP=ON, USE_ROCM=OFF, 

TorchVision: 0.15.2+cu118
OpenCV: 4.7.0
MMEngine: 0.7.4
MMCV: 2.0.0
MMDetection: 3.0.0
MMYOLO: 0.5.0+dc85144

Additional information

No response

[zbf1-999]

我遇到了相同的问题，请问这是为什么，应该如何解决呢

[baconkcode]

change your configs like this model = dict( bbox_head=dict( head_module=dict(num_classes=num_classes)), train_cfg=dict(assigner=dict( type='BatchTaskAlignedAssigner', num_classes=1,#your num_classes use_ciou=True, topk=10, alpha=0.5, beta=6.0, eps=1e-09)) )

[Le0v1n]

After running, a relatively complete config.py file will be generated. Check whether all the class_name related items in this file have been changed to your class_name.