ptrblck / pytorch_misc

Code snippets created for the PyTorch discussion board
546 stars 64 forks source link

UNET + YOLOv6 #11

Open geekdreamer04 opened 2 years ago

geekdreamer04 commented 2 years ago

I am trying to add Unet as a preprocessing layer before the YOLOv6 architecture and when I am trying to do this, I am facing the following error. And how do I combine a Unet architecture with YOLO architecture. Any help from you here will be appreciated. Looking forward to your support as soon as possible. Thank you

ERROR in training steps. ERROR in training loop or eval/save model.

Training completed in 0.000 hours. Traceback (most recent call last): File "tools/train.py", line 112, in main(args) File "tools/train.py", line 102, in main trainer.train() File "/workspace/YOLOv61/yolov6/core/engine.py", line 75, in train self.train_in_loop() File "/workspace/YOLOv61/yolov6/core/engine.py", line 88, in train_in_loop self.train_in_steps() File "/workspace/YOLOv61/yolov6/core/engine.py", line 104, in train_in_steps preds = self.model(images) File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl return forward_call(*input, kwargs) File "/workspace/YOLOv61/yolov6/models/yolo.py", line 39, in forward x = self.detect(x) File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl return forward_call(*input, *kwargs) File "/workspace/YOLOv61/yolov6/models/effidehead.py", line 60, in forward x[i] = self.stemsi File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl return forward_call(input, kwargs) File "/workspace/YOLOv61/yolov6/layers/common.py", line 102, in forward return self.act(self.bn(self.conv(x))) File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl return forward_call(*input, kwargs) File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 446, in forward return self._conv_forward(input, self.weight, self.bias) File "/home/ubuntu/anaconda3/envs/pytorch_p38/lib/python3.8/site-packages/torch/nn/modules/conv.py", line 442, in _conv_forward return F.conv2d(input, weight, bias, self.stride, RuntimeError: Given groups=1, weight of size [256, 256, 1, 1], expected input[8, 128, 20, 20] to have 256 channels, but got 128 channels instead**

ptrblck commented 2 years ago

The error is raised in your YOLO model in:

x = self.detect(x)
...
return self.act(self.bn(self.conv(x)))

which gets an input activation with 128 channels while 256 are expected. Could you check the shape of the UNet output and make sure it's compatible with your YOLO model?