Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Opacus: I made a model-agnostic callback for PyTorch Lightning #524

Open
gianmarcoaversanoenx opened this issue Oct 18, 2022 · 0 comments
Open
Labels
enhancement New feature or request

Comments

@gianmarcoaversanoenx
Copy link

🚀 Feature

I could get Opacus to work with PyTorch Lightning (pl) using a pl.Callback. Note that the callback is model-agnostic, and the model's pl.LightningModule class does not have anything related to Opacus.

Motivation

We need an easy way for PyTorch Lightning users to use Opacus without them having to refactor their LightningModule classes. See below.

Pitch

We need something as follows:
(I could actually implement this for real but it only works only for models with one optimizer)

import pytorch_lightning as pl
from opacus import OpacusCallback
from pl_bolts.models import LitMNIST
from pl_bolts.datamodules import MNISTDataModule

trainer = pl.Trainer(
        callbacks=[
            OpacusCallback(...), # all that is needed for DP-training
        ],
    )
trainer.fit(model=LitMNIST(...), datamodule=MNISTDataModule(...))

Additional context

In my version for OpacusCallback, all I do is call .make_private in the on_train_epoch_start hook:

# --- pseudo code --- #

def on_train_epoch_start(
        self,
        trainer: pl.Trainer,
        pl_module: pl.LightningModule,
    ) -> None:
    optimizers: ty.List[Optimizer] = []
        # for loop: begin
        for i in range(len(trainer.optimizers)):
            optimizer = trainer.optimizers[i]
           # this works
            _, dp_optimizer, _ = self.privacy_engine.make_private( # or make_private_with_epsilon
                    module=pl_module,
                    optimizer=optimizer,
                    data_loader=trainer._data_connector._train_dataloader_source.dataloader(),
                    noise_multiplier=self.noise_multiplier,
                    max_grad_norm=self.max_grad_norm,
                    clipping=self.clipping,  # "flat" or "per_layer" or "adaptive"
                    poisson_sampling=self.poisson_sampling,
                )
            optimizers.append(dp_optimizer)
        ### this will fail
        #  if not hasattr(pl_module, "autograd_grad_sample_hooks"):
        #         pl_module = GradSampleModule(pl_module)
        # dp_optimizer = privacy_engine._prepare_optimizer(
        #         optimizer,
        #         noise_multiplier=self.noise_multiplier,
        #         max_grad_norm=self.max_grad_norm,
        #         expected_batch_size=expected_batch_size,
        #    )
        # for loop: end
    trainer.optimizers = optimizers

What's cool is that this is an EarlyStopping callback, so it will stop training when enough privacy budget has been spent.

@HuanyuZhang HuanyuZhang added enhancement New feature or request labels Sep 12, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

No branches or pull requests

2 participants