Возможно ли модифицировать YOLOv8, чтобы использовать его в качестве извлечителя признаков для других задач?

Question 1

Я изучаю документацию YOLOv8 здесь, но не вижу простого способа сделать то, что я предложил в заголовке. Я хочу загрузить предобученную модель YOLOv8, создать более крупную модель, которая будет содержать YOLOv8 в качестве субмодуля, и изменить функцию прямого распространения YOLOv8, чтобы иметь доступ к функции потерь обнаружения объектов и свёрточным признакам, чтобы их можно было использовать для последующих слоев для других пользовательских задач.

Чтобы сделать это понятнее, вот как yolov8 изначально предназначен для использования согласно https://docs.ultralytics.com/quickstart/#use-with-python:

from ultralytics import YOLO

# Создайте новую модель YOLO с нуля
model = YOLO('yolov8n.yaml')

# Загрузите предобученную модель YOLO (рекомендуется для обучения)
model = YOLO('yolov8n.pt')

# Обучите модель, используя набор данных 'coco128.yaml' в течение 3 эпох
results = model.train(data="coco128.yaml", epochs=3)

# Оцените производительность модели на валидационном наборе
results = model.val()

# Выполните обнаружение объектов на изображении, используя модель
results = model('https://ultralytics.com/images/bus.jpg')

# Экспортируйте модель в формате ONNX
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; думайте об этом как о многозадачном обучении).

Это возможно? Как это может быть сделано?

Question 2

просто несколько идей для обёртки:

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):
        class CustomDataset(torch.utils.data.Dataset):
            def __init__(self, images, gt_boxes):
                self.images = images
                self.gt_boxes = gt_boxes

            def __len__(self):
                return len(self.images)

            def __getitem__(self, idx):
                image = self.images[idx]
                gt_box = self.gt_boxes[idx]
                return (image, gt_box)

        custom_dataloader = torch.utils.data.DataLoader(CustomDataset(images, gt_boxes), batch_size=1)

        with torch.no_grad():
            for batch_idx, (image, gt_box) in enumerate(custom_dataloader):
                features, loss = self.yolov8(image, gt_box)

        logits1 = self.fc1(features)
        logits2 = self.fc2(features)
        logits3 = self.fc3(features)

        return {
            'logits1': logits1,
            'logits2': logits2,
            'logits3': logits3,
            'yolov8_loss': loss,
        }
```

Возможно ли модифицировать YOLOv8, чтобы использовать его в качестве извлечителя признаков для других задач?

Вопрос или проблема

Ответ или решение

Теория (Theory)

Пример (Example)

Применение (Application)