def prepare_for_training()

in flsim/clients/dp_client.py [0:0]

• 25 lines of code
• 3 McCabe index (conditional complexity)

    def prepare_for_training(self, model: IFLModel):
        """
        1- call parent's prepare_for_training
        2- attach the privacy_engine
        """
        model, optimizer, optimizer_scheduler = super().prepare_for_training(model)
        if self.privacy_on:
            batch_size, self.dataset_length = self._get_dataset_stats(model)
            sample_rate = batch_size / self.dataset_length

            self.grad_sample_module = GradSampleModule(model.fl_get_module())

            # pyre-fixme[16]: `DPClient` has no attribute `cfg`.
            if self.cfg.privacy_setting.noise_seed is not None:
                generator = torch.Generator()
                # pyre-fixme[16]
                generator.manual_seed(self.cfg.privacy_setting.noise_seed)
            else:
                generator = None

            optimizer = DPOptimizer(
                optimizer=optimizer,
                noise_multiplier=self.cfg.privacy_setting.noise_multiplier,
                max_grad_norm=self.cfg.privacy_setting.clipping_value,
                expected_batch_size=batch_size,
                generator=generator,
            )

            def accountant_hook(optim: DPOptimizer):
                self.accountant.step(
                    noise_multiplier=optim.noise_multiplier,
                    sample_rate=sample_rate * optim.accumulated_iterations,
                )

            optimizer.attach_step_hook(accountant_hook)

        return model, optimizer, optimizer_scheduler