How to add P2 head on YOLOv8?

Hi, I’m trying to add a P2 head for detecting small targets. Any tips? For now I’m trying to make an analogy from the steps for YOLOv5 (yolov5/yolov5-p2.yaml at master · ultralytics/yolov5 · GitHub), but it would be better to have something more targeted.

I’ve tryied this model yaml:

Ultralytics YOLO :rocket:, GPL-3.0 license

Parameters

nc: 1 # number of classes
depth_multiple: 0.33 # scales module repeats
width_multiple: 0.25 # scales convolution channels

YOLOv8.0n_P2 backbone

backbone:

[from, repeats, module, args]

  • [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
  • [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
  • [-1, 3, C2f, [128, True]]
  • [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
  • [-1, 6, C2f, [256, True]]
  • [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
  • [-1, 6, C2f, [512, True]]
  • [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
  • [-1, 3, C2f, [1024, True]]
  • [-1, 1, SPPF, [1024, 5]] # 9

YOLOv8.0n_P2 head

head:

  • [-1, 1, nn.Upsample, [None, 2, ‘nearest’]]

  • [[-1, 6], 1, Concat, [1]] # cat backbone P4

  • [-1, 3, C2f, [512]] # 12

  • [-1, 1, nn.Upsample, [None, 2, ‘nearest’]]

  • [[-1, 4], 1, Concat, [1]] # cat backbone P3

  • [-1, 3, C2f, [256]] # 15 (P3/8-small)

  • [-1, 1, nn.Upsample, [None, 2, ‘nearest’]]

  • [[-1, 2], 1, Concat, [1]] # cat backbone P2

  • [-1, 1, C2f, [128]] # 18 (P2/4-xsmall)

  • [-1, 1, Conv, [128, 3, 2]]

  • [[-1, 15], 1, Concat, [1]] # cat backbone P3

  • [-1, 3, C2f, [256]] # 21 (P3/8-small)

  • [-1, 1, Conv, [256, 3, 2]]

  • [[-1, 12], 1, Concat, [1]] # cat head P4

  • [-1, 3, C2f, [512]] # 24 (P4/16-medium)

  • [-1, 1, Conv, [512, 3, 2]]

  • [[-1, 9], 1, Concat, [1]] # cat head P5

  • [-1, 3, C2f, [1024]] # 27 (P5/32-large)

  • [[18, 21, 24, 27], 1, Detect, [nc]] # Detect(P2, P3, P4, P5)

But I am getting this bug:

Traceback (most recent call last):
File “/usr/local/bin/yolo”, line 33, in
sys.exit(load_entry_point(‘ultralytics’, ‘console_scripts’, ‘yolo’)())
File “/content/ultralytics/ultralytics/yolo/cfg/init.py”, line 304, in entrypoint
getattr(model, mode)(**overrides) # default args from model
File “/content/ultralytics/ultralytics/yolo/engine/model.py”, line 298, in train
self.trainer = TASK_MAP[self.task]1
KeyError: None

Sorry have you reached to an answer i am looking in same error for a while ?