Update training logic to computes loss over logits

[anwai98]

WIP!

[anwai98]

Hi @constantinpape,

This should be good for a first review from my side!

The intended idea here is simple: 1) compute loss over logit masks (to make this possible, I convert the labels to match the logits dimensions) and 2) work with iterative prompts for logits.

Let me know how it looks!

EDIT: I am gonna run a training for this and see how the results look for LIVECell!

[constantinpape]

Was this never used / not used anymore?

[anwai98]

It was used previously, I removed it now because the samples are logits now. Would you recommend to keep those predictions by upsampling them?

[constantinpape]

What exactly were the samples? Examples for mask predictions? It would maybe be good to keep them, we can discuss later.

[anwai98]

Yes, samples were mask predictions!

[anwai98]

Here's a detailed description on the PR (we see when we can come back to this in future)

Core Idea: Computing loss over predicted logits masks and downsampled ground-truth masks.

The current implementation works, however it does not bring us significant memory advantages (below mentioned are quick try-outs on LIVECell):

# # for default settings
# loss over logits: 47.89 GB
# loss over masks: 49.42 GB

# # for resource-efficient setting (i.e. n_objects=5 for 'vit_b')
# loss over logits: 24.01 GB
# loss over masks: 24.38 GB