diff --git a/opacus/privacy_engine.py b/opacus/privacy_engine.py index 5045e96ca..c9fd8d606 100644 --- a/opacus/privacy_engine.py +++ b/opacus/privacy_engine.py @@ -14,7 +14,7 @@ # limitations under the License. import os import warnings -from typing import IO, Any, BinaryIO, Dict, List, Optional, Tuple, Union +from typing import Any, BinaryIO, Dict, IO, List, Optional, Tuple, Union import torch from opacus.accountants import create_accountant @@ -360,6 +360,21 @@ def make_private( if noise_generator and self.secure_mode: raise ValueError("Passing seed is prohibited in secure mode") + # compare module parameter with optimizer parameters + if not all( + torch.eq(i, j).all() + for i, j in zip( + list(module.parameters()), + sum( + [param_group["params"] for param_group in optimizer.param_groups], + [], + ), + ) + ): + raise ValueError( + "Module parameters are different than optimizer Parameters" + ) + distributed = isinstance(module, (DPDDP, DDP)) module = self._prepare_model( diff --git a/opacus/tests/privacy_engine_test.py b/opacus/tests/privacy_engine_test.py index 224acc875..a26f59f6d 100644 --- a/opacus/tests/privacy_engine_test.py +++ b/opacus/tests/privacy_engine_test.py @@ -29,8 +29,9 @@ from opacus import PrivacyEngine from opacus.layers.dp_multihead_attention import DPMultiheadAttention from opacus.optimizers.optimizer import _generate_noise -from opacus.scheduler import StepNoise, _NoiseScheduler +from opacus.scheduler import _NoiseScheduler, StepNoise from opacus.utils.module_utils import are_state_dict_equal +from opacus.validators import ModuleValidator from opacus.validators.errors import UnsupportedModuleError from torch.utils.data import DataLoader, Dataset, TensorDataset from torchvision import models, transforms @@ -464,6 +465,41 @@ def test_deterministic_run(self): "Model parameters after deterministic run must match", ) + def test_param_equal_module_optimizer(self): + """Test that the privacy engine raises error if nn.Module parameters are not equal to optimizer parameters""" + model = models.densenet121(pretrained=True) + num_ftrs = model.classifier.in_features + model.classifier = nn.Sequential(nn.Linear(num_ftrs, 10), nn.Sigmoid()) + optimizer = torch.optim.SGD( + model.parameters(), lr=0.01, momentum=0, weight_decay=0 + ) + dl = self._init_data() + model = ModuleValidator.fix(model) + privacy_engine = PrivacyEngine() + with self.assertRaisesRegex( + ValueError, "Module parameters are different than optimizer Parameters" + ): + _, _, _ = privacy_engine.make_private( + module=model, + optimizer=optimizer, + data_loader=dl, + noise_multiplier=1.1, + max_grad_norm=1.0, + ) + + # if optimizer is defined after ModuleValidator.fix() then raise no error + optimizer = torch.optim.SGD( + model.parameters(), lr=0.01, momentum=0, weight_decay=0 + ) + _, _, _ = privacy_engine.make_private( + module=model, + optimizer=optimizer, + data_loader=dl, + noise_multiplier=1.1, + max_grad_norm=1.0, + ) + self.assertTrue(1, 1) + @given(noise_scheduler=st.sampled_from([None, StepNoise])) @settings(deadline=None) def test_checkpoints(self, noise_scheduler: Optional[_NoiseScheduler]):