Lr-finder with multiple inputs, outputs and losses

Hello,

Firstly, thank you for this wonderful library. I have a model which expects 2 inputs. I am working with 2 kinds of images, one of size (512, 1536) and the other of size (128, 384). Therefore, my train_loader contains 2 inputs and one target of shape (128, 384, 16). My model has 4 prediction heads and hence is trained using 4 losses for different purposes.

So my collate_fn for the data loader looks like this:

def detection_collate(batch):
    """Custom collate fn for dealing with batches of images that have a different
    number of associated object annotations (bounding boxes).
    Arguments:
        batch: (tuple) A tuple of tensor images and lists of annotations
    Return:
        A tuple containing:
            1) (tensor) batch of images stacked on their 0 dim
            2) (list of tensors) annotations for a given image are stacked on
                                 0 dim
    """
    targets = []
    imgs = []
    deps = []
    for sample in batch:
        imgs.append(sample[0])
        deps.append(sample[1])
        targets.append(sample[2])
    return torch.stack(imgs, 0), torch.stack(deps, 0), torch.stack(targets, 0)

As mentioned, there are 4 different losses: Custom Heatmap (Focal) loss, SmoothL1, SmoothL1, BCE loss.

The forward method of the model expects 2 inputs. A small snippet is shown below:

 def forward(self, x, dep=None, target=None):
        # Backbone: ResNet18, x is image size: (512, 1536)

Here, targets are the labels so to say.

In this case, how do I go about finding the best learning rate using lr-finder? Notably, I can only use batch_size=2 because of the computational limitations.

[YashRunwal]

Hi,

OKay, I will try to create the wrapper. I have to make some changes to the model as well I think. Currently, the model returns losses, and now for this wrapper, it should return the prediction heads which can be passed to this wrapper to calculate losses. I will do this and get back to you tomorrow I think.

You are right, one loss (txty loss) dominates. How do I add some weights to the losses though?

[YashRunwal]

I have used gradient accumulation. I backpropagated the gradients after 64 steps (simulating 64 batch size). But let me check out how to use lr-finder with this. I will get back to you in case I need any help. Thanks for replying promptly. I really appreciate it.