How to pass gradients to backward()

[JayMan91]

In my experiment, I have a loss function, which is not defined by an expression. But I have a formulation of the gradient (formally a subgradient) so I have to pass the gradient manually. In pytorch, I implement it in the following way and it is working fine.

self.optimizer.zero_grad()
y_hat = self.model(x_train)
grad =  compute_grad(y_hat, y)
y_hat.backward(gradient=grad)
self.optimizer.step()

Would the following be a correct implementation in lightning?

def training_step(self, batch, batch_idx):
        opt = self.optimizers()
        x,y = batch
        y_hat =  self(x)
        grad =  compute_grad(y_hat, y)
        opt.zero_grad()
        y_hat.backward(gradient= grad)
        opt.step()

[akihironitta]

@JayMan91 Haven't tried, but you could probably try manual optimization.

def __init__(...):
    ...
    self.automatic_optimization = False  # use manual optimization
    ...

def training_step(...):
    ...
    self.manual_backward(y_hat, gradient=grad)
    ...

manual optimization: https://pytorch-lightning.readthedocs.io/en/latest/common/optimizers.html#manual-optimization
LightningModule.manual_backward: https://pytorch-lightning.readthedocs.io/en/latest/common/lightning_module.html#manual-backward

[JayMan91]

Thanks for your reply. In this case, we have to write optimizer.step() and optimizer.zero_grad(), right?

[akihironitta]

Yes, you need to call them manually in your training step.