Allow compute_dp_sgd_privacy to be called as library function.

PiperOrigin-RevId: 275966906

tensorflow_privacy/privacy/analysis/compute_dp_sgd_privacy.py

@@ -52,11 +52,6 @@ flags.DEFINE_float('noise_multiplier', None, 'Noise multiplier for DP-SGD')
 flags.DEFINE_float('epochs', None, 'Number of epochs (may be fractional)')
 flags.DEFINE_float('delta', 1e-6, 'Target delta')
 
-flags.mark_flag_as_required('N')
-flags.mark_flag_as_required('batch_size')
-flags.mark_flag_as_required('noise_multiplier')
-flags.mark_flag_as_required('epochs')
-
 
 def apply_dp_sgd_analysis(q, sigma, steps, orders, delta):
   """Compute and print results of DP-SGD analysis."""
@@ -79,18 +74,30 @@ def apply_dp_sgd_analysis(q, sigma, steps, orders, delta):
     print('The privacy estimate is likely to be improved by expanding '
           'the set of orders.')
 
+  return eps, opt_order
+
+
+def compute_dp_sgd_privacy(n, batch_size, noise_multiplier, epochs, delta):
+  """Compute epsilon based on the given hyperparameters."""
+  q = batch_size / n  # q - the sampling ratio.
+  if q > 1:
+    raise app.UsageError('n must be larger than the batch size.')
+  orders = ([1.25, 1.5, 1.75, 2., 2.25, 2.5, 3., 3.5, 4., 4.5] +
+            list(range(5, 64)) + [128, 256, 512])
+  steps = int(math.ceil(epochs * n / batch_size))
+
+  return apply_dp_sgd_analysis(q, noise_multiplier, steps, orders, delta)
+
 
 def main(argv):
   del argv  # argv is not used.
 
-  q = FLAGS.batch_size / FLAGS.N  # q - the sampling ratio.
+  assert FLAGS.N is not None, 'Flag N is missing.'
-  if q > 1:
+  assert FLAGS.batch_size is not None, 'Flag batch_size is missing.'
-    raise app.UsageError('N must be larger than the batch size.')
+  assert FLAGS.noise_multiplier is not None, 'Flag noise_multiplier is missing.'
-  orders = ([1.25, 1.5, 1.75, 2., 2.25, 2.5, 3., 3.5, 4., 4.5] +
+  assert FLAGS.epochs is not None, 'Flag epochs is missing.'
-            list(range(5, 64)) + [128, 256, 512])
+  compute_dp_sgd_privacy(FLAGS.N, FLAGS.batch_size, FLAGS.noise_multiplier,
-  steps = int(math.ceil(FLAGS.epochs * FLAGS.N / FLAGS.batch_size))
+                         FLAGS.epochs, FLAGS.delta)
-
-  apply_dp_sgd_analysis(q, FLAGS.noise_multiplier, steps, orders, FLAGS.delta)
 
 
 if __name__ == '__main__':

tensorflow_privacy/privacy/analysis/compute_dp_sgd_privacy_test.py

@@ -0,0 +1,41 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+from tensorflow_privacy.privacy.analysis import compute_dp_sgd_privacy
+
+
+class ComputeDpSgdPrivacyTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      ('Test0', 60000, 150, 1.3, 15, 1e-5, 0.941870567, 19.0),
+      ('Test1', 100000, 100, 1.0, 30, 1e-7, 1.70928734, 13.0),
+      ('Test2', 100000000, 1024, 0.1, 10, 1e-7, 5907984.81339406, 1.25),
+  )
+  def test_compute_dp_sgd_privacy(self, n, batch_size, noise_multiplier, epochs,
+                                  delta, expected_eps, expected_order):
+    eps, order = compute_dp_sgd_privacy.compute_dp_sgd_privacy(
+        n, batch_size, noise_multiplier, epochs, delta)
+    self.assertAlmostEqual(eps, expected_eps)
+    self.assertAlmostEqual(order, expected_order)
+
+if __name__ == '__main__':
+  absltest.main()