Help for adversarial learning with pytorch lighting
Help for adversarial learning with pytorch lighting
What is your question?
Code
the old method for adversarial learning is like this:
fgm = FGM(model)
for batch_input, batch_label in data:
    # normal training
    loss = model(batch_input, batch_label)
    loss.backward() 
    # adversarial training
    fgm.attack() 
    loss_adv = model(batch_input, batch_label)
    loss_adv.backward() 
    fgm.restore()
    optimizer.step()
    model.zero_grad()
but I don't know how to deal it by pytorch lighting
What have you tried?
It seems to work, but I‘m not sure if it is correct or there is a better way to achieve it.
    def backward(self, loss: Tensor, optimizer: Optimizer, optimizer_idx: int, *args, **kwargs) -> None:
        if not self.hparams.do_adv:
            super().backward(loss, optimizer, optimizer_idx, *args, **kwargs)

    def training_step(self, batch, batch_idx):
        outputs = self(**batch)
        loss = outputs[0]
        if self.hparams.do_adv:
            loss.backward()
            self.fgm.attack()
            outputs = self(**batch)
            adv_loss = outputs[0]
            self.fgm.restore()
            adv_loss.backward()
        return loss
What's your environment?

OS: Linux
Packaging conda
Version latest