Trying to find the best way to implement a gradient penalty
❓ Questions and Help
What is your question?
I'm trying to do some custom optimization with Lightning, and I wanted to ask here to see if there are any best practices for doing this sort of thing. It's a bit involved but I'll do my best to explain, bear with me...
I am writing some code with lightning to implement a variety of GAN optimization techniques with gradient penalties, one of which is Symplectic Gradient Adjustment (SGA) from this paper. It involves taking the gradient, and performing back-propagation an additional time to get a hessian-vector product. Put simply, the full loss ends up being the original loss plus the norm of the gradient vector. (Let L be the original loss function, L_new = L + ||∇L||)
Now, I have a class that looks like this (For the sake of simplicity, just pretend that gradient_norm() is a function which takes the gradient of the parameters w.r.t. the norm):
class GAN(LightningModule, ABC):
    ...

    @abstractmethod
    def generator_loss(self, d_z: Tensor) -> Tensor:
        pass

    @abstractmethod
    def discriminator_loss(self, d_x: Tensor, d_z: Tensor) -> Tensor:
        pass
    ...

    @property
    @abstractmethod
    def g_optimizer(self):
        pass

    @property
    @abstractmethod
    def d_optimizer(self):
        pass

    def configure_optimizers(self):
        return {'optimizer': self.g_optimizer(), 'frequency': int(ceil(1. / self.hparams['k']))}, \
               {'optimizer': self.d_optimizer(), 'frequency': int(ceil(self.hparams['k']))}
    ...

class MyGan(GAN):
    def generator_loss(self, d_z: Tensor) -> Tensor:
        return -1 * d_z.mean() + gradient_norm(...)

    def discriminator_loss(self, d_x: Tensor, d_z: Tensor) -> Tensor:
        return d_z.mean() - d_x.mean() + gradient_norm(...)
    ...
This does not work because, as you can see, I have two optimizers. When the generator optimizer is working, the discriminator parameters do not have gradients, and vice versa, which makes this not differentiable unless I write custom optimization code.
My goal is to be able to have some sort of functionality where I can get the update vector from the generator and discriminator optimizers, and then do additional processing in order to get the final update vector to update the weights.
What I'm trying at the moment is a new custom GanOptimizer class, and doing updates manually with the new manual_backward() functionality that was added recently:
class GAN(LightningModule, ABC):
    ...
    def configure_optimizers(self):
        return self.gan_optimizer()

    def optimizers(self) -> GanOptimizer:
        return super(GAN, self).optimizers()

    def training_step(self, batch: Tuple[Tensor, Tensor], *args) -> Dict:
        gan_opt = self.optimizers()

        X, _ = batch

        batch_size = X.shape[0]
        g_X = self.sample_G(batch_size)
        d_z = self.D(g_X)

        # Train generator
        g_loss = self.generator_loss(d_z)

        # Perform backward pass, NO step/zero_grad
        self.manual_backward(g_loss, gan_opt.generator_opt)

        # Train discriminator
        d_x = self.D(X)
        d_loss = self.discriminator_loss(d_x, d_z)

        self.manual_backward(d_loss, gan_opt.discriminator_opt)

        # Perform combined step
        gan_opt.step()
        self.on_before_zero_grad(gan_opt)
        gan_opt.zero_grad()

        return {'g_loss': g_loss, 'd_loss': d_loss}


class GanOptimizer(ABC, Optimizer):
    def __init__(self, g_params: Union[Iterable[Tensor], Iterable[Dict]],
                 d_params: Union[Iterable[Tensor], Iterable[Dict]], defaults: Dict):
        super(GanOptimizer, self).__init__(itertools.chain(g_params, d_params), defaults)
        self._g_params = list(g_params)
        self._d_params = list(d_params)

    @property
    def generator_opt(self) -> Optimizer:
        # "Dummy" optimizer
        return Optimizer(self.g_params, self.defaults)

    @property
    def discriminator_opt(self) -> Optimizer:
        # "Dummy" optimizer
        return Optimizer(self.d_params, self.defaults)

    @property
    def g_params(self):
        return self._g_params

    @property
    def d_params(self):
        return self._d_params
However, this still requires me to have separate optimizers for the generator and discriminator, and custom-implement all of the algorithms (Adam, RMSProp, etc.) and I'm feeling like this is a flawed approach. Like I said before, I'm just wondering if there's a better way to do this, or if it's possible that Lightning could have a feature added which would make this easier. Thanks in advance!
What's your environment?

OS: macOS 10.15
Packaging conda
Version 1.0.4