Extracting YOLOv8 backbone weights

I'm trying to extract the classification backbone from my YOLOv8 model. I saw an issue on the official Ultralytics repository about that but when I tried the method mentioned it didn't work. Here's the link for the issue

What I want to do is extract the weights of the backbone and then use it as if it was a classification model to perform FGSM attack because I tried to do it on the whole model it didn't work even after I extracted class logits from yolo using pytorch hooks

import torch
import torch.nn as nn
from ultralytics import YOLO  

# Step 1: Load trained YOLOv8 model
model = YOLO('/content/best.pt') 
image = '/content/stop.png'

# Step 2: Extract the backbone (CSPDarknet53)
backbone = model.model[0] 
labels = torch.tensor([22])

num_classes = 29  
classify_model = nn.Sequential(
    backbone,  # Use the CSPDarknet53 backbone
    nn.AdaptiveAvgPool2d((1, 1)),  # Global Average Pooling to reduce to (batch_size, channels, 1, 1)
    nn.Flatten(),  # Flatten to (batch_size, channels)
    nn.Linear(in_features=backbone[-1].out_channels, out_features=num_classes)  # Linear layer for classification
)

# Example FGSM attack
def fgsm_attack(model, images, labels, epsilon):
    images.requires_grad = True
    outputs = model(images)
    loss = nn.CrossEntropyLoss()(outputs, labels)
    model.zero_grad()
    loss.backward()  # Compute gradients with respect to input images
    grad_sign = images.grad.data.sign()
    perturbed_image = images + epsilon * grad_sign  # Apply perturbation
    return perturbed_image

# apply FGSM on `classify_model`
perturbed_image = fgsm_attack(classify_model, image, labels, 0.3)

But I get this error 'TypeError: 'DetectionModel' object is not subscriptable'

I tried to do FGSM on the whole model it didn't work even after I extracted class logits from yolo using pytorch hooks. I also tried to extract YOLOv8 backbone and use it as a classification model but this also is not working.

Try this solution to extract the YOLOv8 backbone in your code:

# Step 2: Extract the backbone (CSPDarknet53)
backbone = model.model.model[:10]

First, instead of just model.model go deeper into the model structure to avoid the 'TypeError: 'DetectionModel' object is not subscriptable' error. Let's see in detail what we have and what we need here.

Look for instance at the model = YOLO('yolov8n.pt') object:

model = YOLO('yolov8n.pt')
print(model)

# OUTPUT
'''
YOLO(
  (model): DetectionModel(
    (model): Sequential(
      (0): Conv(
        (conv): Conv2d(3, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        (bn): BatchNorm2d(16, eps=0.001, momentum=0.03, affine=True, track_running_stats=True)
        (act): SiLU(inplace=True)
      )
      (1): Conv(
        (conv): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        (bn): BatchNorm2d(32, eps=0.001, momentum=0.03, affine=True, track_running_stats=True)
        (act): SiLU(inplace=True)
      )
...
        (dfl): DFL(
          (conv): Conv2d(16, 1, kernel_size=(1, 1), stride=(1, 1), bias=False)
        )
      )
    )
  )
)
'''

We need to reach these enumerated modules (0, 1 ...) which are the elements of the Sequential container. So, let's see where it lies:

type(model)

# OUTPUT
ultralytics.models.yolo.model.YOLO
<...>
YOLO (You Only Look Once) object detection model.

type(model.model)

# OUTPUT
ultralytics.nn.tasks.DetectionModel
<...>
YOLOv8 detection model.

type(model.model.model)

# OUTPUT
torch.nn.modules.container.Sequential
A sequential container.

Here it is, so the simplified view of the model is the following, and we need to call model.model.model[:] to reach all the numerated blocks.

YOLO(
    DetectionModel(
        Sequential(
            0: Conv, 
            1: Conv, 
            2, 
            3,
            ...
        )
    )
)

Second, if you look at the model's structure, the backbone is the first 10 layers of the model (0-9), not just the first one. Here comes the final solution backbone = model.model.model[:10].