Weighting different losses based on random initialization of network

[Michael-Geuenich]

I'm working on a model with multiple losses that I add together: an autoencoder reconstruction loss (mse) and a prediction loss (cross entropy).

I now want to weight them based on the reconstruction loss / prediction loss ratio that is derived from a random initialization of the network (prior to training). I've thought about calculating these ratios using the on_train_start hook and then storing them in self. However, as far as I'm aware I do not have access to/cannot pass my dataloaders to on_train_start?

Alternatively, I've also thought about creating a custom method in my model class that calculates the weights, stores them in self and then I would call that method manually after instantiating my model but prior to calling fit. Something like:

class MyModelClass(pl.LightningModule):
    def __init__(self):
        ...
    
    def calculate_loss_weights(self):
        # calculate my ratio
        self.reconstruction_prediction_ratio = ratio

    def training_step(self, loaders):
        # all my training code
        return self.reconstruction_prediction_ratio * reconstruction loss + prediction_loss

data = MyDataClass()
model = MyModelClass()
model.calculate_loss_weights()
trainer = pl.Trainer()
trainer.fit(model, data)

However, I've realized that here too I don't know how to access my data loaders.

Is there any way to implement this, and are there any best practices in terms of implementation I should stick to?

Two things to note:

1. these weights would not change during training but remain fixed
2. I have about 4 different dataloaders at the moment, most of which need to be passed to calculate the loss ratios

[akihironitta]

Hi @Michael-Geuenich, I think #8114 will help.

[Michael-Geuenich]

Hi @akihironitta , Thanks so much for the suggestion! I did indeed get it to work based on the issue you linked.

Specifically

def on_train_start(self):
    loaders = self.trainer.train_dataloader.loaders
    loader1, loader2 = loaders['l1'], loaders['l2']
    X, y = loader1

Got me the loaders. However, the last line threw the following error ValueError: too many values to unpack (expected 2), though I managed to fix that as follows:

def on_train_start(self):
    loaders = self.trainer.train_dataloader.loaders
    loader1, loader2 = loaders['l1'], loaders['l2']
    X, y = next(iter(loader1))

Not sure if this is an issue, but it seems to work.

[akihironitta]

However, the last line threw the following error ValueError: too many values to unpack (expected 2)

@Michael-Geuenich I believe that's because loader1 is a single instance of DataLoader, so it's equivalent to something like:

X, y = DataLoader(...)

I think your updated code looks good to me.

---

If you have any further questions, feel free to ask me here (if they're relevant) or in another GitHub discussion (if they're irrelevant)!