Support DP in PyText (#1366)

Summary:
Pull Request resolved: #1366

Pull Request resolved: #1355

as titled

Reviewed By: snisarg, ashkan-software

Differential Revision: D20844321

fbshipit-source-id: 0825df81462a76b192e06d1e13bcfc8cf64155b8

Author: karthikprasad

docs_requirements.txt

@@ -13,3 +13,4 @@ sentencepiece
 torchtext
 tensorboard==1.14
 pandas
+pytorch-dp

pytext/config/component.py

@@ -32,6 +32,7 @@ class ComponentType(enum.Enum):
     METRIC_REPORTER = "metric_reporter"
     SPARSIFIER = "sparsifier"
     MASKING_FUNCTION = "masking_function"
+    PRIVACY_ENGINE = "privacy_engine"
 
 
 class RegistryError(Exception):
@@ -247,6 +248,12 @@ def create_sparsifier(sparsifier_config, *args, **kwargs):
     )
 
 
+def create_privacy_engine(privacy_engine_config, *args, **kwargs):
+    return create_component(
+        ComponentType.PRIVACY_ENGINE, privacy_engine_config, *args, **kwargs
+    )
+
+
 def create_predictor(predictor_config, *args, **kwargs):
     return create_component(ComponentType.PREDICTOR, predictor_config, *args, **kwargs)
 

pytext/optimizer/__init__.py

@@ -15,5 +15,6 @@
     Optimizer,
     learning_rates,
 )
+from pytext.optimizer.privacy_engine import PrivacyEngine  # noqa
 from pytext.optimizer.radam import RAdam  # noqa
 from pytext.optimizer.swa import StochasticWeightAveraging  # noqa

pytext/optimizer/privacy_engine.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+from typing import List, Optional
+
+import torchdp
+from pytext.config import ConfigBase
+from pytext.config.component import Component, ComponentType
+
+
+class PrivacyEngine(Component):
+    """
+    A wrapper around PrivacyEngine of pytorch-dp
+    """
+
+    __COMPONENT_TYPE__ = ComponentType.PRIVACY_ENGINE
+    __EXPANSIBLE__ = False
+
+    class Config(ConfigBase):
+        noise_multiplier: float
+        max_grad_norm: float
+        batch_size: float
+        dataset_size: float
+        target_delta: Optional[float] = 0.000001
+        alphas: Optional[List[float]] = [1 + x / 10.0 for x in range(1, 100)] + list(
+            range(12, 64)
+        )
+
+    def __init__(
+        self,
+        model,
+        optimizer,
+        noise_multiplier,
+        max_grad_norm,
+        batch_size,
+        dataset_size,
+        target_delta,
+        alphas,
+    ):
+        self.noise_multiplier = noise_multiplier
+        self.max_grad_norm = max_grad_norm
+        self.batch_size = batch_size
+        self.dataset_size = dataset_size
+        self.target_delta = target_delta
+        self.alphas = alphas
+
+        self._privacy_engine = torchdp.PrivacyEngine(
+            model,
+            self.batch_size,
+            self.dataset_size,
+            self.alphas,
+            noise_multiplier=self.noise_multiplier,
+            max_grad_norm=self.max_grad_norm,
+            target_delta=self.target_delta,
+        )
+        self._privacy_engine.attach(optimizer)
+
+    @classmethod
+    def from_config(cls, config: Config, model, optimizer):
+        return cls(
+            model=model,
+            optimizer=optimizer,
+            noise_multiplier=config.noise_multiplier,
+            max_grad_norm=config.max_grad_norm,
+            batch_size=config.batch_size,
+            dataset_size=config.dataset_size,
+            target_delta=config.target_delta,
+            alphas=config.alphas,
+        )
+
+    def attach(self, optimizer):
+        self._privacy_engine.attach(optimizer)
+
+    def detach(self):
+        self._privacy_engine.detach()
+
+    def get_privacy_spent(self):
+        return self._privacy_engine.get_privacy_spent()

pytext/trainers/trainer.py

@@ -13,6 +13,7 @@
     Component,
     ComponentType,
     create_optimizer,
+    create_privacy_engine,
     create_scheduler,
     create_sparsifier,
 )
@@ -21,7 +22,7 @@
 from pytext.metric_reporters import MetricReporter
 from pytext.models.distributed_model import DistributedModel
 from pytext.models.model import Model
-from pytext.optimizer import Adam, Optimizer, learning_rates
+from pytext.optimizer import Adam, Optimizer, PrivacyEngine, learning_rates
 from pytext.optimizer.fp16_optimizer import FP16Optimizer, FP16OptimizerFairseq
 from pytext.optimizer.scheduler import Scheduler
 from pytext.optimizer.sparsifiers.sparsifier import Sparsifier
@@ -119,6 +120,8 @@ class Config(ConfigBase):
         #: backward and master weight will be maintained on original optimizer.
         #: https://arxiv.org/abs/1710.03740
         fp16_args: FP16Optimizer.Config = FP16OptimizerFairseq.Config()
+        # PrivacyEngine related args
+        privacy_engine: Optional[PrivacyEngine.Config] = None
 
     def __init__(self, config: Config, model: torch.nn.Module):
         if config.early_stop_after > 0:
@@ -135,6 +138,11 @@ def __init__(self, config: Config, model: torch.nn.Module):
             self.optimizer: torch.optim.Optimizer = create_optimizer(
                 config.optimizer, model
             )
+        self.privacy_engine: PrivacyEngine = (
+            create_privacy_engine(config.privacy_engine, model, self.optimizer)
+            if config.privacy_engine
+            else None
+        )
 
         self.scheduler: torch.optim.lr_scheduler = (
             create_scheduler(config.scheduler, self.optimizer)
@@ -370,6 +378,7 @@ def train(
             optimizer=self.optimizer,
             scheduler=self.scheduler,
             sparsifier=self.sparsifier,
+            privacy_engine=self.privacy_engine,
             rank=rank,
         )
         return self.train_from_state(

pytext/trainers/training_state.py

@@ -7,7 +7,7 @@
 from pytext.common.constants import Stage
 from pytext.data.tensorizers import Tensorizer
 from pytext.models.model import Model
-from pytext.optimizer import Optimizer
+from pytext.optimizer import Optimizer, PrivacyEngine
 from pytext.optimizer.scheduler import Scheduler
 from pytext.optimizer.sparsifiers.sparsifier import Sparsifier
 
@@ -18,6 +18,7 @@ class TrainingState:
     scheduler: Scheduler
     sparsifier: Sparsifier
     start_time: float
+    privacy_engine: PrivacyEngine
     # epoch counter
     epoch: int = 0
     # step counter: each optimizer.step() increments step_counter

requirements.txt

@@ -10,6 +10,7 @@ onnx>=1.6.0
 python-dateutil==2.8.0
 pandas
 pytorch-pretrained-bert
+pytorch-dp
 regex==2019.11.1
 requests
 scipy