Bounding box regression without a classification task

Question

I am using PyTorch to create a model that detects certain objects in an image. I framed my problem as a regression on bounding boxes, without any classification task whatsoever. The reasoning behind it was that I do not need to know the class, I just need my model to localize the object on an image if it exists.

Each image can have up to 2 objects with their bounding boxes, so my target is 2x2 matrix, or rather a torch tensor with the shape [2,4]

I have 2 blocks of Conv2d and MaxPool2d, after which I have 3 blocks of Linear-ReLU-Dropout, and finally Unflatten to reshape the outputs of the model to fit my loss function calculation. Here is the snippet from the model class

Sequential(
            
  Conv2d(3, 32, kernel_size = (7,7)),
  MaxPool2d(kernel_size = (4,4)),
  Conv2d(32,16 , kernel_size = (3,3)),
  MaxPool2d(kernel_size=(2,2)),
  ReLU(),
  Flatten()    
)
Sequential(
    
   Linear(19600, 256),
   ReLU(),
   Dropout(0.2),
   Linear(256, 64),
   ReLU(),
   Dropout(0.2),
   Linear(64,8),
   Unflatten(1,(2,4))            
        
)

The labels look like this.

1. If there is no objects in the image - it's a zero tensor of shape [2,4]
2. If there is one object in the image - it's a tensor of shape [2,4] where first row is bounding box coordinates in form (xmin,ymin,xmax,ymax) and the second row is zeros] 3.If there are two objects in the image - it's a tensor of shape [2,4] where both rows contain bounding boxes coordinates in form (xmin, ymin, xmax, ymax)

I am using DIoU loss function with mean reduction. What I encounter is the following:

1. My validation loss is constant throughout epochs
2. My training loss descreases up to a certain point where it also becomes constant.
3. Even with such overfitted models, I stilled tried inference to see what kind of predictions I get on test set, and I get bounding box predictions that are very large positive or negative integers that are outside of my image area.

No matter the regularisation technique I try (and I tried Dropout as well as weight decay) or learning rate choice, which I tried setting to a larger value to perhaps avoid stucking in the local minimum, my losses just don't decrease and I can't get quality predictions on inference.

I'm getting the feeling that there is something pretty fundamental that I am missing but just can't seem to find what could it be. I'm sort of new to pytorch, so does anyone have an idea of what I might try to get it to start learning something? I'm happy to provide more code snippets in the comments.

Thank you!

Answer 1

Remove the trailing 0's for no objects and 1 object. Only encode where the object is, not that there isn't another object. This is throwing off the training. It's okay if an image has 0 bounding boxes. So if there are no objects in a given image, just don't put any bounding boxes corresponding to that image in the file.