Hello,
I am pretty new to Yolo models and would really appreciate any kind of help.
So I am trying to create a double head yolov5 model, one for detection and other for classification. I have developed a class for the two headed model. but I am facing a lot of issues, kindly let me know how I could improve it and make a trainable model.
class TwoHead(nn.Module):
def __init__(self, model, ncc = 3):
super(TwoHead, self).__init__()
self.base_model = deepcopy(model)
self.base_model.model = self.base_model.model[:8]
self.detect_head = deepcopy(model)
self.detect_head.model = self.detect_head.model[8:]
self.classification_head = deepcopy(model)
self.classification_head.model = self.classification_head.model[7]
m = self.classification_head.model[-1] # last layer
ch = m.conv.in_channels if hasattr(m, 'conv') else sum([x.in_channels for x in m.m]) # ch into module
c = Classify(ch, ncc) # Classify()
c.i, c.f, c.type = m.i, m.f, 'models.common.Classify' # index, from, type
self.classification_head.model[-1] = c # replace
def forward(self, x):
x = self.base_model(x)
out1 = self.detect_head(x)
out2 = self.classification_head(x)
return out1, out2
twomodel = TwoHead(model) #Yolov5 model with pretrained weights
...
...
...
...
for imgs, targets, classification_targets in pbar:
pred, pred2 = twomodel(imgs) # forward
loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size
loss2 = criterion(pred2, classification_target)
Error that I face is
Traceback (most recent call last):
File "/home/pipu/yolov5/twohead.py", line 664, in <module>
main(opt)
File "/home/pipu/yolov5/twohead.py", line 558, in main
train(opt.hyp, opt, device, callbacks)
File "/home/pipu/yolov5/twohead.py", line 337, in train
pred, pred2 = twomodel(imgs) # forward
File "/home/piyush/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "/home/pipu/yolov5/twohead.py", line 86, in forward
out1 = self.detect_head(x)
File "/home/piyush/miniconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "/home/pipu/yolov5/models/yolo.py", line 211, in forward
return self._forward_once(x, profile, visualize) # single-scale inference, train
File "/home/pipu/yolov5/models/yolo.py", line 120, in _forward_once
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
File "/home/pipu/yolov5/models/yolo.py", line 120, in <listcomp>
x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f] # from earlier layers
IndexError: list index out of range