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!
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.
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
Pl. check if the below code suits your need.
import torch, torch.nn as nn
import torchvision.models as models
from einops import reduce
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
m = Model().cuda()
x = torch.randn(5, 3, 224, 224).cuda()
h1, h2, h3 = m(x)
print(h1.shape, h2.shape, h3.shape)
I don’t think so, I am using the yolov5 network for the execution of this process described above.
thanks for the support!
You could try to change the
self.backbone assignment to appropriate YoloV5 model.
I haven’t used YoloV5 so far.
Hello @InnovArul .
I’m getting some problems with my model Yolov5s aswell. Can’t create a solution for this problem.
Thanks for the help!