$ show meta-pytorch/opacus
Opacus is a production-grade PyTorch library for training deep learning models with differential privacy (DP). It implements the DP-SGD algorithm with per-sample gradient computation, gradient clipping, and noise addition. The library provides a clean API via `PrivacyEngine.make_private()` that wraps any PyTorch model, optimizer, and data loader with differential privacy components. It supports multiple per-sample gradient backends (hooks, functorch, ExpandedWeights, ghost clipping), distributed training (DDP, FSDP, tensor parallelism), and several privacy accounting mechanisms (RDP, Gaussian DP, PRV). Opacus also includes a module validator that automatically detects and replaces incompatible layers (e.g., BatchNorm → GroupNorm, LSTM → DPLSTM), a batch memory manager for virtual batching, and comprehensive scheduling utilities. The codebase is well-tested, actively maintained with recent commits, and includes extensive documentation and tutorials. Although it carries some complexity from multiple backends and experimental features (research folder), the core abstractions are clean, modular, and reusable.
The LLM identified 5 files worth highlighting for reuse:
Implements `UniformWithReplacementSampler` and `DistributedUniformWithReplacementSampler`, providing stateless Poisson subsampling with configurable sample rate and steps. The class is self-contained (no internal imports beyond `torch.utils.data.Sampler` and `torch.distributed`) and can be extracted and reused in any framework needing privacy-preserving data sampling. Demonstrates clean separation of sampling logic from training loops.
Provides `BatchMemoryManager` and `BatchSplittingSampler` for decoupling logical batch size from physical batch size. It splits large batches into smaller micro-batches, coordinates with the optimizer via `signal_skip_step`, and handles gradient accumulation transparently. The implementation is self-contained (imports only `math`, `torch`, and internal `uniform_sampler.py`) and demonstrates a clean context manager pattern for memory-constrained training.
Implements `ModuleValidator` with a registration pattern using `register_module_validator` and `register_module_fixer` decorators. It automatically scans PyTorch modules for DP-incompatible layers (e.g., BatchNorm, LSTM, GRU) and applies fixes (e.g., replacement with DP-safe equivalents). The decorator-based registry allows users to extend the validator with custom checks and fixes, making it a highly reusable pattern for library validation frameworks.
Defines `IAccountant` abstract base class with methods `get_epsilon`, `get_noise_multiplier`, `step`, and `get_tensor`. Concrete implementations (`RDPAccountant`, `GaussianAccountant`, `PRVAccountant`) follow a strategy pattern. The abstraction cleanly separates privacy accounting from the rest of DP training, making it easy to swap or add new accounting mechanisms. The class is self-contained (imports only `abc`, `torch`) and exemplifies good interface design.
Provides `AbstractGradSampleModule` and `AbstractGradSampleHooks` base classes that manage per-sample gradient hooks (forward/backward), `grad_sample` attribute lifecycle, and model cleanup via `to_standard_module()`. The classes are well-designed with clear separation of hook registration, activation storage, and gradient computation. They serve as a reusable foundation for implementing different per-sample gradient backends (hooks, functorch, ExpandedWeights).