PyTorch Lightning Support?

username: James_M
cooked: <p>I’m trying to utilise <code>opacus</code> with the PyTorch Lightning framework which we use as a wrapper around a lot of our models. I can see that there was an effort to integrate this partially into PyTorch Lightning <a href="https://github.com/PyTorchLightning/pytorch-lightning/issues/3305" rel="noopener nofollow ugc">late last year</a> but this seems to have stalled due to lack of bandwidth.</p>
<p>I’ve created <a href="https://gist.github.com/JMMarchant/68168bbc725aac0f616a87618709a7a5" rel="noopener nofollow ugc">a simple MVP</a> but there seems to be a compatibility problem with even this simple model; it throws <code>AttributeError: 'Parameter' object has no attribute 'grad_sample'</code> as soon as it hits the optimization step.</p>
<p>What’s the likely underlying cause of this? I can see on the <code>opacus</code> GitHub that similar errors have been encountered before where it’s been caused by unsupported layers but as the gist shows, this model is incredibly simple so I don’t think it’s any of the layers.</p>
<p>This is with:</p>
<pre><code class="lang-auto">opacus==0.11.0
pytorch-lightning==1.2.1
torch==1.7.1
torchaudio==0.7.2
torchvision==0.8.2
</code></pre>

username: Darktex
cooked: <p>Hi James! Would you have time to file a bug and share a colab so I can take a look? Integrating with Lightning is indeed on our plate <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: James_M
cooked: <p>Sure thing; didn’t want to go straight to a bug report as it didn’t feel right</p>

username: NiWaRe
cooked: <p>Hi <a class="mention" href="/u/darktex">@Darktex</a> <a class="mention" href="/u/james_m">@James_M</a>,<br>
was there some progress made with the interaction with PyTorch Lightning?<br>
I wrote a single integration myself also using the LightningCLI. I’m initializing the PrivacyEngine in the before_fit hook of a custom LightningCLI and attaching it to the configure_optimizer function of a typical LightningModule.<br>
It seems to work, but I would be curious if there’s a best practice way of integration.<br>
Regards</p>

username: sayanghosh
cooked: <p>Hello <a class="mention" href="/u/niware">@NiWaRe</a> <a class="mention" href="/u/james_m">@James_M</a> e have not worked yet on integrating PyTorch Lightning with Opacus. As mentioned, it is on our roadmap. Meanwhile, if you could share your changes either in Google Colab or send out a pull request in Github we would consider your changes when we are ready to start work on this.</p>

username: amin-nejad
cooked: <p>Thanks for describing your solution <a class="mention" href="/u/niware">@NiWaRe</a>, do you have a code snippet you could share?</p>

username: NiWaRe
cooked: <p><a class="mention" href="/u/sayanghosh">@sayanghosh</a> <a class="mention" href="/u/amin-nejad">@amin-nejad</a> sorry was busy last weeks. I can gladly share the snippets, concerning PR should I rather commit to a tutorial in a Jupyter Notebook or think about how to integrate it directly into the framework without the need to overwrite different hooks?</p>
<p>What I have for now (prototyping code, the hparams are the defined params in my PL Model):</p>
<pre><code class="lang-auto">class LightningCLI_Custom(LightningCLI):
[...]
def before_fit(self):
    """Hook to run some code before fit is started"""
    # possible because self.datamodule and self.model are instantiated beforehand
    # in LightningCLI.instantiate_trainer(self) -- see docs

    # TODO: why do I have to call them explictly here 
    #       -- in docs not mentioned (not found in .trainerfit())
    self.datamodule.prepare_data()
    self.datamodule.setup()

    if self.model.hparams.dp:
        if self.model.hparams.dp_tool == "opacus":
            # NOTE: for now the adding to the optimizers is in model.configure_optimizers()
            # because at this point model.configure_optimizers() wasn't called yet. 
            # That's also why we save n_accumulation_steps as a model parameter.
            sample_rate = self.datamodule.batch_size/len(self.datamodule.dataset_train)
            if self.model.hparams.virtual_batch_size &gt;= self.model.hparams.batch_size: 
                self.model.n_accumulation_steps = int(
                    self.model.hparams.virtual_batch_size/self.model.hparams.batch_size
                )
            else: 
                self.model.n_accumulation_steps = 1 # neutral
                print("Virtual batch size has to be bigger than real batch size!")

            # NOTE: For multiple GPU support: see PL code. 
            # For now we only consider shifting to cuda, if there's at least one GPU ('gpus' &gt; 0)
            self.model.privacy_engine = PrivacyEngine(
                self.model.model,
                sample_rate=sample_rate * self.model.n_accumulation_steps,
                target_delta = self.model.hparams.target_delta,
                target_epsilon=self.model.hparams.target_epsilon,
                epochs=self.trainer.max_epochs,
                max_grad_norm=self.model.hparams.L2_clip,
            ).to("cuda:0" if self.trainer.gpus else "cpu")
            # necessary if noise_multiplier is dynamically calculated by opacus
            # in order to ensure that the param is tracked
            self.model.hparams.noise_multiplier = self.model.privacy_engine.noise_multiplier
            print(f"Noise Multiplier: {self.model.privacy_engine.noise_multiplier}")

        else: 
            print("Use either 'opacus' or 'deepee' as DP tool.")

        # self.fit_kwargs is passed to self.trainer.fit() in LightningCLI.fit(self)
        self.fit_kwargs.update({
            'model': self.model
        })

        # in addition to the params saved through the model, save some others from trainer
        important_keys_trainer = ['gpus', 'max_epochs', 'deterministic']
        self.trainer.logger.experiment.config.update(
            {
                important_key:self.config['trainer'][important_key] 
                for important_key in important_keys_trainer
            }
        )
        # the rest is stored as part of the SaveConfigCallbackWandB
        # (too big to store every metric as part of the above config)
    
    # track gradients, etc.
    self.trainer.logger.experiment.watch(self.model)
