repository·main·Indexed 24 days ago
https://github.com/meta-pytorch/opacusA library for training PyTorch models with differential privacy (DP-SGD). It provides the PrivacyEngine to wrap models, optimizers, and data loaders with minimal code changes. Key features include multiple gradient sampling modes (hooks, ew, functorch), non-wrapping mode for model compatibility, and adaptive clipping via PrivacyEngineAdaptiveClipping.
Opacus is a library for training PyTorch models with differential privacy. It is designed to require minimal code changes, maintain high training performance, and allow users to track their privacy budget online during training.
It serves two primary audiences:
The research/ directory contains novel methods and extensions built on top of Opacus. These are community contributions aimed at enhancing DP-SGD (Differentially Private Stochastic Gradient Descent) based on recent research.
Note for users: The code in this directory is provided as-is and is not maintained by the core Opacus team. This means it may experience compatibility issues with future Opacus updates. If you encounter issues, you should contact the specific PR contributor for that method.
Opacus computes and stores per-sample gradients. To prevent memory usage from scaling quadratically with the batch size, Opacus uses virtual batches.
Virtual batches allow you to separate physical steps (the actual computation of gradients) from logical steps (the addition of noise and parameter updates). This enables you to use larger batches for training while maintaining a low memory footprint. Integration is handled via the Batch Memory Manager.
Increasing the batch size increases the sampling rate, which in turn increases the privacy budget (meaning less privacy).
You can counteract this effect by:
Opacus tracks privacy expenditure internally using Rényi Differential Privacy (RDP). The alphas parameter specifies the RDP orders that the privacy engine should use to track this expenditure.
When you call privacy_engine.get_epsilon(delta=delta), the engine searches through the provided alphas to find the optimal order that provides the tightest (epsilon, delta)-DP bound for your given delta.
Best Practices:
alpha returned by get_epsilon is one of the boundary values of your alphas list, you should consider expanding the list.[1 + x / 10.0 for x in range(1, 100)] + list(range(12, 64)).When defining input_shape in config.json, note that parameters shared between the model and the input are listed separately. The actual input shape passed to the layer is constructed as follows:
(batch_size, *input_shape, in_features)(batch_size, in_channels, *input_shape)(batch_size, *input_shape)(input_shape[-D:])(batch_size, num_features, *input_shape)(batch_size, num_channels, *input_shape)(batch_size, *input_shape)not batch_first: (targ_seq_len, batch_size, embed_dim)(batch_size, targ_seq_len, embed_dim)not batch_first: (seq_len, batch_size, input_size)(batch_size, seq_len, input_size)Opacus offers two integration modes for PyTorch models depending on your compatibility needs:
Opacus wraps your model in a GradSampleModule to compute per-sample gradients.
isinstance() type checking to fail (e.g., with HuggingFace Transformers) and adds a _module. prefix to keys in your state_dict.Set wrap_model=False to attach hooks directly to your model instead of wrapping it.
state_dict keys, and offers better compatibility with transformer models.hooks.cleanup() once training is complete.For detailed instructions, refer to the non-wrapping mode tutorial.
The (epsilon, delta) pair quantifies the privacy guarantees of your model:
Calculating Epsilon: Epsilon is computed ex post (after the optimizer run). You can retrieve the current epsilon for a specific delta using:
epsilon = privacy_engine.get_epsilon(delta=delta)Note that (epsilon, delta) provides a conservative upper bound on actual privacy loss; the real loss may be significantly smaller.
Opacus requires that the nn.Module used in your training loop consists of supported components. A module is considered supported if it meets one of the following criteria:
nn.ReLU or nn.Tanh.requires_grad = False.nn.Conv2d (check the grad_sample/ directory for the list of explicitly supported modules).nn.Module that is composed entirely of the supported modules listed above.Note that compatibility depends on the underlying implementation. For example, a module might be logically equivalent to a supported one but use different internal operators that are not yet supported.
class SampleConvNet(nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(1, 16, 8, 2, padding=3)
self.conv2 = nn.Conv2d(16, 32, 4, 2)
self.fc1 = nn.Linear(32 * 4 * 4, 32)
self.fc2 = nn.Linear(32, 10)
def forward(self, x):
# x of shape [B, 1, 28, 28]
x = F.relu(self.conv1(x)) # -> [B, 16, 14, 14]
x = F.max_pool2d(x, 2, 1) # -> [B, 16, 13, 13]
x = F.relu(self.conv2(x)) # -> [B, 32, 5, 5]
x = F.max_pool2d(x, 2, 1) # -> [B, 32, 4, 4]
x = x.view(-1, 32 * 4 * 4) # -> [B, 512]
x = F.relu(self.fc1(x)) # -> [B, 32]
x = self.fc2(x) # -> [B, 10]
return xOpacus provides three different approaches for computing per-sample gradients. Choosing the right one depends on your stability requirements, model architecture, and performance needs:
grad_sample_mode="hooks"): The most stable implementation. It uses backward hooks and requires custom grad sampler methods for trainable layers. Use this if you want a reliable, production-ready experience.grad_sample_mode="ew"): An early beta approach based on PyTorch 1.12+ functionality. It works at the function level rather than the module level, potentially offering better performance (~25% faster) and support for layers that use known operations. Note: This mode requires model wrapping and does not support non-wrapping mode.grad_sample_mode="functorch"): A beta approach using vmap and grad transforms. It is highly flexible and can support models that the hooks-based approach cannot. It also supports non-wrapping mode.Summary Table:
| Feature | Hooks | Expanded Weights | Functorch |
|---|---|---|---|
| Status | Stable | Beta | Beta |
| Non-wrapping mode | Supported | Not supported | Supported |
| Performance | Baseline | ~25% faster | 0-50% slower |
| TorchScript | Not supported | Supported | Not supported |
You can install the latest release of Opacus using pip or conda. For the latest features, you can also install directly from the source.
# Via pip
pip install opacus
# Via conda
conda install -c conda-forge opacus
# From source
git clone https://github.com/pytorch/opacus.git
cd opacus
pip install -e .The Opacus 1.x API introduces a different approach to data handling and privacy configuration. The main changes involve:
model, optimizer, and data_loader. The make_private method returns wrapped versions of these objects.PrivacyEngine constructor no longer accepts training artifacts (like sample_rate or alphas). These are now passed directly to the make_private method.make_private, you must use keyword arguments (e.g., module=model).optimizer.zero_grad() in your training loop, as Opacus no longer clears gradients automatically.convert_batchnorm_modules), use ModuleValidator.fix(model) to automatically apply necessary remediations like BatchNorm -> GroupNorm or LSTM -> DPLSTM.# define your components as usual
model = Net()
optimizer = SGD(model.parameters(), lr=0.05)
data_loader = torch.utils.data.DataLoader(dataset, batch_size=1024)
# enter PrivacyEngine
privacy_engine = PrivacyEngine()
model, optimizer, data_loader = privacy_engine.make_private(
module=model,
optimizer=optimizer,
data_loader=data_loader,
noise_multiplier=1.1,
max_grad_norm=1.0,
)
# Now it's business as usual