Use one backbone with different heads

1

Is it possible to save the results from the backbone and apply them on the heads of the all the other models. My goal was to try to save time by avoiding repeating the backbone part. For example, looking at the diagram, instead of running the 3 complete models (left), only run the backbone 1 time and switch only the heads for the 3 models (right), therefore not repeating executing the backbone every time.

Thank you for the help!

2

Probably you forgot to add the diagram that you are referring to?

In general, running backbone one time and switching between 3 heads is technically possible in pytorch.

1 Like
3

Capturar

The diagram is in this genre. I don’t know how to do it, can you give me some kind of solution?

Thanks for your reply

4

Pl. check if the below code suits your need.

import torch, torch.nn as nn
import torchvision.models as models
from einops import reduce

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = nn.Sequential(*list(models.resnet18(pretrained=True).children())[:-2])
        self.head1 = nn.Linear(512,10)
        self.head2 = nn.Linear(512,10)
        self.head3 = nn.Linear(512,10)

    def forward(self, x):
        # get backbone features
        backbone_feats = self.backbone(x)

        # global average pooling
        backbone_feats = reduce(backbone_feats, 'b c h w -> b c', reduction='mean')

        # pass through heads
        h1 = self.head1(backbone_feats)
        h2 = self.head2(backbone_feats)
        h3 = self.head3(backbone_feats)
        return h1, h2, h3

if __name__=='__main__':
    m = Model().cuda()
    x = torch.randn(5, 3, 224, 224).cuda()
    h1, h2, h3 = m(x)
    print(h1.shape, h2.shape, h3.shape)
1 Like
5

I don’t think so, I am using the yolov5 network for the execution of this process described above.

thanks for the support!

6

You could try to change the self.backbone assignment to appropriate YoloV5 model.
I haven’t used YoloV5 so far.

7

Hello @InnovArul .

I’m getting some problems with my model Yolov5s aswell. Can’t create a solution for this problem.

Thanks for the help!