Custom backward that requires network input

[joecomerisnotavailable]

I am interested in using a loss-balancing class defined here:

https://github.com/facebookresearch/encodec/blob/main/encodec/balancer.py#L31

which has the expected usage

        weights = {'loss_a': 1, 'loss_b': 4}
        balancer = Balancer(weights, ...)
        losses: dict = {}
        losses['loss_a'] = compute_loss_a(x, y)
        losses['loss_b'] = compute_loss_b(x, y)
        if model.training():
            balancer.backward(losses, x)

or in replicating its functionality in a way that cooperates with Pytorch Lightning and preferably still allows for mixed-precision training. The class's function is to re-weight per-loss gradients so that predefined weights for each loss correspond to that loss's proportion of contribution to the norm of the total gradient step.

Since the class replaces the usual backward pass, and requires the model's input as an argument, I'm not sure whether to overwrite the LightningModule's backward, or manual_backward, or bypass both in the closure defined in training_step, or if possibly an alternative implementation of the class's balancing utility that can be called using after_backward hook is required.

My main concern with bypassing or overwriting manual_backward is creating a major slowdown or silently breaking the mixed precision handling.

Thanks in advance for any help with this.

[lukasschmit]

I also want to try this balancer class with Lightning! Did you figure this out?

[harimohanraj]

@joecomerisnotavailable @lukasschmit

I found this issue while also exploring adapting the loss balancer mechanism from Encodec. Initially I thought that you needed to pass the network input as well, but as it turns out, the paper itself says you need to pass the output of the network. See the following:

Then you can use the mechanism by overriding the default training step like usual. You must set self.automatic_optimization = False in your Lightning module initializer. Then, you can do something like the following:

def training_step(self, batch, idx):
        inputs, targets = batch 
        outputs = self.model(inputs)
        losses = {}
        losses['a'] = self.loss_a(outputs, targets)
        losses['b'] = self.loss_b(outputs, targets)

        # We need to setup our own backwards pass to make use of the loss balancer
        opt = self.optimizers()
        opt.zero_grad()
        self.balancer.backward(losses, outputs)
        # If you want to clip gradients, you need to call this function manually:        
        # self.clip_gradients(opt, gradient_clip_val=5.0, gradient_clip_algorithm="norm") 
        opt.step()

This seems to work for me at a first glance, but I'll report back if I encounter any other difficulties.

[DanTremonti]

Hi @harimohanraj , thanks for sharing your implementation. I understand that it has been quite sometime since you posted your thread, but I hope you can help me with a small clarification.

I noticed that the author mentioned

My main concern with bypassing or overwriting manual_backward is creating a major slowdown or silently breaking the mixed precision handling.

and your observations reporting

This seems to work for me at a first glance

Did your implementation have AMP configured? Or did you mean to say that you were able to train without AMP but without any major slowdown?