是否可以修改YOLOv8以将其用作其他任务的特征提取器?

Question

我在这里通读了YOLOv8的文档，但我没有看到一个简单的方法来做我在标题中建议的事情。我想做的是加载一个预训练的YOLOv8模型，创建一个更大的模型，其中将包含YOLOv8作为子模块，并修改YOLOv8的前向函数，以便我可以访问对象检测损失和卷积特征，以便它们可以用于为其他自定义任务提供后续层。

为了让事情更清楚，这是yolov8最初的使用方式，根据https://docs.ultralytics.com/quickstart/#use-with-python:

from ultralytics import YOLO
# Create a new YOLO model from scratch
model = YOLO('yolov8n.yaml')
# Load a pretrained YOLO model (recommended for training)
model = YOLO('yolov8n.pt')
# Train the model using the 'coco128.yaml' dataset for 3 epochs
results = model.train(data='coco128.yaml', epochs=3)
# Evaluate the model's performance on the validation set
results = model.val()
# Perform object detection on an image using the model
results = model('https://ultralytics.com/images/bus.jpg')
# Export the model to ONNX format
success = model.export(format='onnx')

相反，我想这样做:

import torch
import torch.nn as nn
from ultralytics import YOLO
class Yolov8Wrapper(nn.Module):

def __init__(self, yolov8_feature_dim, n1, n2, n3):
super().__init__()
self.yolov8 = YOLO('yolov8n.pt')
self.fc1 = nn.Linear(yolov8_feature_dim, n1)
self.fc2 = nn.Linear(yolov8_feature_dim, n2)
self.fc3 = nn.Linear(yolov8_feature_dim, n3)

def forward(self, images, gt_boxes):
features, loss = self.yolov8(images, gt_boxes)
logits1 = self.fc1(features)
logits2 = self.fc2(features)
logits3 = self.fc3(features)
return {
'logits1': logits1,
'logits2': logits2,
'logits3': logits3,
'yolov8_loss': loss,
}

上面的代码是一个非常简化的草图，当然不会工作，但或多或少就是这个想法。此外，通过创建这个特别的包装器，我将无法使用YOLO库自带的开箱即用的功能来训练、验证和预测，因为它将是一个自定义架构(YOLOv8只是它的一个子模块)。因此，我还需要弄清楚如何编写一个自定义数据加载程序，以便为YOLOv8提供它期望的输入以及我的包装器所需的额外内容(包装器中可能有不同的附加层，根据YOLOv8的特性预测不同的输出，将其视为多任务学习)。

这可能吗?如何做到这一点?

Answer 1

甚至不需要包装器。为了达到这个目的，我做了以下工作:

model = YOLO('yolov8s-cls')
del model.model.model[-1].linear
# output is now 1280