Opacus and Lightining: multple optimizers and automatic optimization

[gianmarcoaversanoenx]

I am successfully running this piece of code:

from torchvision.datasets import MNIST
import torchvision.transforms as tfs
from torch.utils.data import DataLoader

dataloader = DataLoader(
    dataset=MNIST(
        '.data',
        transform = tfs.Compose(
            [
                tfs.ToTensor(),
                tfs.Normalize((0.1307,), (0.3081,)),
            ]
        ),
    ),
    batch_size=8,
)
model = LitSampleConvNetClassifier()

trainer = pl.Trainer(
    logger=False,
    enable_checkpointing=False,
    max_steps=1,
)
trainer.fit(model, train_dataloaders=dataloader)

But only when the definition of LitSampleConvNetClassifier is the following:

from loguru import logger
import os
import warnings

import pytorch_lightning as pl
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.optimizer import Optimizer
import torchmetrics
from opacus import PrivacyEngine, GradSampleModule
from opacus.accountants import RDPAccountant
from opacus.data_loader import DPDataLoader
from opacus.optimizers import DPOptimizer
from opacus.lightning import DPLightningDataModule
from pl_bolts.datamodules import MNISTDataModule
from pytorch_lightning.utilities.cli import LightningCLI


class LitSampleConvNetClassifier(pl.LightningModule):
    def __init__(
        self,
        lr: float = 0.1,
        enable_dp: bool = True,
        delta: float = 1e-5,
        noise_multiplier: float = 1.0,
        max_grad_norm: float = 1.0,
    ):
        """A simple conv-net for classifying MNIST with differential privacy training
        Args:
            lr: Learning rate
            enable_dp: Enables training with privacy guarantees using Opacus (if True), vanilla SGD otherwise
            delta: Target delta for which (eps, delta)-DP is computed
            noise_multiplier: Noise multiplier
            max_grad_norm: Clip per-sample gradients to this norm
        """
        super().__init__()
        # Hyper-parameters
        self.lr = lr
        # Parameters
        self.model = torch.nn.Sequential(
            torch.nn.Conv2d(1, 16, 8, 2, padding=3),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2, 1), 
            torch.nn.Conv2d(16, 32, 4, 2), 
            torch.nn.ReLU(), 
            torch.nn.MaxPool2d(2, 1), 
            torch.nn.Flatten(), 
            torch.nn.Linear(32 * 4 * 4, 32), 
            torch.nn.ReLU(), 
            torch.nn.Linear(32, 10)
        )
        # Metrics
        self.test_accuracy = torchmetrics.Accuracy()
        # Differential privacy
        self.accountant = RDPAccountant()
        self.enable_dp = enable_dp
        self.delta = delta
        self.noise_multiplier = noise_multiplier
        self.max_grad_norm = max_grad_norm
        # Important: This property activates manual optimization.
        self.automatic_optimization = False

    def forward(self, x):
        return self.model(x)

    def configure_optimizers(self):
        print("Configuring optimizers...")
        parameters = list(self.model.parameters())
        params1 = parameters[:3]
        params2 = parameters[3:]
        optimizers = [
            torch.optim.SGD(params1, lr=0.05),
            torch.optim.SGD(params2, lr=0.05),
        ]
        # privacy
        if not isinstance(self.model, GradSampleModule):
            self.model = GradSampleModule(self.model)
        data_loader = self.trainer._data_connector._train_dataloader_source.dataloader()
        sample_rate: float = 1 / len(data_loader)
        dataset_size: int = len(data_loader.dataset)  # type: ignore
        expected_batch_size = int(dataset_size * sample_rate)
        for i, optim in enumerate(optimizers):
            optim = DPOptimizer(
                optimizer=optim,
                noise_multiplier=self.noise_multiplier,
                max_grad_norm=1.0,
                expected_batch_size=expected_batch_size
            )
            optim.attach_step_hook(
                self.accountant.get_optimizer_hook_fn(sample_rate=sample_rate)
            )
            optimizers[i] = optim
        # return
        return optimizers

    def training_step(self, batch, batch_idx):
        optimizers = self.optimizers()
        for optimizer_idx, optimizer in enumerate(optimizers):
            logger.debug(f"Optimizer idx: {optimizer_idx}")
            assert isinstance(optimizer, Optimizer)
            optimizer.zero_grad()
            loss = self.loss(batch)
            self.manual_backward(loss)
            optimizer.step()
        self.log("train/loss", loss, on_step=False, on_epoch=True, prog_bar=True)
        # return loss

    def on_train_epoch_end(self):
        # Logging privacy spent: (epsilon, delta)
        epsilon, best_alpha = self.accountant.get_privacy_spent(delta=self.delta)
        self.log("epsilon", epsilon, on_epoch=True, prog_bar=True)
        print(f"\nepsilon = {epsilon}; best_alpha = {best_alpha}")

    def loss(self, batch):
        data, target = batch
        output = self.model(data)
        loss = F.cross_entropy(output, target)
        return loss

That is, I use manual optimization. If I change the class definition as follows:

# in the __init__() I also do self.automatic_optimization = True
    def training_step(self, batch, batch_idx, optimizer_idx):
        # optimizers = self.optimizers()
        # for optimizer_idx, optimizer in enumerate(optimizers):
            # assert isinstance(optimizer, Optimizer)
            # optimizer.zero_grad()
        logger.debug(f"Optimizer idx: {optimizer_idx}")
        loss = self.loss(batch)
            # self.manual_backward(loss)
            # optimizer.step()
        self.log("train/loss", loss, on_step=False, on_epoch=True, prog_bar=True)
        return loss

I get an error. Why? Independently from what error I get, which is Opacus-related, shouldn't the two pieces of code be equivalent? The documentation says so. If they are, then why does one succeed but the other one fails? How do they differ?

This should be better explained in the documentation.