Returns a callback function which can be used to attach to DPOptimizer Args: sample_rate: Expected sampling rate used for accounting
(
self, sample_rate: float
)
| 67 | pass |
| 68 | |
| 69 | def get_optimizer_hook_fn( |
| 70 | self, sample_rate: float |
| 71 | ) -> Callable[[DPOptimizer], None]: |
| 72 | """ |
| 73 | Returns a callback function which can be used to attach to DPOptimizer |
| 74 | Args: |
| 75 | sample_rate: Expected sampling rate used for accounting |
| 76 | """ |
| 77 | |
| 78 | def hook_fn(optim: DPOptimizer): |
| 79 | # This works for Poisson for both single-node and distributed |
| 80 | # The reason is that the sample rate is the same in both cases (but in |
| 81 | # distributed mode, each node samples among a subset of the data) |
| 82 | self.step( |
| 83 | noise_multiplier=optim.noise_multiplier, |
| 84 | sample_rate=sample_rate * optim.accumulated_iterations, |
| 85 | ) |
| 86 | |
| 87 | return hook_fn |
| 88 | |
| 89 | def state_dict(self, destination: T_state_dict = None) -> T_state_dict: |
| 90 | """ |