1
$\begingroup$

I'm reading through the documentation of YOLOv8 here, but I fail to see an easy way to do what I suggest in the title. What I want to do is to load a pretrained YOLOv8 model, create a bigger model that will contain YOLOv8 as a submodule, and modify the forward function of YOLOv8 so that I may have access to the object detection loss plus the convolutional features, so that they can be used to feed subsequent layers for other custom tasks.

To make things clearer, this is how yolov8 is intended to be used originally according to 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')

Instead, I want to do something like this:

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,
        }

The code above is a very simplified sketch and of course is not going to work, but more or less that's the idea. Moreover, by creating this ad-hoc wrapper I won't be able to use the out-of-the-box functionality to train, validate and predict that comes with the YOLO library, since it would be a custom architecture (YOLOv8 being just a submodule of it). Thus, I also need to figure out how to write a custom dataloader in order to provide YOLOv8 with the input it expects plus the additional stuff required by my wrapper (there could be different additional layers in the wrapper predicting different outputs based on YOLOv8's features, think of it as multitask learning).

Is this possible? How can this be done?

$\endgroup$

1 Answer 1

0
$\begingroup$

just some idea recommendation for wrapper:

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,
        }
```
$\endgroup$
2
  • $\begingroup$ Your answer could be improved with additional supporting information. Please edit to add further details, such as citations or documentation, so that others can confirm that your answer is correct. You can find more information on how to write good answers in the help center. $\endgroup$
    – Community Bot
    Commented Apr 2, 2023 at 16:08
  • $\begingroup$ You note whether or not you tried the code yourself. $\endgroup$
    – Full Array
    Commented Jan 7 at 2:19

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.