</code></pre>
<p>Then in my model:</p>
<pre><code class="lang-auto">class LitModelDP(LightningModule):
    def __init__(...):
      [...]
      # disable auto. backward to be able to add noise and track 
      # the global grad norm (also in the non-dp case, lightning 
      # only does per param grad tracking)
      self.automatic_optimization = False

     # manual training step, eval, etc.

     def configure_optimizers(self):
        optims = {}
        # DeePee: we want params from wrapped mdoel
        # self.paramters() -&gt; self.model.wrapped_model.parameters()
        if self.hparams.optimizer=='sgd':
            optimizer = torch.optim.SGD(
                self.model.parameters(), 
                **self.hparams.opt_kwargs,
            )
        elif self.hparams.optimizer=='adam':
            optimizer = torch.optim.Adam(
                self.model.parameters(), 
                **self.hparams.opt_kwargs,
            )

        if self.hparams.dp_tool=='opacus' and self.hparams.dp:
            self.privacy_engine.attach(optimizer)

        optims.update({'optimizer': optimizer})
        [...]
        return optims
</code></pre>
<p>Calling the LightningCLI at the end</p>
<pre><code class="lang-auto">cli = LightningCLI_Custom(model_class=LitModelDP, datamodule_class=[...])
</code></pre>
<p>I’m very open to feedback and would also gladly help to integrate that into PL. <img src="https://discuss.pytorch.org/images/emoji/apple/+1.png?v=10" title=":+1:" class="emoji" alt=":+1:"></p>

username: karthikprasad
cooked: <p>Hello <a class="mention" href="/u/niware">@NiWaRe</a>,<br>
Apologies for the delay in response. Adding Lightning support to Opacus will certainly electrify it. Thank you so much for working on it.</p>
<p>What took us so long to respond you ask? Well, we have been working on a major refactor to Opacus (v1.0), and this would mean breaking API changes and code movement. Would you mind sending a PR on top of <a href="https://github.com/pytorch/opacus/tree/experimental_v1.0" class="inline-onebox" rel="noopener nofollow ugc">GitHub - pytorch/opacus at experimental_v1.0</a> branch?</p>

username: NiWaRe
cooked: <p>Hi <a class="mention" href="/u/karthikprasad">@karthikprasad</a>,<br>
great yes, I’ll take a look at the new version on the weekend!<br>
Best,<br>
Nicolas</p>

username: Peter_Romov
cooked: <p>Hi everyone!<br>
The problem reported by <a class="mention" href="/u/james_m">@James_M</a> is fixed in this <a href="https://github.com/pytorch/opacus/pull/244" rel="noopener nofollow ugc">pull request</a>. You can now add PrivacyEngine to LightningModel (see the demo in examples/mnist_lightning.py).</p>
<p>As we are working on the new API in experimental_v1.0 branch, I’m exploring ways to implement tighter (ideally seamless) integration between PyTorch Lightning and Opacus. In <a href="https://github.com/pytorch/opacus/pull/250" rel="noopener nofollow ugc">this RFC</a> I’ve shared some thoughts on how this integration can be implemented. Feedback and comments are much appreciated. Actually, I need help there.</p>
<p><a class="mention" href="/u/niware">@NiWaRe</a> Many thanks for your prototype!<br>
In LitModelDP we now can use automatic_optimization=True.<br>
In my demos I simply put all the DP parameters in LightningModule.<strong>init</strong> so that we can enable DP and configure it via --model.xxx cmdline args. Extending LightningCLI is indeed a good next step.</p>

username: NiWaRe
cooked: <p>Hi <a class="mention" href="/u/peter_romov">@Peter_Romov</a> and <a class="mention" href="/u/karthikprasad">@karthikprasad</a>,<br>
I’m sorry I didn’t respond, I was extremely busy with my thesis, which is due in three weeks.<br>
Thanks for implementing this! I’ll take a look latest in three weeks when I’m done with my project <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>