[
  {
    "path": "CONTRIBUTING.md",
    "content": "# How to Contribute\n\nWe'd love to accept your patches and contributions to this project. There are\njust a few small guidelines you need to follow.\n\n## Contributor License Agreement\n\nContributions to this project must be accompanied by a Contributor License\nAgreement. You (or your employer) retain the copyright to your contribution;\nthis simply gives us permission to use and redistribute your contributions as\npart of the project. Head over to <https://cla.developers.google.com/> to see\nyour current agreements on file or to sign a new one.\n\nYou generally only need to submit a CLA once, so if you've already submitted one\n(even if it was for a different project), you probably don't need to do it\nagain.\n\n## Code reviews\n\nAll submissions, including submissions by project members, require review. We\nuse GitHub pull requests for this purpose. Consult\n[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more\ninformation on using pull requests.\n\n## Community Guidelines\n\nThis project follows\n[Google's Open Source Community Guidelines](https://opensource.google.com/conduct/).\n"
  },
  {
    "path": "LICENSE",
    "content": "\n                                 Apache License\n                           Version 2.0, January 2004\n                        http://www.apache.org/licenses/\n\n   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n   1. Definitions.\n\n      \"License\" shall mean the terms and conditions for use, reproduction,\n      and distribution as defined by Sections 1 through 9 of this document.\n\n      \"Licensor\" shall mean the copyright owner or entity authorized by\n      the copyright owner that is granting the License.\n\n      \"Legal Entity\" shall mean the union of the acting entity and all\n      other entities that control, are controlled by, or are under common\n      control with that entity. For the purposes of this definition,\n      \"control\" means (i) the power, direct or indirect, to cause the\n      direction or management of such entity, whether by contract or\n      otherwise, or (ii) ownership of fifty percent (50%) or more of the\n      outstanding shares, or (iii) beneficial ownership of such entity.\n\n      \"You\" (or \"Your\") shall mean an individual or Legal Entity\n      exercising permissions granted by this License.\n\n      \"Source\" form shall mean the preferred form for making modifications,\n      including but not limited to software source code, documentation\n      source, and configuration files.\n\n      \"Object\" form shall mean any form resulting from mechanical\n      transformation or translation of a Source form, including but\n      not limited to compiled object code, generated documentation,\n      and conversions to other media types.\n\n      \"Work\" shall mean the work of authorship, whether in Source or\n      Object form, made available under the License, as indicated by a\n      copyright notice that is included in or attached to the work\n      (an example is provided in the Appendix below).\n\n      \"Derivative Works\" shall mean any work, whether in Source or Object\n      form, that is based on (or derived from) the Work and for which the\n      editorial revisions, annotations, elaborations, or other modifications\n      represent, as a whole, an original work of authorship. For the purposes\n      of this License, Derivative Works shall not include works that remain\n      separable from, or merely link (or bind by name) to the interfaces of,\n      the Work and Derivative Works thereof.\n\n      \"Contribution\" shall mean any work of authorship, including\n      the original version of the Work and any modifications or additions\n      to that Work or Derivative Works thereof, that is intentionally\n      submitted to Licensor for inclusion in the Work by the copyright owner\n      or by an individual or Legal Entity authorized to submit on behalf of\n      the copyright owner. For the purposes of this definition, \"submitted\"\n      means any form of electronic, verbal, or written communication sent\n      to the Licensor or its representatives, including but not limited to\n      communication on electronic mailing lists, source code control systems,\n      and issue tracking systems that are managed by, or on behalf of, the\n      Licensor for the purpose of discussing and improving the Work, but\n      excluding communication that is conspicuously marked or otherwise\n      designated in writing by the copyright owner as \"Not a Contribution.\"\n\n      \"Contributor\" shall mean Licensor and any individual or Legal Entity\n      on behalf of whom a Contribution has been received by Licensor and\n      subsequently incorporated within the Work.\n\n   2. Grant of Copyright License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      copyright license to reproduce, prepare Derivative Works of,\n      publicly display, publicly perform, sublicense, and distribute the\n      Work and such Derivative Works in Source or Object form.\n\n   3. Grant of Patent License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      (except as stated in this section) patent license to make, have made,\n      use, offer to sell, sell, import, and otherwise transfer the Work,\n      where such license applies only to those patent claims licensable\n      by such Contributor that are necessarily infringed by their\n      Contribution(s) alone or by combination of their Contribution(s)\n      with the Work to which such Contribution(s) was submitted. If You\n      institute patent litigation against any entity (including a\n      cross-claim or counterclaim in a lawsuit) alleging that the Work\n      or a Contribution incorporated within the Work constitutes direct\n      or contributory patent infringement, then any patent licenses\n      granted to You under this License for that Work shall terminate\n      as of the date such litigation is filed.\n\n   4. Redistribution. You may reproduce and distribute copies of the\n      Work or Derivative Works thereof in any medium, with or without\n      modifications, and in Source or Object form, provided that You\n      meet the following conditions:\n\n      (a) You must give any other recipients of the Work or\n          Derivative Works a copy of this License; and\n\n      (b) You must cause any modified files to carry prominent notices\n          stating that You changed the files; and\n\n      (c) You must retain, in the Source form of any Derivative Works\n          that You distribute, all copyright, patent, trademark, and\n          attribution notices from the Source form of the Work,\n          excluding those notices that do not pertain to any part of\n          the Derivative Works; and\n\n      (d) If the Work includes a \"NOTICE\" text file as part of its\n          distribution, then any Derivative Works that You distribute must\n          include a readable copy of the attribution notices contained\n          within such NOTICE file, excluding those notices that do not\n          pertain to any part of the Derivative Works, in at least one\n          of the following places: within a NOTICE text file distributed\n          as part of the Derivative Works; within the Source form or\n          documentation, if provided along with the Derivative Works; or,\n          within a display generated by the Derivative Works, if and\n          wherever such third-party notices normally appear. The contents\n          of the NOTICE file are for informational purposes only and\n          do not modify the License. You may add Your own attribution\n          notices within Derivative Works that You distribute, alongside\n          or as an addendum to the NOTICE text from the Work, provided\n          that such additional attribution notices cannot be construed\n          as modifying the License.\n\n      You may add Your own copyright statement to Your modifications and\n      may provide additional or different license terms and conditions\n      for use, reproduction, or distribution of Your modifications, or\n      for any such Derivative Works as a whole, provided Your use,\n      reproduction, and distribution of the Work otherwise complies with\n      the conditions stated in this License.\n\n   5. Submission of Contributions. Unless You explicitly state otherwise,\n      any Contribution intentionally submitted for inclusion in the Work\n      by You to the Licensor shall be under the terms and conditions of\n      this License, without any additional terms or conditions.\n      Notwithstanding the above, nothing herein shall supersede or modify\n      the terms of any separate license agreement you may have executed\n      with Licensor regarding such Contributions.\n\n   6. Trademarks. This License does not grant permission to use the trade\n      names, trademarks, service marks, or product names of the Licensor,\n      except as required for reasonable and customary use in describing the\n      origin of the Work and reproducing the content of the NOTICE file.\n\n   7. Disclaimer of Warranty. Unless required by applicable law or\n      agreed to in writing, Licensor provides the Work (and each\n      Contributor provides its Contributions) on an \"AS IS\" BASIS,\n      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n      implied, including, without limitation, any warranties or conditions\n      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n      PARTICULAR PURPOSE. You are solely responsible for determining the\n      appropriateness of using or redistributing the Work and assume any\n      risks associated with Your exercise of permissions under this License.\n\n   8. Limitation of Liability. In no event and under no legal theory,\n      whether in tort (including negligence), contract, or otherwise,\n      unless required by applicable law (such as deliberate and grossly\n      negligent acts) or agreed to in writing, shall any Contributor be\n      liable to You for damages, including any direct, indirect, special,\n      incidental, or consequential damages of any character arising as a\n      result of this License or out of the use or inability to use the\n      Work (including but not limited to damages for loss of goodwill,\n      work stoppage, computer failure or malfunction, or any and all\n      other commercial damages or losses), even if such Contributor\n      has been advised of the possibility of such damages.\n\n   9. Accepting Warranty or Additional Liability. While redistributing\n      the Work or Derivative Works thereof, You may choose to offer,\n      and charge a fee for, acceptance of support, warranty, indemnity,\n      or other liability obligations and/or rights consistent with this\n      License. However, in accepting such obligations, You may act only\n      on Your own behalf and on Your sole responsibility, not on behalf\n      of any other Contributor, and only if You agree to indemnify,\n      defend, and hold each Contributor harmless for any liability\n      incurred by, or claims asserted against, such Contributor by reason\n      of your accepting any such warranty or additional liability.\n\n   END OF TERMS AND CONDITIONS\n\n   APPENDIX: How to apply the Apache License to your work.\n\n      To apply the Apache License to your work, attach the following\n      boilerplate notice, with the fields enclosed by brackets \"[]\"\n      replaced with your own identifying information. (Don't include\n      the brackets!)  The text should be enclosed in the appropriate\n      comment syntax for the file format. We also recommend that a\n      file or class name and description of purpose be included on the\n      same \"printed page\" as the copyright notice for easier\n      identification within third-party archives.\n\n   Copyright [yyyy] [name of copyright owner]\n\n   Licensed under the Apache License, Version 2.0 (the \"License\");\n   you may not use this file except in compliance with the License.\n   You may obtain a copy of the License at\n\n       http://www.apache.org/licenses/LICENSE-2.0\n\n   Unless required by applicable law or agreed to in writing, software\n   distributed under the License is distributed on an \"AS IS\" BASIS,\n   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n   See the License for the specific language governing permissions and\n   limitations under the License.\n"
  },
  {
    "path": "README.md",
    "content": "# Interval Bound Propagation for Training Verifiably Robust Models\n\nThis repository contains a simple implementation of Interval Bound Propagation\n(IBP) using TensorFlow:\n[https://arxiv.org/abs/1810.12715](https://arxiv.org/abs/1810.12715).\nIt also contains an implementation of CROWN-IBP:\n[https://arxiv.org/abs/1906.06316](https://arxiv.org/abs/1906.06316).\nIt also contains a sentiment analysis example under [`examples/language`](https://github.com/deepmind/interval-bound-propagation/tree/master/examples/language) \nfor [https://arxiv.org/abs/1909.01492](https://arxiv.org/abs/1909.01492).\n\nThis is not an official Google product\n\n## Installation\n\nIBP can be installed with the following command:\n\n```bash\npip install git+https://github.com/deepmind/interval-bound-propagation\n```\n\nIBP will work with both the CPU and GPU version of tensorflow and dm-sonnet, but\nto allow for that it does not list Tensorflow as a requirement, so you need to\ninstall Tensorflow and Sonnet separately if you haven't already done so.\n\n## Usage\n\nThe following command trains a small model on MNIST with epsilon set to 0.3:\n\n```bash\ncd examples\npython train.py --model=small --output_dir=/tmp/small_model\n```\n\n## Pretrained Models\n\nModels trained using IBP and CROWN-IBP can be downloaded\n[here](https://drive.google.com/open?id=1lovI-fUabgs3swMgIe7MLRvHB9KtjzNT).\n\n### IBP models:\n\n| Dataset  | Test epsilon | Model path                 | Clean accuracy | Verified accuracy | Accuracy under attack |\n|----------|--------------|----------------------------|----------------|-------------------|-----------------------|\n| MNIST    | 0.1          | ibp/mnist_0.2_medium       | 98.94%         | 97.08%            | 97.99%                |\n| MNIST    | 0.2          | ibp/mnist_0.4_large_200    | 98.34%         | 95.47%            | 97.06%                |\n| MNIST    | 0.3          | ibp/mnist_0.4_large_200    | 98.34%         | 91.79%            | 96.03%                |\n| MNIST    | 0.4          | ibp/mnist_0.4_large_200    | 98.34%         | 84.99%            | 94.56%                |\n| CIFAR-10 | 2/255        | ibp/cifar_2-255_large_200  | 70.21%         | 44.12%            | 56.53%                |\n| CIFAR-10 | 8/255        | ibp/cifar_8-255_large      | 49.49%         | 31.56%            | 39.53%                |\n\n### CROWN-IBP models:\n\n| Dataset  | Test epsilon | Model path                   | Clean accuracy | Verified accuracy | Accuracy under attack |\n|----------|--------------|------------------------------|----------------|-------------------|-----------------------|\n| MNIST    | 0.1          | crown-ibp/mnist_0.2_large    | 99.03%         | 97.75%            | 98.34%                |\n| MNIST    | 0.2          | crown-ibp/mnist_0.4_large    | 98.38%         | 96.13%            | 97.28%                |\n| MNIST    | 0.3          | crown-ibp/mnist_0.4_large    | 98.38%         | 93.32%            | 96.38%                |\n| MNIST    | 0.4          | crown-ibp/mnist_0.4_large    | 98.38%         | 87.51%            | 94.95%                |\n| CIFAR-10 | 2/255        | crown-ibp/cifar_2-255_large  | 71.52%         | 53.97%            | 59.72%                |\n| CIFAR-10 | 8/255        | crown-ibp/cifar_8-255_large  | 47.14%         | 33.30%            | 36.81%                |\n| CIFAR-10 | 16/255       | crown-ibp/cifar_16-255_large | 34.19%         | 23.08%            | 26.55%                |\n\nIn these tables, we evaluated the verified accuracy using IBP only.\nWe evaluted the accuracy under attack using a 20-step untargeted PGD attack.\nYou can evaluate these models yourself using `eval.py`, for example:\n\n```bash\ncd examples\npython eval.py --model_dir pretrained_models/ibp/mnist_0.4_large_200/ \\\n  --epsilon 0.3\n```\n\nNote that we evaluated the CIFAR-10 2/255 CROWN-IBP model using CROWN-IBP\n(instead of pure IBP). You can do so yourself by setting the flag\n`--bound_method=crown-ibp`:\n\n```bash\npython eval.py --model_dir pretrained_models/crown-ibp/cifar_2-255_large/ \\\n  --epsilon 0.00784313725490196 --bound_method=crown-ibp\n```\n\n## Giving credit\n\nIf you use this code in your work, we ask that you cite this paper:\n\nSven Gowal, Krishnamurthy Dvijotham, Robert Stanforth, Rudy Bunel, Chongli Qin,\nJonathan Uesato, Relja Arandjelovic, Timothy Mann, and Pushmeet Kohli.\n\"On the Effectiveness of Interval Bound Propagation for Training Verifiably\nRobust Models.\" _arXiv preprint arXiv:1810.12715 (2018)_.\n\nIf you use CROWN-IBP, we also ask that you cite:\n\nHuan Zhang, Hongge Chen, Chaowei Xiao, Sven Gowal, Robert Stanforth, Bo Li,\nDuane Boning, Cho-Jui Hsieh.\n\"Towards Stable and Efficient Training of Verifiably Robust Neural Networks.\"\n_arXiv preprint arXiv:1906.06316 (2019)_.\n\nIf you use the sentiment analysis example, please cite:\n\nPo-Sen Huang, Robert Stanforth, Johannes Welbl, Chris Dyer, Dani Yogatama, Sven Gowal, Krishnamurthy Dvijotham, Pushmeet Kohli.\n\"Achieving Verified Robustness to Symbol Substitutions via Interval Bound Propagation.\"\n_EMNLP 2019_.\n\n\n## Acknowledgements\n\nIn addition to the people involved in the original IBP publication, we would\nlike to thank Huan Zhang, Sumanth Dathathri and Johannes Welbl for their\ncontributions.\n\n"
  },
  {
    "path": "examples/eval.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Evaluates a verifiable model on Mnist or CIFAR-10.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nimport interval_bound_propagation as ibp\nimport tensorflow.compat.v1 as tf\n\n\nFLAGS = flags.FLAGS\nflags.DEFINE_enum('dataset', 'auto', ['auto', 'mnist', 'cifar10'], 'Dataset '\n                  '(\"auto\", \"mnist\" or \"cifar10\"). When set to \"auto\", '\n                  'the dataset is inferred from the model directory path.')\nflags.DEFINE_enum('model', 'auto', ['auto', 'tiny', 'small', 'medium',\n                                    'large_200', 'large'], 'Model size. '\n                  'When set to \"auto\", the model name is inferred from the '\n                  'model directory path.')\nflags.DEFINE_string('model_dir', None, 'Model checkpoint directory.')\nflags.DEFINE_enum('bound_method', 'ibp', ['ibp', 'crown-ibp'],\n                  'Bound progataion method. For models trained with CROWN-IBP '\n                  'and beta_final=1 (e.g., CIFAR 2/255), use \"crown-ibp\". '\n                  'Otherwise use \"ibp\".')\nflags.DEFINE_integer('batch_size', 200, 'Batch size.')\nflags.DEFINE_float('epsilon', .3, 'Target epsilon.')\n\n\ndef layers(model_size):\n  \"\"\"Returns the layer specification for a given model name.\"\"\"\n  if model_size == 'tiny':\n    return (\n        ('linear', 100),\n        ('activation', 'relu'))\n  elif model_size == 'small':\n    return (\n        ('conv2d', (4, 4), 16, 'VALID', 2),\n        ('activation', 'relu'),\n        ('conv2d', (4, 4), 32, 'VALID', 1),\n        ('activation', 'relu'),\n        ('linear', 100),\n        ('activation', 'relu'))\n  elif model_size == 'medium':\n    return (\n        ('conv2d', (3, 3), 32, 'VALID', 1),\n        ('activation', 'relu'),\n        ('conv2d', (4, 4), 32, 'VALID', 2),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 64, 'VALID', 1),\n        ('activation', 'relu'),\n        ('conv2d', (4, 4), 64, 'VALID', 2),\n        ('activation', 'relu'),\n        ('linear', 512),\n        ('activation', 'relu'),\n        ('linear', 512),\n        ('activation', 'relu'))\n  elif model_size == 'large_200':\n    # Some old large checkpoints have 200 hidden neurons in the last linear\n    # layer.\n    return (\n        ('conv2d', (3, 3), 64, 'SAME', 1),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 64, 'SAME', 1),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 128, 'SAME', 2),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 128, 'SAME', 1),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 128, 'SAME', 1),\n        ('activation', 'relu'),\n        ('linear', 200),\n        ('activation', 'relu'))\n  elif model_size == 'large':\n    return (\n        ('conv2d', (3, 3), 64, 'SAME', 1),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 64, 'SAME', 1),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 128, 'SAME', 2),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 128, 'SAME', 1),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 128, 'SAME', 1),\n        ('activation', 'relu'),\n        ('linear', 512),\n        ('activation', 'relu'))\n  else:\n    raise ValueError('Unknown model: \"{}\"'.format(model_size))\n\n\ndef show_metrics(metric_values, bound_method='ibp'):\n  if bound_method == 'crown-ibp':\n    verified_accuracy = metric_values.crown_ibp_verified_accuracy\n  else:\n    verified_accuracy = metric_values.verified_accuracy\n  print('nominal accuracy = {:.2f}%, '\n        'verified accuracy = {:.2f}%, '\n        'accuracy under PGD attack = {:.2f}%'.format(\n            metric_values.nominal_accuracy * 100.,\n            verified_accuracy* 100.,\n            metric_values.attack_accuracy * 100.))\n\n\ndef main(unused_args):\n  dataset = FLAGS.dataset\n  if FLAGS.dataset == 'auto':\n    if 'mnist' in FLAGS.model_dir:\n      dataset = 'mnist'\n    elif 'cifar' in FLAGS.model_dir:\n      dataset = 'cifar10'\n    else:\n      raise ValueError('Cannot guess the dataset name. Please specify '\n                       '--dataset manually.')\n\n  model_name = FLAGS.model\n  if FLAGS.model == 'auto':\n    model_names = ['large_200', 'large', 'medium', 'small', 'tiny']\n    for name in model_names:\n      if name in FLAGS.model_dir:\n        model_name = name\n        logging.info('Using guessed model name \"%s\".', model_name)\n        break\n    if model_name == 'auto':\n      raise ValueError('Cannot guess the model name. Please specify --model '\n                       'manually.')\n\n  checkpoint_path = tf.train.latest_checkpoint(FLAGS.model_dir)\n  if checkpoint_path is None:\n    raise OSError('Cannot find a valid checkpoint in {}.'.format(\n        FLAGS.model_dir))\n\n  # Dataset.\n  input_bounds = (0., 1.)\n  num_classes = 10\n  if dataset == 'mnist':\n    data_train, data_test = tf.keras.datasets.mnist.load_data()\n  else:\n    assert dataset == 'cifar10', (\n        'Unknown dataset \"{}\"'.format(dataset))\n    data_train, data_test = tf.keras.datasets.cifar10.load_data()\n    data_train = (data_train[0], data_train[1].flatten())\n    data_test = (data_test[0], data_test[1].flatten())\n\n  # Base predictor network.\n  original_predictor = ibp.DNN(num_classes, layers(model_name))\n  predictor = original_predictor\n  if dataset == 'cifar10':\n    mean = (0.4914, 0.4822, 0.4465)\n    std = (0.2023, 0.1994, 0.2010)\n    predictor = ibp.add_image_normalization(original_predictor, mean, std)\n  if FLAGS.bound_method == 'crown-ibp':\n    predictor = ibp.crown.VerifiableModelWrapper(predictor)\n  else:\n    predictor = ibp.VerifiableModelWrapper(predictor)\n\n  # Test using while loop.\n  def get_test_metrics(batch_size, attack_builder=ibp.UntargetedPGDAttack):\n    \"\"\"Returns the test metrics.\"\"\"\n    num_test_batches = len(data_test[0]) // batch_size\n    assert len(data_test[0]) % batch_size == 0, (\n        'Test data is not a multiple of batch size.')\n\n    def cond(i, *unused_args):\n      return i < num_test_batches\n\n    def body(i, metrics):\n      \"\"\"Compute the sum of all metrics.\"\"\"\n      test_data = ibp.build_dataset(data_test, batch_size=batch_size,\n                                    sequential=True)\n      predictor(test_data.image, override=True, is_training=False)\n      input_interval_bounds = ibp.IntervalBounds(\n          tf.maximum(test_data.image - FLAGS.epsilon, input_bounds[0]),\n          tf.minimum(test_data.image + FLAGS.epsilon, input_bounds[1]))\n      predictor.propagate_bounds(input_interval_bounds)\n      test_specification = ibp.ClassificationSpecification(\n          test_data.label, num_classes)\n      test_attack = attack_builder(predictor, test_specification, FLAGS.epsilon,\n                                   input_bounds=input_bounds,\n                                   optimizer_builder=ibp.UnrolledAdam)\n\n      # Use CROWN-IBP bound or IBP bound.\n      if FLAGS.bound_method == 'crown-ibp':\n        test_losses = ibp.crown.Losses(predictor, test_specification,\n                                       test_attack, use_crown_ibp=True,\n                                       crown_bound_schedule=tf.constant(1.))\n      else:\n        test_losses = ibp.Losses(predictor, test_specification, test_attack)\n\n      test_losses(test_data.label)\n      new_metrics = []\n      for m, n in zip(metrics, test_losses.scalar_metrics):\n        new_metrics.append(m + n)\n      return i + 1, new_metrics\n\n    if FLAGS.bound_method == 'crown-ibp':\n      metrics = ibp.crown.ScalarMetrics\n    else:\n      metrics = ibp.ScalarMetrics\n    total_count = tf.constant(0, dtype=tf.int32)\n    total_metrics = [tf.constant(0, dtype=tf.float32)\n                     for _ in range(len(metrics._fields))]\n    total_count, total_metrics = tf.while_loop(\n        cond,\n        body,\n        loop_vars=[total_count, total_metrics],\n        back_prop=False,\n        parallel_iterations=1)\n    total_count = tf.cast(total_count, tf.float32)\n    test_metrics = []\n    for m in total_metrics:\n      test_metrics.append(m / total_count)\n    return metrics(*test_metrics)\n\n  test_metrics = get_test_metrics(\n      FLAGS.batch_size, ibp.UntargetedPGDAttack)\n\n  # Prepare to load the pretrained-model.\n  saver = tf.compat.v1.train.Saver(original_predictor.get_variables())\n\n  # Run everything.\n  tf_config = tf.ConfigProto()\n  tf_config.gpu_options.allow_growth = True\n  with tf.train.SingularMonitoredSession(config=tf_config) as sess:\n    logging.info('Restoring from checkpoint \"%s\".', checkpoint_path)\n    saver.restore(sess, checkpoint_path)\n    logging.info('Evaluating at epsilon = %f.', FLAGS.epsilon)\n    metric_values = sess.run(test_metrics)\n    show_metrics(metric_values, FLAGS.bound_method)\n\n\nif __name__ == '__main__':\n  flags.mark_flag_as_required('model_dir')\n  app.run(main)\n"
  },
  {
    "path": "examples/language/README.md",
    "content": "# Achieving Verified Robustness to Symbol Substitutions via Interval Bound Propagation\n\nHere contains an implementation of\n[Achieving Verified Robustness to Symbol Substitutions via Interval Bound \nPropagation](https://arxiv.org/abs/1909.01492).\n\n## Installation\n\nThe installation can be done with the following commands:\n\n```bash\npip3 install \"tensorflow-gpu<2\" \"dm-sonnet<2\" \"tensorflow-probability==0.7.0\" \"tensorflow-datasets\" \"absl-py\"\npip3 install git+https://github.com/deepmind/interval-bound-propagation\n```\n\n\n## Usage\n\nThe following command reproduces the [SST](https://nlp.stanford.edu/sentiment/) \ncharacter level experiments using perturbation radius of 3:\n\n```bash\ncd examples/language\npython3 robust_train.py\n```\n\nYou should expect to see the following at the end of training\n(note we only use SST dev set only for evaluation here).\n\n```bash\nstep: 149900, train loss: 0.392112, verifiable train loss: 0.826042,\ntrain accuracy: 0.850000, dev accuracy: 0.747619, test accuracy: 0.747619,\nTrain Bound = -0.42432, train verified: 0.800,\ndev verified: 0.695, test verified: 0.695\nbest dev acc 0.780952        best test acc   0.780952\nbest verified dev acc        0.716667        best verified test acc  0.716667\n```\n\nWe can verify the model in \n`config['model_location']='/tmp/robust_model/checkpoint/final'` using IBP.\n\nFor example, after changing `config['delta']=1.`, we can evaluate the IBP \nverified accuracy with perturbation radius of 1:\n\n```bash\npython3 robust_train.py --analysis --batch_size=1\n```\n\nWe expect to see results like the following:\n\n```bash\ntest final correct: 0.748, verified: 0.722\n{'datasplit': 'test', 'nominal': 0.7477064220183486,\n'verify': 0.7224770642201835, 'delta': 1.0, \n'num_perturbations': 268,\n'model_location': '/tmp/robust_model/checkpoint/final', 'final': True}\n```\n\nWe can also exhaustively search all valid perturbations to exhaustively verify\nthe models.\n\n```bash\npython3 exhaustive_verification.py --num_examples=0\n```\n\nWe should expect the following results\n\n```bash\nverified_proportion: 0.7350917431192661\n{'delta': 1, 'character_level': True, 'mode': 'validation', 'checkpoint_path': '/tmp/robust_model/checkpoint/final', 'verified_proportion': 0.7350917431192661}\n```\n\nThe IBP verified accuracy ` 0.7224770642201835` is a lower bound of the\nexhaustive verification results, `0.7350917431192661`.\n\nFurthermore, we can also align the predictions between the IBP verification \nand exhaustive verification. There should not be cases where IBP can verify \n(no attack can change the predictions) and exhaustive verification cannot \nverify (there exist an attack that can change the predictions), since IBP \nprovides a lower bound on the true robustness accuracy (via exhaustive search).\n\n\n## Reference\n\nIf you use this code in your work, please cite the accompanying paper:\n\n```\n@inproceedings{huang-2019-achieving,\n    title = \"Achieving Verified Robustness to Symbol Substitutions via Interval Bound Propagation\",\n    author = \"Po-Sen Huang and\n      Robert Stanforth and\n      Johannes Welbl and\n      Chris Dyer and\n      Dani Yogatama and\n      Sven Gowal  and\n      Krishnamurthy Dvijotham and\n      Pushmeet Kohli\",\n    booktitle = \"Empirical Methods in Natural Language Processing (EMNLP)\",\n    year = \"2019\",\n    pages = \"4081--4091\",\n}\n```\n\n## Disclaimer\n\nThis is not an official Google product.\n"
  },
  {
    "path": "examples/language/config.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Configuration parameters for sentence representation models.\"\"\"\n\n\ndef get_config():\n  \"\"\"Returns the default configuration as a dict.\"\"\"\n\n  config = {}\n\n  config['dataset'] = 'sst'\n  # Convolutional architecture.\n  # Format: Tuple/List for a Conv layer (filters, kernel_size, pooling_size)\n  # Otherwise, nonlinearity.\n  config['conv_architecture'] = ((100, 5, 1), 'relu')\n\n  # Fully connected layer 1 hidden sizes (0 means no layer).\n  config['conv_fc1'] = 0\n\n  # Fully connected layer 2 hidden sizes (0 means no layer).\n  config['conv_fc2'] = 0\n\n  # Number of allowable perturbations.\n  # (delta specifies the budget, i.e., how many may be used at once.)\n  config['delta'] = 3.0\n\n  # Allow each character to be changed to another character.\n  config['synonym_filepath'] = 'data/character_substitution_enkey_sub1.json'\n  config['max_padded_length'] = 268\n  # (~1*268) Max num_perturbations.\n  # seqlen * max_number_synonyms (total number of elementary perturbations)\n  config['num_perturbations'] = 268\n\n  config['vocab_filename'] = 'data/sst_binary_character_vocabulary_sorted.txt'\n  # Need to add pad for analysis (which is what is used after\n  # utils.get_merged_vocabulary_file).\n  config['vocab_filename_pad'] = (\n      'data/sst_binary_character_vocabulary_sorted_pad.txt')\n\n  config['embedding_dim'] = 150\n\n  config['delta_schedule'] = True\n  config['verifiable_loss_schedule'] = True\n\n  # Ratio between the task loss and verifiable loss.\n  config['verifiable_loss_ratio'] = 0.75\n\n  # Aggregrated loss of the verifiable training objective\n  # (among softmax, mean, max).\n  config['verifiable_training_aggregation'] = 'softmax'\n\n  config['data_id'] = 1\n\n  config['model_location'] = '/tmp/robust_model/checkpoint/final'\n\n  return config\n"
  },
  {
    "path": "examples/language/data/character_substitution_enkey_sub1.json",
    "content": "{\"z\": [\"x\"], \"y\": [\"t\"], \"x\": [\"s\"], \"w\": [\"d\"], \"v\": [\"c\"], \"u\": [\"8\"], \"t\": [\"f\"], \"s\": [\"e\"], \"r\": [\"g\"], \"q\": [\"s\"], \"p\": [\";\"], \"o\": [\"k\"], \"n\": [\"m\"], \"m\": [\"j\"], \"l\": [\"p\"], \"k\": [\".\"], \"j\": [\"i\"], \"i\": [\"u\"], \"h\": [\"n\"], \"g\": [\"v\"], \"f\": [\"c\"], \"e\": [\"r\"], \"d\": [\"f\"], \"c\": [\"d\"], \"b\": [\"g\"], \"a\": [\"x\"]}\n"
  },
  {
    "path": "examples/language/data/sst_binary_character_vocabulary_sorted.txt",
    "content": " \n!\n#\n$\n%\n&\n'\n(\n)\n*\n+\n,\n-\n.\n/\n0\n1\n2\n3\n4\n5\n6\n7\n8\n9\n:\n;\n=\n?\n`\na\nb\nc\nd\ne\nf\ng\nh\ni\nj\nk\nl\nm\nn\no\np\nq\nr\ns\nt\nu\nv\nw\nx\ny\nz\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n"
  },
  {
    "path": "examples/language/data/sst_binary_character_vocabulary_sorted_pad.txt",
    "content": "<PAD>\n \n!\n#\n$\n%\n&\n'\n(\n)\n*\n+\n,\n-\n.\n/\n0\n1\n2\n3\n4\n5\n6\n7\n8\n9\n:\n;\n=\n?\n`\na\nb\nc\nd\ne\nf\ng\nh\ni\nj\nk\nl\nm\nn\no\np\nq\nr\ns\nt\nu\nv\nw\nx\ny\nz\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n"
  },
  {
    "path": "examples/language/exhaustive_verification.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Functionality for exhaustive adversarial attacks on synonym perturbations.\n\nModels restored from checkpoint can be tested w.r.t their robustness to\nexhaustive-search adversaries, which have a fixed perturbation budget with which\nthey can flip words to synonyms.\n\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport collections\nimport copy\nimport imp\nimport json\nimport pprint\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nimport numpy as np\nimport tensorflow.compat.v1 as tf\nimport tensorflow_datasets as tfds\nimport tqdm\n\nimport interactive_example\n\n\nflags.DEFINE_boolean('character_level', True, 'Character level model.')\nflags.DEFINE_boolean('debug_mode', False, 'Debug mode.')\nflags.DEFINE_string('checkpoint_path', '/tmp/robust_model/checkpoint/final',\n                    'Checkpoint path.')\nflags.DEFINE_string('dataset', 'sst', 'Dataset name. train, dev, or test.')\nflags.DEFINE_string('mode', 'validation', 'Dataset part. train, dev, or test.')\nflags.DEFINE_string('config_path', './config.py',\n                    'Path to training configuration file.')\nflags.DEFINE_string('task', 'sst', 'One of snli, mnli, sick, sst.')\nflags.DEFINE_integer('batch_size', 30, 'Batch size.')\nflags.DEFINE_string('pooling', 'average', 'One of averge, sum, max, last.')\nflags.DEFINE_boolean('fine_tune_embeddings', True, 'Finetune embeddings.')\nflags.DEFINE_integer('num_oov_buckets', 1, 'Number of out-of-vocab buckets.')\n\nflags.DEFINE_integer('delta', 1, 'Maximum perturbation radius')\nflags.DEFINE_integer('skip_batches', 0, 'Skip this number of batches'\n                     ' for analysis.')\nflags.DEFINE_integer('num_examples', 100, 'Analyze this number of examples. '\n                     ' 0 suggest the whole dataset.')\nflags.DEFINE_integer('truncated_len', 0, 'truncated sentence length. '\n                     ' 0 suggest the whole sentence.')\nflags.DEFINE_integer('max_padded_length', 0, 'max_padded_length. '\n                     ' 0 suggest no change.')\nflags.DEFINE_integer('num_perturbations', 0, 'num_perturbations. '\n                     ' 0 suggest no change.')\n\nFLAGS = flags.FLAGS\n\n\ndef load_synonyms(synonym_filepath=None):\n  \"\"\"Loads synonym dictionary. Returns as defaultdict(list).\"\"\"\n  with tf.gfile.Open(synonym_filepath) as f:\n    synonyms = json.load(f)\n  synonyms_ = collections.defaultdict(list)\n  synonyms_.update(synonyms)\n  return synonyms_\n\n\ndef load_dataset(mode='validation', character_level=False):\n  \"\"\"Loads SST dataset.\n\n  Takes data from disk/cns if it exists, otherwise out of tensorflow graph.\n  Args:\n    mode: string. Either train, dev, or test.\n    character_level: bool. Whether to return character-level, or token level\n      inputs.\n  Returns:\n    List of (input, output) pairs, where input is a list of strings (tokens),\n      and output is an integer (categorical label in [0,1]).\n  \"\"\"\n  message = 'Loading SST {}, character_level {}'.format(mode,\n                                                        str(character_level))\n  logging.info(message)\n  dataset = tfds.load(name='glue/sst2', split=mode)\n  minibatch = dataset.batch(1).make_one_shot_iterator().get_next()\n  label_list, input_list = [], []\n  with tf.train.SingularMonitoredSession() as session:\n    while True:\n      output_nodes = (minibatch['label'], minibatch['sentence'])\n      label, sentence = session.run(output_nodes)\n      label_list.append(label[0])\n      input_list.append([chr(i) for i in sentence[0]])\n\n  # zip together.\n  dataset = [(in_, out_) for (in_, out_) in zip(input_list, label_list)]\n  return dataset\n\n\ndef expand_by_one_perturbation(original_tokenized_sentence,\n                               tokenized_sentence, synonym_dict):\n  \"\"\"Expands given sentence by all possible synonyms.\n\n  Note that only a single synonym replacement is applied, and it is applied\n  everywhere, i.e. for every mention of the word with the synonym.\n  Args:\n    original_tokenized_sentence: List[str]. List of tokens.\n    tokenized_sentence: List[str]. List of tokens.\n    synonym_dict: dict, mapping words (str) to lists of synonyms (list of str)\n  Returns:\n    new_sentences_list: List[List[str]]. Outer list is across different synonym\n      replacements. Inner list is over (str) tokens.\n  \"\"\"\n  new_sentences_list = []\n  for i_outer, (original_token, _) in enumerate(zip(\n      original_tokenized_sentence, tokenized_sentence)):\n    synonyms = synonym_dict[original_token]\n    for synonym in synonyms:  # replace only one particular mention\n      new_sentence = copy.copy(tokenized_sentence)\n      new_sentence[i_outer] = synonym\n      new_sentences_list.append(new_sentence)\n\n  return new_sentences_list\n\n\ndef find_up_to_depth_k_perturbations(\n    original_tokenized_sentence, tokenized_sentence, synonym_dict, k):\n  \"\"\"Takes sentence, finds all sentences reachable using k token perturbations.\n\n  Args:\n    original_tokenized_sentence: List[str]. List of tokens.\n    tokenized_sentence: List[str]. List of tokens.\n    synonym_dict: dict, mapping words (str) to lists of synonyms (list of str)\n    k: int. perturbation depth parameter.\n  Returns:\n    output_sentences: List[List[str]]. List of tokenised sentences.\n  \"\"\"\n  # Case: recursion ends - no further perturbations.\n  if k == 0:\n    return [tokenized_sentence]\n  else:\n    # Expand by one level.\n    expanded_sentences = expand_by_one_perturbation(original_tokenized_sentence,\n                                                    tokenized_sentence,\n                                                    synonym_dict)\n\n    # Call recursive function one level deeper for each expanded sentence.\n    expanded_sentences_deeper = []\n    for sentence in expanded_sentences:\n      new_sentences = find_up_to_depth_k_perturbations(\n          original_tokenized_sentence, sentence, synonym_dict, k-1)\n      expanded_sentences_deeper.extend(new_sentences)\n\n  output_sentences = expanded_sentences + expanded_sentences_deeper\n  output_sentences = remove_duplicates(output_sentences)\n  return output_sentences\n\n\ndef remove_duplicates(list_of_list_of_tokens):\n  # Convert list of str to str.\n  sentences = ['|'.join(s) for s in list_of_list_of_tokens]\n  sentences = set(sentences)  # Now hashable -> remove duplicates.\n  sentences = [s.split('|') for s in sentences]  # Convert to original format.\n  return sentences\n\n\ndef verify_exhaustively(sample, synonym_dict, sst_model, delta,\n                        truncated_len=0):\n  \"\"\"Returns True if a sample can be verified, False otherwise.\n\n  Args:\n    sample: a 2-tuple (x,y), where x is a tokenised sentence (List[str]), and y\n      is a label (int).\n    synonym_dict: str -> List[str]. Keys are words, values are word lists with\n      synonyms for the key word.\n    sst_model: InteractiveSentimentPredictor instance. Used to make predictions.\n    delta: int. How many synonym perturbations to maximally allow.\n    truncated_len: int. Truncate sentence to truncated_len. 0 for unchanged.\n  Returns:\n    verified: bool. Whether all possible perturbed version of input sentence x\n      up to perturbation radius delta have the correct prediction.\n  \"\"\"\n  (x, y) = sample\n  counter_example = None\n  counter_prediction = None\n  # Create (potentially long) list of perturbed sentences from x.\n  if truncated_len > 0:\n    x = x[: truncated_len]\n  # Add original sentence.\n  altered_sentences = find_up_to_depth_k_perturbations(x, x, synonym_dict,\n                                                       delta)\n  altered_sentences = altered_sentences + [x]\n  # Form batches of these altered sentences.\n  batch = []\n  num_forward_passes = len(altered_sentences)\n\n  for sentence in altered_sentences:\n    any_prediction_wrong = False\n    batch.append(sentence)\n    # When batch_size is reached, make predictions, break if any label flip\n    if len(batch) == sst_model.batch_size:\n      # np array of size [batch_size]\n      predictions, _ = sst_model.batch_predict_sentiment(\n          batch, is_tokenised=True)\n\n      # Check any prediction that is different from the true label.\n      any_prediction_wrong = np.any(predictions != y)\n      if any_prediction_wrong:\n        wrong_index = np.where(predictions != y)[0].tolist()[0]\n        counter_example = ' '.join([str(c) for c in batch[wrong_index]])\n        if FLAGS.debug_mode:\n          logging.info('\\nOriginal example: %s, prediction: %d',\n                       ' '.join([str(c) for c in sentence]), y)\n          logging.info('\\ncounter example:  %s, prediction: %s',\n                       counter_example, predictions[wrong_index].tolist())\n        counter_prediction = predictions[wrong_index]\n        # Break. No need to evaluate further.\n        return False, counter_example, counter_prediction, num_forward_passes\n\n      # Start filling up the next batch.\n      batch = []\n\n  if not batch:\n    # No remainder, not previously broken the loop.\n    return True, None, None, num_forward_passes\n  else:\n    # Remainder -- what didn't fit into a full batch of size batch_size.\n    # We use the first altered_sentence to pad.\n    batch += [altered_sentences[0]]*(sst_model.batch_size-len(batch))\n    assert len(batch) == sst_model.batch_size\n    predictions, _ = sst_model.batch_predict_sentiment(batch, is_tokenised=True)\n    any_prediction_wrong = np.any(predictions != y)\n    if any_prediction_wrong:\n      wrong_index = np.where(predictions != y)[0].tolist()[0]\n      counter_example = ' '.join([str(c) for c in batch[wrong_index]])\n      if FLAGS.debug_mode:\n        logging.info('\\nOriginal example: %s, prediction: %d',\n                     ' '.join([str(c) for c in sentence]), y)  # pylint: disable=undefined-loop-variable\n        logging.info('\\ncounter example:  %s, prediction: %s', counter_example,\n                     predictions[wrong_index].tolist())\n      counter_prediction = predictions[wrong_index]\n    return (not any_prediction_wrong, counter_example,\n            counter_prediction, num_forward_passes)\n\n\ndef verify_dataset(dataset, config_dict, model_location, synonym_dict, delta):\n  \"\"\"Tries to verify against perturbation attacks up to delta.\"\"\"\n  sst_model = interactive_example.InteractiveSentimentPredictor(\n      config_dict, model_location,\n      max_padded_length=FLAGS.max_padded_length,\n      num_perturbations=FLAGS.num_perturbations)\n  verified_list = []  # Holds boolean entries, across dataset.\n  samples = []\n  labels = []\n  counter_examples = []\n  counter_predictions = []\n  total_num_forward_passes = []\n  logging.info('dataset size: %d', len(dataset))\n  num_examples = FLAGS.num_examples if FLAGS.num_examples else len(dataset)\n  logging.info('skip_batches: %d', FLAGS.skip_batches)\n  logging.info('num_examples: %d', num_examples)\n  logging.info('new dataset size: %d',\n               len(dataset[FLAGS.skip_batches:FLAGS.skip_batches+num_examples]))\n  for i, sample in tqdm.tqdm(enumerate(\n      dataset[FLAGS.skip_batches:FLAGS.skip_batches+num_examples])):\n    if FLAGS.debug_mode:\n      logging.info('index: %d', i)\n      (verified_bool, counter_example, counter_prediction, num_forward_passes\n      ) = verify_exhaustively(\n          sample, synonym_dict, sst_model, delta, FLAGS.truncated_len)\n      samples.append(''.join(sample[0]))\n      labels.append(sample[1])\n      counter_examples.append(counter_example)\n      counter_predictions.append(counter_prediction)\n      total_num_forward_passes.append(num_forward_passes)\n    else:\n      verified_bool, _, _, num_forward_passes = verify_exhaustively(\n          sample, synonym_dict, sst_model, delta, FLAGS.truncated_len)\n    verified_list.append(verified_bool)\n\n  verified_proportion = np.mean(verified_list)\n  assert len(verified_list) == len(\n      dataset[FLAGS.skip_batches:FLAGS.skip_batches+num_examples])\n  return (verified_proportion, verified_list, samples, counter_examples,\n          counter_predictions, total_num_forward_passes)\n\n\ndef example(synonym_dict, dataset, k=2):\n  \"\"\"Example usage of functions above.\"\"\"\n\n  # The below example x has these synonyms.\n  # 'decree' --> [edict, order],\n  # 'tubes' --> 'pipes';\n  # 'refrigerated' --> ['cooled', 'chilled']\n  x = ['the', 'refrigerated', 'decree', 'tubes']\n\n  # Example: 1 perturbation.\n  new_x = expand_by_one_perturbation(x, x, synonym_dict)\n  pprint.pprint(sorted(new_x))\n\n  # Example: up to k perturbations.\n  new_x = find_up_to_depth_k_perturbations(x, x, synonym_dict, k)\n  pprint.pprint(sorted(new_x))\n\n  # Statistics: how large is the combinatorial space of perturbations?\n  total_x = []\n  size_counter = collections.Counter()\n  for (x, _) in tqdm.tqdm(dataset):\n    new_x = find_up_to_depth_k_perturbations(x, x, synonym_dict, k)\n    size_counter[len(new_x)] += 1\n    total_x.extend(new_x)\n\n  # Histogram for perturbation space size, computed across dataset.\n  pprint.pprint([x for x in sorted(size_counter.items(), key=lambda xx: xx[0])])\n\n  # Total number of inputs for forward pass if comprehensively evaluated.\n  pprint.pprint(len(total_x))\n\n\ndef main(args):\n  del args\n\n  # Read the config file into a new ad-hoc module.\n  with open(FLAGS.config_path, 'r') as config_file:\n    config_code = config_file.read()\n    config_module = imp.new_module('config')\n    exec(config_code, config_module.__dict__)  # pylint: disable=exec-used\n  config = config_module.get_config()\n\n  config_dict = {'task': FLAGS.task,\n                 'batch_size': FLAGS.batch_size,\n                 'pooling': FLAGS.pooling,\n                 'learning_rate': 0.,\n                 'config': config,\n                 'embedding_dim': config['embedding_dim'],\n                 'fine_tune_embeddings': FLAGS.fine_tune_embeddings,\n                 'num_oov_buckets': FLAGS.num_oov_buckets,\n                 'max_grad_norm': 0.}\n\n  # Maximum verification range.\n  delta = FLAGS.delta\n  character_level = FLAGS.character_level\n  mode = FLAGS.mode\n  model_location = FLAGS.checkpoint_path\n\n  # Load synonyms.\n  synonym_filepath = config['synonym_filepath']\n  synonym_dict = load_synonyms(synonym_filepath)\n\n  # Load data.\n  dataset = load_dataset(mode, character_level)\n\n  # Compute verifiable accuracy on dataset.\n  (verified_proportion, _, _, _, _, _) = verify_dataset(dataset, config_dict,\n                                                        model_location,\n                                                        synonym_dict, delta)\n  logging.info('verified_proportion:')\n  logging.info(str(verified_proportion))\n  logging.info({\n      'delta': FLAGS.delta,\n      'character_level': FLAGS.character_level,\n      'mode': FLAGS.mode,\n      'checkpoint_path': FLAGS.checkpoint_path,\n      'verified_proportion': verified_proportion\n  })\n\n\nif __name__ == '__main__':\n  logging.set_stderrthreshold('info')\n  app.run(main)\n"
  },
  {
    "path": "examples/language/interactive_example.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Minimum code to interact with a pretrained Stanford Sentiment Treebank model.\n\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport collections\n\nimport numpy as np\nfrom six.moves import range\nimport tensorflow.compat.v1 as tf\n\nimport robust_model\n\n\nSparseTensorValue = collections.namedtuple(\n    'SparseTensorValue', ['indices', 'values', 'dense_shape'])\n\n\nclass InteractiveSentimentPredictor(object):\n  \"\"\"Can be used to interact with a trained sentiment analysis model.\"\"\"\n\n  def __init__(self, config_dict, model_location, max_padded_length=0,\n               num_perturbations=0):\n    self.graph_tensor_producer = robust_model.RobustModel(**config_dict)\n\n    self.batch_size = self.graph_tensor_producer.batch_size\n    if max_padded_length:\n      self.graph_tensor_producer.config.max_padded_length = max_padded_length\n    if num_perturbations:\n      self.graph_tensor_producer.config.num_perturbations = num_perturbations\n    self.graph_tensors = self.graph_tensor_producer()\n\n    network_saver = tf.train.Saver(self.graph_tensor_producer.variables)\n    self.open_session = tf.Session()\n    self.open_session.run(tf.tables_initializer())\n    network_saver.restore(self.open_session, model_location)\n\n  def batch_predict_sentiment(self, list_of_sentences, is_tokenised=True):\n    \"\"\"Computes sentiment predictions for a batch of sentences.\n\n    Note: the model batch size is usually hard-coded in the model (e.g. at 64).\n    We require that len(list_of_sentences)==self.batch_size.\n    If padding is necessary to reach as many sentences, this should happen\n    outside of this function.\n\n    Important: we assume that each sentence has the same number of tokens.\n    Args:\n      list_of_sentences: List[str] in case is_tokenised is False, or\n        List[List[str]] in case is_tokenised is True. Holds inputs whose\n        sentiment is to be classified.\n      is_tokenised: bool. Whether sentences are already tokenised. If not,\n        naive whitespace splitting tokenisation is applied.\n    Returns:\n      batch_label_predictions: np.array of shape [self.batch_size] holding\n        integers, representing model predictions for each input.\n    \"\"\"\n\n    # Prepare inputs.\n    tokenised_sentence_list = []\n    for sentence in list_of_sentences:\n      if not is_tokenised:\n        tokenised_sentence = sentence.lower().split(' ')\n      else:\n        tokenised_sentence = sentence\n      tokenised_sentence_list.append(tokenised_sentence)\n    length = len(tokenised_sentence_list[0])\n    assert all([len(x) == length for x in tokenised_sentence_list])\n    assert len(tokenised_sentence_list) == self.batch_size\n\n    # Construct sparse tensor holding token information.\n    indices = np.zeros([self.batch_size*length, 2])\n    dense_shape = [self.batch_size, length]\n    # Loop over words. All sentences have the same length.\n    for j, _ in enumerate(tokenised_sentence_list[0]):\n      for i in range(self.batch_size):  # Loop over samples.\n        offset = i*length + j\n        indices[offset, 0] = i\n        indices[offset, 1] = j\n\n    # Define sparse tensor values.\n    tokenised_sentence_list = [word for sentence in tokenised_sentence_list  # pylint:disable=g-complex-comprehension\n                               for word in sentence]\n    values = np.array(tokenised_sentence_list)\n    mb_tokens = SparseTensorValue(indices=indices, values=values,\n                                  dense_shape=dense_shape)\n    mb_num_tokens = np.array([length]*self.batch_size)\n\n    # Fill feed_dict with input token information.\n    feed_dict = {}\n    feed_dict[self.graph_tensors['dev']['tokens']] = mb_tokens\n    feed_dict[self.graph_tensors['dev']['num_tokens']] = mb_num_tokens\n\n    # Generate model predictions [batch_size x n_labels].\n    logits = self.open_session.run(self.graph_tensors['dev']['predictions'],\n                                   feed_dict)\n    batch_label_predictions = np.argmax(logits, axis=1)\n\n    return batch_label_predictions, logits\n\n  def predict_sentiment(self, sentence, tokenised=False):\n    \"\"\"Computes sentiment of a sentence.\"\"\"\n    # Create inputs to tensorflow graph.\n    if tokenised:\n      inputstring_tokenised = sentence\n    else:\n      assert isinstance(sentence, str)\n      # Simple tokenisation.\n      inputstring_tokenised = sentence.lower().split(' ')\n    length = len(inputstring_tokenised)\n\n    # Construct inputs to sparse tensor holding token information.\n    indices = np.zeros([self.batch_size*length, 2])\n    dense_shape = [self.batch_size, length]\n    for j, _ in enumerate(inputstring_tokenised):\n      for i in range(self.batch_size):\n        offset = i*length + j\n        indices[offset, 0] = i\n        indices[offset, 1] = j\n    values = inputstring_tokenised*self.batch_size\n    mb_tokens = SparseTensorValue(indices=indices, values=np.array(values),\n                                  dense_shape=dense_shape)\n    mb_num_tokens = np.array([length]*self.batch_size)\n\n    # Fill feeddict with input token information.\n    feed_dict = {}\n    feed_dict[self.graph_tensors['dev']['tokens']] = mb_tokens\n    feed_dict[self.graph_tensors['dev']['num_tokens']] = mb_num_tokens\n    # Generate predictions.\n    logits = self.open_session.run(self.graph_tensors['dev']['predictions'],\n                                   feed_dict)\n    predicted_label = np.argmax(logits, axis=1)\n    final_prediction = predicted_label[0]\n    # Check that prediction same everywhere (had batch of identical inputs).\n    assert np.all(predicted_label == final_prediction)\n    return final_prediction, logits\n"
  },
  {
    "path": "examples/language/models.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Models for sentence representation.\"\"\"\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\ndef _max_pool_1d(x, pool_size=2, name='max_pool_1d'):\n  with tf.name_scope(name, 'MaxPool1D', [x, pool_size]):\n    return tf.squeeze(\n        tf.nn.max_pool(tf.expand_dims(x, 1),\n                       [1, 1, pool_size, 1],\n                       [1, 1, pool_size, 1],\n                       'VALID'),\n        axis=1)\n\n\nclass SentenceRepresenterConv(snt.AbstractModule):\n  \"\"\"Use stacks of 1D Convolutions to build a sentence representation.\"\"\"\n\n  def __init__(self,\n               config,\n               keep_prob=1.,\n               pooling='max',\n               name='sentence_rep_conv'):\n    super(SentenceRepresenterConv, self).__init__(name=name)\n    self._config = config\n    self._pooling = pooling\n    self._keep_prob = keep_prob\n\n  def _build(self, padded_word_embeddings, length):\n    x = padded_word_embeddings\n    for layer in self._config['conv_architecture']:\n      if isinstance(layer, tuple) or isinstance(layer, list):\n        filters, kernel_size, pooling_size = layer\n        conv = snt.Conv1D(\n            output_channels=filters,\n            kernel_shape=kernel_size)\n        x = conv(x)\n        if pooling_size and pooling_size > 1:\n          x = _max_pool_1d(x, pooling_size)\n      elif layer == 'relu':\n        x = tf.nn.relu(x)\n        if self._keep_prob < 1:\n          x = tf.nn.dropout(x, keep_prob=self._keep_prob)\n      else:\n        raise RuntimeError('Bad layer type {} in conv'.format(layer))\n    # Final layer pools over the remaining sequence length to get a\n    # fixed sized vector.\n    if self._pooling == 'max':\n      x = tf.reduce_max(x, axis=1)\n    elif self._pooling == 'average':\n      x = tf.reduce_sum(x, axis=1)\n      lengths = tf.expand_dims(tf.cast(length, tf.float32), axis=1)\n      x = x / lengths\n\n    if self._config['conv_fc1']:\n      fc1_layer = snt.Linear(output_size=self._config['conv_fc1'])\n      x = tf.nn.relu(fc1_layer(x))\n      if self._keep_prob < 1:\n        x = tf.nn.dropout(x, keep_prob=self._keep_prob)\n    if self._config['conv_fc2']:\n      fc2_layer = snt.Linear(output_size=self._config['conv_fc2'])\n      x = tf.nn.relu(fc2_layer(x))\n      if self._keep_prob < 1:\n        x = tf.nn.dropout(x, keep_prob=self._keep_prob)\n\n    return x\n\n\n\n"
  },
  {
    "path": "examples/language/robust_model.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Train verifiable robust models.\"\"\"\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport collections\n\nfrom absl import logging\nimport interval_bound_propagation as ibp\nimport numpy as np\nimport six\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\nimport tensorflow_datasets as tfds\nimport tensorflow_probability as tfp\n\nfrom tensorflow.contrib import lookup as contrib_lookup\nimport models\nimport utils\n\n\nEmbeddedDataset = collections.namedtuple(\n    'EmbeddedDataset',\n    ['embedded_inputs', 'length', 'input_tokens', 'sentiment'])\n\nDataset = collections.namedtuple(\n    'Dataset',\n    ['tokens', 'num_tokens', 'sentiment'])\n\nPerturbation = collections.namedtuple(\n    'Perturbation',\n    ['positions', 'tokens'])\n\n\ndef _pad_fixed(x, axis, padded_length):\n  \"\"\"Pads a tensor to a fixed size (rather than batch-specific).\"\"\"\n  pad_shape = x.shape.as_list()\n  pad_shape[axis] = tf.maximum(padded_length - tf.shape(x)[axis], 0)\n  # Pad zero as in utils.get_padded_indexes.\n  padded = tf.concat([x, tf.zeros(dtype=x.dtype, shape=pad_shape)], axis=axis)\n  assert axis == 1\n  padded = padded[:, :padded_length]\n\n  padded_shape = padded.shape.as_list()\n  padded_shape[axis] = padded_length\n  padded.set_shape(padded_shape)\n  return padded\n\n\nclass GeneratedDataset(snt.AbstractModule):\n  \"\"\"A dataset wrapper for data_gen such that it behaves like sst_binary.\"\"\"\n\n  def __init__(self, data_gen, batch_size, mode='train',\n               num_examples=0,\n               dataset_name='glue/sst2',\n               name='generated_dataset'):\n    super(GeneratedDataset, self).__init__(name=name)\n    self._data_gen = data_gen\n    self._batch_size = batch_size\n    self._mode = mode\n    self._shuffle = True if mode == 'train' else False\n    self._num_examples = num_examples\n    self._dataset_name = dataset_name\n\n  def get_row_lengths(self, sparse_tensor_input):\n    # sparse_tensor_input is a tf.SparseTensor\n    # In RaggedTensor, row_lengths is a vector with shape `[nrows]`,\n    # which specifies the length of each row.\n    rt = tf.RaggedTensor.from_sparse(sparse_tensor_input)\n    return rt.row_lengths()\n\n  def _build(self):\n    dataset = tfds.load(name=self._dataset_name, split=self._mode)\n    minibatch = dataset.map(parse).repeat()\n\n    if self._shuffle:\n      minibatch = minibatch.shuffle(self._batch_size*100)\n    minibatch = minibatch.batch(\n        self._batch_size).make_one_shot_iterator().get_next()\n    minibatch['sentiment'].set_shape([self._batch_size])\n    minibatch['sentence'] = tf.SparseTensor(\n        indices=minibatch['sentence'].indices,\n        values=minibatch['sentence'].values,\n        dense_shape=[self._batch_size, minibatch['sentence'].dense_shape[1]])\n    # minibatch.sentence sparse tensor with dense shape\n    # [batch_size x seq_length], length: [batch_size]\n    return Dataset(\n        tokens=minibatch['sentence'],\n        num_tokens=self.get_row_lengths(minibatch['sentence']),\n        sentiment=minibatch['sentiment'],\n    )\n\n  @property\n  def num_examples(self):\n    return self._num_examples\n\n\ndef parse(data_dict):\n  \"\"\"Parse dataset from _data_gen into the same format as sst_binary.\"\"\"\n  sentiment = data_dict['label']\n  sentence = data_dict['sentence']\n  dense_chars = tf.decode_raw(sentence, tf.uint8)\n  dense_chars.set_shape((None,))\n  chars = tfp.math.dense_to_sparse(dense_chars)\n  if six.PY3:\n    safe_chr = lambda c: '?' if c >= 128 else chr(c)\n  else:\n    safe_chr = chr\n  to_char = np.vectorize(safe_chr)\n  chars = tf.SparseTensor(indices=chars.indices,\n                          values=tf.py_func(to_char, [chars.values], tf.string),\n                          dense_shape=chars.dense_shape)\n  return {'sentiment': sentiment,\n          'sentence': chars}\n\n\nclass RobustModel(snt.AbstractModule):\n  \"\"\"Model for applying sentence representations for different tasks.\"\"\"\n\n  def __init__(self,\n               task,\n               batch_size,\n               pooling,\n               learning_rate,\n               config,\n               embedding_dim,\n               fine_tune_embeddings=False,\n               num_oov_buckets=1000,\n               max_grad_norm=5.0,\n               name='robust_model'):\n    super(RobustModel, self).__init__(name=name)\n    self.config = config\n    self.task = task\n    self.batch_size = batch_size\n    self.pooling = pooling\n    self.learning_rate = learning_rate\n    self.embedding_dim = embedding_dim\n    self.fine_tune_embeddings = fine_tune_embeddings\n    self.num_oov_buckets = num_oov_buckets\n    self.max_grad_norm = max_grad_norm\n    self.linear_classifier = None\n\n  def add_representer(self, vocab_filename, padded_token=None):\n    \"\"\"Add sentence representer to the computation graph.\n\n    Args:\n      vocab_filename: the name of vocabulary files.\n      padded_token: padded_token to the vocabulary.\n    \"\"\"\n\n    self.embed_pad = utils.EmbedAndPad(\n        self.batch_size,\n        [self._lines_from_file(vocab_filename)],\n        embedding_dim=self.embedding_dim,\n        num_oov_buckets=self.num_oov_buckets,\n        fine_tune_embeddings=self.fine_tune_embeddings,\n        padded_token=padded_token)\n\n    self.keep_prob = tf.placeholder(tf.float32, shape=None, name='keep_prob')\n\n    # Model to get a sentence representation from embeddings.\n    self.sentence_representer = models.SentenceRepresenterConv(\n        self.config, keep_prob=self.keep_prob, pooling=self.pooling)\n\n  def add_dataset(self):\n    \"\"\"Add datasets.\n\n    Returns:\n      train_data, dev_data, test_data, num_classes\n    \"\"\"\n    if self.config.get('dataset', '') == 'sst':\n      train_data = GeneratedDataset(None, self.batch_size, mode='train',\n                                    num_examples=67349)\n      dev_data = GeneratedDataset(None, self.batch_size, mode='validation',\n                                  num_examples=872)\n      test_data = GeneratedDataset(None, self.batch_size, mode='validation',\n                                   num_examples=872)\n      num_classes = 2\n      return train_data, dev_data, test_data, num_classes\n    else:\n      raise ValueError('Not supported dataset')\n\n  def get_representation(self, tokens, num_tokens):\n    if tokens.dtype == tf.float32:\n      return self.sentence_representer(tokens, num_tokens)\n    else:  # dtype == tf.string\n      return self.sentence_representer(self.embed_pad(tokens), num_tokens)\n\n  def add_representation(self, minibatch):\n    \"\"\"Compute sentence representations.\n\n    Args:\n      minibatch: a minibatch of sequences of embeddings.\n    Returns:\n      joint_rep: representation of sentences or concatenation of\n        sentence vectors.\n    \"\"\"\n\n    joint_rep = self.get_representation(minibatch.tokens, minibatch.num_tokens)\n    result = {'representation1': joint_rep}\n    return joint_rep, result\n\n  def add_train_ops(self,\n                    num_classes,\n                    joint_rep,\n                    minibatch):\n    \"\"\"Add ops for training in the computation graph.\n\n    Args:\n      num_classes: number of classes to predict in the task.\n      joint_rep: the joint sentence representation if the input is sentence\n        pairs or the representation for the sentence if the input is a single\n        sentence.\n      minibatch: a minibatch of sequences of embeddings.\n    Returns:\n      train_accuracy: the accuracy on the training dataset\n      loss: training loss.\n      opt_step: training op.\n    \"\"\"\n    if self.linear_classifier is None:\n      classifier_layers = []\n      classifier_layers.append(snt.Linear(num_classes))\n      self.linear_classifier = snt.Sequential(classifier_layers)\n    logits = self.linear_classifier(joint_rep)\n    # Losses and optimizer.\n    def get_loss(logits, labels):\n      return tf.reduce_mean(\n          tf.nn.sparse_softmax_cross_entropy_with_logits(\n              labels=labels, logits=logits))\n\n    loss = get_loss(logits, minibatch.sentiment)\n    train_accuracy = utils.get_accuracy(logits, minibatch.sentiment)\n    opt_step = self._add_optimize_op(loss)\n    return train_accuracy, loss, opt_step\n\n  def create_perturbation_ops(self, minibatch, synonym_values, vocab_table):\n    \"\"\"Perturb data_batch using synonym_values.\"\"\"\n    data_batch = _pad_fixed(\n        utils.get_padded_indexes(vocab_table, minibatch.tokens,\n                                 self.batch_size), axis=1,\n        padded_length=self.config['max_padded_length'])\n\n    # synonym_values: [vocab_size x max_num_synonyms]\n    # data_batch: [batch_size x seq_length]\n    # [batch_size x seq_length x max_num_synonyms] - synonyms for each token.\n    # Defaults to same word in case of no other synonyms.\n    synonym_ids = tf.gather(synonym_values, data_batch, axis=0)\n\n    # Split along batchsize. Elements shape: [seq_length x max_num_synonyms].\n    synonym_ids_per_example = tf.unstack(synonym_ids, axis=0)\n\n    # Loop across batch.\n    # synonym_ids_this_example shape: [seq_length x max_num_synonyms]\n    sequence_positions_across_batch, values_across_batch = [], []\n    for i_sample, synonym_ids_this_example in enumerate(\n        synonym_ids_per_example):\n      # [num_nonzero, 2]. The rows are pairs of (t,s), where t is an index for\n      # a time step, and s is an index into the max_num_synonyms dimension.\n      nonzero_indices = tf.where(synonym_ids_this_example)\n\n      # shape [num_nonzero]. Corresponding to the entries at nonzero_indices\n      synonym_tokens = tf.gather_nd(params=synonym_ids_this_example,\n                                    indices=nonzero_indices)\n\n      # [num_nonzero] - Of the (t,s) pairs in nonzero_indices, pick only the\n      # time dimension (t), corresponding to perturbation positions in the\n      # sequence.\n      perturbation_positions_this_example = nonzero_indices[:, 0]\n\n      # The main logic is done. Now follows padding to a fixed length of\n      # num_perturbations. However, this cannot be done with 0-padding, as it\n      # would introduce a new (zero) vertex. Instead, we duplicate existing\n      # tokens as perturbations (which have no effect), until we have reached a\n      # total of num_perturbations perturbations. In this case, the padded\n      # tokens are the original tokens from the data_batch. The padded positions\n      # are all the positions (using range) corresponding to the padded tokens.\n\n      # How often seq-length fits into maximum num perturbations\n      padding_multiplier = tf.floordiv(self.config['num_perturbations'],\n                                       tf.cast(minibatch.num_tokens[i_sample],\n                                               tf.int32)) + 1\n\n      # original tokens  # [seq_length]\n      original_tokens = data_batch[i_sample, :minibatch.num_tokens[i_sample]]\n      # [padding_multiplier * seq_length]. Repeat several times, use as padding.\n      padding_tokens = tf.tile(original_tokens, multiples=[padding_multiplier])\n      synonym_tokens_padded = tf.concat([synonym_tokens, tf.cast(padding_tokens,\n                                                                 dtype=tf.int64)\n                                        ], axis=0)\n      # Crop at exact num_perturbations size.\n      synonym_tokens_padded = synonym_tokens_padded[\n          :self.config['num_perturbations']]\n\n      # [seq_length] padding sequence positions with tiles of range()\n      pad_positions = tf.range(minibatch.num_tokens[i_sample], delta=1)\n      # [padding_multiplier*seq_length]\n      padding_positions = tf.tile(pad_positions, multiples=[padding_multiplier])\n      perturbation_positions_this_example_padded = tf.concat(\n          [perturbation_positions_this_example, tf.cast(padding_positions,\n                                                        dtype=tf.int64)],\n          axis=0)\n      # Crop at exact size num_perturbations.\n      sequence_positions_padded = perturbation_positions_this_example_padded[\n          :self.config['num_perturbations']]\n\n      # Collect across the batch for tf.stack later.\n      sequence_positions_across_batch.append(sequence_positions_padded)\n      values_across_batch.append(synonym_tokens_padded)\n\n    # Both [batch_size x max_n_perturbations]\n    perturbation_positions = tf.stack(sequence_positions_across_batch, axis=0)\n    perturbation_tokens = tf.stack(values_across_batch, axis=0)\n\n    # Explicitly setting the shape to self.config['num_perturbations']\n    perturbation_positions_shape = perturbation_positions.shape.as_list()\n    perturbation_positions_shape[1] = self.config['num_perturbations']\n    perturbation_positions.set_shape(perturbation_positions_shape)\n    perturbation_tokens_shape = perturbation_tokens.shape.as_list()\n    perturbation_tokens_shape[1] = self.config['num_perturbations']\n    perturbation_tokens.set_shape(perturbation_tokens_shape)\n\n    return Perturbation(\n        positions=perturbation_positions,\n        tokens=perturbation_tokens)\n\n  def _add_optimize_op(self, loss):\n    \"\"\"Add ops for training.\"\"\"\n    global_step = tf.Variable(0, trainable=False)\n    learning_rate = tf.Variable(self.learning_rate, trainable=False)\n    tvars = tf.trainable_variables()\n    grads, _ = tf.clip_by_global_norm(tf.gradients(loss, tvars),\n                                      self.max_grad_norm)\n    opt = tf.train.AdamOptimizer(learning_rate)\n    opt_step = opt.apply_gradients(zip(grads, tvars),\n                                   global_step=global_step)\n    return opt_step\n\n  def embed_dataset(self, minibatch, vocab_table):\n    return EmbeddedDataset(\n        embedded_inputs=_pad_fixed(\n            self.embed_pad(minibatch.tokens),\n            axis=1,\n            padded_length=self.config['max_padded_length']),\n        input_tokens=_pad_fixed(\n            utils.get_padded_indexes(vocab_table, minibatch.tokens,\n                                     self.batch_size),\n            axis=1,\n            padded_length=self.config['max_padded_length']),\n        length=tf.minimum(self.config['max_padded_length'],\n                          tf.cast(minibatch.num_tokens, tf.int32)),\n        sentiment=minibatch.sentiment)\n\n  def compute_mask_vertices(self, data_batch, perturbation):\n    \"\"\"Compute perturbation masks and perbuted vertices.\n\n    Args:\n      data_batch: EmbeddedDataset object.\n      perturbation: Perturbation object.\n\n    Returns:\n      masks: Positions where there are perturbations.\n      vertices: The resulting embeddings of the perturbed inputs.\n    \"\"\"\n    # The following are all shaped (after broadcasting) as:\n    # (batch_size, num_perturbations, seq_length, embedding_size).\n    embedding = self.embed_pad._embeddings  # pylint: disable=protected-access\n    # (batch_size, 1, seq_length, emb_dim)\n    original_vertices = tf.expand_dims(data_batch.embedded_inputs, axis=1)\n    # (batch_size, num_perturbation, 1, emb_dim])\n    perturbation_vertices = tf.gather(\n        embedding, tf.expand_dims(perturbation.tokens, axis=2))\n    # (batch_size, num_perturbations, seq_length, 1)\n    mask = tf.expand_dims(\n        tf.one_hot(perturbation.positions,\n                   depth=self.config['max_padded_length']), axis=3)\n    # (batch_size, num_perturbations, seq_length, embedding_size)\n    vertices = (1 - mask) * original_vertices + mask * perturbation_vertices\n    return mask, vertices\n\n  def preprocess_databatch(self, minibatch, vocab_table, perturbation):\n    data_batch = self.embed_dataset(minibatch, vocab_table)\n    mask, vertices = self.compute_mask_vertices(data_batch, perturbation)\n    return data_batch, mask, vertices\n\n  def add_verifiable_objective(self,\n                               minibatch,\n                               vocab_table,\n                               perturbation,\n                               stop_gradient=False):\n    # pylint: disable=g-missing-docstring\n    data_batch = self.embed_dataset(minibatch, vocab_table)\n    _, vertices = self.compute_mask_vertices(data_batch, perturbation)\n\n    def classifier(embedded_inputs):\n      representation = self.sentence_representer(embedded_inputs,\n                                                 data_batch.length)\n      return self.linear_classifier(representation)\n\n    # Verification graph.\n    network = ibp.VerifiableModelWrapper(classifier)\n    network(data_batch.embedded_inputs)\n\n    input_bounds = ibp.SimplexBounds(\n        vertices=vertices,\n        nominal=data_batch.embedded_inputs,\n        r=(self.delta if not stop_gradient else self.config['delta']))\n    network.propagate_bounds(input_bounds)\n\n    # Calculate the verifiable objective.\n    verifiable_obj = verifiable_objective(\n        network, data_batch.sentiment, margin=1.)\n\n    return verifiable_obj\n\n  def run_classification(self, inputs, labels, length):\n    prediction = self.run_prediction(inputs, length)\n    correct = tf.cast(tf.equal(labels, tf.argmax(prediction, 1)),\n                      dtype=tf.float32)\n    return correct\n\n  def compute_verifiable_loss(self, verifiable_obj, labels):\n    \"\"\"Compute verifiable training objective.\n\n    Args:\n      verifiable_obj: Verifiable training objective.\n      labels: Ground truth labels.\n    Returns:\n      verifiable_loss: Aggregrated loss of the verifiable training objective.\n    \"\"\"\n    # Three options: reduce max, reduce mean, and softmax.\n    if self.config['verifiable_training_aggregation'] == 'mean':\n      verifiable_loss = tf.reduce_mean(\n          verifiable_obj)  # average across all target labels\n    elif self.config['verifiable_training_aggregation'] == 'max':\n      # Worst target label only.\n      verifiable_loss = tf.reduce_mean(tf.reduce_max(verifiable_obj, axis=0))\n    elif self.config['verifiable_training_aggregation'] == 'softmax':\n      # This assumes that entries in verifiable_obj belonging to the true class\n      # are set to a (large) negative value, so to not affect the softmax much.\n\n      # [batch_size]. Compute x-entropy against one-hot distrib. for true label.\n      verifiable_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(\n          logits=tf.transpose(verifiable_obj), labels=labels)\n\n      verifiable_loss = tf.reduce_mean(\n          verifiable_loss)  # aggregation across batch\n    else:\n      logging.info(self.config['verifiable_training_aggregation'])\n      raise ValueError(\n          'Bad input argument for verifiable_training_aggregation used.')\n\n    return verifiable_loss\n\n  def compute_verifiable_verified(self, verifiable_obj):\n    # Overall upper bound is maximum over all incorrect target classes.\n    bound = tf.reduce_max(verifiable_obj, axis=0)\n    verified = tf.cast(bound <= 0, dtype=tf.float32)\n    return bound, verified\n\n  def run_prediction(self, inputs, length):\n    representation = self.sentence_representer(inputs, length)\n    prediction = self.linear_classifier(representation)\n    return prediction\n\n  def sentiment_accuracy_op(self, minibatch):\n    \"\"\"Compute accuracy of dev/test set on the task of sentiment analysis.\n\n    Args:\n      minibatch: a batch of sequences of embeddings.\n    Returns:\n      num_correct: the number of examples that are predicted correctly on the\n        given dataset.\n    \"\"\"\n\n    rep = self.get_representation(minibatch.tokens, minibatch.num_tokens)\n    logits = self.linear_classifier(rep)\n    num_correct = utils.get_num_correct_predictions(logits,\n                                                    minibatch.sentiment)\n    return num_correct\n\n  def add_dev_eval_ops(self, minibatch):\n    \"\"\"Add ops for evaluating on the dev/test set.\n\n    Args:\n      minibatch: a batch of sequence of embeddings.\n    Returns:\n      num_correct: the number of examples that are predicted correctly.\n    \"\"\"\n    num_correct = self.sentiment_accuracy_op(minibatch)\n    return num_correct\n\n  def _build(self):\n    \"\"\"Build the computation graph.\n\n    Returns:\n      graph_tensors: list of ops that are to be executed during\n        training/evaluation.\n    \"\"\"\n    train_data, dev_data, test_data, num_classes = self.add_dataset()\n    train_minibatch = train_data()\n    dev_minibatch = dev_data()\n    test_minibatch = test_data()\n\n    # Load the vocab without padded_token and add it to the add_representer\n    # later. Otherwise, it will be sorted.\n    vocab_filename = self.config['vocab_filename']\n    self.add_representer(vocab_filename, padded_token=b'<PAD>')\n\n    graph_tensors = self._build_graph_with_datasets(\n        train_minibatch, dev_minibatch, test_minibatch, num_classes)\n    graph_tensors['dev_num_examples'] = dev_data.num_examples\n    graph_tensors['test_num_examples'] = test_data.num_examples\n    return graph_tensors\n\n  def _build_graph_with_datasets(self,\n                                 train_minibatch,\n                                 dev_minibatch,\n                                 test_minibatch,\n                                 num_classes):\n    \"\"\"Returns the training/evaluation ops.\"\"\"\n    self.keep_prob = 1.  # Using literal 1 (not placeholder) skips dropout op.\n    self.sentence_representer._keep_prob = 1.  # pylint:disable=protected-access\n\n    # Build the graph as per the base class.\n    (train_joint_rep, _) = self.add_representation(train_minibatch)\n\n    (train_accuracy,\n     loss,\n     opt_step) = self.add_train_ops(num_classes, train_joint_rep,\n                                    train_minibatch)\n\n    dev_num_correct = self.add_dev_eval_ops(dev_minibatch)\n    test_num_correct = self.add_dev_eval_ops(test_minibatch)\n    graph_tensors = {\n        'loss': loss,\n        'train_op': opt_step,\n        'train_accuracy': train_accuracy,\n        'dev_num_correct': dev_num_correct,\n        'test_num_correct': test_num_correct,\n        'keep_prob': self.keep_prob\n    }\n\n    vocab_table = self.embed_pad.vocab_table\n    vocab_size = self.embed_pad.vocab_size\n\n    verifiable_loss_ratio = tf.constant(\n        self.config['verifiable_loss_ratio'],\n        dtype=tf.float32,\n        name='verifiable_loss_ratio')\n    self.delta = tf.constant(self.config['delta'],\n                             dtype=tf.float32, name='delta')\n\n    lookup_token = tf.placeholder(tf.string, shape=None, name='lookup_token')\n    indices = vocab_table.lookup(lookup_token)\n\n    self.vocab_list = contrib_lookup.index_to_string_table_from_file(\n        self.config['vocab_filename_pad'])\n    lookup_token_index = tf.placeholder(tf.int64, shape=None,\n                                        name='lookup_token_index')\n    lookup_token_string = self.vocab_list.lookup(lookup_token_index)\n\n    synonym_values = tf.placeholder(tf.int64, shape=[None, None],\n                                    name='synonym_values')\n    synonym_counts = tf.placeholder(tf.int64, shape=[None],\n                                    name='synonym_counts')\n\n    train_perturbation = self.create_perturbation_ops(\n        train_minibatch, synonym_values, vocab_table)\n\n    train_data_batch, _, _ = self.preprocess_databatch(\n        train_minibatch, vocab_table, train_perturbation)\n\n    train_words = self.vocab_list.lookup(train_data_batch.input_tokens)\n\n    # [num_targets x batchsize]\n    verifiable_obj = self.add_verifiable_objective(\n        train_minibatch, vocab_table, train_perturbation, stop_gradient=False)\n\n    train_nominal = self.run_classification(train_data_batch.embedded_inputs,\n                                            train_data_batch.sentiment,\n                                            train_data_batch.length)\n    train_bound, train_verified = self.compute_verifiable_verified(\n        verifiable_obj)\n    verifiable_loss = self.compute_verifiable_loss(verifiable_obj,\n                                                   train_minibatch.sentiment)\n\n    if (self.config['verifiable_loss_ratio']) > 1.0:\n      raise ValueError('Loss ratios sum up to more than 1.0')\n\n    total_loss = (1 - verifiable_loss_ratio) * graph_tensors['loss']\n    if self.config['verifiable_loss_ratio'] != 0:\n      total_loss += verifiable_loss_ratio * verifiable_loss\n\n    # Attack on dev/test set.\n    dev_perturbation = self.create_perturbation_ops(\n        dev_minibatch, synonym_values, vocab_table)\n    # [num_targets x batchsize]\n    dev_verifiable_obj = self.add_verifiable_objective(\n        dev_minibatch, vocab_table, dev_perturbation, stop_gradient=True)\n    dev_bound, dev_verified = self.compute_verifiable_verified(\n        dev_verifiable_obj)\n\n    dev_data_batch, _, _ = self.preprocess_databatch(\n        dev_minibatch, vocab_table, dev_perturbation)\n\n    test_perturbation = self.create_perturbation_ops(\n        test_minibatch, synonym_values, vocab_table)\n    # [num_targets x batchsize]\n    test_verifiable_obj = self.add_verifiable_objective(\n        test_minibatch, vocab_table, test_perturbation, stop_gradient=True)\n    test_bound, test_verified = self.compute_verifiable_verified(\n        test_verifiable_obj)\n\n    test_data_batch, _, _ = self.preprocess_databatch(\n        test_minibatch, vocab_table, test_perturbation)\n\n    dev_words = self.vocab_list.lookup(dev_data_batch.input_tokens)\n    test_words = self.vocab_list.lookup(test_data_batch.input_tokens)\n    dev_nominal = self.run_classification(dev_data_batch.embedded_inputs,\n                                          dev_data_batch.sentiment,\n                                          dev_data_batch.length)\n    test_nominal = self.run_classification(test_data_batch.embedded_inputs,\n                                           test_data_batch.sentiment,\n                                           test_data_batch.length)\n\n    dev_predictions = self.run_prediction(dev_data_batch.embedded_inputs,\n                                          dev_data_batch.length)\n    test_predictions = self.run_prediction(test_data_batch.embedded_inputs,\n                                           test_data_batch.length)\n\n    with tf.control_dependencies([train_verified, test_verified, dev_verified]):\n      opt_step = self._add_optimize_op(total_loss)\n\n    graph_tensors['total_loss'] = total_loss\n    graph_tensors['verifiable_loss'] = verifiable_loss\n    graph_tensors['train_op'] = opt_step\n    graph_tensors['indices'] = indices\n    graph_tensors['lookup_token_index'] = lookup_token_index\n    graph_tensors['lookup_token_string'] = lookup_token_string\n    graph_tensors['lookup_token'] = lookup_token\n    graph_tensors['vocab_size'] = vocab_size\n    graph_tensors['synonym_values'] = synonym_values\n    graph_tensors['synonym_counts'] = synonym_counts\n    graph_tensors['verifiable_loss_ratio'] = verifiable_loss_ratio\n    graph_tensors['delta'] = self.delta\n\n    graph_tensors['train'] = {\n        'bound': train_bound,\n        'verified': train_verified,\n        'words': train_words,\n        'sentiment': train_minibatch.sentiment,\n        'correct': train_nominal,\n    }\n    graph_tensors['dev'] = {\n        'predictions': dev_predictions,\n        'data_batch': dev_data_batch,\n        'tokens': dev_minibatch.tokens,\n        'num_tokens': dev_minibatch.num_tokens,\n        'minibatch': dev_minibatch,\n        'bound': dev_bound,\n        'verified': dev_verified,\n        'words': dev_words,\n        'sentiment': dev_minibatch.sentiment,\n        'correct': dev_nominal,\n    }\n    graph_tensors['test'] = {\n        'predictions': test_predictions,\n        'data_batch': test_data_batch,\n        'tokens': test_minibatch.tokens,\n        'num_tokens': test_minibatch.num_tokens,\n        'minibatch': test_minibatch,\n        'bound': test_bound,\n        'verified': test_verified,\n        'words': test_words,\n        'sentiment': test_minibatch.sentiment,\n        'correct': test_nominal,\n    }\n\n    return graph_tensors\n\n  def _lines_from_file(self, filename):\n    with open(filename, 'rb') as f:\n      return f.read().splitlines()\n\n\ndef verifiable_objective(network, labels, margin=0.):\n  \"\"\"Computes the verifiable objective.\n\n  Args:\n    network: `ibp.VerifiableModelWrapper` for the network to verify.\n    labels: 1D integer tensor of shape (batch_size) of labels for each\n      input example.\n    margin: Verifiable objective values for correct class will be forced to\n      `-margin`, thus disregarding large negative bounds when maximising. By\n      default this is set to 0.\n\n  Returns:\n    2D tensor of shape (num_classes, batch_size) containing verifiable objective\n      for each target class, for each example.\n  \"\"\"\n  last_layer = network.output_module\n\n  # Objective, elided with final linear layer.\n  obj_w, obj_b = targeted_objective(\n      last_layer.module.w, last_layer.module.b, labels)\n\n  # Relative bounds on the objective.\n  per_neuron_objective = tf.maximum(\n      obj_w * last_layer.input_bounds.lower_offset,\n      obj_w * last_layer.input_bounds.upper_offset)\n  verifiable_obj = tf.reduce_sum(\n      per_neuron_objective,\n      axis=list(range(2, per_neuron_objective.shape.ndims)))\n\n  # Constant term (objective layer bias).\n  verifiable_obj += tf.reduce_sum(\n      obj_w * last_layer.input_bounds.nominal,\n      axis=list(range(2, obj_w.shape.ndims)))\n  verifiable_obj += obj_b\n\n  # Filter out cases in which the target class is the correct class.\n  # Using `margin` makes the irrelevant cases of target=correct return\n  # a large negative value, which will be ignored by the reduce_max.\n  num_classes = last_layer.output_bounds.shape[-1]\n  verifiable_obj = filter_correct_class(\n      verifiable_obj, num_classes, labels, margin=margin)\n\n  return verifiable_obj\n\n\ndef targeted_objective(final_w, final_b, labels):\n  \"\"\"Determines final layer weights for attacks targeting each class.\n\n  Args:\n    final_w: 2D tensor of shape (last_hidden_layer_size, num_classes)\n      containing the weights for the final linear layer.\n    final_b: 1D tensor of shape (num_classes) containing the biases for the\n      final hidden layer.\n    labels: 1D integer tensor of shape (batch_size) of labels for each\n      input example.\n\n  Returns:\n    obj_w: Tensor of shape (num_classes, batch_size, last_hidden_layer_size)\n      containing weights (to use in place of final linear layer weights)\n      for targeted attacks.\n    obj_b: Tensor of shape (num_classes, batch_size) containing bias\n      (to use in place of final linear layer biases) for targeted attacks.\n  \"\"\"\n  # Elide objective with final linear layer.\n  final_wt = tf.transpose(final_w)\n  obj_w = tf.expand_dims(final_wt, axis=1) - tf.gather(final_wt, labels, axis=0)\n  obj_b = tf.expand_dims(final_b, axis=1) - tf.gather(final_b, labels, axis=0)\n  return obj_w, obj_b\n\n\ndef filter_correct_class(verifiable_obj, num_classes, labels, margin):\n  \"\"\"Filters out the objective when the target class contains the true label.\n\n  Args:\n    verifiable_obj: 2D tensor of shape (num_classes, batch_size) containing\n      verifiable objectives.\n    num_classes: number of target classes.\n    labels: 1D tensor of shape (batch_size) containing the labels for each\n      example in the batch.\n    margin: Verifiable objective values for correct class will be forced to\n      `-margin`, thus disregarding large negative bounds when maximising.\n\n  Returns:\n   2D tensor of shape (num_classes, batch_size) containing the corrected\n   verifiable objective values for each (class, example).\n  \"\"\"\n  targets_to_filter = tf.expand_dims(\n      tf.range(num_classes, dtype=labels.dtype), axis=1)\n  neq = tf.not_equal(targets_to_filter, labels)\n  verifiable_obj = tf.where(neq, verifiable_obj, -margin *\n                            tf.ones_like(verifiable_obj))\n  return verifiable_obj\n"
  },
  {
    "path": "examples/language/robust_train.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Train verifiably robust models.\"\"\"\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport imp\nimport json\nimport os\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\n\nimport numpy as np\nfrom six.moves import range\nimport tensorflow.compat.v1 as tf\n\nimport robust_model\n\n\nflags.DEFINE_string('config_path', 'config.py',\n                    'Path to training configuration file.')\nflags.DEFINE_integer('batch_size', 40, 'Batch size.')\nflags.DEFINE_integer('num_train_steps', 150000, 'Number of training steps.')\nflags.DEFINE_integer('num_oov_buckets', 1,\n                     'Number of out of vocabulary buckets.')\nflags.DEFINE_integer('report_every', 100,\n                     'Report test loss every N batches.')\nflags.DEFINE_float('schedule_ratio', 0.8,\n                   'The final delta and verifiable_loss_ratio are reached when '\n                   'the number of steps equals schedule_ratio * '\n                   'num_train_steps.')\nflags.DEFINE_float('learning_rate', 0.001, 'Learning rate.')\nflags.DEFINE_float('max_grad_norm', 5.0, 'Maximum norm of gradients.')\nflags.DEFINE_boolean('fine_tune_embeddings', True, 'Finetune embeddings.')\nflags.DEFINE_string('task', 'sst', 'One of snli, mnli, sick, sst.')\nflags.DEFINE_string('pooling', 'average', 'One of averge, sum, max, last.')\nflags.DEFINE_boolean('analysis', False, 'Analysis mode.')\nflags.DEFINE_string('analysis_split', 'test', 'Analysis dataset split.')\nflags.DEFINE_string('experiment_root',\n                    '/tmp/robust_model/',\n                    'Path to save trained models.')\nflags.DEFINE_string(\n    'tensorboard_dir', None,\n    'Tensorboard folder. If not specified, set under experiment_root')\nFLAGS = flags.FLAGS\n\n\ndef load_synonyms(synonym_filepath=None):\n  synonyms = None\n  with open(synonym_filepath) as f:\n    synonyms = json.load(f)\n  return synonyms\n\n\ndef construct_synonyms(synonym_filepath):\n  synonyms = load_synonyms(synonym_filepath)\n  synonym_keys = list(synonyms.keys())\n  synonym_values = [synonyms[k] for k in synonym_keys]\n  max_synoynm_counts = max([len(s) for s in synonym_values])\n  synonym_value_lens = [len(x) for x in synonym_values]\n  # Add 0 for the first starting point.\n  synonym_value_lens_cum = np.cumsum([0] + synonym_value_lens)\n  synonym_values_list = [word for val in synonym_values for word in val]  # pylint: disable=g-complex-comprehension\n  return synonym_keys, max_synoynm_counts, synonym_value_lens_cum, synonym_values_list\n\n\ndef linear_schedule(step, init_step, final_step, init_value, final_value):\n  \"\"\"Linear schedule.\"\"\"\n  assert final_step >= init_step\n  if init_step == final_step:\n    return final_value\n  rate = np.float32(step - init_step) / float(final_step - init_step)\n  linear_value = rate * (final_value - init_value) + init_value\n  return np.clip(linear_value, min(init_value, final_value),\n                 max(init_value, final_value))\n\n\ndef config_train_summary(task, train_accuracy, loss):\n  \"\"\"Add ops for summary in the computation graph.\n\n  Args:\n    task: string name of task being trained for.\n    train_accuracy: training accuracy.\n    loss: training loss.\n\n  Returns:\n    train_summary: summary for training.\n    saver: tf.saver, used to save the checkpoint with the best dev accuracy.\n  \"\"\"\n  train_acc_summ = tf.summary.scalar(('%s_train_accuracy' % task),\n                                     train_accuracy)\n  loss_summ = tf.summary.scalar('loss', loss)\n  train_summary = tf.summary.merge([train_acc_summ, loss_summ])\n  return train_summary\n\n\ndef write_tf_summary(writer, step, tag, value):\n  summary = tf.Summary()\n  summary.value.add(tag=tag, simple_value=value)\n  writer.add_summary(summary, step)\n\n\ndef train(config_dict, synonym_filepath,\n          batch_size, num_train_steps, schedule_ratio, report_every,\n          checkpoint_path, tensorboard_dir):\n  \"\"\"Model training.\"\"\"\n  graph_tensor_producer = robust_model.RobustModel(**config_dict)\n  graph_tensors = graph_tensor_producer()\n\n  synonym_keys, max_synoynm_counts, synonym_value_lens_cum, \\\n      synonym_values_list = construct_synonyms(synonym_filepath)\n  train_summary = config_train_summary(config_dict['task'],\n                                       graph_tensors['train_accuracy'],\n                                       graph_tensors['loss'])\n\n  tf.gfile.MakeDirs(checkpoint_path)\n\n  best_dev_accuracy = 0.0\n  best_test_accuracy = 0.0\n  best_verified_dev_accuracy = 0.0\n  best_verified_test_accuracy = 0.0\n\n  network_saver = tf.train.Saver(graph_tensor_producer.variables)\n  with tf.train.SingularMonitoredSession() as session:\n    logging.info('Initialize parameters...')\n    writer = tf.summary.FileWriter(tensorboard_dir, session.graph)\n    input_feed = {}\n\n    # Tokenize synonyms.\n    tokenize_synonyms = [[] for _ in range(graph_tensors['vocab_size'])]\n    lookup_indices_keys = session.run(graph_tensors['indices'],\n                                      feed_dict={graph_tensors['lookup_token']:\n                                                 synonym_keys})\n    lookup_indices_values = session.run(graph_tensors['indices'],\n                                        feed_dict={\n                                            graph_tensors['lookup_token']:\n                                            synonym_values_list})\n    for i, key_index in enumerate(lookup_indices_keys):\n      tokenize_synonyms[key_index] = lookup_indices_values[\n          synonym_value_lens_cum[i]:synonym_value_lens_cum[i+1]].tolist()\n\n    synonym_values_np = np.zeros([graph_tensors['vocab_size'],\n                                  max_synoynm_counts])\n    for i in range(graph_tensors['vocab_size']):\n      # False-safe case. No perturbations. Set it as itself.\n      synonym_values_np[i][0] = i\n      for j in range(len(tokenize_synonyms[i])):\n        synonym_values_np[i][j] = tokenize_synonyms[i][j]\n    synonym_counts_np = [len(s) for s in tokenize_synonyms]\n    input_feed[graph_tensors['synonym_values']] = synonym_values_np\n    input_feed[graph_tensors['synonym_counts']] = synonym_counts_np\n\n    warmup_steps = 0\n    for step in range(num_train_steps):\n      config = config_dict['config']\n      if config['delta'] > 0.0 and config['delta_schedule']:\n        delta = linear_schedule(\n            step, 0., schedule_ratio * num_train_steps,\n            0., config['delta'])\n        input_feed[graph_tensors['delta']] = delta\n\n      if (config['verifiable_loss_ratio'] > 0.0 and\n          config['verifiable_loss_schedule']):\n        if delta > 0.0 and warmup_steps == 0:\n          warmup_steps = step\n        if delta > 0.0:\n          verifiable_loss_ratio = linear_schedule(\n              step, warmup_steps, schedule_ratio * num_train_steps,\n              0., config['verifiable_loss_ratio'])\n        else:\n          verifiable_loss_ratio = 0.0\n        input_feed[\n            graph_tensors['verifiable_loss_ratio']] = verifiable_loss_ratio\n\n      total_loss_np, loss_np, verifiable_loss_np, train_accuracy_np, \\\n          train_bound, train_verified, \\\n          verifiable_loss_ratio_val, delta_val, \\\n          train_summary_py, _ = session.run(\n              [graph_tensors['total_loss'],\n               graph_tensors['loss'],\n               graph_tensors['verifiable_loss'],\n               graph_tensors['train_accuracy'],\n               graph_tensors['train']['bound'],\n               graph_tensors['train']['verified'],\n               graph_tensors['verifiable_loss_ratio'],\n               graph_tensors['delta'],\n               train_summary,\n               graph_tensors['train_op']], input_feed)\n\n      writer.add_summary(train_summary_py, step)\n      if step % report_every == 0 or step == num_train_steps - 0:\n        dev_total_num_correct = 0.0\n        test_total_num_correct = 0.0\n        dev_verified_count = 0.0\n        test_verified_count = 0.0\n        dev_num_batches = graph_tensors['dev_num_examples'] // batch_size\n        test_num_batches = graph_tensors['test_num_examples'] // batch_size\n        dev_total_num_examples = dev_num_batches * batch_size\n        test_total_num_examples = test_num_batches * batch_size\n        for _ in range(dev_num_batches):\n          correct, verified = session.run(\n              [graph_tensors['dev_num_correct'],\n               graph_tensors['dev']['verified']], input_feed)\n          dev_total_num_correct += correct\n          dev_verified_count += np.sum(verified)\n        for _ in range(test_num_batches):\n          correct, verified = session.run(\n              [graph_tensors['test_num_correct'],\n               graph_tensors['test']['verified']], input_feed)\n          test_total_num_correct += correct\n          test_verified_count += np.sum(verified)\n        dev_accuracy = dev_total_num_correct / dev_total_num_examples\n        test_accuracy = test_total_num_correct / test_total_num_examples\n        dev_verified_accuracy = dev_verified_count / dev_total_num_examples\n        test_verified_accuracy = test_verified_count / test_total_num_examples\n\n        write_tf_summary(writer, step, tag='dev_accuracy', value=dev_accuracy)\n        write_tf_summary(writer, step, tag='test_accuracy', value=test_accuracy)\n        write_tf_summary(writer, step, tag='train_bound_summary',\n                         value=np.mean(train_bound))\n        write_tf_summary(writer, step, tag='train_verified_summary',\n                         value=np.mean(train_verified))\n        write_tf_summary(writer, step, tag='dev_verified_summary',\n                         value=np.mean(dev_verified_accuracy))\n        write_tf_summary(writer, step, tag='test_verified_summary',\n                         value=np.mean(test_verified_accuracy))\n        write_tf_summary(writer, step, tag='total_loss_summary',\n                         value=total_loss_np)\n        write_tf_summary(writer, step, tag='verifiable_train_loss_summary',\n                         value=verifiable_loss_np)\n\n        logging.info('verifiable_loss_ratio: %f, delta: %f',\n                     verifiable_loss_ratio_val, delta_val)\n        logging.info('step: %d, '\n                     'train loss: %f, '\n                     'verifiable train loss: %f, '\n                     'train accuracy: %f, '\n                     'dev accuracy: %f, '\n                     'test accuracy: %f, ', step, loss_np,\n                     verifiable_loss_np, train_accuracy_np,\n                     dev_accuracy, test_accuracy)\n        dev_verified_accuracy_mean = np.mean(dev_verified_accuracy)\n        test_verified_accuracy_mean = np.mean(test_verified_accuracy)\n        logging.info('Train Bound = %.05f, train verified: %.03f, '\n                     'dev verified: %.03f, test verified: %.03f',\n                     np.mean(train_bound),\n                     np.mean(train_verified), dev_verified_accuracy_mean,\n                     test_verified_accuracy_mean)\n        if dev_accuracy > best_dev_accuracy:\n          # Store most accurate model so far.\n          network_saver.save(session.raw_session(),\n                             os.path.join(checkpoint_path, 'best'))\n          best_dev_accuracy = dev_accuracy\n          best_test_accuracy = test_accuracy\n        logging.info('best dev acc\\t%f\\tbest test acc\\t%f',\n                     best_dev_accuracy, best_test_accuracy)\n        if dev_verified_accuracy_mean > best_verified_dev_accuracy:\n          # Store model with best verified accuracy so far.\n          network_saver.save(session.raw_session(),\n                             os.path.join(checkpoint_path, 'best_verified'))\n          best_verified_dev_accuracy = dev_verified_accuracy_mean\n          best_verified_test_accuracy = test_verified_accuracy_mean\n        logging.info('best verified dev acc\\t%f\\tbest verified test acc\\t%f',\n                     best_verified_dev_accuracy, best_verified_test_accuracy)\n\n        network_saver.save(session.raw_session(),\n                           os.path.join(checkpoint_path, 'model'))\n        writer.flush()\n\n    # Store model at end of training.\n    network_saver.save(session.raw_session(),\n                       os.path.join(checkpoint_path, 'final'))\n\n\ndef analysis(config_dict, synonym_filepath,\n             model_location, batch_size, batch_offset=0,\n             total_num_batches=0, datasplit='test', delta=3.0,\n             num_perturbations=5, max_padded_length=0):\n  \"\"\"Run analysis.\"\"\"\n  tf.reset_default_graph()\n  if datasplit not in ['train', 'dev', 'test']:\n    raise ValueError('Invalid datasplit: %s' % datasplit)\n  logging.info('model_location: %s', model_location)\n  logging.info('num_perturbations: %d', num_perturbations)\n  logging.info('delta: %f', delta)\n\n  logging.info('Run analysis, datasplit: %s, batch %d', datasplit, batch_offset)\n  synonym_keys, max_synoynm_counts, synonym_value_lens_cum, \\\n      synonym_values_list = construct_synonyms(synonym_filepath)\n\n  graph_tensor_producer = robust_model.RobustModel(**config_dict)\n  # Use new batch size.\n  graph_tensor_producer.batch_size = batch_size\n  # Overwrite the config originally in the saved checkpoint.\n  logging.info('old delta %f, old num_perturbations: %d',\n               graph_tensor_producer.config['delta'],\n               graph_tensor_producer.config['num_perturbations'])\n  graph_tensor_producer.config['delta'] = delta\n  graph_tensor_producer.config['num_perturbations'] = num_perturbations\n  if max_padded_length > 0:\n    graph_tensor_producer.config['max_padded_length'] = max_padded_length\n\n  logging.info('new delta %f, num_perturbations: %d, max_padded_length: %d',\n               graph_tensor_producer.config['delta'],\n               graph_tensor_producer.config['num_perturbations'],\n               graph_tensor_producer.config['max_padded_length'])\n  logging.info('graph_tensors.config: %s', graph_tensor_producer.config)\n\n  graph_tensors = graph_tensor_producer()\n  network_saver = tf.train.Saver(graph_tensor_producer.variables)\n  with tf.train.SingularMonitoredSession() as session:\n    network_saver.restore(session.raw_session(), model_location)\n\n    for _ in range(batch_offset):\n      # Seek to the correct batch.\n      session.run(graph_tensors[datasplit]['sentiment'])\n\n    input_feed = {}\n    # Tokenize synonyms.\n    tokenize_synonyms = [[] for _ in range(graph_tensors['vocab_size'])]\n    lookup_indices_keys = session.run(graph_tensors['indices'],\n                                      feed_dict={graph_tensors['lookup_token']:\n                                                     synonym_keys})\n    lookup_indices_values = session.run(graph_tensors['indices'],\n                                        feed_dict={\n                                            graph_tensors['lookup_token']:\n                                            synonym_values_list})\n    for i, key_index in enumerate(lookup_indices_keys):\n      tokenize_synonyms[key_index] = lookup_indices_values[\n          synonym_value_lens_cum[i]:synonym_value_lens_cum[i+1]].tolist()\n\n    synonym_values_np = np.zeros([graph_tensors['vocab_size'],\n                                  max_synoynm_counts])\n    for i in range(graph_tensors['vocab_size']):\n      # False-safe case. No perturbations. Set it as itself.\n      synonym_values_np[i][0] = i\n      for j in range(len(tokenize_synonyms[i])):\n        synonym_values_np[i][j] = tokenize_synonyms[i][j]\n    synonym_counts_np = [len(s) for s in tokenize_synonyms]\n    input_feed[graph_tensors['synonym_values']] = synonym_values_np\n    input_feed[graph_tensors['synonym_counts']] = synonym_counts_np\n\n    total_num_batches = (\n        graph_tensors['%s_num_examples' % datasplit] //\n        batch_size) if total_num_batches == 0 else total_num_batches\n    total_num_examples = total_num_batches * batch_size\n    logging.info('total number of examples  %d', total_num_examples)\n    logging.info('total number of batches  %d', total_num_batches)\n\n    total_correct, total_verified = 0.0, 0.0\n    for ibatch in range(total_num_batches):\n      results = session.run(graph_tensors[datasplit], input_feed)\n      logging.info('batch: %d, %s bound = %.05f, verified: %.03f,'\n                   ' nominally correct: %.03f',\n                   ibatch, datasplit, np.mean(results['bound']),\n                   np.mean(results['verified']),\n                   np.mean(results['correct']))\n      total_correct += sum(results['correct'])\n      total_verified += sum(results['verified'])\n\n    total_correct /= total_num_examples\n    total_verified /= total_num_examples\n    logging.info('%s final correct: %.03f, verified: %.03f',\n                 datasplit, total_correct, total_verified)\n    logging.info({\n        'datasplit': datasplit,\n        'nominal': total_correct,\n        'verify': total_verified,\n        'delta': delta,\n        'num_perturbations': num_perturbations,\n        'model_location': model_location,\n        'final': True\n    })\n\n\ndef main(_):\n  # Read the config file into a new ad-hoc module.\n  with open(FLAGS.config_path, 'r') as config_file:\n    config_code = config_file.read()\n    config_module = imp.new_module('config')\n    exec(config_code, config_module.__dict__)  # pylint: disable=exec-used\n  config = config_module.get_config()\n\n  config_dict = {'task': FLAGS.task,\n                 'batch_size': FLAGS.batch_size,\n                 'pooling': FLAGS.pooling,\n                 'learning_rate': FLAGS.learning_rate,\n                 'config': config,\n                 'embedding_dim': config['embedding_dim'],\n                 'fine_tune_embeddings': FLAGS.fine_tune_embeddings,\n                 'num_oov_buckets': FLAGS.num_oov_buckets,\n                 'max_grad_norm': FLAGS.max_grad_norm}\n\n  if FLAGS.analysis:\n    logging.info('Analyze model location: %s', config['model_location'])\n    base_batch_offset = 0\n    analysis(config_dict, config['synonym_filepath'], config['model_location'],\n             FLAGS.batch_size, base_batch_offset,\n             0, datasplit=FLAGS.analysis_split,\n             delta=config['delta'],\n             num_perturbations=config['num_perturbations'],\n             max_padded_length=config['max_padded_length'])\n\n  else:\n    checkpoint_path = os.path.join(FLAGS.experiment_root, 'checkpoint')\n\n    if FLAGS.tensorboard_dir is None:\n      tensorboard_dir = os.path.join(FLAGS.experiment_root, 'tensorboard')\n    else:\n      tensorboard_dir = FLAGS.tensorboard_dir\n\n    train(config_dict, config['synonym_filepath'],\n          FLAGS.batch_size,\n          num_train_steps=FLAGS.num_train_steps,\n          schedule_ratio=FLAGS.schedule_ratio,\n          report_every=FLAGS.report_every,\n          checkpoint_path=checkpoint_path,\n          tensorboard_dir=tensorboard_dir)\n\n\nif __name__ == '__main__':\n  app.run(main)\n"
  },
  {
    "path": "examples/language/utils.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utilities for sentence representation.\"\"\"\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport tempfile\n\nfrom absl import logging\nimport sonnet as snt\nimport tensorflow as tf\nfrom tensorflow.contrib import lookup as contrib_lookup\n\n\ndef get_padded_embeddings(embeddings,\n                          vocabulary_table,\n                          tokens, batch_size,\n                          token_indexes=None):\n  \"\"\"Reshapes and pads 'raw' word embeddings.\n\n  Say we have batch of B tokenized sentences, of variable length, with a total\n  of W tokens. For example, B = 2 and W = 3 + 4 = 7:\n  [['The',   'cat', 'eats'],\n   [  'A', 'black',  'cat', 'jumps']]\n\n  Since rows have variable length, this cannot be represented as a tf.Tensor.\n  It is represented as a tf.SparseTensor, with 7 values & indexes:\n  indices: [[0,0], [0,1], [0,2], [1,0], [1,1], [1,2], [1,3]]\n  values: ['The', 'cat', 'eats', 'A', 'black', 'cat', 'jumps']\n\n  We have also built a vocabulary table:\n  vocabulary table: ['cat', 'The', 'A', 'black', 'eats', 'jumps']\n\n  We also have the embeddings, a WxD matrix of floats\n  representing each word in the vocabulary table as a normal tf.Tensor.\n\n  For example, with D=3, embeddings could be:\n  [[0.4, 0.5, -0.6],  # This is the embedding for word 0 = 'cat'\n   [0.1, -0.3, 0.6],  # This is the embedding for word 1 = 'The''\n   [0.7, 0.8, -0.9],  # This is the embedding for word 2 = 'A'\n   [-0.1, 0.9, 0.7],  # This is the embedding for word 3 = 'black'\n   [-0.2, 0.4, 0.7],  # This is the embedding for word 4 = 'eats\n   [0.3, -0.5, 0.2]]  # This is the embedding for word 5 = 'jumps'\n\n  This function builds a normal tf.Tensor containing the embeddings for the\n  tokens provided, in the correct order, with appropriate 0 padding.\n\n  In our example, the returned tensor would be:\n  [[[0.1, -0.3, 0.6], [0.4, 0.5, -0.6], [-0.2, 0.4, 0.7], [0.0, 0.0, 0.0]],\n   [[0.7, 0.8, -0.9], [-0.1, 0.9, 0.7], [0.4, 0.5, -0.6], [0.3, -0.5, 0.2]]]\n\n  Note that since the first sentence has only 3 words, the 4th embedding gets\n  replaced by a D-dimensional vector of 0.\n\n  Args:\n    embeddings: [W, D] Tensor of floats, containing the embeddings, initialized\n        with the same vocabulary file as vocabulary_table.\n    vocabulary_table: a tf.contrib.lookup.LookupInterface,\n        containing the vocabulary, initialized with the same vocabulary file as\n        embeddings.\n    tokens: [B, ?] SparseTensor of strings, the tokens.\n    batch_size: Python integer.\n    token_indexes: A Boolean, indicating whether the input tokens are\n        token ids or string.\n\n  Returns:\n    [B, L, D] Tensor of floats: the embeddings in the correct order,\n    appropriately padded with 0.0, where L = max(num_tokens) and B = batch_size\n  \"\"\"\n  embedding_dim = embeddings.get_shape()[1].value  # D in docstring above.\n  num_tokens_in_batch = tf.shape(tokens.indices)[0]  # W in the docstring above.\n  max_length = tokens.dense_shape[1]  # This is L in the docstring above.\n\n  # Get indices of tokens in vocabulary_table.\n  if token_indexes is not None:\n    indexes = token_indexes\n  else:\n    indexes = vocabulary_table.lookup(tokens.values)\n\n  # Get word embeddings.\n  tokens_embeddings = tf.gather(embeddings, indexes)\n\n  # Shape of the return tensor.\n  new_shape = tf.cast(\n      tf.stack([batch_size, max_length, embedding_dim], axis=0), tf.int32)\n\n  # Build the vector of indices for the return Tensor.\n  # In the example above, indices_final would be:\n  # [[[0,0,0], [0,0,1], [0,0,2]],\n  #  [[0,1,0], [0,1,1], [0,1,2]],\n  #  [[0,2,0], [0,2,1], [0,2,2]],\n  #  [[1,0,0], [1,0,1], [1,0,2]],\n  #  [[1,1,0], [1,1,1], [1,1,2]],\n  #  [[1,2,0], [1,2,1], [1,2,2]],\n  #  [[1,3,0], [1,3,1], [1,3,2]]]\n  tiled = tf.tile(tokens.indices, [1, embedding_dim])\n  indices_tiled = tf.cast(\n      tf.reshape(tiled, [num_tokens_in_batch * embedding_dim, 2]), tf.int32)\n  indices_linear = tf.expand_dims(\n      tf.tile(tf.range(0, embedding_dim), [num_tokens_in_batch]), axis=1)\n  indices_final = tf.concat([indices_tiled, indices_linear], axis=1)\n\n  # Build the dense Tensor.\n  embeddings_padded = tf.sparse_to_dense(\n      sparse_indices=indices_final,\n      output_shape=new_shape,\n      sparse_values=tf.reshape(tokens_embeddings,\n                               [num_tokens_in_batch * embedding_dim]))\n  embeddings_padded.set_shape((batch_size, None, embedding_dim))\n\n  return embeddings_padded\n\n\ndef get_padded_indexes(vocabulary_table,\n                       tokens, batch_size,\n                       token_indexes=None):\n  \"\"\"Get the indices of tokens from vocabulary table.\n\n  Args:\n    vocabulary_table: a tf.contrib.lookup.LookupInterface,\n        containing the vocabulary, initialized with the same vocabulary file as\n        embeddings.\n    tokens: [B, ?] SparseTensor of strings, the tokens.\n    batch_size: Python integer.\n    token_indexes: A Boolean, indicating whether the input tokens are\n        token ids or string.\n\n  Returns:\n    [B, L] Tensor of integers: indices of tokens in the correct order,\n    appropriately padded with 0, where L = max(num_tokens) and B = batch_size\n  \"\"\"\n  num_tokens_in_batch = tf.shape(tokens.indices)[0]\n  max_length = tokens.dense_shape[1]\n\n  # Get indices of tokens in vocabulary_table.\n  if token_indexes is not None:\n    indexes = token_indexes\n  else:\n    indexes = vocabulary_table.lookup(tokens.values)\n\n  # Build the dense Tensor.\n  indexes_padded = tf.sparse_to_dense(\n      sparse_indices=tokens.indices,\n      output_shape=[batch_size, max_length],\n      sparse_values=tf.reshape(indexes,\n                               [num_tokens_in_batch]))\n  indexes_padded.set_shape((batch_size, None))\n\n  return indexes_padded\n\n\nclass EmbedAndPad(snt.AbstractModule):\n  \"\"\"Embed and pad tokenized words.\n\n  This class primary functionality is similar to get_padded_embeddings.\n  It stores references to the embeddings and vocabulary table for convenience,\n  so that the user does not have to keep and pass them around.\n  \"\"\"\n\n  def __init__(self,\n               batch_size,\n               vocabularies,\n               embedding_dim,\n               num_oov_buckets=1000,\n               fine_tune_embeddings=False,\n               padded_token=None,\n               name='embed_and_pad'):\n    super(EmbedAndPad, self).__init__(name=name)\n    self._batch_size = batch_size\n    vocab_file, vocab_size = get_merged_vocabulary_file(vocabularies,\n                                                        padded_token)\n    self._vocab_size = vocab_size\n    self._num_oov_buckets = num_oov_buckets\n\n    # Load vocabulary table for index lookup.\n    self._vocabulary_table = contrib_lookup.index_table_from_file(\n        vocabulary_file=vocab_file,\n        num_oov_buckets=num_oov_buckets,\n        vocab_size=self._vocab_size)\n\n    def create_initializer(initializer_range=0.02):\n      \"\"\"Creates a `truncated_normal_initializer` with the given range.\"\"\"\n      # The default value is chosen from language/bert/modeling.py.\n      return tf.truncated_normal_initializer(stddev=initializer_range)\n\n    self._embeddings = tf.get_variable('embeddings_matrix',\n                                       [self._vocab_size + num_oov_buckets,\n                                        embedding_dim],\n                                       trainable=fine_tune_embeddings,\n                                       initializer=create_initializer())\n\n  def _build(self, tokens):\n    padded_embeddings = get_padded_embeddings(\n        self._embeddings, self._vocabulary_table, tokens, self._batch_size)\n    return padded_embeddings\n\n  @property\n  def vocab_table(self):\n    return self._vocabulary_table\n\n  @property\n  def vocab_size(self):\n    return self._vocab_size + self._num_oov_buckets\n\n\ndef get_accuracy(logits, labels):\n  \"\"\"Top 1 accuracy from logits and labels.\"\"\"\n  return tf.reduce_mean(tf.cast(tf.nn.in_top_k(logits, labels, 1), tf.float32))\n\n\ndef get_num_correct_predictions(logits, labels):\n  \"\"\"Get the number of correct predictions over a batch.\"\"\"\n  predictions = tf.cast(tf.argmax(logits, axis=1), tf.int64)\n  evals = tf.equal(predictions, labels)\n  num_correct = tf.reduce_sum(tf.cast(evals, tf.float64))\n  return num_correct\n\n\ndef get_merged_vocabulary_file(vocabularies, padded_token=None):\n  \"\"\"Merges several vocabulary files into one temporary file.\n\n  The TF object that loads the embedding expects a vocabulary file, to know\n  which embeddings it should load.\n  See tf.contrib.embedding.load_embedding_initializer.\n\n  When we want to train/test on several datasets simultaneously we need to merge\n  their vocabulary files into a single file.\n\n  Args:\n    vocabularies: Iterable of vocabularies. Each vocabulary should be\n        a list of tokens.\n    padded_token: If not None, add the padded_token to the first index.\n  Returns:\n    outfilename: Name of the merged file. Contains the union of all tokens in\n        filenames, without duplicates, one token per line.\n    vocabulary_size: Count of tokens in the merged file.\n  \"\"\"\n  uniques = [set(vocabulary) for vocabulary in vocabularies]\n  unique_merged = frozenset().union(*uniques)\n  unique_merged_sorted = sorted(unique_merged)\n  if padded_token is not None:\n    # Add padded token as 0 index.\n    unique_merged_sorted = [padded_token] + unique_merged_sorted\n  vocabulary_size = len(unique_merged_sorted)\n  outfile = tempfile.NamedTemporaryFile(delete=False)\n  outfile.write(b'\\n'.join(unique_merged_sorted))\n  outfilename = outfile.name\n  logging.info('Merged vocabulary file with %d tokens: %s', vocabulary_size,\n               outfilename)\n  outfile.close()\n  return outfilename, vocabulary_size\n"
  },
  {
    "path": "examples/train.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Trains a verifiable model on Mnist or CIFAR-10.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport os\n\nfrom absl import app\nfrom absl import flags\nfrom absl import logging\nimport interval_bound_propagation as ibp\nimport tensorflow.compat.v1 as tf\n\n\nFLAGS = flags.FLAGS\nflags.DEFINE_enum('dataset', 'mnist', ['mnist', 'cifar10'],\n                  'Dataset (either \"mnist\" or \"cifar10\").')\nflags.DEFINE_enum('model', 'tiny', ['tiny', 'small', 'medium', 'large'],\n                  'Model size.')\nflags.DEFINE_string('output_dir', '/tmp/ibp_model', 'Output directory.')\n\n# Options.\nflags.DEFINE_integer('steps', 60001, 'Number of steps in total.')\nflags.DEFINE_integer('test_every_n', 2000,\n                     'Number of steps between testing iterations.')\nflags.DEFINE_integer('warmup_steps', 2000, 'Number of warm-up steps.')\nflags.DEFINE_integer('rampup_steps', 10000, 'Number of ramp-up steps.')\nflags.DEFINE_integer('batch_size', 200, 'Batch size.')\nflags.DEFINE_float('epsilon', .3, 'Target epsilon.')\nflags.DEFINE_float('epsilon_train', .33, 'Train epsilon.')\nflags.DEFINE_string('learning_rate', '1e-3,1e-4@15000,1e-5@25000',\n                    'Learning rate schedule of the form: '\n                    'initial_learning_rate[,learning:steps]*. E.g., \"1e-3\" or '\n                    '\"1e-3,1e-4@15000,1e-5@25000\".')\nflags.DEFINE_float('nominal_xent_init', 1.,\n                   'Initial weight for the nominal cross-entropy.')\nflags.DEFINE_float('nominal_xent_final', .5,\n                   'Final weight for the nominal cross-entropy.')\nflags.DEFINE_float('verified_xent_init', 0.,\n                   'Initial weight for the verified cross-entropy.')\nflags.DEFINE_float('verified_xent_final', .5,\n                   'Final weight for the verified cross-entropy.')\nflags.DEFINE_float('crown_bound_init', 0.,\n                   'Initial weight for mixing the CROWN bound with the IBP '\n                   'bound in the verified cross-entropy.')\nflags.DEFINE_float('crown_bound_final', 0.,\n                   'Final weight for mixing the CROWN bound with the IBP '\n                   'bound in the verified cross-entropy.')\nflags.DEFINE_float('attack_xent_init', 0.,\n                   'Initial weight for the attack cross-entropy.')\nflags.DEFINE_float('attack_xent_final', 0.,\n                   'Initial weight for the attack cross-entropy.')\n\n\ndef show_metrics(step_value, metric_values, loss_value=None):\n  print('{}: {}nominal accuracy = {:.2f}%, '\n        'verified = {:.2f}%, attack = {:.2f}%'.format(\n            step_value,\n            'loss = {}, '.format(loss_value) if loss_value is not None else '',\n            metric_values.nominal_accuracy * 100.,\n            metric_values.verified_accuracy * 100.,\n            metric_values.attack_accuracy * 100.))\n\n\ndef layers(model_size):\n  \"\"\"Returns the layer specification for a given model name.\"\"\"\n  if model_size == 'tiny':\n    return (\n        ('linear', 100),\n        ('activation', 'relu'))\n  elif model_size == 'small':\n    return (\n        ('conv2d', (4, 4), 16, 'VALID', 2),\n        ('activation', 'relu'),\n        ('conv2d', (4, 4), 32, 'VALID', 1),\n        ('activation', 'relu'),\n        ('linear', 100),\n        ('activation', 'relu'))\n  elif model_size == 'medium':\n    return (\n        ('conv2d', (3, 3), 32, 'VALID', 1),\n        ('activation', 'relu'),\n        ('conv2d', (4, 4), 32, 'VALID', 2),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 64, 'VALID', 1),\n        ('activation', 'relu'),\n        ('conv2d', (4, 4), 64, 'VALID', 2),\n        ('activation', 'relu'),\n        ('linear', 512),\n        ('activation', 'relu'),\n        ('linear', 512),\n        ('activation', 'relu'))\n  elif model_size == 'large':\n    return (\n        ('conv2d', (3, 3), 64, 'SAME', 1),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 64, 'SAME', 1),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 128, 'SAME', 2),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 128, 'SAME', 1),\n        ('activation', 'relu'),\n        ('conv2d', (3, 3), 128, 'SAME', 1),\n        ('activation', 'relu'),\n        ('linear', 512),\n        ('activation', 'relu'))\n  else:\n    raise ValueError('Unknown model: \"{}\"'.format(model_size))\n\n\ndef main(unused_args):\n  logging.info('Training IBP on %s...', FLAGS.dataset.upper())\n  step = tf.train.get_or_create_global_step()\n\n  # Learning rate.\n  learning_rate = ibp.parse_learning_rate(step, FLAGS.learning_rate)\n\n  # Dataset.\n  input_bounds = (0., 1.)\n  num_classes = 10\n  if FLAGS.dataset == 'mnist':\n    data_train, data_test = tf.keras.datasets.mnist.load_data()\n  else:\n    assert FLAGS.dataset == 'cifar10', (\n        'Unknown dataset \"{}\"'.format(FLAGS.dataset))\n    data_train, data_test = tf.keras.datasets.cifar10.load_data()\n    data_train = (data_train[0], data_train[1].flatten())\n    data_test = (data_test[0], data_test[1].flatten())\n  data = ibp.build_dataset(data_train, batch_size=FLAGS.batch_size,\n                           sequential=False)\n  if FLAGS.dataset == 'cifar10':\n    data = data._replace(image=ibp.randomize(\n        data.image, (32, 32, 3), expand_shape=(40, 40, 3),\n        crop_shape=(32, 32, 3), vertical_flip=True))\n\n  # Base predictor network.\n  original_predictor = ibp.DNN(num_classes, layers(FLAGS.model))\n  predictor = original_predictor\n  if FLAGS.dataset == 'cifar10':\n    mean = (0.4914, 0.4822, 0.4465)\n    std = (0.2023, 0.1994, 0.2010)\n    predictor = ibp.add_image_normalization(original_predictor, mean, std)\n  if FLAGS.crown_bound_init > 0 or FLAGS.crown_bound_final > 0:\n    logging.info('Using CROWN-IBP loss.')\n    model_wrapper = ibp.crown.VerifiableModelWrapper\n    loss_helper = ibp.crown.create_classification_losses\n  else:\n    model_wrapper = ibp.VerifiableModelWrapper\n    loss_helper = ibp.create_classification_losses\n  predictor = model_wrapper(predictor)\n\n  # Training.\n  train_losses, train_loss, _ = loss_helper(\n      step,\n      data.image,\n      data.label,\n      predictor,\n      FLAGS.epsilon_train,\n      loss_weights={\n          'nominal': {\n              'init': FLAGS.nominal_xent_init,\n              'final': FLAGS.nominal_xent_final,\n              'warmup': FLAGS.verified_xent_init + FLAGS.nominal_xent_init\n          },\n          'attack': {\n              'init': FLAGS.attack_xent_init,\n              'final': FLAGS.attack_xent_final\n          },\n          'verified': {\n              'init': FLAGS.verified_xent_init,\n              'final': FLAGS.verified_xent_final,\n              'warmup': 0.\n          },\n          'crown_bound': {\n              'init': FLAGS.crown_bound_init,\n              'final': FLAGS.crown_bound_final,\n              'warmup': 0.\n          },\n      },\n      warmup_steps=FLAGS.warmup_steps,\n      rampup_steps=FLAGS.rampup_steps,\n      input_bounds=input_bounds)\n  saver = tf.train.Saver(original_predictor.get_variables())\n  optimizer = tf.train.AdamOptimizer(learning_rate)\n  update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)\n  with tf.control_dependencies(update_ops):\n    train_op = optimizer.minimize(train_loss, step)\n\n  # Test using while loop.\n  def get_test_metrics(batch_size, attack_builder=ibp.UntargetedPGDAttack):\n    \"\"\"Returns the test metrics.\"\"\"\n    num_test_batches = len(data_test[0]) // batch_size\n    assert len(data_test[0]) % batch_size == 0, (\n        'Test data is not a multiple of batch size.')\n\n    def cond(i, *unused_args):\n      return i < num_test_batches\n\n    def body(i, metrics):\n      \"\"\"Compute the sum of all metrics.\"\"\"\n      test_data = ibp.build_dataset(data_test, batch_size=batch_size,\n                                    sequential=True)\n      predictor(test_data.image, override=True, is_training=False)\n      input_interval_bounds = ibp.IntervalBounds(\n          tf.maximum(test_data.image - FLAGS.epsilon, input_bounds[0]),\n          tf.minimum(test_data.image + FLAGS.epsilon, input_bounds[1]))\n      predictor.propagate_bounds(input_interval_bounds)\n      test_specification = ibp.ClassificationSpecification(\n          test_data.label, num_classes)\n      test_attack = attack_builder(predictor, test_specification, FLAGS.epsilon,\n                                   input_bounds=input_bounds,\n                                   optimizer_builder=ibp.UnrolledAdam)\n      test_losses = ibp.Losses(predictor, test_specification, test_attack)\n      test_losses(test_data.label)\n      new_metrics = []\n      for m, n in zip(metrics, test_losses.scalar_metrics):\n        new_metrics.append(m + n)\n      return i + 1, new_metrics\n\n    total_count = tf.constant(0, dtype=tf.int32)\n    total_metrics = [tf.constant(0, dtype=tf.float32)\n                     for _ in range(len(ibp.ScalarMetrics._fields))]\n    total_count, total_metrics = tf.while_loop(\n        cond,\n        body,\n        loop_vars=[total_count, total_metrics],\n        back_prop=False,\n        parallel_iterations=1)\n    total_count = tf.cast(total_count, tf.float32)\n    test_metrics = []\n    for m in total_metrics:\n      test_metrics.append(m / total_count)\n    return ibp.ScalarMetrics(*test_metrics)\n\n  test_metrics = get_test_metrics(\n      FLAGS.batch_size, ibp.UntargetedPGDAttack)\n  summaries = []\n  for f in test_metrics._fields:\n    summaries.append(\n        tf.summary.scalar(f, getattr(test_metrics, f)))\n  test_summaries = tf.summary.merge(summaries)\n  test_writer = tf.summary.FileWriter(os.path.join(FLAGS.output_dir, 'test'))\n\n  # Run everything.\n  tf_config = tf.ConfigProto()\n  tf_config.gpu_options.allow_growth = True\n  with tf.train.SingularMonitoredSession(config=tf_config) as sess:\n    for _ in range(FLAGS.steps):\n      iteration, loss_value, _ = sess.run(\n          [step, train_losses.scalar_losses.nominal_cross_entropy, train_op])\n      if iteration % FLAGS.test_every_n == 0:\n        metric_values, summary = sess.run([test_metrics, test_summaries])\n        test_writer.add_summary(summary, iteration)\n        show_metrics(iteration, metric_values, loss_value=loss_value)\n    saver.save(sess._tf_sess(),  # pylint: disable=protected-access\n               os.path.join(FLAGS.output_dir, 'model'),\n               global_step=FLAGS.steps - 1)\n\n\nif __name__ == '__main__':\n  app.run(main)\n"
  },
  {
    "path": "interval_bound_propagation/__init__.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Library to train verifiably robust neural networks.\n\nFor more details see paper: On the Effectiveness of Interval Bound Propagation\nfor Training Verifiably Robust Models.\n\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom interval_bound_propagation.src.attacks import MemoryEfficientMultiTargetedPGDAttack\nfrom interval_bound_propagation.src.attacks import MultiTargetedPGDAttack\nfrom interval_bound_propagation.src.attacks import pgd_attack\nfrom interval_bound_propagation.src.attacks import RestartedAttack\nfrom interval_bound_propagation.src.attacks import UnrolledAdam\nfrom interval_bound_propagation.src.attacks import UnrolledFGSMDescent\nfrom interval_bound_propagation.src.attacks import UnrolledGradientDescent\nfrom interval_bound_propagation.src.attacks import UnrolledSPSAAdam\nfrom interval_bound_propagation.src.attacks import UnrolledSPSAFGSMDescent\nfrom interval_bound_propagation.src.attacks import UnrolledSPSAGradientDescent\nfrom interval_bound_propagation.src.attacks import UntargetedAdaptivePGDAttack\nfrom interval_bound_propagation.src.attacks import UntargetedPGDAttack\nfrom interval_bound_propagation.src.attacks import UntargetedTop5PGDAttack\nfrom interval_bound_propagation.src.bounds import AbstractBounds\nfrom interval_bound_propagation.src.bounds import IntervalBounds\nimport interval_bound_propagation.src.crown as crown\nfrom interval_bound_propagation.src.fastlin import RelativeSymbolicBounds\nfrom interval_bound_propagation.src.fastlin import SymbolicBounds\nimport interval_bound_propagation.src.layer_utils as layer_utils\nfrom interval_bound_propagation.src.layers import BatchNorm\nfrom interval_bound_propagation.src.layers import ImageNorm\nfrom interval_bound_propagation.src.loss import Losses\nfrom interval_bound_propagation.src.loss import ScalarLosses\nfrom interval_bound_propagation.src.loss import ScalarMetrics\nfrom interval_bound_propagation.src.model import DNN\nfrom interval_bound_propagation.src.model import StandardModelWrapper\nfrom interval_bound_propagation.src.model import VerifiableModelWrapper\nfrom interval_bound_propagation.src.relative_bounds import RelativeIntervalBounds\nfrom interval_bound_propagation.src.simplex_bounds import SimplexBounds\nfrom interval_bound_propagation.src.specification import ClassificationSpecification\nfrom interval_bound_propagation.src.specification import LeastLikelyClassificationSpecification\nfrom interval_bound_propagation.src.specification import LinearSpecification\nfrom interval_bound_propagation.src.specification import RandomClassificationSpecification\nfrom interval_bound_propagation.src.specification import Specification\nfrom interval_bound_propagation.src.specification import TargetedClassificationSpecification\nfrom interval_bound_propagation.src.utils import add_image_normalization\nfrom interval_bound_propagation.src.utils import build_dataset\nfrom interval_bound_propagation.src.utils import create_attack\nfrom interval_bound_propagation.src.utils import create_classification_losses\nfrom interval_bound_propagation.src.utils import create_specification\nfrom interval_bound_propagation.src.utils import get_attack_builder\nfrom interval_bound_propagation.src.utils import linear_schedule\nfrom interval_bound_propagation.src.utils import parse_learning_rate\nfrom interval_bound_propagation.src.utils import randomize\nfrom interval_bound_propagation.src.utils import smooth_schedule\nfrom interval_bound_propagation.src.verifiable_wrapper import BatchFlattenWrapper\nfrom interval_bound_propagation.src.verifiable_wrapper import BatchNormWrapper\nfrom interval_bound_propagation.src.verifiable_wrapper import BatchReshapeWrapper\nfrom interval_bound_propagation.src.verifiable_wrapper import ConstWrapper\nfrom interval_bound_propagation.src.verifiable_wrapper import ImageNormWrapper\nfrom interval_bound_propagation.src.verifiable_wrapper import IncreasingMonotonicWrapper\nfrom interval_bound_propagation.src.verifiable_wrapper import LinearConv1dWrapper\nfrom interval_bound_propagation.src.verifiable_wrapper import LinearConv2dWrapper\nfrom interval_bound_propagation.src.verifiable_wrapper import LinearConvWrapper\nfrom interval_bound_propagation.src.verifiable_wrapper import LinearFCWrapper\nfrom interval_bound_propagation.src.verifiable_wrapper import ModelInputWrapper\nfrom interval_bound_propagation.src.verifiable_wrapper import PiecewiseMonotonicWrapper\nfrom interval_bound_propagation.src.verifiable_wrapper import VerifiableWrapper\n\n\n__version__ = '1.10'\n"
  },
  {
    "path": "interval_bound_propagation/src/__init__.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Library to train verifiably robust neural networks.\"\"\"\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n"
  },
  {
    "path": "interval_bound_propagation/src/attacks.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utilities to define attacks.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport abc\nimport collections\n\nimport six\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\nnest = tf.nest\n\n\n@six.add_metaclass(abc.ABCMeta)\nclass UnrolledOptimizer(object):\n  \"\"\"In graph optimizer to be used in tf.while_loop.\"\"\"\n\n  def __init__(self, colocate_gradients_with_ops=False):\n    self._colocate_gradients_with_ops = colocate_gradients_with_ops\n\n  @abc.abstractmethod\n  def minimize(self, loss, x, optim_state):\n    \"\"\"Compute a new value of `x` to minimize `loss`.\n\n    Args:\n      loss: A scalar Tensor, the value to be minimized. `loss` should be a\n        continuous function of `x` which supports gradients, `loss = f(x)`.\n      x: A list of Tensors, the values to be updated. This is analogous to the\n        `var_list` argument in standard TF Optimizer.\n      optim_state: A (possibly nested) dict, containing any state info needed\n        for the optimizer.\n\n    Returns:\n      new_x: A list of Tensors, the same length as `x`, which are updated\n      new_optim_state: A new state, with the same structure as `optim_state`,\n        which have been updated.\n    \"\"\"\n\n  @abc.abstractmethod\n  def init_state(self, x):\n    \"\"\"Returns the initial state of the optimizer.\n\n    Args:\n      x: A list of Tensors, which will be optimized.\n\n    Returns:\n      Any structured output.\n    \"\"\"\n\n\nclass UnrolledGradientDescent(UnrolledOptimizer):\n  \"\"\"Vanilla gradient descent optimizer.\"\"\"\n\n  _State = collections.namedtuple('State', ['iteration'])  # pylint: disable=invalid-name\n\n  def __init__(self, lr=.1, lr_fn=None, fgsm=False,\n               colocate_gradients_with_ops=False):\n    super(UnrolledGradientDescent, self).__init__(\n        colocate_gradients_with_ops=colocate_gradients_with_ops)\n    self._lr_fn = (lambda i: lr) if lr_fn is None else lr_fn\n    self._fgsm = fgsm\n\n  def init_state(self, unused_x):\n    return self._State(tf.constant(0, dtype=tf.int64))\n\n  def minimize(self, loss, x, optim_state):\n    \"\"\"Refer to parent class documentation.\"\"\"\n    lr = self._lr_fn(optim_state.iteration)\n    grads = self.gradients(loss, x)\n    if self._fgsm:\n      grads = [tf.sign(g) for g in grads]\n    new_x = [None] * len(x)\n    for i in range(len(x)):\n      new_x[i] = x[i] - lr * grads[i]\n    new_optim_state = self._State(optim_state.iteration + 1)\n    return new_x, new_optim_state\n\n  def gradients(self, loss, x):\n    return tf.gradients(\n        loss, x, colocate_gradients_with_ops=self._colocate_gradients_with_ops)\n\n\n# Syntactic sugar.\nclass UnrolledFGSMDescent(UnrolledGradientDescent):\n  \"\"\"Identical to UnrolledGradientDescent but forces FGM steps.\"\"\"\n\n  def __init__(self, lr=.1, lr_fn=None,\n               colocate_gradients_with_ops=False):\n    super(UnrolledFGSMDescent, self).__init__(\n        lr, lr_fn, True, colocate_gradients_with_ops)\n\n\nclass UnrolledAdam(UnrolledOptimizer):\n  \"\"\"The Adam optimizer defined in https://arxiv.org/abs/1412.6980.\"\"\"\n\n  _State = collections.namedtuple('State', ['t', 'm', 'u'])  # pylint: disable=invalid-name\n\n  def __init__(self, lr=0.1, lr_fn=None, beta1=0.9, beta2=0.999, epsilon=1e-9,\n               colocate_gradients_with_ops=False):\n    super(UnrolledAdam, self).__init__(\n        colocate_gradients_with_ops=colocate_gradients_with_ops)\n    self._lr_fn = (lambda i: lr) if lr_fn is None else lr_fn\n    self._beta1 = beta1\n    self._beta2 = beta2\n    self._epsilon = epsilon\n\n  def init_state(self, x):\n    return self._State(\n        t=tf.constant(0, dtype=tf.int64),\n        m=[tf.zeros_like(v) for v in x],\n        u=[tf.zeros_like(v) for v in x])\n\n  def _apply_gradients(self, grads, x, optim_state):\n    \"\"\"Applies gradients.\"\"\"\n    lr = self._lr_fn(optim_state.t)\n    new_optim_state = self._State(\n        t=optim_state.t + 1,\n        m=[None] * len(x),\n        u=[None] * len(x))\n    t = tf.cast(new_optim_state.t, tf.float32)\n    new_x = [None] * len(x)\n    for i in range(len(x)):\n      g = grads[i]\n      m_old = optim_state.m[i]\n      u_old = optim_state.u[i]\n      new_optim_state.m[i] = self._beta1 * m_old + (1. - self._beta1) * g\n      new_optim_state.u[i] = self._beta2 * u_old + (1. - self._beta2) * g * g\n      m_hat = new_optim_state.m[i] / (1. - tf.pow(self._beta1, t))\n      u_hat = new_optim_state.u[i] / (1. - tf.pow(self._beta2, t))\n      new_x[i] = x[i] - lr * m_hat / (tf.sqrt(u_hat) + self._epsilon)\n    return new_x, new_optim_state\n\n  def minimize(self, loss, x, optim_state):\n    grads = self.gradients(loss, x)\n    return self._apply_gradients(grads, x, optim_state)\n\n  def gradients(self, loss, x):\n    return tf.gradients(\n        loss, x, colocate_gradients_with_ops=self._colocate_gradients_with_ops)\n\n\ndef _spsa_gradients(loss_fn, x, delta=0.01, num_samples=16, num_iterations=4):\n  \"\"\"Compute gradient estimates using SPSA.\n\n  Args:\n    loss_fn: Callable that takes a single argument of shape [batch_size, ...]\n      and returns the loss contribution of each element of the batch as a\n      tensor of shape [batch_size].\n    x: List of tensors with a single element. We only support computation of\n      the gradient of the loss with respect to x[0]. We take a list as input to\n      keep the same API call as tf.gradients.\n    delta: The gradients are computed by computing the loss within x - delta and\n      x + delta.\n    num_samples: The total number of random samples used to compute the gradient\n      is `num_samples` times `num_iterations`. `num_samples` contributes to the\n      gradient by tiling `x` `num_samples` times.\n    num_iterations: The total number of random samples used to compute the\n      gradient is `num_samples` times `num_iterations`. `num_iterations`\n      contributes to the gradient by iterating using a `tf.while_loop`.\n\n  Returns:\n    List of tensors with a single element corresponding to the gradient of\n    loss_fn(x[0]) with respect to x[0].\n  \"\"\"\n\n  if len(x) != 1:\n    raise NotImplementedError('SPSA gradients with respect to multiple '\n                              'variables is not supported.')\n  # loss_fn takes a single argument.\n  tensor = x[0]\n\n  def _get_delta(x):\n    return delta * tf.sign(\n        tf.random_uniform(tf.shape(x), minval=-1., maxval=1., dtype=x.dtype))\n\n  # Process batch_size samples at a time.\n  def cond(i, *_):\n    return tf.less(i, num_iterations)\n\n  def loop_body(i, total_grad):\n    \"\"\"Compute gradient estimate.\"\"\"\n    batch_size = tf.shape(tensor)[0]\n    # The tiled tensor has shape [num_samples, batch_size, ...]\n    tiled_tensor = tf.expand_dims(tensor, axis=0)\n    tiled_tensor = tf.tile(tiled_tensor,\n                           [num_samples] + [1] * len(tensor.shape))\n    # The tiled tensor has now shape [2, num_samples, batch_size, ...].\n    delta = _get_delta(tiled_tensor)\n    tiled_tensor = tf.stack(\n        [tiled_tensor + delta, tiled_tensor - delta], axis=0)\n    # Compute loss with shape [2, num_samples, batch_size].\n    losses = loss_fn(\n        tf.reshape(tiled_tensor,\n                   [2 * num_samples, batch_size] + tensor.shape.as_list()[1:]))\n    losses = tf.reshape(losses, [2, num_samples, batch_size])\n\n    # Compute approximate gradient using broadcasting.\n    shape = losses.shape.as_list() + [1] * (len(tensor.shape) - 1)\n    shape = [(s or -1) for s in shape]  # Remove None.\n    losses = tf.reshape(losses, shape)\n    g = tf.reduce_mean((losses[0] - losses[1]) / (2. * delta), axis=0)\n    return [i + 1, g / num_iterations + total_grad]\n\n  _, g = tf.while_loop(\n      cond,\n      loop_body,\n      loop_vars=[tf.constant(0.), tf.zeros_like(tensor)],\n      parallel_iterations=1,\n      back_prop=False)\n  return [g]\n\n\n@six.add_metaclass(abc.ABCMeta)\nclass UnrolledSPSA(object):\n  \"\"\"Abstract class that represents an optimizer based on SPSA.\"\"\"\n\n\nclass UnrolledSPSAGradientDescent(UnrolledGradientDescent, UnrolledSPSA):\n  \"\"\"Optimizer for gradient-free attacks in https://arxiv.org/abs/1802.05666.\n\n  Gradients estimates are computed using Simultaneous Perturbation Stochastic\n  Approximation (SPSA).\n  \"\"\"\n\n  def __init__(self, lr=0.1, lr_fn=None, fgsm=False,\n               colocate_gradients_with_ops=False, delta=0.01, num_samples=32,\n               num_iterations=4, loss_fn=None):\n    super(UnrolledSPSAGradientDescent, self).__init__(\n        lr, lr_fn, fgsm, colocate_gradients_with_ops)\n    assert num_samples % 2 == 0, 'Number of samples must be even'\n    self._delta = delta\n    self._num_samples = num_samples // 2  # Since we mirror +/- delta later.\n    self._num_iterations = num_iterations\n    assert loss_fn is not None, 'loss_fn must be specified.'\n    self._loss_fn = loss_fn\n\n  def gradients(self, loss, x):\n    return _spsa_gradients(self._loss_fn, x, self._delta, self._num_samples,\n                           self._num_iterations)\n\n\n# Syntactic sugar.\nclass UnrolledSPSAFGSMDescent(UnrolledSPSAGradientDescent):\n  \"\"\"Identical to UnrolledSPSAGradientDescent but forces FGSM steps.\"\"\"\n\n  def __init__(self, lr=.1, lr_fn=None,\n               colocate_gradients_with_ops=False, delta=0.01, num_samples=32,\n               num_iterations=4, loss_fn=None):\n    super(UnrolledSPSAFGSMDescent, self).__init__(\n        lr, lr_fn, True, colocate_gradients_with_ops, delta, num_samples,\n        num_iterations, loss_fn)\n\n\nclass UnrolledSPSAAdam(UnrolledAdam, UnrolledSPSA):\n  \"\"\"Optimizer for gradient-free attacks in https://arxiv.org/abs/1802.05666.\n\n  Gradients estimates are computed using Simultaneous Perturbation Stochastic\n  Approximation (SPSA), combined with the ADAM update rule.\n  \"\"\"\n\n  def __init__(self, lr=0.1, lr_fn=None, beta1=0.9, beta2=0.999, epsilon=1e-9,\n               colocate_gradients_with_ops=False, delta=0.01, num_samples=32,\n               num_iterations=4, loss_fn=None):\n    super(UnrolledSPSAAdam, self).__init__(lr, lr_fn, beta1, beta2, epsilon,\n                                           colocate_gradients_with_ops)\n    assert num_samples % 2 == 0, 'Number of samples must be even'\n    self._delta = delta\n    self._num_samples = num_samples // 2  # Since we mirror +/- delta later.\n    self._num_iterations = num_iterations\n    assert loss_fn is not None, 'loss_fn must be specified.'\n    self._loss_fn = loss_fn\n\n  def gradients(self, loss, x):\n    return _spsa_gradients(self._loss_fn, x, self._delta, self._num_samples,\n                           self._num_iterations)\n\n\ndef _is_spsa_optimizer(cls):\n  return issubclass(cls, UnrolledSPSA)\n\n\ndef wrap_optimizer(cls, **default_kwargs):\n  \"\"\"Wraps an optimizer such that __init__ uses the specified kwargs.\"\"\"\n\n  class WrapperUnrolledOptimizer(cls):\n\n    def __init__(self, *args, **kwargs):\n      new_kwargs = default_kwargs.copy()\n      new_kwargs.update(kwargs)\n      super(WrapperUnrolledOptimizer, self).__init__(*args, **new_kwargs)\n  return WrapperUnrolledOptimizer\n\n\ndef _project_perturbation(perturbation, epsilon, input_image, image_bounds):\n  \"\"\"Project `perturbation` onto L-infinity ball of radius `epsilon`.\"\"\"\n  clipped_perturbation = tf.clip_by_value(perturbation, -epsilon, epsilon)\n  new_image = tf.clip_by_value(input_image + clipped_perturbation,\n                               image_bounds[0], image_bounds[1])\n  return new_image - input_image\n\n\ndef pgd_attack(loss_fn, input_image, epsilon, num_steps,\n               optimizer=UnrolledGradientDescent(),\n               project_perturbation=_project_perturbation,\n               image_bounds=None, random_init=1.):\n  \"\"\"Projected gradient descent for generating adversarial images.\n\n  Args:\n    loss_fn: A callable which takes `input_image` and `label` as arguments, and\n      returns the loss, a scalar Tensor, we will be minimized\n    input_image: Tensor, a batch of images\n    epsilon: float, the L-infinity norm of the maximum allowable perturbation\n    num_steps: int, the number of steps of gradient descent\n    optimizer: An `UnrolledOptimizer` object\n    project_perturbation: A function, which will be used to enforce some\n      constraint. It should have the same signature as `_project_perturbation`.\n      Note that if you use a custom projection function, you should double-check\n      your implementation, since an incorrect implementation will not error,\n      and will appear to work fine.\n    image_bounds: A pair of floats: minimum and maximum pixel value. If None\n      (default), the bounds are assumed to be 0 and 1.\n    random_init: Probability of starting from random location rather than\n      nominal input image.\n\n  Returns:\n    adversarial version of `input_image`, with L-infinity difference less than\n      epsilon, which tries to minimize loss_fn.\n  \"\"\"\n  image_bounds = image_bounds or (0., 1.)\n  random_shape = [tf.shape(input_image)[0]] + [1] * (len(input_image.shape) - 1)\n  use_random_init = tf.cast(\n      tf.random_uniform(random_shape) < float(random_init), tf.float32)\n  init_perturbation = use_random_init * tf.random_uniform(\n      tf.shape(input_image), minval=-epsilon, maxval=epsilon)\n  init_perturbation = project_perturbation(init_perturbation,\n                                           epsilon, input_image, image_bounds)\n  init_optim_state = optimizer.init_state([init_perturbation])\n\n  def loop_body(i, perturbation, flat_optim_state):\n    \"\"\"Update perturbation to input image.\"\"\"\n    optim_state = nest.pack_sequence_as(structure=init_optim_state,\n                                        flat_sequence=flat_optim_state)\n    loss = loss_fn(input_image + perturbation)\n    new_perturbation_list, new_optim_state = optimizer.minimize(\n        loss, [perturbation], optim_state)\n    projected_perturbation = project_perturbation(\n        new_perturbation_list[0], epsilon, input_image, image_bounds)\n    return i + 1, projected_perturbation, nest.flatten(new_optim_state)\n\n  def cond(i, *_):\n    return tf.less(i, num_steps)\n\n  flat_init_optim_state = nest.flatten(init_optim_state)\n  _, final_perturbation, _ = tf.while_loop(\n      cond,\n      loop_body,\n      loop_vars=[tf.constant(0.), init_perturbation, flat_init_optim_state],\n      parallel_iterations=1,\n      back_prop=False)\n\n  adversarial_image = input_image + final_perturbation\n  return tf.stop_gradient(adversarial_image)\n\n\n@six.add_metaclass(abc.ABCMeta)\nclass Attack(snt.AbstractModule):\n  \"\"\"Defines an attack as a Sonnet module.\"\"\"\n\n  def __init__(self, predictor, specification, name, predictor_kwargs=None):\n    super(Attack, self).__init__(name=name)\n    self._predictor = predictor\n    self._specification = specification\n    if predictor_kwargs is None:\n      self._kwargs = {'intermediate': {}, 'final': {}}\n    else:\n      self._kwargs = predictor_kwargs\n    self._forced_mode = None\n    self._target_class = None\n\n  def _eval_fn(self, x, mode='intermediate'):\n    \"\"\"Runs the logits corresponding to `x`.\n\n    Args:\n      x: input to the predictor network.\n      mode: Either \"intermediate\" or \"final\". Selects the desired predictor\n        arguments.\n\n    Returns:\n      Tensor of logits.\n    \"\"\"\n    if self._forced_mode is not None:\n      mode = self._forced_mode\n    return self._predictor(x, **self._kwargs[mode])\n\n  @abc.abstractmethod\n  def _build(self, inputs, labels):\n    \"\"\"Returns the adversarial attack around inputs.\"\"\"\n\n  @abc.abstractproperty\n  def logits(self):\n    \"\"\"Returns the logits corresponding to the best attack.\"\"\"\n\n  @abc.abstractproperty\n  def attack(self):\n    \"\"\"Returns the best attack.\"\"\"\n\n  @abc.abstractproperty\n  def success(self):\n    \"\"\"Returns whether the attack was successful.\"\"\"\n\n  def force_mode(self, mode):\n    \"\"\"Only used by RestartedAttack to force the evaluation mode.\"\"\"\n    self._forced_mode = mode\n\n  @property\n  def target_class(self):\n    \"\"\"Returns the target class if this attack is a targeted attacks.\"\"\"\n    return self._target_class\n\n  @target_class.setter\n  def target_class(self, t):\n    self._target_class = t\n\n\n@six.add_metaclass(abc.ABCMeta)\nclass PGDAttack(Attack):\n  \"\"\"Defines a PGD attack.\"\"\"\n\n  def __init__(self, predictor, specification, epsilon, lr=.1, lr_fn=None,\n               num_steps=20, num_restarts=1, input_bounds=(0., 1.),\n               random_init=1., optimizer_builder=UnrolledGradientDescent,\n               project_perturbation=_project_perturbation,\n               predictor_kwargs=None):\n    super(PGDAttack, self).__init__(predictor, specification, name='pgd',\n                                    predictor_kwargs=predictor_kwargs)\n    self._num_steps = num_steps\n    self._num_restarts = num_restarts\n    self._epsilon = epsilon\n    self._lr = lr\n    self._lr_fn = lr_fn\n    self._input_bounds = input_bounds\n    self._random_init = random_init\n    self._optimizer_builder = optimizer_builder\n    self._project_perturbation = project_perturbation\n\n  # Helper functions.\n  def prepare_inputs(self, inputs):\n    \"\"\"Tiles inputs according to number of restarts.\"\"\"\n    batch_size = tf.shape(inputs)[0]\n    input_shape = list(inputs.shape.as_list()[1:])\n    duplicated_inputs = tf.expand_dims(inputs, axis=0)\n    # Shape is [num_restarts, batch_size, ...]\n    duplicated_inputs = tf.tile(\n        duplicated_inputs,\n        [self._num_restarts, 1] + [1] * len(input_shape))\n    # Shape is [num_restarts * batch_size, ...]\n    duplicated_inputs = tf.reshape(\n        duplicated_inputs, [self._num_restarts * batch_size] + input_shape)\n    return batch_size, input_shape, duplicated_inputs\n\n  def prepare_labels(self, labels):\n    \"\"\"Tiles labels according to number of restarts.\"\"\"\n    return tf.tile(labels, [self._num_restarts])\n\n  def find_worst_attack(self, objective_fn, adversarial_input, batch_size,\n                        input_shape):\n    \"\"\"Returns the attack that maximizes objective_fn.\"\"\"\n    adversarial_objective = objective_fn(adversarial_input)\n    adversarial_objective = tf.reshape(adversarial_objective, [-1, batch_size])\n    adversarial_input = tf.reshape(adversarial_input,\n                                   [-1, batch_size] + input_shape)\n    i = tf.argmax(adversarial_objective, axis=0)\n    j = tf.cast(tf.range(tf.shape(adversarial_objective)[1]), i.dtype)\n    ij = tf.stack([i, j], axis=1)\n    return tf.gather_nd(adversarial_input, ij)\n\n\ndef _maximize_margin(bounds):\n  # Bounds has shape [num_restarts, batch_size, num_specs].\n  return tf.reduce_max(bounds, axis=-1)\n\n\ndef _any_greater(bounds):\n  # Bounds has shape [batch_size, num_specs].\n  bounds = tf.reduce_max(bounds, axis=-1)\n  return bounds > 0.\n\n\ndef _maximize_topk_hinge_margin(bounds, k=5, margin=.1):\n  # Bounds has shape [num_restarts, batch_size, num_specs].\n  b = tf.nn.top_k(bounds, k=k, sorted=False).values\n  return tf.reduce_sum(tf.minimum(b, margin), axis=-1)\n\n\ndef _topk_greater(bounds, k=5):\n  # Bounds has shape [batch_size, num_specs].\n  b = tf.nn.top_k(bounds, k=k, sorted=False).values\n  return tf.reduce_min(b, axis=-1) > 0.\n\n\nclass UntargetedPGDAttack(PGDAttack):\n  \"\"\"Defines an untargeted PGD attack.\"\"\"\n\n  def __init__(self, predictor, specification, epsilon, lr=.1, lr_fn=None,\n               num_steps=20, num_restarts=1, input_bounds=(0., 1.),\n               random_init=1., optimizer_builder=UnrolledGradientDescent,\n               project_perturbation=_project_perturbation,\n               objective_fn=_maximize_margin, success_fn=_any_greater,\n               predictor_kwargs=None):\n    super(UntargetedPGDAttack, self).__init__(\n        predictor, specification, epsilon, lr, lr_fn, num_steps, num_restarts,\n        input_bounds, random_init, optimizer_builder, project_perturbation,\n        predictor_kwargs)\n    self._objective_fn = objective_fn\n    self._success_fn = success_fn\n\n  def _build(self, inputs, labels):\n    batch_size, input_shape, duplicated_inputs = self.prepare_inputs(inputs)\n    duplicated_labels = self.prepare_labels(labels)\n\n    # Define objectives.\n    def objective_fn(x):\n      model_logits = self._eval_fn(x)  # [restarts * batch_size, output].\n      model_logits = tf.reshape(\n          model_logits, [self._num_restarts, batch_size, -1])\n      bounds = self._specification.evaluate(model_logits)\n      # Output has dimension [num_restarts, batch_size].\n      return self._objective_fn(bounds)\n\n    # Only used for SPSA.\n    # The input to this loss is the perturbation (not the image).\n    # The first dimension corresponds to the number of SPSA samples.\n    # Shape of perturbations is [num_samples, restarts * batch_size, ...]\n    def spsa_loss_fn(perturbation):\n      \"\"\"Computes the loss per SPSA sample.\"\"\"\n      x = tf.reshape(\n          perturbation + tf.expand_dims(duplicated_inputs, axis=0),\n          [-1] + duplicated_inputs.shape.as_list()[1:])\n      model_logits = self._eval_fn(x)\n      num_outputs = tf.shape(model_logits)[1]\n      model_logits = tf.reshape(\n          model_logits, [-1, batch_size, num_outputs])\n      bounds = self._specification.evaluate(model_logits)\n      losses = -self._objective_fn(bounds)\n      return tf.reshape(losses, [-1])\n\n    def reduced_loss_fn(x):\n      # Pick worse attack, output has shape [num_restarts, batch_size].\n      return -tf.reduce_sum(objective_fn(x))\n\n    # Use targeted attacks as specified by the specification.\n    if _is_spsa_optimizer(self._optimizer_builder):\n      optimizer = self._optimizer_builder(lr=self._lr, lr_fn=self._lr_fn,\n                                          loss_fn=spsa_loss_fn)\n    else:\n      optimizer = self._optimizer_builder(lr=self._lr, lr_fn=self._lr_fn)\n\n    adversarial_input = pgd_attack(\n        reduced_loss_fn, duplicated_inputs, epsilon=self._epsilon,\n        num_steps=self._num_steps, image_bounds=self._input_bounds,\n        random_init=self._random_init, optimizer=optimizer,\n        project_perturbation=self._project_perturbation)\n    adversarial_input = self.adapt(duplicated_inputs, adversarial_input,\n                                   duplicated_labels)\n\n    self._attack = self.find_worst_attack(objective_fn, adversarial_input,\n                                          batch_size, input_shape)\n    self._logits = self._eval_fn(self._attack, mode='final')\n    self._success = self._success_fn(self._specification.evaluate(self._logits))\n    return self._attack\n\n  @property\n  def logits(self):\n    self._ensure_is_connected()\n    return self._logits\n\n  @property\n  def attack(self):\n    self._ensure_is_connected()\n    return self._attack\n\n  @property\n  def success(self):\n    self._ensure_is_connected()\n    return self._success\n\n  def adapt(self, original_inputs, adversarial_inputs, labels):\n    \"\"\"Function called after PGD to adapt adversarial examples.\"\"\"\n    return adversarial_inputs\n\n\nclass UntargetedTop5PGDAttack(UntargetedPGDAttack):\n  \"\"\"Defines an untargeted PGD attack on top-5.\"\"\"\n\n  def __init__(self, predictor, specification, epsilon, lr=.1, lr_fn=None,\n               num_steps=20, num_restarts=1, input_bounds=(0., 1.),\n               random_init=1., optimizer_builder=UnrolledGradientDescent,\n               project_perturbation=_project_perturbation,\n               objective_fn=_maximize_topk_hinge_margin, predictor_kwargs=None):\n    super(UntargetedTop5PGDAttack, self).__init__(\n        predictor, specification, epsilon, lr=lr, lr_fn=lr_fn,\n        num_steps=num_steps, num_restarts=num_restarts,\n        input_bounds=input_bounds, random_init=random_init,\n        optimizer_builder=optimizer_builder,\n        project_perturbation=project_perturbation, objective_fn=objective_fn,\n        success_fn=_topk_greater, predictor_kwargs=predictor_kwargs)\n\n\nclass UntargetedAdaptivePGDAttack(UntargetedPGDAttack):\n  \"\"\"Uses an adaptive scheme to pick attacks that are just strong enough.\"\"\"\n\n  def adapt(self, original_inputs, adversarial_inputs, labels):\n    \"\"\"Runs binary search to find the first misclassified input.\"\"\"\n    batch_size = tf.shape(original_inputs)[0]\n    binary_search_iterations = 10\n\n    def cond(i, *_):\n      return tf.less(i, binary_search_iterations)\n\n    def get(m):\n      m = tf.reshape(m, [batch_size] + [1] * (len(original_inputs.shape) - 1))\n      return (adversarial_inputs - original_inputs) * m + original_inputs\n\n    def is_attack_successful(m):\n      logits = self._eval_fn(get(m))\n      return self._success_fn(self._specification.evaluate(logits))\n\n    def loop_body(i, lower, upper):\n      m = (lower + upper) * .5\n      success = is_attack_successful(m)\n      new_lower = tf.where(success, lower, m)\n      new_upper = tf.where(success, m, upper)\n      return i + 1, new_lower, new_upper\n\n    lower = tf.zeros(shape=[batch_size])\n    upper = tf.ones(shape=[batch_size])\n    _, lower, upper = tf.while_loop(\n        cond,\n        loop_body,\n        loop_vars=[tf.constant(0.), lower, upper],\n        parallel_iterations=1,\n        back_prop=False)\n    # If lower is incorrectly classified, pick lower; otherwise pick upper.\n    success = is_attack_successful(lower)\n    return get(tf.where(success, lower, upper))\n\n\nclass MultiTargetedPGDAttack(PGDAttack):\n  \"\"\"Runs targeted attacks for each specification.\"\"\"\n\n  def __init__(self, predictor, specification, epsilon, lr=.1, lr_fn=None,\n               num_steps=20, num_restarts=1, input_bounds=(0., 1.),\n               random_init=1., optimizer_builder=UnrolledGradientDescent,\n               project_perturbation=_project_perturbation,\n               max_specifications=0, random_specifications=False,\n               predictor_kwargs=None):\n    super(MultiTargetedPGDAttack, self).__init__(\n        predictor, specification, epsilon, lr=lr, lr_fn=lr_fn,\n        num_steps=num_steps, num_restarts=num_restarts,\n        input_bounds=input_bounds, random_init=random_init,\n        optimizer_builder=optimizer_builder,\n        project_perturbation=project_perturbation,\n        predictor_kwargs=predictor_kwargs)\n    self._max_specifications = max_specifications\n    self._random_specifications = random_specifications\n\n  def _build(self, inputs, labels):\n    batch_size = tf.shape(inputs)[0]\n    num_specs = self._specification.num_specifications\n    if self._max_specifications > 0 and self._max_specifications < num_specs:\n      model_logits = self._eval_fn(inputs)\n      bounds = self._specification.evaluate(model_logits)\n      _, idx = tf.math.top_k(bounds, k=self._max_specifications, sorted=False)\n      if self._random_specifications:\n        idx = tf.random.uniform(shape=tf.shape(idx),\n                                maxval=self._specification.num_specifications,\n                                dtype=idx.dtype)\n      idx = tf.tile(tf.expand_dims(idx, 0), [self._num_restarts, 1, 1])\n      select_fn = lambda x: tf.gather(x, idx, batch_dims=len(idx.shape) - 1)\n    else:\n      select_fn = lambda x: x\n\n    input_shape = list(inputs.shape.as_list()[1:])\n    duplicated_inputs = tf.expand_dims(inputs, axis=0)\n    # Shape is [num_restarts * num_specifications, batch_size, ...]\n    duplicated_inputs = tf.tile(\n        duplicated_inputs,\n        [self._num_restarts * num_specs, 1] + [1] * len(input_shape))\n    # Shape is [num_restarts * num_specifications * batch_size, ...]\n    duplicated_inputs = tf.reshape(duplicated_inputs, [-1] + input_shape)\n\n    def objective_fn(x):\n      # Output has shape [restarts * num_specs * batch_size, output].\n      model_logits = self._eval_fn(x)\n      model_logits = tf.reshape(\n          model_logits, [self._num_restarts, num_specs, batch_size, -1])\n      # Output has shape [num_restarts, batch_size, num_specs].\n      return self._specification.evaluate(model_logits)\n\n    def reduced_loss_fn(x):\n      # Negate as we minimize.\n      return -tf.reduce_sum(select_fn(objective_fn(x)))\n\n    # Use targeted attacks as specified by the specification.\n    if _is_spsa_optimizer(self._optimizer_builder):\n      raise ValueError('\"UnrolledSPSA*\" unsupported in '\n                       'MultiTargetedPGDAttack')\n    optimizer = self._optimizer_builder(lr=self._lr, lr_fn=self._lr_fn)\n    adversarial_input = pgd_attack(\n        reduced_loss_fn, duplicated_inputs,\n        epsilon=self._epsilon, num_steps=self._num_steps,\n        image_bounds=self._input_bounds, random_init=self._random_init,\n        optimizer=optimizer, project_perturbation=self._project_perturbation)\n    # Get best attack.\n    adversarial_objective = objective_fn(adversarial_input)\n    adversarial_objective = tf.transpose(adversarial_objective, [0, 2, 1])\n    adversarial_objective = tf.reshape(adversarial_objective, [-1, batch_size])\n    adversarial_input = tf.reshape(adversarial_input,\n                                   [-1, batch_size] + input_shape)\n    i = tf.argmax(adversarial_objective, axis=0)\n    j = tf.cast(tf.range(tf.shape(adversarial_objective)[1]), i.dtype)\n    ij = tf.stack([i, j], axis=1)\n    self._attack = tf.gather_nd(adversarial_input, ij)\n    self._logits = self._eval_fn(self._attack, mode='final')\n    # Count the number of sample that violate any specification.\n    bounds = tf.reduce_max(self._specification.evaluate(self._logits), axis=1)\n    self._success = (bounds > 0.)\n    return self._attack\n\n  @property\n  def logits(self):\n    self._ensure_is_connected()\n    return self._logits\n\n  @property\n  def attack(self):\n    self._ensure_is_connected()\n    return self._attack\n\n  @property\n  def success(self):\n    self._ensure_is_connected()\n    return self._success\n\n\nclass MemoryEfficientMultiTargetedPGDAttack(PGDAttack):\n  \"\"\"Defines a targeted PGD attack for each specification using while_loop.\"\"\"\n\n  def __init__(self, predictor, specification, epsilon, lr=.1, lr_fn=None,\n               num_steps=20, num_restarts=1, input_bounds=(0., 1.),\n               random_init=1., optimizer_builder=UnrolledGradientDescent,\n               project_perturbation=_project_perturbation,\n               max_specifications=0, random_specifications=False,\n               predictor_kwargs=None):\n    super(MemoryEfficientMultiTargetedPGDAttack, self).__init__(\n        predictor, specification, epsilon, lr=lr, lr_fn=lr_fn,\n        num_steps=num_steps, num_restarts=num_restarts,\n        input_bounds=input_bounds, random_init=random_init,\n        optimizer_builder=optimizer_builder,\n        project_perturbation=project_perturbation,\n        predictor_kwargs=predictor_kwargs)\n    self._max_specifications = max_specifications\n    self._random_specifications = random_specifications\n\n  def _build(self, inputs, labels):\n    batch_size, input_shape, duplicated_inputs = self.prepare_inputs(inputs)\n    if (self._max_specifications > 0 and\n        self._max_specifications < self._specification.num_specifications):\n      num_specs = self._max_specifications\n      model_logits = self._eval_fn(inputs)\n      bounds = self._specification.evaluate(model_logits)\n      _, idx = tf.math.top_k(bounds, k=num_specs, sorted=False)\n      if self._random_specifications:\n        idx = tf.random.uniform(shape=tf.shape(idx),\n                                maxval=self._specification.num_specifications,\n                                dtype=idx.dtype)\n      idx = tf.tile(tf.expand_dims(idx, 0), [self._num_restarts, 1, 1])\n      def select_fn(x, i):\n        return tf.squeeze(\n            tf.gather(x, tf.expand_dims(idx[:, :, i], -1),\n                      batch_dims=len(idx.shape) - 1),\n            axis=-1)\n    else:\n      num_specs = self._specification.num_specifications\n      select_fn = lambda x, i: x[:, :, i]\n\n    def objective_fn(x):\n      model_logits = self._eval_fn(x)  # [restarts * batch_size, output].\n      model_logits = tf.reshape(\n          model_logits, [self._num_restarts, batch_size, -1])\n      # Output has dimension [num_restarts, batch_size, num_specifications].\n      return self._specification.evaluate(model_logits)\n\n    def flat_objective_fn(x):\n      return _maximize_margin(objective_fn(x))\n\n    def build_loss_fn(idx):\n      def _reduced_loss_fn(x):\n        # Pick worse attack, output has shape [num_restarts, batch_size].\n        return -tf.reduce_sum(select_fn(objective_fn(x), idx))\n      return _reduced_loss_fn\n\n    if _is_spsa_optimizer(self._optimizer_builder):\n      raise ValueError('\"UnrolledSPSA*\" unsupported in '\n                       'MultiTargetedPGDAttack')\n    optimizer = self._optimizer_builder(lr=self._lr, lr_fn=self._lr_fn)\n\n    # Run a separate PGD attack for each specification.\n    def cond(spec_idx, unused_attack, success):\n      # If we are already successful, we break.\n      return tf.logical_and(spec_idx < num_specs,\n                            tf.logical_not(tf.reduce_all(success)))\n\n    def body(spec_idx, attack, success):\n      \"\"\"Runs a separate PGD attack for each specification.\"\"\"\n      adversarial_input = pgd_attack(\n          build_loss_fn(spec_idx), duplicated_inputs,\n          epsilon=self._epsilon, num_steps=self._num_steps,\n          image_bounds=self._input_bounds, random_init=self._random_init,\n          optimizer=optimizer, project_perturbation=self._project_perturbation)\n      new_attack = self.find_worst_attack(flat_objective_fn, adversarial_input,\n                                          batch_size, input_shape)\n      new_logits = self._eval_fn(new_attack)\n      # Count the number of sample that violate any specification.\n      new_success = _any_greater(self._specification.evaluate(new_logits))\n      # The first iteration always sets the attack and logits.\n      use_new_values = tf.logical_or(tf.equal(spec_idx, 0), new_success)\n      print_op = tf.print('Processed specification #', spec_idx)\n      with tf.control_dependencies([print_op]):\n        new_spec_idx = spec_idx + 1\n      return (new_spec_idx,\n              tf.where(use_new_values, new_attack, attack),\n              tf.logical_or(success, new_success))\n\n    _, self._attack, self._success = tf.while_loop(\n        cond, body, back_prop=False, parallel_iterations=1,\n        loop_vars=[\n            tf.constant(0, dtype=tf.int32),\n            inputs,\n            tf.zeros([tf.shape(inputs)[0]], dtype=tf.bool),\n        ])\n    self._logits = self._eval_fn(self._attack, mode='final')\n    return self._attack\n\n  @property\n  def logits(self):\n    self._ensure_is_connected()\n    return self._logits\n\n  @property\n  def attack(self):\n    self._ensure_is_connected()\n    return self._attack\n\n  @property\n  def success(self):\n    self._ensure_is_connected()\n    return self._success\n\n\nclass RestartedAttack(Attack):\n  \"\"\"Wraps an attack to run it multiple times using a tf.while_loop.\"\"\"\n\n  def __init__(self, inner_attack, num_restarts=1):\n    super(RestartedAttack, self).__init__(\n        inner_attack._predictor,  # pylint: disable=protected-access\n        inner_attack._specification,  # pylint: disable=protected-access\n        name='restarted_' + inner_attack.module_name,\n        predictor_kwargs=inner_attack._kwargs)  # pylint: disable=protected-access\n    self._inner_attack = inner_attack\n    self._num_restarts = num_restarts\n    # Prevent the inner attack from updating batch normalization statistics.\n    self._inner_attack.force_mode('intermediate')\n\n  def _build(self, inputs, labels):\n\n    def cond(i, unused_attack, success):\n      # If we are already successful, we break.\n      return tf.logical_and(i < self._num_restarts,\n                            tf.logical_not(tf.reduce_all(success)))\n\n    def body(i, attack, success):\n      new_attack = self._inner_attack(inputs, labels)\n      new_success = self._inner_attack.success\n      # The first iteration always sets the attack.\n      use_new_values = tf.logical_or(tf.equal(i, 0), new_success)\n      return (i + 1,\n              tf.where(use_new_values, new_attack, attack),\n              tf.logical_or(success, new_success))\n\n    _, self._attack, self._success = tf.while_loop(\n        cond, body, back_prop=False, parallel_iterations=1,\n        loop_vars=[\n            tf.constant(0, dtype=tf.int32),\n            inputs,\n            tf.zeros([tf.shape(inputs)[0]], dtype=tf.bool),\n        ])\n    self._logits = self._eval_fn(self._attack, mode='final')\n    return self._attack\n\n  @property\n  def logits(self):\n    self._ensure_is_connected()\n    return self._logits\n\n  @property\n  def attack(self):\n    self._ensure_is_connected()\n    return self._attack\n\n  @property\n  def success(self):\n    self._ensure_is_connected()\n    return self._success\n"
  },
  {
    "path": "interval_bound_propagation/src/bounds.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Definition of input bounds to each layer.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport abc\nimport itertools\n\nimport six\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\n@six.add_metaclass(abc.ABCMeta)\nclass AbstractBounds(object):\n  \"\"\"Abstract bounds class.\"\"\"\n\n  def __init__(self):\n    self._update_cache_op = None\n\n  @classmethod\n  @abc.abstractmethod\n  def convert(cls, bounds):\n    \"\"\"Converts another bound type to this type.\"\"\"\n\n  @abc.abstractproperty\n  def shape(self):\n    \"\"\"Returns shape (as list) of the tensor, including batch dimension.\"\"\"\n\n  def concretize(self):\n    return self\n\n  def _raise_not_implemented(self, name):\n    raise NotImplementedError(\n        '{} modules are not supported by \"{}\".'.format(\n            name, self.__class__.__name__))\n\n  def apply_linear(self, wrapper, w, b):  # pylint: disable=unused-argument\n    self._raise_not_implemented('snt.Linear')\n\n  def apply_conv1d(self, wrapper, w, b, padding, stride):  # pylint: disable=unused-argument\n    self._raise_not_implemented('snt.Conv1D')\n\n  def apply_conv2d(self, wrapper, w, b, padding, strides):  # pylint: disable=unused-argument\n    self._raise_not_implemented('snt.Conv2D')\n\n  def apply_increasing_monotonic_fn(self, wrapper, fn, *args, **parameters):  # pylint: disable=unused-argument\n    self._raise_not_implemented(fn.__name__)\n\n  def apply_piecewise_monotonic_fn(self, wrapper, fn, boundaries, *args):  # pylint: disable=unused-argument\n    self._raise_not_implemented(fn.__name__)\n\n  def apply_batch_norm(self, wrapper, mean, variance, scale, bias, epsilon):  # pylint: disable=unused-argument\n    self._raise_not_implemented('ibp.BatchNorm')\n\n  def apply_batch_reshape(self, wrapper, shape):  # pylint: disable=unused-argument\n    self._raise_not_implemented('snt.BatchReshape')\n\n  def apply_softmax(self, wrapper):  # pylint: disable=unused-argument\n    self._raise_not_implemented('tf.nn.softmax')\n\n  @property\n  def update_cache_op(self):\n    \"\"\"TF op to update cached bounds for re-use across session.run calls.\"\"\"\n    if self._update_cache_op is None:\n      raise ValueError('Bounds not cached: enable_caching() not called.')\n    return self._update_cache_op\n\n  def enable_caching(self):\n    \"\"\"Enables caching the bounds for re-use across session.run calls.\"\"\"\n    if self._update_cache_op is not None:\n      raise ValueError('Bounds already cached: enable_caching() called twice.')\n    self._update_cache_op = self._set_up_cache()\n\n  def _set_up_cache(self):\n    \"\"\"Replace fields with cached versions.\n\n    Returns:\n      TensorFlow op to update the cache.\n    \"\"\"\n    return tf.no_op()  # By default, don't cache.\n\n  def _cache_with_update_op(self, tensor):\n    \"\"\"Creates non-trainable variable to cache the tensor across sess.run calls.\n\n    Args:\n      tensor: Tensor to cache.\n\n    Returns:\n      cached_tensor: Non-trainable variable to contain the cached value\n        of `tensor`.\n      update_op: TensorFlow op to re-evaluate `tensor` and assign the result\n        to `cached_tensor`.\n    \"\"\"\n    cached_tensor = tf.get_variable(\n        tensor.name.replace(':', '__') + '_ibp_cache',\n        shape=tensor.shape, dtype=tensor.dtype, trainable=False)\n    update_op = tf.assign(cached_tensor, tensor)\n    return cached_tensor, update_op\n\n\nclass IntervalBounds(AbstractBounds):\n  \"\"\"Axis-aligned bounding box.\"\"\"\n\n  def __init__(self, lower, upper):\n    super(IntervalBounds, self).__init__()\n    self._lower = lower\n    self._upper = upper\n\n  @property\n  def lower(self):\n    return self._lower\n\n  @property\n  def upper(self):\n    return self._upper\n\n  @property\n  def shape(self):\n    return self.lower.shape.as_list()\n\n  def __iter__(self):\n    yield self.lower\n    yield self.upper\n\n  @classmethod\n  def convert(cls, bounds):\n    if isinstance(bounds, tf.Tensor):\n      return cls(bounds, bounds)\n    bounds = bounds.concretize()\n    if not isinstance(bounds, cls):\n      raise ValueError('Cannot convert \"{}\" to \"{}\"'.format(bounds,\n                                                            cls.__name__))\n    return bounds\n\n  def apply_linear(self, wrapper, w, b):\n    return self._affine(w, b, tf.matmul)\n\n  def apply_conv1d(self, wrapper, w, b, padding, stride):\n    return self._affine(w, b, tf.nn.conv1d, padding=padding, stride=stride)\n\n  def apply_conv2d(self, wrapper, w, b, padding, strides):\n    return self._affine(w, b, tf.nn.convolution,\n                        padding=padding, strides=strides)\n\n  def _affine(self, w, b, fn, **kwargs):\n    c = (self.lower + self.upper) / 2.\n    r = (self.upper - self.lower) / 2.\n    c = fn(c, w, **kwargs)\n    if b is not None:\n      c = c + b\n    r = fn(r, tf.abs(w), **kwargs)\n    return IntervalBounds(c - r, c + r)\n\n  def apply_increasing_monotonic_fn(self, wrapper, fn, *args, **parameters):\n    args_lower = [self.lower] + [a.lower for a in args]\n    args_upper = [self.upper] + [a.upper for a in args]\n    return IntervalBounds(fn(*args_lower), fn(*args_upper))\n\n  def apply_piecewise_monotonic_fn(self, wrapper, fn, boundaries, *args):\n    valid_values = []\n    for a in [self] + list(args):\n      vs = []\n      vs.append(a.lower)\n      vs.append(a.upper)\n      for b in boundaries:\n        vs.append(\n            tf.maximum(a.lower, tf.minimum(a.upper, b * tf.ones_like(a.lower))))\n      valid_values.append(vs)\n    outputs = []\n    for inputs in itertools.product(*valid_values):\n      outputs.append(fn(*inputs))\n    outputs = tf.stack(outputs, axis=-1)\n    return IntervalBounds(tf.reduce_min(outputs, axis=-1),\n                          tf.reduce_max(outputs, axis=-1))\n\n  def apply_batch_norm(self, wrapper, mean, variance, scale, bias, epsilon):\n    # Element-wise multiplier.\n    multiplier = tf.rsqrt(variance + epsilon)\n    if scale is not None:\n      multiplier *= scale\n    w = multiplier\n    # Element-wise bias.\n    b = -multiplier * mean\n    if bias is not None:\n      b += bias\n    b = tf.squeeze(b, axis=0)\n    # Because the scale might be negative, we need to apply a strategy similar\n    # to linear.\n    c = (self.lower + self.upper) / 2.\n    r = (self.upper - self.lower) / 2.\n    c = tf.multiply(c, w) + b\n    r = tf.multiply(r, tf.abs(w))\n    return IntervalBounds(c - r, c + r)\n\n  def apply_batch_reshape(self, wrapper, shape):\n    return IntervalBounds(snt.BatchReshape(shape)(self.lower),\n                          snt.BatchReshape(shape)(self.upper))\n\n  def apply_softmax(self, wrapper):\n    ub = self.upper\n    lb = self.lower\n    # Keep diagonal and take opposite bound for non-diagonals.\n    lbs = tf.matrix_diag(lb) + tf.expand_dims(ub, axis=-2) - tf.matrix_diag(ub)\n    ubs = tf.matrix_diag(ub) + tf.expand_dims(lb, axis=-2) - tf.matrix_diag(lb)\n    # Get diagonal entries after softmax operation.\n    ubs = tf.matrix_diag_part(tf.nn.softmax(ubs))\n    lbs = tf.matrix_diag_part(tf.nn.softmax(lbs))\n    return IntervalBounds(lbs, ubs)\n\n  def _set_up_cache(self):\n    self._lower, update_lower_op = self._cache_with_update_op(self._lower)\n    self._upper, update_upper_op = self._cache_with_update_op(self._upper)\n    return tf.group([update_lower_op, update_upper_op])\n\n\n"
  },
  {
    "path": "interval_bound_propagation/src/crown.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"CROWN-IBP implementation.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport collections\n\nfrom absl import logging\nfrom interval_bound_propagation.src import bounds\nfrom interval_bound_propagation.src import fastlin\nfrom interval_bound_propagation.src import loss\nfrom interval_bound_propagation.src import model\nfrom interval_bound_propagation.src import specification as specification_lib\nfrom interval_bound_propagation.src import utils\nfrom interval_bound_propagation.src import verifiable_wrapper\nimport tensorflow.compat.v1 as tf\n\n\nclass BackwardBounds(bounds.AbstractBounds):\n  \"\"\"Implementation of backward bound propagation used by CROWN.\"\"\"\n\n  def __init__(self, lower, upper):\n    super(BackwardBounds, self).__init__()\n    # Setting \"lower\" or \"upper\" to None will avoid creating the computation\n    # graph for CROWN lower or upper bounds. For verifiable training, only the\n    # upper bound is necessary.\n    self._lower = lower\n    self._upper = upper\n\n  @property\n  def lower(self):\n    return self._lower\n\n  @property\n  def upper(self):\n    return self._upper\n\n  @property\n  def shape(self):\n    return self.lower.shape.as_list()\n\n  def concretize(self):\n    \"\"\"Returns lower and upper interval bounds.\"\"\"\n    lb = ub = None\n    if self.lower is not None:\n      lb = (\n          tf.einsum('nsi,ni->ns',\n                    self._reshape_to_rank(tf.maximum(self.lower.w, 0), 3),\n                    self._reshape_to_rank(self.lower.lower, 2)) +\n          tf.einsum('nsi,ni->ns',\n                    self._reshape_to_rank(tf.minimum(self.lower.w, 0), 3),\n                    self._reshape_to_rank(self.lower.upper, 2)))\n      lb += self.lower.b\n    if self.upper is not None:\n      ub = (\n          tf.einsum('nsi,ni->ns',\n                    self._reshape_to_rank(tf.maximum(self.upper.w, 0), 3),\n                    self._reshape_to_rank(self.upper.upper, 2)) +\n          tf.einsum('nsi,ni->ns',\n                    self._reshape_to_rank(tf.minimum(self.upper.w, 0), 3),\n                    self._reshape_to_rank(self.upper.lower, 2)))\n      ub += self.upper.b\n    return bounds.IntervalBounds(lb, ub)\n\n  @classmethod\n  def convert(cls, other_bounds):\n    if isinstance(other_bounds, cls):\n      return other_bounds\n    raise RuntimeError('BackwardBounds does not support conversion from any '\n                       'other bound type.')\n\n  def apply_linear(self, wrapper, w, b):\n    \"\"\"Propagate CROWN bounds backward through a linear layer.\"\"\"\n    def _linear_propagate(bound):\n      \"\"\"Propagate one side of the bound.\"\"\"\n      new_bound_w = tf.einsum('nsk,lk->nsl', bound.w, w)\n      if b is not None:\n        bias = tf.tensordot(bound.w, b, axes=1)\n      return fastlin.LinearExpression(w=new_bound_w, b=bias + bound.b,\n                                      lower=wrapper.input_bounds.lower,\n                                      upper=wrapper.input_bounds.upper)\n    ub_expr = _linear_propagate(self.upper) if self.upper else None\n    lb_expr = _linear_propagate(self.lower) if self.lower else None\n    return BackwardBounds(lb_expr, ub_expr)\n\n  def apply_conv2d(self, wrapper, w, b, padding, strides):\n    \"\"\"Propagate CROWN bounds backward through a convolution layer.\"\"\"\n    def _conv2d_propagate(bound):\n      \"\"\"Propagate one side of the bound.\"\"\"\n      s = tf.shape(bound.w)\n      # Variable bound.w has shape (batch_size, num_specs, H, W, C),\n      # resize it to (batch_size * num_specs, H, W, C) for batch processing.\n      effective_batch_size = tf.reshape(s[0] * s[1], [1])\n      batched_shape = tf.concat([effective_batch_size, s[2:]], 0)\n      # The output of a deconvolution is the input shape of the corresponding\n      # convolution.\n      output_shape = wrapper.input_bounds.lower.shape\n      batched_output_shape = tf.concat([effective_batch_size, output_shape[1:]],\n                                       0)\n      # Batched transpose convolution for efficiency.\n      bound_batch = tf.nn.conv2d_transpose(tf.reshape(bound.w, batched_shape),\n                                           filter=w,\n                                           output_shape=batched_output_shape,\n                                           strides=[1] + list(strides) + [1],\n                                           padding=padding)\n      # Reshape results to (batch_size, num_specs, new_H, new_W, new_C).\n      new_shape = tf.concat(\n          [tf.reshape(s[0], [1]), tf.reshape(s[1], [1]), output_shape[1:]], 0)\n      new_bound_w = tf.reshape(bound_batch, new_shape)\n      # If this convolution has bias, multiplies it with current w.\n      bias = 0\n      if b is not None:\n        # Variable bound.w has dimension (batch_size, num_specs, H, W, C),\n        # accumulate H and W, and do a dot product for each channel C.\n        bias = tf.tensordot(tf.reduce_sum(bound.w, [2, 3]), b, axes=1)\n      return fastlin.LinearExpression(w=new_bound_w, b=bias + bound.b,\n                                      lower=wrapper.input_bounds.lower,\n                                      upper=wrapper.input_bounds.upper)\n    ub_expr = _conv2d_propagate(self.upper) if self.upper else None\n    lb_expr = _conv2d_propagate(self.lower) if self.lower else None\n    return BackwardBounds(lb_expr, ub_expr)\n\n  def _get_monotonic_fn_bound(self, wrapper, fn):\n    \"\"\"Compute CROWN upper and lower linear bounds for a given function fn.\"\"\"\n    # Get lower and upper bounds from forward IBP pass.\n    lb, ub = wrapper.input_bounds.lower, wrapper.input_bounds.upper\n    if fn.__name__ == 'relu':\n      # CROWN upper and lower linear bounds for ReLU.\n      f_lb = tf.minimum(lb, 0)\n      f_ub = tf.maximum(ub, 0)\n      # When both ub and lb are very close to 0 we might have NaN issue,\n      # so we have to avoid this happening.\n      f_ub = tf.maximum(f_ub, f_lb + 1e-8)\n      # CROWN upper/lower scaling matrices and biases.\n      ub_scaling_matrix = f_ub / (f_ub - f_lb)\n      ub_bias = -f_lb * ub_scaling_matrix\n      # Expand dimension for using broadcast later.\n      ub_scaling_matrix = tf.expand_dims(ub_scaling_matrix, 1)\n      lb_scaling_matrix = tf.cast(tf.greater(ub_scaling_matrix, .5),\n                                  dtype=tf.float32)\n      lb_bias = 0.\n    # For 'apply' fn we need to differentiate them through the wrapper.\n    elif isinstance(wrapper, verifiable_wrapper.ImageNormWrapper):\n      inner_module = wrapper.inner_module\n      ub_scaling_matrix = lb_scaling_matrix = inner_module.scale\n      ub_bias = - inner_module.offset * inner_module.scale\n      lb_bias = ub_bias\n    else:\n      raise NotImplementedError('monotonic fn {} is not supported '\n                                'by BackwardBounds'.format(fn.__name__))\n    return ub_scaling_matrix, lb_scaling_matrix, ub_bias, lb_bias\n\n  def apply_increasing_monotonic_fn(self, wrapper, fn, *args):\n    \"\"\"Propagate CROWN bounds backward through a increasing monotonic fn.\"\"\"\n    # Function _get_monotonic_fn_bound returns matrix and bias term for linear\n    # relaxation.\n    (ub_scaling_matrix, lb_scaling_matrix,\n     ub_bias, lb_bias) = self._get_monotonic_fn_bound(wrapper, fn)\n    def _propagate_monotonic_fn(bound, ub_mult, lb_mult):\n      # Matrix multiplication by a diagonal matrix.\n      new_bound_w = ub_mult * ub_scaling_matrix + lb_mult * lb_scaling_matrix\n      # Matrix vector product for the bias term. ub_bias or lb_bias might be 0\n      # or a constant, or need broadcast. They will be handled optimally.\n      b = self._matvec(ub_mult, ub_bias) + self._matvec(lb_mult, lb_bias)\n      return fastlin.LinearExpression(w=new_bound_w, b=bound.b + b,\n                                      lower=wrapper.input_bounds.lower,\n                                      upper=wrapper.input_bounds.upper)\n    # Multiplies w to upper or lower scaling terms according to its sign.\n    ub_expr = _propagate_monotonic_fn(\n        self.upper, tf.maximum(self.upper.w, 0),\n        tf.minimum(self.upper.w, 0)) if self.upper else None\n    lb_expr = _propagate_monotonic_fn(\n        self.lower, tf.minimum(self.lower.w, 0),\n        tf.maximum(self.lower.w, 0)) if self.lower else None\n    return BackwardBounds(lb_expr, ub_expr)\n\n  def apply_batch_reshape(self, wrapper, shape):\n    \"\"\"Propagate CROWN bounds backward through a reshape layer.\"\"\"\n    input_shape = wrapper.input_bounds.lower.shape[1:]\n    def _propagate_batch_flatten(bound):\n      new_bound_w = tf.reshape(\n          bound.w, tf.concat([tf.shape(bound.w)[:2], input_shape], 0))\n      return fastlin.LinearExpression(w=new_bound_w, b=bound.b,\n                                      lower=wrapper.input_bounds.lower,\n                                      upper=wrapper.input_bounds.upper)\n    ub_expr = _propagate_batch_flatten(self.upper) if self.upper else None\n    lb_expr = _propagate_batch_flatten(self.lower) if self.lower else None\n    return BackwardBounds(lb_expr, ub_expr)\n\n  @staticmethod\n  def _reshape_to_rank(a, rank):\n    \"\"\"Reshapes to the given rank while keeping the first (rank-1) dims.\"\"\"\n    shape = tf.concat([tf.shape(a)[0:(rank - 1)], [-1]], axis=-1)\n    return tf.reshape(a, shape)\n\n  @staticmethod\n  def _matvec(a, b):\n    \"\"\"Specialized matvec detecting the case where b is 0 or constant.\"\"\"\n    if isinstance(b, int) or isinstance(b, float):\n      if b == 0:\n        # For efficiency we directly return constant 0, no graph generated.\n        return 0\n      else:\n        # Broadcasting a constant.\n        return a * b\n    elif len(b.shape) == 1:\n      # Need to broadcast against all examples in the batch. This can be done\n      # using an einsum \"tf.einsum('ns...c,c->ns', a, b)\" but it currently\n      # triggers a compiler bug on TPUs, thus we use the following instead.\n      return tf.einsum('nsc,c->ns', tf.reduce_sum(a, [2, 3]), b)\n    else:\n      # Normal 1D or 3D mat-vec product.\n      return tf.einsum('nsi,ni->ns',\n                       BackwardBounds._reshape_to_rank(a, 3),\n                       BackwardBounds._reshape_to_rank(b, 2))\n\n\nScalarMetrics = collections.namedtuple('ScalarMetrics', [\n    'nominal_accuracy',\n    # Verified accuracy using pure IBP bounds.\n    'verified_accuracy',\n    # Verified accuracy using CROWN and IBP mixture.\n    'crown_ibp_verified_accuracy',\n    'attack_accuracy',\n    'attack_success'])\n\nScalarLosses = collections.namedtuple('ScalarLosses', [\n    'nominal_cross_entropy',\n    'attack_cross_entropy',\n    'verified_loss'])\n\n\nclass Losses(loss.Losses):\n  \"\"\"Helper to compute CROWN-IBP losses.\"\"\"\n\n  def __init__(self, predictor, specification=None, pgd_attack=None,\n               interval_bounds_loss_type='xent',\n               interval_bounds_hinge_margin=10.,\n               label_smoothing=0.,\n               use_crown_ibp=False,\n               crown_bound_schedule=None):\n    super(Losses, self).__init__(predictor, specification, pgd_attack,\n                                 interval_bounds_loss_type,\n                                 interval_bounds_hinge_margin,\n                                 label_smoothing)\n    self._use_crown_ibp = use_crown_ibp\n    self._crown_bound_schedule = crown_bound_schedule\n\n  def _get_specification_bounds(self):\n    \"\"\"Get upper bounds on specification. Used for building verified loss.\"\"\"\n    ibp_bounds = self._specification(self._predictor.modules)\n    # Compute verified accuracy using IBP bounds.\n    v = tf.reduce_max(ibp_bounds, axis=1)\n    self._interval_bounds_accuracy = tf.reduce_mean(\n        tf.cast(v <= 0., tf.float32))\n    # CROWN-IBP bounds.\n    if self._use_crown_ibp:\n      logging.info('CROWN-IBP active')\n      def _build_crown_ibp_bounds():\n        \"\"\"Create the computationally expensive CROWN bounds for tf.cond.\"\"\"\n        predictor = self._predictor\n        # CROWN is computed backwards so we need to start with a\n        # initial bound related to the specification.\n        init_crown_bounds = create_initial_backward_bounds(self._specification,\n                                                           predictor.modules)\n        # Now propagate the specification matrix layer by layer;\n        # we only need the CROWN upper bound, do not need lower bound.\n        crown_bound = predictor.propagate_bound_backward(init_crown_bounds,\n                                                         compute_upper=True,\n                                                         compute_lower=False)\n        # A linear mixture of the two bounds with a schedule.\n        return self._crown_bound_schedule * crown_bound.upper + \\\n               (1. - self._crown_bound_schedule) * ibp_bounds\n      # If the coefficient for CROWN bound is close to 0, compute IBP only.\n      mixture_bounds = tf.cond(self._crown_bound_schedule < 1e-6,\n                               lambda: ibp_bounds, _build_crown_ibp_bounds)\n      v = tf.reduce_max(mixture_bounds, axis=1)\n      self._crown_ibp_accuracy = tf.reduce_mean(tf.cast(v <= 0., tf.float32))\n    else:\n      mixture_bounds = ibp_bounds\n      self._crown_ibp_accuracy = tf.constant(0.)\n    return mixture_bounds\n\n  @property\n  def scalar_metrics(self):\n    self._ensure_is_connected()\n    return ScalarMetrics(self._nominal_accuracy,\n                         self._interval_bounds_accuracy,\n                         self._crown_ibp_accuracy,\n                         self._attack_accuracy,\n                         self._attack_success)\n\n  @property\n  def scalar_losses(self):\n    self._ensure_is_connected()\n    return ScalarLosses(self._cross_entropy,\n                        self._attack_cross_entropy,\n                        self._verified_loss)\n\n\nclass VerifiableModelWrapper(model.VerifiableModelWrapper):\n  \"\"\"Model wrapper with CROWN-IBP backward bound propagation.\"\"\"\n\n  def _propagate(self, current_module, current_bounds):\n    \"\"\"Propagate CROWN bounds in a backwards manner.\"\"\"\n    # Construct bounds for this layer.\n    if isinstance(current_module, verifiable_wrapper.ModelInputWrapper):\n      if current_module.index != 0:\n        raise NotImplementedError('CROWN backpropagation does not support '\n                                  'multiple inputs.')\n      return current_bounds\n    # Propagate the bounds through the current layer.\n    new_bounds = current_module.propagate_bounds(current_bounds)\n    prev_modules = self._module_depends_on[current_module]\n    # We assume that each module only depends on one module.\n    if len(prev_modules) != 1:\n      raise NotImplementedError('CROWN for non-sequential networks is not '\n                                'implemented.')\n    return self._propagate(prev_modules[0], new_bounds)\n\n  def propagate_bound_backward(self, initial_bound,\n                               compute_upper=True, compute_lower=False):\n    \"\"\"Propagates CROWN bounds backward through the network.\n\n    This function assumes that we have obtained bounds for all intermediate\n    layers using IBP. Currently only sequential networks are implemented.\n\n    Args:\n\n      initial_bound: A BackwardBounds object containing the initial matrices\n        and biases to start bound propagation.\n      compute_upper: Set to True to construct the computation graph for the\n        CROWN upper bound. For verified training, only the upper bound is\n        needed. Default is True.\n      compute_lower: Set to True to construct the computation graph for the\n        CROWN lower bound. Default is False.\n\n    Returns:\n      IntervalBound instance corresponding to bounds on the specification.\n    \"\"\"\n    if (not compute_upper) and (not compute_lower):\n      raise ValueError('At least one of \"compute_upper\" or \"compute_lower\" '\n                       'needs to be True')\n    self._ensure_is_connected()\n    # We start bound propagation from the logit layer.\n    logit_layer = self._produced_by[self._logits.name]\n    # If only one of ub or lb is needed, we set the unnecessary one to None.\n    ub = initial_bound.upper if compute_upper else None\n    lb = initial_bound.lower if compute_lower else None\n    bound = BackwardBounds(lb, ub)\n    crown_bound = self._propagate(logit_layer, bound)\n    return crown_bound.concretize()\n\n\ndef create_initial_backward_bounds(spec, modules):\n  \"\"\"Create the initial BackwardBounds according to specification.\"\"\"\n  last_bounds = bounds.IntervalBounds.convert(modules[-1].input_bounds)\n  if isinstance(spec, specification_lib.ClassificationSpecification):\n    c_correct = tf.expand_dims(\n        tf.one_hot(spec.correct_idx[:, 1], spec.num_specifications + 1), 1)\n    c_wrong = tf.one_hot(spec.wrong_idx[:, :, 1], spec.num_specifications + 1)\n    c = c_wrong - c_correct\n    b = tf.zeros(spec.num_specifications)\n    lb = ub = fastlin.LinearExpression(w=c, b=b, lower=last_bounds.lower,\n                                       upper=last_bounds.upper)\n  elif isinstance(spec, specification_lib.LinearSpecification):\n    b = spec.d if spec.d is not None else tf.zeros(spec.num_specifications)\n    lb = ub = fastlin.LinearExpression(w=spec.c, b=b, lower=last_bounds.lower,\n                                       upper=last_bounds.upper)\n  else:\n    raise ValueError('Unknown specification class type \"{}\"'.format(str(spec)))\n  return BackwardBounds(lb, ub)\n\n\ndef create_classification_losses(\n    global_step, inputs, label, predictor_network, epsilon, loss_weights,\n    warmup_steps=0, rampup_steps=-1, input_bounds=(0., 1.), options=None):\n  \"\"\"Create the training loss for CROWN-IBP.\"\"\"\n  def _is_loss_active(init, final, warmup=None):\n    return init > 0. or final > 0. or (warmup is not None and warmup > 0.)\n  if 'crown_bound' in loss_weights:\n    schedule = utils.build_loss_schedule(global_step, warmup_steps,\n                                         rampup_steps,\n                                         **loss_weights.get('crown_bound'))\n    use_crown_ibp = _is_loss_active(**loss_weights.get('crown_bound'))\n  else:\n    schedule = None\n    use_crown_ibp = False\n  # Use the loss builder for CROWN-IBP with additional kwargs.\n  def _loss_builder(*args, **kwargs):\n    kwargs.update(dict(use_crown_ibp=use_crown_ibp,\n                       crown_bound_schedule=schedule))\n    return Losses(*args, **kwargs)\n  return utils.create_classification_losses(\n      global_step, inputs, label, predictor_network, epsilon,\n      loss_weights, warmup_steps, rampup_steps, input_bounds,\n      loss_builder=_loss_builder, options=options)\n"
  },
  {
    "path": "interval_bound_propagation/src/fastlin.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Fast-Lin symbolic bound calculation for common neural network layers.\n\nThe Fast-Lin algorithm expresses lower and upper bounds of each layer of\na neural network as a symbolic linear expression in the input neurons,\nrelaxing the ReLU layers to retain linearity at the expense of tightness.\n\nReference: \"Towards Fast Computation of Certified Robustness for ReLU Networks\",\nhttps://arxiv.org/pdf/1804.09699.pdf.\n\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport collections\n\nfrom absl import logging\nfrom interval_bound_propagation.src import bounds as basic_bounds\nfrom interval_bound_propagation.src import relative_bounds\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\n# Holds the linear expressions serving as bounds.\n# w: [batch_size, input_size, output_shape] storing the weights.\n# b: [batch_size, output_shape] storing the bias.\n# lower: [batch_size, input_size] storing the lower bounds on inputs.\n# upper: [batch_size, input_size] storing the upper bounds on inputs.\n# `lower` and `upper` tensors are always flattened representations of the\n# original inputs.\nLinearExpression = collections.namedtuple(\n    'LinearExpression', ['w', 'b', 'lower', 'upper'])\n\n\nclass SymbolicBounds(basic_bounds.AbstractBounds):\n  \"\"\"Fast-Lin bounds (https://arxiv.org/abs/1804.09699).\"\"\"\n\n  def __init__(self, lower, upper):\n    super(SymbolicBounds, self).__init__()\n    self._lower = lower\n    self._upper = upper\n    self._prior_bounds = None\n    self._concretized = None\n\n  @property\n  def lower(self):\n    return self._lower\n\n  @property\n  def upper(self):\n    return self._upper\n\n  @property\n  def shape(self):\n    return self.lower.b.shape.as_list()\n\n  def concretize(self):\n    \"\"\"Returns lower and upper interval bounds.\"\"\"\n    if self._concretized is None:\n      # Construct once and cache.\n      lb, ub = self._concretize_bounds(self.lower, self.upper)\n\n      # Apply intersections with prior runs.\n      if self._prior_bounds is not None:\n        lb = tf.maximum(lb, self._prior_bounds.lower)\n        ub = tf.minimum(ub, self._prior_bounds.upper)\n      self._concretized = basic_bounds.IntervalBounds(lb, ub)\n\n    return self._concretized\n\n  def with_priors(self, existing_bounds):\n    if existing_bounds is not None:\n      self._prior_bounds = existing_bounds.concretize()\n      # These priors are applied the next time concretize() is called.\n      self._concretized = None\n    return self\n\n  @classmethod\n  def convert(cls, bounds):\n    if isinstance(bounds, cls):\n      return bounds\n\n    if isinstance(bounds, tf.Tensor):\n      bounds = basic_bounds.IntervalBounds(bounds, bounds)\n    bounds = bounds.concretize()\n    if not isinstance(bounds, basic_bounds.IntervalBounds):\n      raise ValueError('Cannot convert \"{}\" to \"SymbolicBounds\"'.format(bounds))\n\n    lower, upper = cls._initial_symbolic_bounds(bounds.lower, bounds.upper)\n    return cls(lower, upper)\n\n  def apply_linear(self, wrapper, w, b):\n    w_pos = tf.maximum(w, 0)\n    w_neg = tf.minimum(w, 0)\n    lb = self._add_expression(\n        self._scale_expression(self.lower, w_pos),\n        self._scale_expression(self.upper, w_neg)\n    )\n    lb = self._add_bias(lb, b)\n    ub = self._add_expression(\n        self._scale_expression(self.lower, w_neg),\n        self._scale_expression(self.upper, w_pos)\n    )\n    ub = self._add_bias(ub, b)\n    return SymbolicBounds(lb, ub).with_priors(wrapper.output_bounds)\n\n  def apply_conv1d(self, wrapper, w, b, padding, stride):\n    w_pos = tf.maximum(w, 0)\n    w_neg = tf.minimum(w, 0)\n    lb = self._add_expression(\n        self._conv1d_expression(self.lower, w_pos, padding, stride),\n        self._conv1d_expression(self.upper, w_neg, padding, stride))\n    lb = self._add_bias(lb, b)\n\n    ub = self._add_expression(\n        self._conv1d_expression(self.upper, w_pos, padding, stride),\n        self._conv1d_expression(self.lower, w_neg, padding, stride))\n    ub = self._add_bias(ub, b)\n    return SymbolicBounds(lb, ub).with_priors(wrapper.output_bounds)\n\n  def apply_conv2d(self, wrapper, w, b, padding, strides):\n    w_pos = tf.maximum(w, 0)\n    w_neg = tf.minimum(w, 0)\n    lb = self._add_expression(\n        self._conv2d_expression(self.lower, w_pos, padding, strides),\n        self._conv2d_expression(self.upper, w_neg, padding, strides))\n    lb = self._add_bias(lb, b)\n\n    ub = self._add_expression(\n        self._conv2d_expression(self.upper, w_pos, padding, strides),\n        self._conv2d_expression(self.lower, w_neg, padding, strides))\n    ub = self._add_bias(ub, b)\n    return SymbolicBounds(lb, ub).with_priors(wrapper.output_bounds)\n\n  def apply_increasing_monotonic_fn(self, wrapper, fn, *args, **parameters):\n    if fn.__name__ != 'relu':\n      # Fallback to regular interval bound propagation for unsupported\n      # operations.\n      logging.warn('\"%s\" is not supported by SymbolicBounds. '\n                   'Fallback on IntervalBounds.', fn.__name__)\n      interval_bounds = basic_bounds.IntervalBounds.convert(self)\n      converted_args = [basic_bounds.IntervalBounds.convert(b) for b in args]\n      interval_bounds = interval_bounds._increasing_monotonic_fn(  # pylint: disable=protected-access\n          fn, *converted_args)\n      return self.convert(interval_bounds)\n\n    concrete = self.concretize()\n    lb, ub = concrete.lower, concrete.upper\n    is_ambiguous = tf.logical_and(ub > 0, lb < 0)\n    # Ensure denominator is always positive, even when not needed.\n    ambiguous_denom = tf.where(is_ambiguous, ub - lb, tf.ones_like(ub))\n    scale = tf.where(\n        is_ambiguous, ub / ambiguous_denom,\n        tf.where(lb >= 0, tf.ones_like(lb), tf.zeros_like(lb)))\n    bias = tf.where(is_ambiguous, -lb, tf.zeros_like(lb))\n    lb_out = LinearExpression(\n        w=tf.expand_dims(scale, 1) * self.lower.w,\n        b=scale * self.lower.b,\n        lower=self.lower.lower, upper=self.lower.upper)\n    ub_out = LinearExpression(\n        w=tf.expand_dims(scale, 1) * self.upper.w,\n        b=scale * (self.upper.b + bias),\n        lower=self.upper.lower, upper=self.upper.upper)\n    return SymbolicBounds(lb_out, ub_out).with_priors(wrapper.output_bounds)\n\n  def apply_batch_reshape(self, wrapper, shape):\n    return SymbolicBounds(self._batch_reshape_expression(self.lower, shape),\n                          self._batch_reshape_expression(self.upper, shape)\n                         ).with_priors(wrapper.output_bounds)\n\n  # Helper methods.\n  @staticmethod\n  def _add_bias(expr, b):\n    \"\"\"Add bias b to a linear expression.\"\"\"\n    if b is None:\n      return expr\n    return LinearExpression(w=expr.w, b=expr.b + b,\n                            lower=expr.lower, upper=expr.upper)\n\n  @staticmethod\n  def _add_expression(expr_a, expr_b):\n    \"\"\"Add two expression together.\"\"\"\n    return LinearExpression(w=expr_a.w + expr_b.w, b=expr_a.b + expr_b.b,\n                            lower=expr_a.lower, upper=expr_b.upper)\n\n  @staticmethod\n  def _scale_expression(expr, w):\n    \"\"\"Scale a linear expression by w.\"\"\"\n    b = tf.matmul(expr.b, w)\n    w = tf.tensordot(expr.w, w, axes=1)\n    return LinearExpression(w=w, b=b, lower=expr.lower, upper=expr.upper)\n\n  @staticmethod\n  def _conv1d_expression(expr, w, padding, stride):\n    \"\"\"Scale a linear expression by w (through a convolutional layer).\"\"\"\n    b = tf.nn.conv1d(expr.b, w, padding=padding, stride=stride)\n    shape = tf.concat([[tf.reduce_prod(tf.shape(expr.w)[:2])],\n                       tf.shape(expr.w)[2:]], axis=0)\n    w = tf.nn.conv1d(tf.reshape(expr.w, shape), w, padding=padding,\n                     stride=stride)\n    shape = tf.concat([tf.shape(expr.w)[:2], tf.shape(w)[1:]], axis=0)\n    w = tf.reshape(w, shape)\n    return LinearExpression(w=w, b=b, lower=expr.lower, upper=expr.upper)\n\n  @staticmethod\n  def _conv2d_expression(expr, w, padding, strides):\n    \"\"\"Scale a linear expression by w (through a convolutional layer).\"\"\"\n    b = tf.nn.convolution(expr.b, w, padding=padding, strides=strides)\n    shape = tf.concat([[tf.reduce_prod(tf.shape(expr.w)[:2])],\n                       tf.shape(expr.w)[2:]], axis=0)\n    w = tf.nn.convolution(tf.reshape(expr.w, shape), w, padding=padding,\n                          strides=strides)\n    shape = tf.concat([tf.shape(expr.w)[:2], tf.shape(w)[1:]], axis=0)\n    w = tf.reshape(w, shape)\n    return LinearExpression(w=w, b=b, lower=expr.lower, upper=expr.upper)\n\n  @staticmethod\n  def _batch_reshape_expression(expr, shape):\n    w = snt.BatchReshape(shape, preserve_dims=2)(expr.w)\n    b = snt.BatchReshape(shape)(expr.b)\n    return LinearExpression(w=w, b=b, lower=expr.lower, upper=expr.upper)\n\n  @staticmethod\n  def _concretize_bounds(lower, upper):\n    \"\"\"Returns lower and upper interval bounds.\"\"\"\n    if len(lower.b.shape) == 2:\n      equation = 'ijk,ij->ik'\n    elif len(lower.b.shape) == 3:\n      equation = 'ijnc,ij->inc'\n    elif len(lower.b.shape) == 4:\n      equation = 'ijhwc,ij->ihwc'\n    else:\n      raise NotImplementedError('Shape unsupported: {}'.format(lower.b.shape))\n\n    lb = (tf.einsum(equation, tf.maximum(lower.w, 0), lower.lower) +\n          tf.einsum(equation, tf.minimum(lower.w, 0), lower.upper) +\n          lower.b)\n    ub = (tf.einsum(equation, tf.maximum(upper.w, 0), upper.upper) +\n          tf.einsum(equation, tf.minimum(upper.w, 0), upper.lower) +\n          upper.b)\n    return lb, ub\n\n  @staticmethod\n  def _initial_symbolic_bounds(lb, ub):\n    \"\"\"Returns symbolic bounds for the given interval bounds.\"\"\"\n    batch_size = tf.shape(lb)[0]\n    input_shape = lb.shape[1:]\n    zero = tf.zeros_like(lb)\n    lb = snt.BatchFlatten()(lb)\n    ub = snt.BatchFlatten()(ub)\n    input_size = tf.shape(lb)[1]\n    output_shape = tf.concat([[input_size], input_shape], axis=0)\n    identity = tf.reshape(tf.eye(input_size), output_shape)\n    identity = tf.expand_dims(identity, 0)\n    identity = tf.tile(identity, [batch_size] + [1] * (len(input_shape) + 1))\n    expr = LinearExpression(w=identity, b=zero,\n                            lower=lb, upper=ub)\n    return expr, expr\n\n\nclass RelativeSymbolicBounds(SymbolicBounds):\n  \"\"\"Relative-to-nominal variant of Fast-Lin bounds.\"\"\"\n\n  def __init__(self, lower_offset, upper_offset, nominal):\n    super(RelativeSymbolicBounds, self).__init__(lower_offset, upper_offset)\n    self._nominal = nominal\n\n  def concretize(self):\n    \"\"\"Returns lower and upper interval bounds.\"\"\"\n    if self._concretized is None:\n      # Construct once and cache.\n      lb_offset, ub_offset = self._concretize_bounds(self.lower, self.upper)\n\n      # Apply intersections with prior runs.\n      if self._prior_bounds is not None:\n        lb_offset = tf.maximum(lb_offset, self._prior_bounds.lower_offset)\n        ub_offset = tf.minimum(ub_offset, self._prior_bounds.upper_offset)\n      self._concretized = relative_bounds.RelativeIntervalBounds(\n          lb_offset, ub_offset, self._nominal)\n\n    return self._concretized\n\n  @classmethod\n  def convert(cls, bounds):\n    if isinstance(bounds, cls):\n      return bounds\n\n    if isinstance(bounds, tf.Tensor):\n      bounds = relative_bounds.RelativeIntervalBounds(\n          tf.zeros_like(bounds), tf.zeros_like(bounds), bounds)\n    bounds = bounds.concretize()\n    if not isinstance(bounds, relative_bounds.RelativeIntervalBounds):\n      raise ValueError(\n          'Cannot convert \"{}\" to \"RelativeSymbolicBounds\"'.format(bounds))\n\n    lower, upper = cls._initial_symbolic_bounds(bounds.lower_offset,\n                                                bounds.upper_offset)\n    return cls(lower, upper, bounds.nominal)\n\n  def apply_linear(self, wrapper, w, b):\n    bounds_out = super(RelativeSymbolicBounds, self).apply_linear(\n        wrapper, w, b=None)\n\n    nominal_out = tf.matmul(self._nominal, w)\n    if b is not None:\n      nominal_out += b\n\n    return RelativeSymbolicBounds(\n        bounds_out.lower, bounds_out.upper, nominal_out).with_priors(\n            wrapper.output_bounds)\n\n  def apply_conv1d(self, wrapper, w, b, padding, stride):\n    bounds_out = super(RelativeSymbolicBounds, self).apply_conv1d(\n        wrapper, w, b=None, padding=padding, stride=stride)\n\n    nominal_out = tf.nn.conv1d(self._nominal, w,\n                               padding=padding, stride=stride)\n    if b is not None:\n      nominal_out += b\n\n    return RelativeSymbolicBounds(\n        bounds_out.lower, bounds_out.upper, nominal_out).with_priors(\n            wrapper.output_bounds)\n\n  def apply_conv2d(self, wrapper, w, b, padding, strides):\n    bounds_out = super(RelativeSymbolicBounds, self).apply_conv2d(\n        wrapper, w, b=None, padding=padding, strides=strides)\n\n    nominal_out = tf.nn.convolution(self._nominal, w,\n                                    padding=padding, strides=strides)\n    if b is not None:\n      nominal_out += b\n\n    return RelativeSymbolicBounds(\n        bounds_out.lower, bounds_out.upper, nominal_out).with_priors(\n            wrapper.output_bounds)\n\n  def apply_increasing_monotonic_fn(self, wrapper, fn, *args, **parameters):\n    if fn.__name__ != 'relu':\n      # Fallback to regular interval bound propagation for unsupported\n      # operations.\n      logging.warn('\"%s\" is not supported by RelativeSymbolicBounds. '\n                   'Fallback on RelativeIntervalBounds.', fn.__name__)\n      interval_bounds = relative_bounds.RelativeIntervalBounds.convert(self)\n      converted_args = [relative_bounds.RelativeIntervalBounds.convert(b)\n                        for b in args]\n      interval_bounds = interval_bounds._increasing_monotonic_fn(  # pylint: disable=protected-access\n          fn, *converted_args)\n      return self.convert(interval_bounds)\n\n    concrete = self.concretize()\n    lb, ub = concrete.lower_offset, concrete.upper_offset\n    is_ambiguous = tf.logical_and(ub > -self._nominal, lb < -self._nominal)\n    # Ensure denominator is always positive, even when not needed.\n    ambiguous_denom = tf.where(is_ambiguous, ub - lb, tf.ones_like(ub))\n    scale = tf.where(\n        is_ambiguous, (self._nominal + ub) / ambiguous_denom,\n        tf.where(lb >= -self._nominal, tf.ones_like(lb), tf.zeros_like(lb)))\n    scale_complement = tf.where(\n        is_ambiguous, -(self._nominal + lb) / ambiguous_denom,\n        tf.where(lb >= -self._nominal, tf.zeros_like(lb), tf.ones_like(lb)))\n    # Need lb_out.b = scale * (nom_in + lb_in.b) - nom_out\n    # and ub_out.b = scale * (nom_in + ub_in.b - min(nom_in + lb, 0)) - nom_out\n    lower_bias = (scale * (tf.minimum(self._nominal, 0.)) +\n                  scale_complement * tf.minimum(-self._nominal, 0.))\n    upper_bias = (scale * tf.maximum(tf.minimum(-self._nominal, 0.) - lb,\n                                     tf.minimum(self._nominal, 0.)) +\n                  scale_complement * tf.minimum(-self._nominal, 0.))\n    lb_out = LinearExpression(\n        w=tf.expand_dims(scale, 1) * self.lower.w,\n        b=scale * self.lower.b + lower_bias,\n        lower=self.lower.lower, upper=self.lower.upper)\n    ub_out = LinearExpression(\n        w=tf.expand_dims(scale, 1) * self.upper.w,\n        b=scale * self.upper.b + upper_bias,\n        lower=self.upper.lower, upper=self.upper.upper)\n\n    nominal_out = tf.nn.relu(self._nominal)\n    return RelativeSymbolicBounds(\n        lb_out, ub_out, nominal_out).with_priors(wrapper.output_bounds)\n\n  def apply_batch_reshape(self, wrapper, shape):\n    bounds_out = super(RelativeSymbolicBounds, self).apply_batch_reshape(\n        wrapper, shape)\n    nominal_out = snt.BatchReshape(shape)(self._nominal)\n    return RelativeSymbolicBounds(\n        bounds_out.lower, bounds_out.upper, nominal_out).with_priors(\n            wrapper.output_bounds)\n"
  },
  {
    "path": "interval_bound_propagation/src/layer_utils.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Graph construction for dual verification.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom interval_bound_propagation.src import layers\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\ndef conv_output_shape(input_shape, w, padding, strides):\n  \"\"\"Calculates the output shape of the given N-D convolution.\n\n  Args:\n    input_shape: Integer list of length N+1 specifying the non-batch dimensions\n      of the inputs: [input_height, input_width, input_channels].\n    w: (N+2)D tensor of shape (kernel_height, kernel_width, input_channels,\n      output_channels) containing weights for the convolution.\n    padding: `\"VALID\"` or `\"SAME\"`, the convolution's padding algorithm.\n    strides: Integer list of length N: `[vertical_stride, horizontal_stride]`.\n\n  Returns:\n    Integer list of length N+1 specifying the non-batch dimensions\n      of the outputs: [output_height, output_width, output_channels].\n\n  Raises:\n    ValueError: if an unsupported convolution dimensionality is encountered.\n  \"\"\"\n  # Connect a convolution (never to be run) to infer the output's\n  # spatial structure.\n  dummy_inputs = tf.zeros(dtype=w.dtype, shape=([1] + input_shape))\n  if len(w.shape) == 4:\n    dummy_outputs = tf.nn.convolution(dummy_inputs,\n                                      w, padding=padding, strides=strides)\n  elif len(w.shape) == 3:\n    dummy_outputs = tf.nn.conv1d(dummy_inputs,\n                                 w, padding=padding, stride=strides[0])\n  else:\n    raise ValueError()\n  return dummy_outputs.shape.as_list()[1:]\n\n\ndef materialise_conv(w, b, input_shape, padding, strides):\n  \"\"\"Converts an N-D convolution to an equivalent linear layer.\n\n  Args:\n    w: (N+2)D tensor of shape (kernel_height, kernel_width, input_channels,\n      output_channels) containing the convolution weights.\n    b: 1D tensor of shape (output_channels) containing the convolution biases,\n      or `None` if no biases.\n    input_shape: Integer list of length N+1 specifying the non-batch dimensions\n      of the inputs: [input_height, input_width, input_channels].\n    padding: `\"VALID\"` or `\"SAME\"`, the convolution's padding algorithm.\n    strides: Integer list of length N: `[vertical_stride, horizontal_stride]`.\n\n  Returns:\n    w: 2D tensor of shape (input_height * input_width * input_channels,\n      output_height * output_width * output_channels) containing weights.\n    b: 1D tensor of shape (output_height * output_width * output_channels)\n      containing biases, or `None` if no biases.\n\n  Raises:\n    ValueError: if an unsupported convolution dimensionality is encountered.\n  \"\"\"\n  if len(input_shape) == 3:\n    return _materialise_conv2d(w, b, input_shape[0], input_shape[1],\n                               padding, strides)\n  elif len(input_shape) == 2:\n    return _materialise_conv1d(w, b, input_shape[0], padding, strides[0])\n  else:\n    raise ValueError()\n\n\ndef _materialise_conv2d(w, b, input_height, input_width, padding, strides):\n  \"\"\"Converts a convolution to an equivalent linear layer.\n\n  Args:\n    w: 4D tensor of shape (kernel_height, kernel_width, input_channels,\n      output_channels) containing the convolution weights.\n    b: 1D tensor of shape (output_channels) containing the convolution biases,\n      or `None` if no biases.\n    input_height: height of the input tensor.\n    input_width: width of the input tensor.\n    padding: `\"VALID\"` or `\"SAME\"`, the convolution's padding algorithm.\n    strides: Integer list of `[vertical_stride, horizontal_stride]`.\n\n  Returns:\n    w: 2D tensor of shape (input_height * input_width * input_channels,\n      output_height * output_width * output_channels) containing weights.\n    b: 1D tensor of shape (output_height * output_width * output_channels)\n      containing biases, or `None` if no biases.\n  \"\"\"\n  kernel_height = w.shape[0].value\n  kernel_width = w.shape[1].value\n  input_channels = w.shape[2].value\n  output_channels = w.shape[3].value\n\n  # Temporarily move the input_channels dimension to output_channels.\n  w = tf.reshape(w, shape=(kernel_height, kernel_width, 1,\n                           input_channels * output_channels))\n  # Apply the convolution to elementary (i.e. one-hot) inputs.\n  diagonal_input = tf.reshape(\n      tf.eye(input_height * input_width, dtype=w.dtype),\n      shape=[input_height * input_width, input_height, input_width, 1])\n  conv = tf.nn.convolution(\n      diagonal_input, w,\n      padding=padding, strides=strides)\n  output_height = conv.shape[1].value\n  output_width = conv.shape[2].value\n  # conv is of shape (input_height * input_width, output_height, output_width,\n  #                   input_channels * output_channels).\n  # Reshape it to (input_height * input_width * input_channels,\n  #                output_height * output_width * output_channels).\n  w = tf.reshape(conv, shape=(\n      [input_height * input_width,\n       output_height, output_width,\n       input_channels, output_channels]))\n  w = tf.transpose(w, perm=[0, 3, 1, 2, 4])\n  w = tf.reshape(w, shape=(\n      [input_height * input_width * input_channels,\n       output_height * output_width * output_channels]))\n\n  # Broadcast b over spatial dimensions.\n  b = tf.tile(b, [output_height * output_width]) if b is not None else None\n\n  return w, b\n\n\ndef _materialise_conv1d(w, b, input_length, padding, stride):\n  \"\"\"Converts a convolution to an equivalent linear layer.\n\n  Args:\n    w: 3D tensor of shape (kernel_length, input_channels,\n      output_channels) containing the convolution weights.\n    b: 1D tensor of shape (output_channels) containing the convolution biases,\n      or `None` if no biases.\n    input_length: length of the input tensor.\n    padding: `\"VALID\"` or `\"SAME\"`, the convolution's padding algorithm.\n    stride: Integer stride.\n\n  Returns:\n    w: 2D tensor of shape (input_length * input_channels,\n      output_length * output_channels) containing weights.\n    b: 1D tensor of shape (output_length * output_channels)\n      containing biases, or `None` if no biases.\n  \"\"\"\n  kernel_length = w.shape[0].value\n  input_channels = w.shape[1].value\n  output_channels = w.shape[2].value\n\n  # Temporarily move the input_channels dimension to output_channels.\n  w = tf.reshape(w, shape=(kernel_length, 1,\n                           input_channels * output_channels))\n  # Apply the convolution to elementary (i.e. one-hot) inputs.\n  diagonal_input = tf.reshape(\n      tf.eye(input_length, dtype=w.dtype),\n      shape=[input_length, input_length, 1])\n  conv = tf.nn.conv1d(\n      diagonal_input, w,\n      padding=padding, stride=stride)\n  output_length = conv.shape[1].value\n  # conv is of shape (input_length, output_length,\n  #                   input_channels * output_channels).\n  # Reshape it to (input_length * input_channels,\n  #                output_length * output_channels).\n  w = tf.reshape(conv, shape=(\n      [input_length,\n       output_length,\n       input_channels, output_channels]))\n  w = tf.transpose(w, perm=[0, 2, 1, 3])\n  w = tf.reshape(w, shape=(\n      [input_length * input_channels,\n       output_length * output_channels]))\n\n  # Broadcast b over spatial dimensions.\n  b = tf.tile(b, [output_length]) if b is not None else None\n\n  return w, b\n\n\ndef decode_batchnorm(batchnorm_module):\n  \"\"\"Calculates the neuron-wise multipliers and biases of the batch norm layer.\n\n  Note that, in the case of a convolution, the returned bias will have\n  spatial dimensions.\n\n  Args:\n    batchnorm_module: `snt.BatchNorm` module.\n\n  Returns:\n    w: 1D tensor of shape (output_size) or 3D tensor of shape\n      (output_height, output_width, output_channels) containing\n      neuron-wise multipliers for the batch norm layer.\n    b: 1D tensor of shape (output_size) or 3D tensor of shape\n      (output_height, output_width, output_channels) containing\n      neuron-wise biases for the batch norm layer.\n  \"\"\"\n  if isinstance(batchnorm_module, layers.BatchNorm):\n    mean = batchnorm_module.mean\n    variance = batchnorm_module.variance\n    variance_epsilon = batchnorm_module.epsilon\n    scale = batchnorm_module.scale\n    offset = batchnorm_module.bias\n\n  else:\n    assert isinstance(batchnorm_module, snt.BatchNorm)\n    mean = batchnorm_module.moving_mean\n    variance = batchnorm_module.moving_variance\n    variance_epsilon = batchnorm_module._eps  # pylint: disable=protected-access\n    try:\n      scale = batchnorm_module.gamma\n    except snt.Error:\n      scale = None\n    try:\n      offset = batchnorm_module.beta\n    except snt.Error:\n      offset = None\n\n  w = tf.rsqrt(variance + variance_epsilon)\n  if scale is not None:\n    w *= scale\n\n  b = -w * mean\n  if offset is not None:\n    b += offset\n\n  # Batchnorm vars have a redundant leading dim.\n  w = tf.squeeze(w, axis=0)\n  b = tf.squeeze(b, axis=0)\n  return w, b\n\n\ndef combine_with_batchnorm(w, b, batchnorm_module):\n  \"\"\"Combines a linear layer and a batch norm into a single linear layer.\n\n  Calculates the weights and biases of the linear layer formed by\n  applying the specified linear layer followed by the batch norm.\n\n  Note that, in the case of a convolution, the returned bias will have\n  spatial dimensions.\n\n  Args:\n    w: 2D tensor of shape (input_size, output_size) or 4D tensor of shape\n      (kernel_height, kernel_width, input_channels, output_channels) containing\n      weights for the linear layer.\n    b: 1D tensor of shape (output_size) or (output_channels) containing biases\n      for the linear layer, or `None` if no bias.\n    batchnorm_module: `snt.BatchNorm` module.\n\n  Returns:\n    w: 2D tensor of shape (input_size, output_size) or 4D tensor of shape\n      (kernel_height, kernel_width, input_channels, output_channels) containing\n      weights for the combined layer.\n    b: 1D tensor of shape (output_size) or 3D tensor of shape\n      (output_height, output_width, output_channels) containing\n      biases for the combined layer.\n  \"\"\"\n  if b is None:\n    b = tf.zeros(dtype=w.dtype, shape=())\n\n  w_bn, b_bn = decode_batchnorm(batchnorm_module)\n  return w * w_bn, b * w_bn + b_bn\n"
  },
  {
    "path": "interval_bound_propagation/src/layers.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Additional Sonnet modules.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\n# Slightly altered version of snt.BatchNorm that allows to easily grab which\n# mean and variance are currently in use (whether the last _build was\n# invoked with is_training=True or False).\n# Modifications include:\n# - Removing fused option (which we do not support).\n# - Removing test_local_stats (which we do not support).\n# - Providing a mean and variance property.\n# - Provides scale, bias properties that return None if there are none.\nclass BatchNorm(snt.BatchNorm):\n  \"\"\"Batch normalization module, including optional affine transformation.\"\"\"\n\n  def __init__(self, axis=None, offset=True, scale=False,\n               decay_rate=0.999, eps=1e-3, initializers=None,\n               partitioners=None, regularizers=None,\n               update_ops_collection=None, name='batch_norm'):\n    \"\"\"Constructs a BatchNorm module. See original code for more details.\"\"\"\n    super(BatchNorm, self).__init__(\n        axis=axis, offset=offset, scale=scale, decay_rate=decay_rate, eps=eps,\n        initializers=initializers, partitioners=partitioners,\n        regularizers=regularizers, fused=False,\n        update_ops_collection=update_ops_collection, name=name)\n\n  def _build_statistics(self, input_batch, axis, use_batch_stats, stat_dtype):\n    \"\"\"Builds the statistics part of the graph when using moving variance.\"\"\"\n    self._mean, self._variance = super(BatchNorm, self)._build_statistics(\n        input_batch, axis, use_batch_stats, stat_dtype)\n    return self._mean, self._variance\n\n  def _build(self, input_batch, is_training=True, test_local_stats=False,\n             reuse=False):\n    \"\"\"Connects the BatchNorm module into the graph.\n\n    Args:\n      input_batch: A Tensor of arbitrary dimension. By default, the final\n        dimension is not reduced over when computing the minibatch statistics.\n      is_training: A boolean to indicate if the module should be connected in\n        training mode, meaning the moving averages are updated. Can be a Tensor.\n      test_local_stats: A boolean to indicate if the statistics should be from\n        the local batch. When is_training is True, test_local_stats is not used.\n      reuse: If True, the statistics computed by previous call to _build\n        are used and is_training is ignored. Otherwise, behaves like a normal\n        batch normalization layer.\n\n    Returns:\n      A tensor with the same shape as `input_batch`.\n\n    Raises:\n      ValueError: If `axis` is not valid for the\n        input shape or has negative entries.\n    \"\"\"\n    if reuse:\n      self._ensure_is_connected()\n      return tf.nn.batch_normalization(\n          input_batch, self._mean, self._variance, self._beta, self._gamma,\n          self._eps, name='batch_norm')\n    else:\n      return super(BatchNorm, self)._build(input_batch, is_training,\n                                           test_local_stats=test_local_stats)\n\n  @property\n  def scale(self):\n    self._ensure_is_connected()\n    return tf.stop_gradient(self._gamma) if self._gamma is not None else None\n\n  @property\n  def bias(self):\n    self._ensure_is_connected()\n    return tf.stop_gradient(self._beta) if self._beta is not None else None\n\n  @property\n  def mean(self):\n    self._ensure_is_connected()\n    return tf.stop_gradient(self._mean)\n\n  @property\n  def variance(self):\n    self._ensure_is_connected()\n    return tf.stop_gradient(self._variance)\n\n  @property\n  def epsilon(self):\n    self._ensure_is_connected()\n    return self._eps\n\n\nclass ImageNorm(snt.AbstractModule):\n  \"\"\"Module that does per channel normalization.\"\"\"\n\n  def __init__(self, mean, std, name='image_norm'):\n    \"\"\"Constructs a module that does (x[:, :, c] - mean[c]) / std[c].\"\"\"\n    super(ImageNorm, self).__init__(name=name)\n    if isinstance(mean, float):\n      mean = [mean]\n    if isinstance(std, float):\n      std = [std]\n    scale = []\n    for s in std:\n      if s <= 0.:\n        raise ValueError('Cannot use negative standard deviations.')\n      scale.append(1. / s)\n    with self._enter_variable_scope():\n      # Using broadcasting.\n      self._scale = tf.constant(scale, dtype=tf.float32)\n      self._offset = tf.constant(mean, dtype=tf.float32)\n\n  def _build(self, inputs):\n    return self.apply(inputs)\n\n  @property\n  def scale(self):\n    return self._scale\n\n  @property\n  def offset(self):\n    return self._offset\n\n  # Provide a function that allows to use the IncreasingMonotonicWrapper.\n  def apply(self, inputs):\n    return (inputs - self._offset) * self._scale\n"
  },
  {
    "path": "interval_bound_propagation/src/loss.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Helper to keep track of the different losses.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport collections\n\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n# Used to pick the least violated specification.\n_BIG_NUMBER = 1e25\n\n\nScalarMetrics = collections.namedtuple('ScalarMetrics', [\n    'nominal_accuracy',\n    'verified_accuracy',\n    'attack_accuracy',\n    'attack_success'])\n\n\nScalarLosses = collections.namedtuple('ScalarLosses', [\n    'nominal_cross_entropy',\n    'attack_cross_entropy',\n    'verified_loss'])\n\n\nclass Losses(snt.AbstractModule):\n  \"\"\"Helper to compute our losses.\"\"\"\n\n  def __init__(self, predictor, specification=None, pgd_attack=None,\n               interval_bounds_loss_type='xent',\n               interval_bounds_hinge_margin=10.,\n               label_smoothing=0.):\n    super(Losses, self).__init__(name='losses')\n    self._predictor = predictor\n    self._specification = specification\n    self._attack = pgd_attack\n    # Loss type can be any combination of:\n    #   xent: cross-entropy loss\n    #   hinge: hinge loss\n    #   softplus: softplus loss\n    # with\n    #   all: using all specifications.\n    #   most: using only the specification that is the most violated.\n    #   least: using only the specification that is the least violated.\n    #   random_n: using a random subset of the specifications.\n    # E.g.: \"xent_max\" or \"hinge_random_3\".\n    tokens = interval_bounds_loss_type.split('_', 1)\n    if len(tokens) == 1:\n      loss_type, loss_mode = tokens[0], 'all'\n    else:\n      loss_type, loss_mode = tokens\n      if loss_mode.startswith('random'):\n        loss_mode, num_samples = loss_mode.split('_', 1)\n        self._interval_bounds_loss_n = int(num_samples)\n    if loss_type not in ('xent', 'hinge', 'softplus'):\n      raise ValueError('interval_bounds_loss_type must be either \"xent\", '\n                       '\"hinge\" or \"softplus\".')\n    if loss_mode not in ('all', 'most', 'random', 'least'):\n      raise ValueError('interval_bounds_loss_type must be followed by either '\n                       '\"all\", \"most\", \"random_N\" or \"least\".')\n    self._interval_bounds_loss_type = loss_type\n    self._interval_bounds_loss_mode = loss_mode\n    self._interval_bounds_hinge_margin = interval_bounds_hinge_margin\n    self._label_smoothing = label_smoothing\n\n  def _build(self, labels):\n    self._build_nominal_loss(labels)\n    self._build_verified_loss(labels)\n    self._build_attack_loss(labels)\n\n  def _build_nominal_loss(self, labels):\n    \"\"\"Build natural cross-entropy loss on clean data.\"\"\"\n    # Cross-entropy.\n    nominal_logits = self._predictor.logits\n    if self._label_smoothing > 0:\n      num_classes = nominal_logits.shape[1].value\n      one_hot_labels = tf.one_hot(labels, num_classes)\n      smooth_positives = 1. - self._label_smoothing\n      smooth_negatives = self._label_smoothing / num_classes\n      one_hot_labels = one_hot_labels * smooth_positives + smooth_negatives\n      nominal_cross_entropy = tf.nn.softmax_cross_entropy_with_logits_v2(\n          labels=one_hot_labels, logits=nominal_logits)\n      self._one_hot_labels = one_hot_labels\n    else:\n      nominal_cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(\n          labels=labels, logits=nominal_logits)\n    self._cross_entropy = tf.reduce_mean(nominal_cross_entropy)\n    # Accuracy.\n    nominal_correct_examples = tf.equal(labels, tf.argmax(nominal_logits, 1))\n    self._nominal_accuracy = tf.reduce_mean(\n        tf.cast(nominal_correct_examples, tf.float32))\n\n  def _get_specification_bounds(self):\n    \"\"\"Get upper bounds on specification. Used for building verified loss.\"\"\"\n    ibp_bounds = self._specification(self._predictor.modules)\n    # Compute verified accuracy using IBP bounds.\n    v = tf.reduce_max(ibp_bounds, axis=1)\n    self._interval_bounds_accuracy = tf.reduce_mean(\n        tf.cast(v <= 0., tf.float32))\n    return ibp_bounds\n\n  def _build_verified_loss(self, labels):\n    \"\"\"Build verified loss using an upper bound on specification.\"\"\"\n    if not self._specification:\n      self._verified_loss = tf.constant(0.)\n      self._interval_bounds_accuracy = tf.constant(0.)\n      return\n    # Interval bounds.\n    bounds = self._get_specification_bounds()\n    # Select specifications.\n    if self._interval_bounds_loss_mode == 'all':\n      pass  # Keep bounds the way it is.\n    elif self._interval_bounds_loss_mode == 'most':\n      bounds = tf.reduce_max(bounds, axis=1, keepdims=True)\n    elif self._interval_bounds_loss_mode == 'random':\n      idx = tf.random.uniform(\n          [tf.shape(bounds)[0], self._interval_bounds_loss_n],\n          0, tf.shape(bounds)[1], dtype=tf.int32)\n      bounds = tf.batch_gather(bounds, idx)\n    else:\n      assert self._interval_bounds_loss_mode == 'least'\n      # This picks the least violated contraint.\n      mask = tf.cast(bounds < 0., tf.float32)\n      smallest_violation = tf.reduce_min(\n          bounds + mask * _BIG_NUMBER, axis=1, keepdims=True)\n      has_violations = tf.less(\n          tf.reduce_sum(mask, axis=1, keepdims=True) + .5,\n          tf.cast(tf.shape(bounds)[1], tf.float32))\n      largest_bounds = tf.reduce_max(bounds, axis=1, keepdims=True)\n      bounds = tf.where(has_violations, smallest_violation, largest_bounds)\n\n    if self._interval_bounds_loss_type == 'xent':\n      v = tf.concat(\n          [bounds, tf.zeros([tf.shape(bounds)[0], 1], dtype=bounds.dtype)],\n          axis=1)\n      l = tf.concat(\n          [tf.zeros_like(bounds),\n           tf.ones([tf.shape(bounds)[0], 1], dtype=bounds.dtype)],\n          axis=1)\n      self._verified_loss = tf.reduce_mean(\n          tf.nn.softmax_cross_entropy_with_logits_v2(\n              labels=tf.stop_gradient(l), logits=v))\n    elif self._interval_bounds_loss_type == 'softplus':\n      self._verified_loss = tf.reduce_mean(\n          tf.nn.softplus(bounds + self._interval_bounds_hinge_margin))\n    else:\n      assert self._interval_bounds_loss_type == 'hinge'\n      self._verified_loss = tf.reduce_mean(\n          tf.maximum(bounds, -self._interval_bounds_hinge_margin))\n\n  def _build_attack_loss(self, labels):\n    \"\"\"Build adversarial loss using PGD attack.\"\"\"\n    # PGD attack.\n    if not self._attack:\n      self._attack_accuracy = tf.constant(0.)\n      self._attack_success = tf.constant(1.)\n      self._attack_cross_entropy = tf.constant(0.)\n      return\n    if not isinstance(self._predictor.inputs, tf.Tensor):\n      raise ValueError('Multiple inputs is not supported.')\n    self._attack(self._predictor.inputs, labels)\n    correct_examples = tf.equal(labels, tf.argmax(self._attack.logits, 1))\n    self._attack_accuracy = tf.reduce_mean(\n        tf.cast(correct_examples, tf.float32))\n    self._attack_success = tf.reduce_mean(\n        tf.cast(self._attack.success, tf.float32))\n    if self._label_smoothing > 0:\n      attack_cross_entropy = tf.nn.softmax_cross_entropy_with_logits_v2(\n          labels=self._one_hot_labels, logits=self._attack.logits)\n    else:\n      attack_cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(\n          labels=labels, logits=self._attack.logits)\n    self._attack_cross_entropy = tf.reduce_mean(attack_cross_entropy)\n\n  @property\n  def scalar_metrics(self):\n    self._ensure_is_connected()\n    return ScalarMetrics(self._nominal_accuracy,\n                         self._interval_bounds_accuracy,\n                         self._attack_accuracy,\n                         self._attack_success)\n\n  @property\n  def scalar_losses(self):\n    self._ensure_is_connected()\n    return ScalarLosses(self._cross_entropy,\n                        self._attack_cross_entropy,\n                        self._verified_loss)\n"
  },
  {
    "path": "interval_bound_propagation/src/model.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Sonnet modules that represent the predictor network.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport collections\n\nfrom absl import logging\nfrom interval_bound_propagation.src import layers\nfrom interval_bound_propagation.src import verifiable_wrapper\nimport numpy as np\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\n# Set of supported activations. Must be monotonic and attributes of `tf.nn`.\n_ALLOWED_ACTIVATIONS = set([\n    'elu',\n    'leaky_relu',\n    'relu',\n    'relu6',\n    'selu',\n    'sigmoid',\n    'softplus',\n    'softsign',\n    'tanh',\n])\n\n# Mapping between graph node ops and their TensorFlow function.\n_MONOTONIC_NODE_OPS = {\n    'Elu': tf.nn.elu,\n    'LeakyRelu': tf.nn.leaky_relu,\n    'Relu': tf.nn.relu,\n    'Relu6': tf.nn.relu6,\n    'Selu': tf.nn.selu,\n    'Sigmoid': tf.nn.sigmoid,\n    'Softplus': tf.nn.softplus,\n    'Softsign': tf.nn.softsign,\n    'Tanh': tf.nn.tanh,\n}\n\n\nclass VerifiableModelWrapper(snt.AbstractModule):\n  \"\"\"Wraps a predictor network.\"\"\"\n\n  def __init__(self, net_builder, name='verifiable_predictor'):\n    \"\"\"Constructor for the verifiable model.\n\n    Args:\n      net_builder: A callable that returns output logits from an input.\n        net_builder must accept two arguments: the input (as the first\n        argument) and is_training (as the second).\n      name: Sonnet module name.\n    \"\"\"\n    super(VerifiableModelWrapper, self).__init__(name=name)\n    self._net_builder = net_builder\n\n  @property\n  def wrapped_network(self):\n    return self._net_builder\n\n  @property\n  def output_size(self):\n    self._ensure_is_connected()\n    return self._num_classes\n\n  @property\n  def logits(self):\n    self._ensure_is_connected()\n    return self._logits\n\n  @property\n  def inputs(self):\n    self._ensure_is_connected()\n    return self._inputs\n\n  @property\n  def input_wrappers(self):\n    self._ensure_is_connected()\n    return self._model_inputs\n\n  @property\n  def modules(self):\n    self._ensure_is_connected()\n    return self._modules\n\n  def dependencies(self, module):\n    self._ensure_is_connected()\n    return self._module_depends_on[module]\n\n  @property\n  def output_module(self):\n    self._ensure_is_connected()\n    return self._produced_by[self._logits.name]\n\n  def fanout_of(self, node):\n    \"\"\"Looks up fan-out for a given node.\n\n    Args:\n      node: `ibp.VerifiableWrapper` occurring in the network either as an\n        operation, or as the initial input.\n\n    Returns:\n      Number of times `node` occurs as the input of another operation within\n        the network, or 1 if `node` is the overall output.\n    \"\"\"\n    return self._fanouts[node]\n\n  def _build(self, *z0, **kwargs):\n    \"\"\"Outputs logits from input z0.\n\n    Args:\n      *z0: inputs as `Tensor`.\n      **kwargs: Other arguments passed directly to the _build() function of the\n        wrapper model. Assumes the possible presence of `override` (defaults to\n        False). However, if False, this function does not update any internal\n        state and reuses any components computed by a previous call to _build().\n        If there were no previous calls to _build(), behaves as if it was set to\n        True.\n\n    Returns:\n      logits resulting from using z0 as inputs.\n    \"\"\"\n    override = not self.is_connected\n    if 'override' in kwargs:\n      override = kwargs['override'] or override\n      del kwargs['override']\n    if override:\n      self._inputs = z0[0] if len(z0) == 1 else z0\n      # Build underlying verifiable modules.\n      self._model_inputs = []\n      self._modules = []\n      self._produced_by = {}  # Connection graph.\n      self._fanouts = collections.Counter()\n      for i, z in enumerate(z0):\n        self._model_inputs.append(verifiable_wrapper.ModelInputWrapper(i))\n        self._produced_by[z.name] = self._model_inputs[-1]\n      self._module_depends_on = collections.defaultdict(list)\n      self._output_by_module = {}\n      with snt.observe_connections(self._observer):\n        logits = self._net_builder(*z0, **kwargs)\n      # Logits might be produced by a non-Sonnet module.\n      self._backtrack(logits, max_depth=100)\n      # Log analysis.\n      for m in self._modules:\n        logging.info('Found: %s', m)\n        output_shape = self._output_by_module[m].shape.as_list()[1:]\n        logging.info('  Output shape: %s => %d units', output_shape,\n                     np.prod(output_shape))\n        for depends in self._module_depends_on[m]:\n          logging.info('  Depends on: %s', depends)\n      logging.info('Final logits produced by: %s',\n                   self._produced_by[logits.name])\n      self._logits = logits\n      self._num_classes = logits.shape[-1].value\n    else:\n      # Must have been connected once before.\n      self._ensure_is_connected()\n      logits = self._net_builder(*z0, **kwargs)\n    return logits\n\n  def _observer(self, subgraph):\n    input_nodes = self._inputs_for_observed_module(subgraph)\n    if input_nodes is None:\n      # We do not fail as we want to allow higher-level Sonnet components.\n      # In practice, the rest of the logic will fail if we are unable to\n      # connect all low-level modules.\n      logging.warn('Unprocessed module \"%s\"', str(subgraph.module))\n      return\n    if subgraph.outputs in input_nodes:\n      # The Sonnet module is just returning its input as its output.\n      # This may happen with a reshape in which the shape does not change.\n      return\n\n    self._add_module(self._wrapper_for_observed_module(subgraph),\n                     subgraph.outputs, *input_nodes)\n\n  def _inputs_for_observed_module(self, subgraph):\n    \"\"\"Extracts input tensors from a connected Sonnet module.\n\n    This default implementation supports common layer types, but should be\n    overridden if custom layer types are to be supported.\n\n    Args:\n      subgraph: `snt.ConnectedSubGraph` specifying the Sonnet module being\n        connected, and its inputs and outputs.\n\n    Returns:\n      List of input tensors, or None if not a supported Sonnet module.\n    \"\"\"\n    m = subgraph.module\n    # Only support a few operations for now.\n    if not (isinstance(m, snt.BatchReshape) or\n            isinstance(m, snt.Linear) or\n            isinstance(m, snt.Conv1D) or\n            isinstance(m, snt.Conv2D) or\n            isinstance(m, snt.BatchNorm) or\n            isinstance(m, layers.ImageNorm)):\n      return None\n\n    if isinstance(m, snt.BatchNorm):\n      return subgraph.inputs['input_batch'],\n    else:\n      return subgraph.inputs['inputs'],\n\n  def _wrapper_for_observed_module(self, subgraph):\n    \"\"\"Creates a wrapper for a connected Sonnet module.\n\n    This default implementation supports common layer types, but should be\n    overridden if custom layer types are to be supported.\n\n    Args:\n      subgraph: `snt.ConnectedSubGraph` specifying the Sonnet module being\n        connected, and its inputs and outputs.\n\n    Returns:\n      `ibp.VerifiableWrapper` for the Sonnet module.\n    \"\"\"\n    m = subgraph.module\n    if isinstance(m, snt.BatchReshape):\n      shape = subgraph.outputs.get_shape()[1:].as_list()\n      return verifiable_wrapper.BatchReshapeWrapper(m, shape)\n    elif isinstance(m, snt.Linear):\n      return verifiable_wrapper.LinearFCWrapper(m)\n    elif isinstance(m, snt.Conv1D):\n      return verifiable_wrapper.LinearConv1dWrapper(m)\n    elif isinstance(m, snt.Conv2D):\n      return verifiable_wrapper.LinearConv2dWrapper(m)\n    elif isinstance(m, layers.ImageNorm):\n      return verifiable_wrapper.ImageNormWrapper(m)\n    else:\n      assert isinstance(m, snt.BatchNorm)\n      return verifiable_wrapper.BatchNormWrapper(m)\n\n  def _backtrack(self, node, max_depth=100):\n    if node.name not in self._produced_by:\n      if max_depth <= 0:\n        raise ValueError('Unable to backtrack through the graph. '\n                         'Consider using more basic Sonnet modules.')\n      self._wrap_node(node, max_depth=(max_depth - 1))\n    self._fanouts[self._produced_by[node.name]] += 1\n\n  def _wrap_node(self, node, **kwargs):\n    \"\"\"Adds an IBP wrapper for the node, and backtracks through its inputs.\n\n    This default implementation supports common layer types, but should be\n    overridden if custom layer types are to be supported.\n\n    Implementations should create a `ibp.VerifiableWrapper` and then invoke\n    `self._add_module(wrapper, node, *input_node, **kwargs)`.\n\n    Args:\n      node: TensorFlow graph node to wrap for IBP.\n      **kwargs: Context to pass to `self._add_module`.\n    \"\"\"\n    # Group all unary monotonic ops at the end.\n    if node.op.type in ('Add', 'AddV2', 'Mul', 'Sub', 'Maximum', 'Minimum'):\n      input_node0 = node.op.inputs[0]\n      input_node1 = node.op.inputs[1]\n      if node.op.type in ('Add', 'AddV2'):\n        w = verifiable_wrapper.IncreasingMonotonicWrapper(tf.add)\n      elif node.op.type == 'Mul':\n        w = verifiable_wrapper.PiecewiseMonotonicWrapper(tf.multiply)\n      elif node.op.type == 'Sub':\n        w = verifiable_wrapper.PiecewiseMonotonicWrapper(tf.subtract)\n      elif node.op.type == 'Maximum':\n        w = verifiable_wrapper.IncreasingMonotonicWrapper(tf.maximum)\n      elif node.op.type == 'Minimum':\n        w = verifiable_wrapper.IncreasingMonotonicWrapper(tf.minimum)\n\n      self._add_module(w, node, input_node0, input_node1, **kwargs)\n      return\n    elif node.op.type == 'ConcatV2':\n      num_inputs = node.op.get_attr('N')\n      assert num_inputs == len(node.op.inputs) - 1\n      inputs = node.op.inputs[:num_inputs]\n      axis = node.op.inputs[num_inputs]\n      def concat(*args):\n        return tf.concat(args, axis=axis)\n      self._add_module(\n          verifiable_wrapper.IncreasingMonotonicWrapper(concat, axis=axis),\n          node, *inputs, **kwargs)\n      return\n    elif node.op.type == 'Softmax':\n      input_node = node.op.inputs[0]\n      self._add_module(verifiable_wrapper.SoftmaxWrapper(), node, input_node,\n                       **kwargs)\n      return\n    elif node.op.type == 'Const':\n      self._add_module(verifiable_wrapper.ConstWrapper(node), node, **kwargs)\n      return\n\n    # The rest are all unary monotonic ops.\n    parameters = dict()\n    if node.op.type in _MONOTONIC_NODE_OPS:\n      input_node = node.op.inputs[0]\n      # Leaky ReLUs are a special case since they have a second argument.\n      if node.op.type == 'LeakyRelu':\n        parameters = dict(alpha=node.op.get_attr('alpha'))\n        # Use function definition instead of lambda for clarity.\n        def leaky_relu(x):\n          return tf.nn.leaky_relu(x, **parameters)\n        fn = leaky_relu\n      else:\n        fn = _MONOTONIC_NODE_OPS[node.op.type]\n\n    elif node.op.type in ('Mean', 'Max', 'Sum', 'Min'):\n      # reduce_mean/max have two inputs. The first one should be produced by a\n      # upstream node, while the two one should represent the axis.\n      input_node = node.op.inputs[0]\n      parameters = dict(axis=node.op.inputs[1],\n                        keep_dims=node.op.get_attr('keep_dims'))\n      # Use function definition instead of lambda for clarity.\n      def reduce_max(x):\n        return tf.reduce_max(x, **parameters)\n      def reduce_mean(x):\n        return tf.reduce_mean(x, **parameters)\n      def reduce_min(x):\n        return tf.reduce_min(x, **parameters)\n      def reduce_sum(x):\n        return tf.reduce_sum(x, **parameters)\n      fn = dict(\n          Max=reduce_max, Mean=reduce_mean, Sum=reduce_sum,\n          Min=reduce_min)[node.op.type]\n\n    elif node.op.type == 'ExpandDims':\n      input_node = node.op.inputs[0]\n      parameters = dict(axis=node.op.inputs[1])\n      def expand_dims(x):\n        return tf.expand_dims(x, **parameters)\n      fn = expand_dims\n\n    elif node.op.type == 'Transpose':\n      input_node = node.op.inputs[0]\n      parameters = dict(perm=node.op.inputs[1])\n      def transpose(x):\n        return tf.transpose(x, **parameters)\n      fn = transpose\n\n    elif node.op.type == 'Squeeze':\n      input_node = node.op.inputs[0]\n      parameters = dict(axis=node.op.get_attr('squeeze_dims'))\n      def squeeze(x):\n        return tf.squeeze(x, **parameters)\n      fn = squeeze\n\n    elif node.op.type == 'Pad':\n      input_node = node.op.inputs[0]\n      parameters = dict(paddings=node.op.inputs[1])\n      def pad(x):\n        return tf.pad(x, **parameters)\n      fn = pad\n\n    elif node.op.type in ('MaxPool', 'AvgPool'):\n      input_node = node.op.inputs[0]\n      parameters = dict(\n          ksize=node.op.get_attr('ksize'),\n          strides=node.op.get_attr('strides'),\n          padding=node.op.get_attr('padding'),\n          data_format=node.op.get_attr('data_format'),\n      )\n      if node.op.type == 'MaxPool':\n        def max_pool(x):\n          return tf.nn.max_pool(x, **parameters)\n        fn = max_pool\n      elif node.op.type == 'AvgPool':\n        def avg_pool(x):\n          return tf.nn.avg_pool(x, **parameters)\n        fn = avg_pool\n\n    elif node.op.type == 'Reshape':\n      input_node = node.op.inputs[0]\n      parameters = dict(shape=node.op.inputs[1])\n      def reshape(x):\n        return tf.reshape(x, **parameters)\n      fn = reshape\n\n    elif node.op.type == 'Identity':\n      input_node = node.op.inputs[0]\n      def identity(x):\n        return tf.identity(x)\n      fn = identity\n\n    elif node.op.type == 'MatrixDiag':\n      input_node = node.op.inputs[0]\n      def matrix_diag(x):\n        return tf.matrix_diag(x)\n      fn = matrix_diag\n\n    elif node.op.type == 'Slice':\n      input_node = node.op.inputs[0]\n      parameters = dict(\n          begin=node.op.inputs[1],\n          size=node.op.inputs[2],\n      )\n      def regular_slice(x):\n        return tf.slice(x, **parameters)\n      fn = regular_slice\n\n    elif node.op.type == 'StridedSlice':\n      input_node = node.op.inputs[0]\n      parameters = dict(\n          begin=node.op.inputs[1],\n          end=node.op.inputs[2],\n          strides=node.op.inputs[3],\n          begin_mask=node.op.get_attr('begin_mask'),\n          end_mask=node.op.get_attr('end_mask'),\n          ellipsis_mask=node.op.get_attr('ellipsis_mask'),\n          new_axis_mask=node.op.get_attr('new_axis_mask'),\n          shrink_axis_mask=node.op.get_attr('shrink_axis_mask'),\n      )\n      def strided_slice(x):\n        return tf.strided_slice(x, **parameters)\n      fn = strided_slice\n\n    elif node.op.type == 'Fill':\n      input_node = node.op.inputs[1]  # Shape is the first argument.\n      dims = node.op.inputs[0]\n      parameters = dict(dims=dims)\n      def fill(x):\n        return tf.fill(dims, x)\n      fn = fill\n\n    elif node.op.type == 'RealDiv':\n      # The denominator is assumed to be constant but is permitted to be\n      # example-dependent, for example a sequence's length prior to padding.\n      input_node = node.op.inputs[0]\n      denom = node.op.inputs[1]\n      parameters = dict(denom=denom)\n      def quotient(x):\n        return x / denom\n      fn = quotient\n\n    else:\n      raise NotImplementedError(\n          'Unsupported operation: \"{}\" with\\n{}.'.format(node.op.type, node.op))\n\n    self._add_module(\n        verifiable_wrapper.IncreasingMonotonicWrapper(fn, **parameters),\n        node, input_node, **kwargs)\n\n  def _add_module(self, wrapper, node, *input_nodes, **kwargs):\n    \"\"\"Adds the given node wrapper, first backtracking through its inputs.\n\n    Args:\n      wrapper: `ibp.VerifiableWrapper` for the node.\n      node: TensorFlow graph node.\n      *input_nodes: Input nodes for `node`.\n      **kwargs: Contains the `max_depth` argument for recursive _backtrack call.\n    \"\"\"\n    for input_node in input_nodes:\n      self._backtrack(input_node, **kwargs)\n    self._modules.append(wrapper)\n    self._produced_by[node.name] = self._modules[-1]\n    self._module_depends_on[self._modules[-1]].extend(\n        [self._produced_by[input_node.name] for input_node in input_nodes])\n    self._output_by_module[self._modules[-1]] = node\n\n  def propagate_bounds(self, *input_bounds):\n    \"\"\"Propagates input bounds through the network.\n\n    Args:\n      *input_bounds: `AbstractBounds` instance corresponding to z0.\n\n    Returns:\n      The final output bounds corresponding to the output logits.\n    \"\"\"\n    self._ensure_is_connected()\n    def _get_bounds(input_module):\n      \"\"\"Retrieves the bounds corresponding to a module.\"\"\"\n      # All bounds need to be canonicalized to the same type. In particular, we\n      # need to handle the case of constant bounds specially. We convert them\n      # to the same type as input_bounds.\n      if isinstance(input_module, verifiable_wrapper.ConstWrapper):\n        return input_bounds[0].convert(input_module.output_bounds)\n      return input_module.output_bounds\n\n    # Initialise inputs' bounds.\n    for model_input in self._model_inputs:\n      model_input.output_bounds = input_bounds[model_input.index]\n    # By construction, this list is topologically sorted.\n    for m in self._modules:\n      # Construct combined input bounds.\n      upstream_bounds = [_get_bounds(b) for b in self._module_depends_on[m]]\n      m.propagate_bounds(*upstream_bounds)\n    # We assume that the last module is the final output layer.\n    return self._produced_by[self._logits.name].output_bounds\n\n\nclass StandardModelWrapper(snt.AbstractModule):\n  \"\"\"Wraps a predictor network that keeps track of inputs and logits.\"\"\"\n\n  def __init__(self, net_builder, name='verifiable_predictor'):\n    \"\"\"Constructor for a non-verifiable model.\n\n    This wrapper can be used to seamlessly use loss.py and utils.py without\n    IBP verification.\n\n    Args:\n      net_builder: A callable that returns output logits from an input.\n        net_builder must accept two arguments: the input (as the first\n        argument) and is_training (as the second).\n      name: Sonnet module name.\n    \"\"\"\n    super(StandardModelWrapper, self).__init__(name=name)\n    self._net_builder = net_builder\n\n  @property\n  def wrapped_network(self):\n    return self._net_builder\n\n  @property\n  def output_size(self):\n    self._ensure_is_connected()\n    return self._num_classes\n\n  @property\n  def logits(self):\n    self._ensure_is_connected()\n    return self._logits\n\n  @property\n  def inputs(self):\n    self._ensure_is_connected()\n    return self._inputs\n\n  @property\n  def modules(self):\n    raise RuntimeError('Model is not wrapped by a VerifiableModelWrapper. '\n                       'Bounds cannot be propagated.')\n\n  def propagate_bounds(self, *input_bounds):\n    raise RuntimeError('Model is not wrapped by a VerifiableModelWrapper. '\n                       'Bounds cannot be propagated.')\n\n  def _build(self, *z0, **kwargs):\n    \"\"\"Outputs logits from input z0.\n\n    Args:\n      *z0: inputs as `Tensor`.\n      **kwargs: Other arguments passed directly to the _build() function of the\n        wrapper model. Assumes the possible presence of `override` (defaults to\n        False). However, if False, this function does not update any internal\n        state and reuses any components computed by a previous call to _build().\n        If there were no previous calls to _build(), behaves as if it was set to\n        True.\n\n    Returns:\n      logits resulting from using z0 as inputs.\n    \"\"\"\n    override = not self.is_connected\n    if 'override' in kwargs:\n      override = kwargs['override'] or override\n      del kwargs['override']\n    if override:\n      self._inputs = z0[0] if len(z0) == 1 else z0\n      logits = self._net_builder(*z0, **kwargs)\n      self._logits = logits\n      self._num_classes = logits.shape[-1].value\n    else:\n      # Must have been connected once before.\n      self._ensure_is_connected()\n      logits = self._net_builder(*z0, **kwargs)\n    return logits\n\n\nclass DNN(snt.AbstractModule):\n  \"\"\"Simple feed-forward neural network.\"\"\"\n\n  def __init__(self, num_classes, layer_types, l2_regularization_scale=0.,\n               name='predictor'):\n    \"\"\"Constructor for the DNN.\n\n    Args:\n      num_classes: Output size.\n      layer_types: Iterable of tuples. Each tuple must be one of the following:\n        * ('conv2d', (kernel_height, width), channels, padding, stride)\n        * ('linear', output_size)\n        * ('batch_normalization',)\n        * ('activation', activation)\n        Convolutional layers must precede all linear layers.\n      l2_regularization_scale: Scale of the L2 regularization on the weights\n        of each layer.\n      name: Sonnet module name.\n    \"\"\"\n    super(DNN, self).__init__(name=name)\n    self._layer_types = list(layer_types)\n    self._layer_types.append(('linear', num_classes))\n    if l2_regularization_scale > 0.:\n      regularizer = tf.keras.regularizers.l2(l=0.5*l2_regularization_scale)\n      self._regularizers = {'w': regularizer}\n    else:\n      self._regularizers = None\n    # The following allows to reuse previous batch norm statistics.\n    self._batch_norms = {}\n\n  def _build(self, z0, is_training=True, test_local_stats=False, reuse=False):\n    \"\"\"Outputs logits.\"\"\"\n    zk = z0\n    conv2d_id = 0\n    linear_id = 0\n    name = None\n    for spec in self._layer_types:\n      if spec[0] == 'conv2d':\n        if linear_id > 0:\n          raise ValueError('Convolutional layers must precede fully connected '\n                           'layers.')\n        name = 'conv2d_{}'.format(conv2d_id)\n        conv2d_id += 1\n        (_, (kernel_height, kernel_width), channels, padding, stride) = spec\n        m = snt.Conv2D(output_channels=channels,\n                       kernel_shape=(kernel_height, kernel_width),\n                       padding=padding, stride=stride, use_bias=True,\n                       regularizers=self._regularizers,\n                       initializers=_create_conv2d_initializer(\n                           zk.get_shape().as_list()[1:], channels,\n                           (kernel_height, kernel_width)),\n                       name=name)\n        zk = m(zk)\n      elif spec[0] == 'linear':\n        must_flatten = (linear_id == 0 and len(zk.shape) > 2)\n        if must_flatten:\n          zk = snt.BatchFlatten()(zk)\n        name = 'linear_{}'.format(linear_id)\n        linear_id += 1\n        output_size = spec[1]\n        m = snt.Linear(output_size,\n                       regularizers=self._regularizers,\n                       initializers=_create_linear_initializer(\n                           np.prod(zk.get_shape().as_list()[1:]), output_size),\n                       name=name)\n        zk = m(zk)\n      elif spec[0] == 'batch_normalization':\n        if name is None:\n          raise ValueError('Batch normalization only supported after linear '\n                           'layers.')\n        name += '_batch_norm'\n        m = layers.BatchNorm(name=name)\n        if reuse:\n          if m.scope_name not in self._batch_norms:\n            raise ValueError('Cannot set reuse to True without connecting the '\n                             'module once before.')\n          m = self._batch_norms[m.scope_name]\n        else:\n          self._batch_norms[m.scope_name] = m\n        zk = m(zk, is_training=is_training, test_local_stats=test_local_stats,\n               reuse=reuse)\n      elif spec[0] == 'activation':\n        if spec[1] not in _ALLOWED_ACTIVATIONS:\n          raise NotImplementedError(\n              'Only the following activations are supported {}'.format(\n                  list(_ALLOWED_ACTIVATIONS)))\n        name = None\n        m = getattr(tf.nn, spec[1])\n        zk = m(zk)\n    return zk\n\n\ndef _create_conv2d_initializer(\n    input_shape, output_channels, kernel_shape, dtype=tf.float32):  # pylint: disable=unused-argument\n  \"\"\"Returns a default initializer for the weights of a convolutional module.\"\"\"\n  return {\n      'w': tf.orthogonal_initializer(),\n      'b': tf.zeros_initializer(dtype=dtype),\n  }\n\n\ndef _create_linear_initializer(input_size, output_size, dtype=tf.float32):  # pylint: disable=unused-argument\n  \"\"\"Returns a default initializer for the weights of a linear module.\"\"\"\n  return {\n      'w': tf.orthogonal_initializer(),\n      'b': tf.zeros_initializer(dtype=dtype),\n  }\n"
  },
  {
    "path": "interval_bound_propagation/src/relative_bounds.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Interval bounds expressed relative to a nominal value.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom interval_bound_propagation.src import bounds as basic_bounds\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\nclass RelativeIntervalBounds(basic_bounds.AbstractBounds):\n  \"\"\"Upper and lower bounds, as a delta relative to nominal values.\"\"\"\n\n  def __init__(self, lower_offset, upper_offset, nominal):\n    super(RelativeIntervalBounds, self).__init__()\n    self._lower_offset = lower_offset\n    self._upper_offset = upper_offset\n    self._nominal = nominal\n\n  @property\n  def lower_offset(self):\n    \"\"\"Returns lower bounds, expressed relative to nominal values.\"\"\"\n    return self._lower_offset\n\n  @property\n  def upper_offset(self):\n    \"\"\"Returns upper bounds, expressed relative to nominal values.\"\"\"\n    return self._upper_offset\n\n  @property\n  def nominal(self):\n    return self._nominal\n\n  @property\n  def lower(self):\n    \"\"\"Returns absolute lower bounds.\"\"\"\n    return self.nominal + self.lower_offset\n\n  @property\n  def upper(self):\n    \"\"\"Returns absolute upper bounds.\"\"\"\n    return self.nominal + self.upper_offset\n\n  @property\n  def shape(self):\n    return self.lower_offset.shape.as_list()\n\n  @classmethod\n  def convert(cls, bounds):\n    if isinstance(bounds, tf.Tensor):\n      return cls(tf.zeros_like(bounds), tf.zeros_like(bounds), bounds)\n    bounds = bounds.concretize()\n    if not isinstance(bounds, cls):\n      raise ValueError('Cannot convert \"{}\" to \"{}\"'.format(bounds,\n                                                            cls.__name__))\n    return bounds\n\n  def apply_batch_reshape(self, wrapper, shape):\n    \"\"\"Propagates the bounds through a reshape.\n\n    Args:\n      wrapper: Contains prior bounds from a previous iteration.\n      shape: output shape, excluding the batch dimension.\n\n    Returns:\n      Output bounds.\n    \"\"\"\n    reshape = snt.BatchReshape(shape)\n    return RelativeIntervalBounds(\n        reshape(self.lower_offset),\n        reshape(self.upper_offset),\n        reshape(self.nominal))\n\n  def apply_linear(self, wrapper, w, b):\n    \"\"\"Propagates the bounds through a linear layer.\n\n    Args:\n      wrapper: Contains prior bounds from a previous iteration.\n      w: 2D tensor of shape (input_size, output_size) containing\n        weights for the linear layer.\n      b: 1D tensor of shape (output_size) containing biases for the linear\n        layer, or `None` if no bias.\n\n    Returns:\n      Output bounds.\n    \"\"\"\n    w_pos = tf.maximum(w, 0)\n    w_neg = tf.minimum(w, 0)\n    lb = (tf.matmul(self.lower_offset, w_pos) +\n          tf.matmul(self.upper_offset, w_neg))\n    ub = (tf.matmul(self.upper_offset, w_pos) +\n          tf.matmul(self.lower_offset, w_neg))\n\n    nominal_out = tf.matmul(self.nominal, w)\n    if b is not None:\n      nominal_out += b\n\n    return RelativeIntervalBounds(lb, ub, nominal_out)\n\n  def apply_conv1d(self, wrapper, w, b, padding, stride):\n    \"\"\"Propagates the bounds through a 1D convolution layer.\n\n    Args:\n      wrapper: Contains prior bounds from a previous iteration.\n      w: 3D tensor of shape (kernel_length, input_channels, output_channels)\n        containing weights for the convolution.\n      b: 1D tensor of shape (output_channels) containing biases for the\n        convolution, or `None` if no bias.\n      padding: `\"VALID\"` or `\"SAME\"`, the convolution's padding algorithm.\n      stride: Integer stride.\n\n    Returns:\n      Output bounds.\n    \"\"\"\n    w_pos = tf.maximum(w, 0)\n    w_neg = tf.minimum(w, 0)\n    lb = (tf.nn.conv1d(self.lower_offset, w_pos,\n                       padding=padding, stride=stride) +\n          tf.nn.conv1d(self.upper_offset, w_neg,\n                       padding=padding, stride=stride))\n    ub = (tf.nn.conv1d(self.upper_offset, w_pos,\n                       padding=padding, stride=stride) +\n          tf.nn.conv1d(self.lower_offset, w_neg,\n                       padding=padding, stride=stride))\n\n    nominal_out = tf.nn.conv1d(self.nominal, w,\n                               padding=padding, stride=stride)\n    if b is not None:\n      nominal_out += b\n\n    return RelativeIntervalBounds(lb, ub, nominal_out)\n\n  def apply_conv2d(self, wrapper, w, b, padding, strides):\n    \"\"\"Propagates the bounds through a 2D convolution layer.\n\n    Args:\n      wrapper: Contains prior bounds from a previous iteration.\n      w: 4D tensor of shape (kernel_height, kernel_width, input_channels,\n        output_channels) containing weights for the convolution.\n      b: 1D tensor of shape (output_channels) containing biases for the\n        convolution, or `None` if no bias.\n      padding: `\"VALID\"` or `\"SAME\"`, the convolution's padding algorithm.\n      strides: Integer list of length N: `[vertical_stride, horizontal_stride]`.\n\n    Returns:\n      Output bounds.\n    \"\"\"\n    w_pos = tf.maximum(w, 0)\n    w_neg = tf.minimum(w, 0)\n    lb = (tf.nn.convolution(self.lower_offset, w_pos,\n                            padding=padding, strides=strides) +\n          tf.nn.convolution(self.upper_offset, w_neg,\n                            padding=padding, strides=strides))\n    ub = (tf.nn.convolution(self.upper_offset, w_pos,\n                            padding=padding, strides=strides) +\n          tf.nn.convolution(self.lower_offset, w_neg,\n                            padding=padding, strides=strides))\n\n    nominal_out = tf.nn.convolution(self.nominal, w,\n                                    padding=padding, strides=strides)\n    if b is not None:\n      nominal_out += b\n\n    return RelativeIntervalBounds(lb, ub, nominal_out)\n\n  def apply_increasing_monotonic_fn(self, wrapper, fn, *args, **parameters):\n    \"\"\"Propagates the bounds through a non-linear activation layer or `add` op.\n\n    Args:\n      wrapper: Contains prior bounds from a previous iteration.\n      fn: String specifying non-linear activation function.\n        May be one of: sig, relu, tanh, elu, leaky_relu.\n        Anything else denotes identity.\n      *args: Other inputs' bounds, for a multi-input node (e.g. Add).\n      **parameters: Optional parameters if activation is parameterised, e.g.\n        `{'alpha': 0.2}` for leaky ReLu.\n\n    Returns:\n      Output bounds.\n    \"\"\"\n    if fn.__name__ in ('add', 'reduce_mean', 'reduce_sum', 'avg_pool'):\n      return RelativeIntervalBounds(\n          fn(self.lower_offset, *[bounds.lower_offset for bounds in args]),\n          fn(self.upper_offset, *[bounds.upper_offset for bounds in args]),\n          fn(self.nominal, *[bounds.nominal for bounds in args]))\n    else:\n      assert not args, 'unary function expected'\n      nominal_out = fn(self.nominal)\n      if fn.__name__ == 'reduce_max':\n        lb, ub = _maxpool_bounds(fn, None, None,\n                                 self.lower_offset, self.upper_offset,\n                                 nominal_in=self.nominal,\n                                 nominal_out=nominal_out)\n      elif fn.__name__ == 'max_pool':\n        lb, ub = _maxpool_bounds(fn,\n                                 parameters['ksize'][1:-1],\n                                 parameters['strides'][1:-1],\n                                 self.lower_offset, self.upper_offset,\n                                 nominal_in=self.nominal,\n                                 nominal_out=nominal_out)\n      else:\n        lb, ub = _activation_bounds(fn, self.lower_offset, self.upper_offset,\n                                    nominal_in=self.nominal,\n                                    parameters=parameters)\n      return RelativeIntervalBounds(lb, ub, nominal_out)\n\n  def apply_batch_norm(self, wrapper, mean, variance, scale, bias, epsilon):\n    \"\"\"Propagates the bounds through a batch norm layer.\n\n    Args:\n      wrapper: Contains prior bounds from a previous iteration.\n      mean: Learnt batch mean.\n      variance: Learnt batch variance.\n      scale: Trained component-wise scale variable.\n      bias: Trained component-wise bias variable.\n      epsilon: Epsilon for avoiding instability when `variance` is very small.\n\n    Returns:\n      Output bounds.\n    \"\"\"\n    lb = tf.nn.batch_normalization(self.lower_offset,\n                                   tf.zeros_like(mean), variance,\n                                   None, scale, epsilon)\n    ub = tf.nn.batch_normalization(self.upper_offset,\n                                   tf.zeros_like(mean), variance,\n                                   None, scale, epsilon)\n    # It's just possible that the batchnorm's scale is negative.\n    lb, ub = tf.minimum(lb, ub), tf.maximum(lb, ub)\n\n    nominal_out = tf.nn.batch_normalization(self.nominal,\n                                            mean, variance,\n                                            bias, scale, epsilon)\n    return RelativeIntervalBounds(lb, ub, nominal_out)\n\n  def _set_up_cache(self):\n    self._lower_offset, update_lower = self._cache_with_update_op(\n        self._lower_offset)\n    self._upper_offset, update_upper = self._cache_with_update_op(\n        self._upper_offset)\n    return tf.group([update_lower, update_upper])\n\n\ndef _maxpool_bounds(module, kernel_shape, strides, lb_in, ub_in,\n                    nominal_in, nominal_out):\n  \"\"\"Calculates naive bounds on output of an N-D max pool layer.\n\n  Args:\n    module: Callable for max-pool operation.\n    kernel_shape: Integer list of `[kernel_height, kernel_width]`,\n      or `None` to aggregate over the layer`s entire spatial extent.\n    strides: Integer list of `[vertical_stride, horizontal_stride]`.\n    lb_in: (N+2)D tensor of shape (batch_size, input_height, input_width,\n      layer_channels) containing lower bounds on the inputs to the\n      max pool layer.\n    ub_in: (N+2)D tensor of shape (batch_size, input_height, input_width,\n      layer_channels) containing upper bounds on the inputs to the\n      max pool layer.\n    nominal_in: (N+2)D tensor of shape (batch_size, input_height, input_width,\n      layer_channels) containing nominal input values.\n      Inputs bounds are interpreted relative to this.\n    nominal_out: (N+2)D tensor of shape (batch_size, output_height,output_width,\n      layer_channels) containing nominal input values.\n      The returned output bounds are expressed relative to this.\n\n  Returns:\n    lb_out: (N+2)D tensor of shape (batch_size, output_height, output_width,\n      layer_channels) with lower bounds on the outputs of the max pool layer.\n    ub_out: (N+2)D tensor of shape (batch_size, output_height, output_width,\n      layer_channels) with upper bounds on the outputs of the max pool layer.\n  \"\"\"\n  if kernel_shape is None:\n    nominal_out = tf.reduce_max(nominal_in,\n                                axis=list(range(1, nominal_in.shape.ndims-1)),\n                                keepdims=True)\n    return (module((nominal_in - nominal_out) + lb_in),\n            module((nominal_in - nominal_out) + ub_in))\n  else:\n    # Must perform the max on absolute bounds, as the kernels may overlap.\n    # TODO(stanforth) investigate a more numerically stable implementation\n    del strides\n    return (module(nominal_in + lb_in) - nominal_out,\n            module(nominal_in + ub_in) - nominal_out)\n\n\ndef _activation_bounds(nl_fun, lb_in, ub_in, nominal_in, parameters=None):\n  \"\"\"Calculates naive bounds on output of an activation layer.\n\n  Inputs bounds are interpreted relative to `nominal_in`, and the returned\n  output bounds are expressed relative to `nominal_out=nl(nominal_in)`.\n\n  Args:\n    nl_fun: Callable implementing the activation function itself.\n    lb_in: (N+2)D tensor of shape (batch_size, layer_height, layer_width,\n      layer_channels) containing lower bounds on the pre-activations.\n    ub_in: (N+2)D tensor of shape (batch_size, layer_height, layer_width,\n      layer_channels) containing upper bounds on the pre-activations.\n    nominal_in: (N+2)D tensor of shape (batch_size, input_height, input_width,\n      layer_channels) containing nominal input values.\n    parameters: Optional parameter dict if activation is parameterised, e.g.\n      `{'alpha': 0.2}` for leaky ReLu.\n\n  Returns:\n    lb_out: 2D tensor of shape (batch_size, layer_size) or\n      4D tensor of shape (batch_size, layer_height, layer_width, layer_channels)\n      with lower bounds on the activations.\n    ub_out: 2D tensor of shape (batch_size, layer_size) or\n      4D tensor of shape (batch_size, layer_height, layer_width, layer_channels)\n      with upper bounds on the activations.\n  \"\"\"\n  if nl_fun.__name__ == 'relu':\n    return (\n        tf.maximum(tf.minimum(nominal_in, 0.) + lb_in,\n                   tf.minimum(-nominal_in, 0.)),  # pylint:disable=invalid-unary-operand-type\n        tf.maximum(tf.minimum(nominal_in, 0.) + ub_in,\n                   tf.minimum(-nominal_in, 0.)))  # pylint:disable=invalid-unary-operand-type\n  elif nl_fun.__name__ == 'leaky_relu':\n    alpha = parameters['alpha']\n    return (\n        tf.maximum(\n            lb_in + tf.minimum(nominal_in, 0.) * (1. - alpha),\n            alpha * lb_in + tf.minimum(-nominal_in, 0.) * (1. - alpha)),  # pylint:disable=invalid-unary-operand-type\n        tf.maximum(\n            ub_in + tf.minimum(nominal_in, 0.) * (1. - alpha),\n            alpha * ub_in + tf.minimum(-nominal_in, 0.) * (1. - alpha)))  # pylint:disable=invalid-unary-operand-type\n  else:\n    nominal_out = nl_fun(nominal_in)\n    return (nl_fun(nominal_in + lb_in) - nominal_out,\n            nl_fun(nominal_in + ub_in) - nominal_out)\n"
  },
  {
    "path": "interval_bound_propagation/src/simplex_bounds.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Naive bound calculation for common neural network layers.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom interval_bound_propagation.src import bounds as basic_bounds\nfrom interval_bound_propagation.src import relative_bounds\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\nclass SimplexBounds(basic_bounds.AbstractBounds):\n  \"\"\"Specifies a bounding simplex within an embedding space.\"\"\"\n\n  def __init__(self, vertices, nominal, r):\n    \"\"\"Initialises the simplex bounds.\n\n    Args:\n      vertices: Tensor of shape (num_vertices, *input_shape)\n        or of shape (batch_size, num_vertices, *input_shape)\n        containing the vertices in embedding space.\n      nominal: Tensor of shape (batch_size, *input_shape) specifying\n        the unperturbed inputs in embedding space, where `*input_shape`\n        denotes either (embedding_size,) for flat input (e.g. bag-of-words)\n        or (input_length, embedding_channels) for sequence input.\n      r: Scalar specifying the dilation factor of the simplex. The dilated\n        simplex will have vertices `nominal + r * (vertices-nominal)`.\n    \"\"\"\n    super(SimplexBounds, self).__init__()\n    self._vertices = vertices\n    self._nominal = nominal\n    self._r = r\n\n  @property\n  def vertices(self):\n    return self._vertices\n\n  @property\n  def nominal(self):\n    return self._nominal\n\n  @property\n  def r(self):\n    return self._r\n\n  @property\n  def shape(self):\n    return self.nominal.shape.as_list()\n\n  @classmethod\n  def convert(cls, bounds):\n    if not isinstance(bounds, cls):\n      raise ValueError('Cannot convert \"{}\" to \"{}\"'.format(bounds,\n                                                            cls.__name__))\n    return bounds\n\n  def apply_batch_reshape(self, wrapper, shape):\n    reshape = snt.BatchReshape(shape)\n    if self.vertices.shape.ndims == self.nominal.shape.ndims:\n      reshape_vertices = reshape\n    else:\n      reshape_vertices = snt.BatchReshape(shape, preserve_dims=2)\n    return SimplexBounds(reshape_vertices(self.vertices),\n                         reshape(self.nominal),\n                         self.r)\n\n  def apply_linear(self, wrapper, w, b):\n    mapped_centres = tf.matmul(self.nominal, w)\n    mapped_vertices = tf.tensordot(self.vertices, w, axes=1)\n\n    lb, ub = _simplex_bounds(mapped_vertices, mapped_centres, self.r, -2)\n\n    nominal_out = tf.matmul(self.nominal, w)\n    if b is not None:\n      nominal_out += b\n\n    return relative_bounds.RelativeIntervalBounds(lb, ub, nominal_out)\n\n  def apply_conv1d(self, wrapper, w, b, padding, stride):\n    mapped_centres = tf.nn.conv1d(self.nominal, w,\n                                  padding=padding, stride=stride)\n    if self.vertices.shape.ndims == 3:\n      # `self.vertices` has no batch dimension; its shape is\n      # (num_vertices, input_length, embedding_channels).\n      mapped_vertices = tf.nn.conv1d(self.vertices, w,\n                                     padding=padding, stride=stride)\n    elif self.vertices.shape.ndims == 4:\n      # `self.vertices` has shape\n      # (batch_size, num_vertices, input_length, embedding_channels).\n      # Vertices are different for each example in the batch,\n      # e.g. for word perturbations.\n      mapped_vertices = snt.BatchApply(\n          lambda x: tf.nn.conv1d(x, w, padding=padding, stride=stride))(\n              self.vertices)\n    else:\n      raise ValueError('\"vertices\" must have either 3 or 4 dimensions.')\n\n    lb, ub = _simplex_bounds(mapped_vertices, mapped_centres, self.r, -3)\n\n    nominal_out = tf.nn.conv1d(self.nominal, w,\n                               padding=padding, stride=stride)\n    if b is not None:\n      nominal_out += b\n\n    return relative_bounds.RelativeIntervalBounds(lb, ub, nominal_out)\n\n  def apply_conv2d(self, wrapper, w, b, padding, strides):\n    mapped_centres = tf.nn.convolution(self.nominal, w,\n                                       padding=padding, strides=strides)\n    if self.vertices.shape.ndims == 4:\n      # `self.vertices` has no batch dimension; its shape is\n      # (num_vertices, input_height, input_width, input_channels).\n      mapped_vertices = tf.nn.convolution(self.vertices, w,\n                                          padding=padding, strides=strides)\n    elif self.vertices.shape.ndims == 5:\n      # `self.vertices` has shape\n      # (batch_size, num_vertices, input_height, input_width, input_channels).\n      # Vertices are different for each example in the batch.\n      mapped_vertices = snt.BatchApply(\n          lambda x: tf.nn.convolution(x, w, padding=padding, strides=strides))(\n              self.vertices)\n    else:\n      raise ValueError('\"vertices\" must have either 4 or 5 dimensions.')\n\n    lb, ub = _simplex_bounds(mapped_vertices, mapped_centres, self.r, -4)\n\n    nominal_out = tf.nn.convolution(self.nominal, w,\n                                    padding=padding, strides=strides)\n    if b is not None:\n      nominal_out += b\n\n    return relative_bounds.RelativeIntervalBounds(lb, ub, nominal_out)\n\n  def apply_increasing_monotonic_fn(self, wrapper, fn, *args, **parameters):\n    if fn.__name__ in ('add', 'reduce_mean', 'reduce_sum', 'avg_pool'):\n      if self.vertices.shape.ndims == self.nominal.shape.ndims:\n        vertices_fn = fn\n      else:\n        vertices_fn = snt.BatchApply(fn, n_dims=2)\n      return SimplexBounds(\n          vertices_fn(self.vertices, *[bounds.vertices for bounds in args]),\n          fn(self.nominal, *[bounds.nominal for bounds in args]),\n          self.r)\n\n    elif fn.__name__ == 'quotient':\n      return SimplexBounds(\n          self.vertices / tf.expand_dims(parameters['denom'], axis=1),\n          fn(self.nominal),\n          self.r)\n\n    else:\n      return super(SimplexBounds, self).apply_increasing_monotonic_fn(\n          wrapper, fn, *args, **parameters)\n\n\ndef _simplex_bounds(mapped_vertices, mapped_centres, r, axis):\n  \"\"\"Calculates naive bounds on the given layer-mapped vertices.\n\n  Args:\n    mapped_vertices: Tensor of shape (num_vertices, *output_shape)\n      or of shape (batch_size, num_vertices, *output_shape)\n      containing the vertices in the layer's output space.\n    mapped_centres: Tensor of shape (batch_size, *output_shape)\n      containing the layer's nominal outputs.\n    r: Scalar in [0, 1) specifying the radius (in vocab space) of the simplex.\n    axis: Index of the `num_vertices` dimension of `mapped_vertices`.\n\n  Returns:\n    lb_out: Tensor of shape (batch_size, *output_shape) with lower bounds\n      on the outputs of the affine layer.\n    ub_out: Tensor of shape (batch_size, *output_shape) with upper bounds\n      on the outputs of the affine layer.\n  \"\"\"\n  # Use the negative of r, instead of the complement of r, as\n  # we're shifting the input domain to be centred at the origin.\n  lb_out = -r * mapped_centres + r * tf.reduce_min(mapped_vertices, axis=axis)\n  ub_out = -r * mapped_centres + r * tf.reduce_max(mapped_vertices, axis=axis)\n  return lb_out, ub_out\n\n"
  },
  {
    "path": "interval_bound_propagation/src/specification.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Defines the output specifications.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport abc\n\nfrom absl import logging\n\nfrom interval_bound_propagation.src import bounds as bounds_lib\nfrom interval_bound_propagation.src import verifiable_wrapper\nimport six\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\n@six.add_metaclass(abc.ABCMeta)\nclass Specification(snt.AbstractModule):\n  \"\"\"Defines a specification.\"\"\"\n\n  def __init__(self, name, collapse=True):\n    super(Specification, self).__init__(name=name)\n    self._collapse = collapse\n\n  @abc.abstractmethod\n  def _build(self, modules):\n    \"\"\"Computes the worst-case specification value.\"\"\"\n\n  @abc.abstractmethod\n  def evaluate(self, logits):\n    \"\"\"Computes the specification value.\n\n    Args:\n      logits: The logits Tensor can have different shapes, i.e.,\n        [batch_size, num_classes]: The output should be [batch_size, num_specs].\n        [num_restarts, batch_size, num_classes]: The output should be\n          [num_restarts, batch_size, num_specs]. Used by UntargetedPGDAttack.\n        [num_restarts, num_specs, batch_size, num_classes]: The output should\n          be [num_restarts, batch_size, num_specs]. For this case, the\n          specifications must be evaluated individually for each column\n          (axis = 1). Used by MultiTargetedPGDAttack.\n\n    Returns:\n      The specification values evaluated at the network output.\n    \"\"\"\n\n  @abc.abstractproperty\n  def num_specifications(self):\n    \"\"\"Returns the number of specifications.\"\"\"\n\n  @property\n  def collapse(self):\n    return self._collapse\n\n\nclass LinearSpecification(Specification):\n  \"\"\"Linear specifications: c^T * z_K + d <= 0.\"\"\"\n\n  def __init__(self, c, d=None, prune_irrelevant=True, collapse=True):\n    \"\"\"Builds a linear specification module.\"\"\"\n    super(LinearSpecification, self).__init__(name='specs', collapse=collapse)\n    # c has shape [batch_size, num_specifications, num_outputs]\n    # d has shape [batch_size, num_specifications]\n    # Some specifications may be irrelevant (not a function of the output).\n    # We automatically remove them for clarity. We expect the number of\n    # irrelevant specs to be equal for all elements of a batch.\n    # Shape is [batch_size, num_specifications]\n    if prune_irrelevant:\n      irrelevant = tf.equal(tf.reduce_sum(\n          tf.cast(tf.abs(c) > 1e-6, tf.int32), axis=-1, keepdims=True), 0)\n      batch_size = tf.shape(c)[0]\n      num_outputs = tf.shape(c)[2]\n      irrelevant = tf.tile(irrelevant, [1, 1, num_outputs])\n      self._c = tf.reshape(\n          tf.boolean_mask(c, tf.logical_not(irrelevant)),\n          [batch_size, -1, num_outputs])\n    else:\n      self._c = c\n    self._d = d\n\n  def _build(self, modules):\n    \"\"\"Outputs specification value.\"\"\"\n    # inputs have shape [batch_size, num_outputs].\n    if not (self.collapse and\n            isinstance(modules[-1], verifiable_wrapper.LinearFCWrapper)):\n      logging.info('Elision of last layer disabled.')\n      bounds = modules[-1].output_bounds\n      w = self._c\n      b = self._d\n    else:\n      logging.info('Elision of last layer active.')\n      # Collapse the last layer.\n      bounds = modules[-1].input_bounds\n      w = modules[-1].module.w\n      b = modules[-1].module.b\n      w = tf.einsum('ijk,lk->ijl', self._c, w)\n      b = tf.einsum('ijk,k->ij', self._c, b)\n      if self._d is not None:\n        b += self._d\n\n    # Maximize z * w + b s.t. lower <= z <= upper.\n    bounds = bounds_lib.IntervalBounds.convert(bounds)\n    c = (bounds.lower + bounds.upper) / 2.\n    r = (bounds.upper - bounds.lower) / 2.\n    c = tf.einsum('ij,ikj->ik', c, w)\n    if b is not None:\n      c += b\n    r = tf.einsum('ij,ikj->ik', r, tf.abs(w))\n\n    # output has shape [batch_size, num_specifications].\n    return c + r\n\n  def evaluate(self, logits):\n    if len(logits.shape) == 2:\n      output = tf.einsum('ij,ikj->ik', logits, self._c)\n    elif len(logits.shape) == 3:\n      output = tf.einsum('rij,ikj->rik', logits, self._c)\n    else:\n      assert len(logits.shape) == 4\n      output = tf.einsum('rsbo,bso->rbs', logits, self._c)\n    if self._d is not None:\n      output += self._d\n    return output\n\n  @property\n  def num_specifications(self):\n    return tf.shape(self._c)[1]\n\n  @property\n  def c(self):\n    return self._c\n\n  @property\n  def d(self):\n    return self._d\n\n\nclass ClassificationSpecification(Specification):\n  \"\"\"Creates a linear specification that corresponds to a classification.\n\n  This class is not a standard LinearSpecification as it does not materialize\n  the c and d tensors.\n  \"\"\"\n\n  def __init__(self, label, num_classes, collapse=True):\n    super(ClassificationSpecification, self).__init__(name='specs',\n                                                      collapse=collapse)\n    self._label = label\n    self._num_classes = num_classes\n    # Precompute indices.\n    with self._enter_variable_scope():\n      indices = []\n      for i in range(self._num_classes):\n        indices.append(list(range(i)) + list(range(i + 1, self._num_classes)))\n      indices = tf.constant(indices, dtype=tf.int32)\n      self._correct_idx, self._wrong_idx = self._build_indices(label, indices)\n\n  def _build(self, modules):\n    if not (self.collapse and\n            isinstance(modules[-1], verifiable_wrapper.LinearFCWrapper)):\n      logging.info('Elision of last layer disabled.')\n      bounds = modules[-1].output_bounds\n      bounds = bounds_lib.IntervalBounds.convert(bounds)\n      correct_class_logit = tf.gather_nd(bounds.lower, self._correct_idx)\n      wrong_class_logits = tf.gather_nd(bounds.upper, self._wrong_idx)\n      return wrong_class_logits - tf.expand_dims(correct_class_logit, 1)\n\n    logging.info('Elision of last layer active.')\n    bounds = modules[-1].input_bounds\n    bounds = bounds_lib.IntervalBounds.convert(bounds)\n    batch_size = tf.shape(bounds.lower)[0]\n    w = modules[-1].module.w\n    b = modules[-1].module.b\n    w_t = tf.tile(tf.expand_dims(tf.transpose(w), 0), [batch_size, 1, 1])\n    b_t = tf.tile(tf.expand_dims(b, 0), [batch_size, 1])\n    w_correct = tf.expand_dims(tf.gather_nd(w_t, self._correct_idx), -1)\n    b_correct = tf.expand_dims(tf.gather_nd(b_t, self._correct_idx), 1)\n    w_wrong = tf.transpose(tf.gather_nd(w_t, self._wrong_idx), [0, 2, 1])\n    b_wrong = tf.gather_nd(b_t, self._wrong_idx)\n    w = w_wrong - w_correct\n    b = b_wrong - b_correct\n    # Maximize z * w + b s.t. lower <= z <= upper.\n    c = (bounds.lower + bounds.upper) / 2.\n    r = (bounds.upper - bounds.lower) / 2.\n    c = tf.einsum('ij,ijk->ik', c, w)\n    if b is not None:\n      c += b\n    r = tf.einsum('ij,ijk->ik', r, tf.abs(w))\n    return c + r\n\n  def evaluate(self, logits):\n    if len(logits.shape) == 2:\n      correct_class_logit = tf.gather_nd(logits, self._correct_idx)\n      correct_class_logit = tf.expand_dims(correct_class_logit, -1)\n      wrong_class_logits = tf.gather_nd(logits, self._wrong_idx)\n    elif len(logits.shape) == 3:\n      # [num_restarts, batch_size, num_classes] to\n      # [num_restarts, batch_size, num_specs]\n      logits = tf.transpose(logits, [1, 2, 0])  # Put restart dimension last.\n      correct_class_logit = tf.gather_nd(logits, self._correct_idx)\n      correct_class_logit = tf.transpose(correct_class_logit)\n      correct_class_logit = tf.expand_dims(correct_class_logit, -1)\n      wrong_class_logits = tf.gather_nd(logits, self._wrong_idx)\n      wrong_class_logits = tf.transpose(wrong_class_logits, [2, 0, 1])\n    else:\n      assert len(logits.shape) == 4\n      # [num_restarts, num_specs, batch_size, num_classes] to\n      # [num_restarts, batch_size, num_specs].\n      logits = tf.transpose(logits, [2, 3, 1, 0])\n      correct_class_logit = tf.gather_nd(logits, self._correct_idx)\n      correct_class_logit = tf.transpose(correct_class_logit, [2, 0, 1])\n      batch_size = tf.shape(logits)[0]\n      wrong_idx = tf.concat([\n          self._wrong_idx,\n          tf.tile(tf.reshape(tf.range(self.num_specifications, dtype=tf.int32),\n                             [1, self.num_specifications, 1]),\n                  [batch_size, 1, 1])], axis=-1)\n      wrong_class_logits = tf.gather_nd(logits, wrong_idx)\n      wrong_class_logits = tf.transpose(wrong_class_logits, [2, 0, 1])\n    return wrong_class_logits - correct_class_logit\n\n  @property\n  def num_specifications(self):\n    return self._num_classes - 1\n\n  @property\n  def correct_idx(self):\n    return self._correct_idx\n\n  @property\n  def wrong_idx(self):\n    return self._wrong_idx\n\n  def _build_indices(self, label, indices):\n    batch_size = tf.shape(label)[0]\n    i = tf.range(batch_size, dtype=tf.int32)\n    correct_idx = tf.stack([i, tf.cast(label, tf.int32)], axis=1)\n    wrong_idx = tf.stack([\n        tf.tile(tf.reshape(i, [batch_size, 1]), [1, self._num_classes - 1]),\n        tf.gather(indices, label),\n    ], axis=2)\n    return correct_idx, wrong_idx\n\n\nclass TargetedClassificationSpecification(ClassificationSpecification):\n  \"\"\"Defines a specification that compares the true class with another.\"\"\"\n\n  def __init__(self, label, num_classes, target_class, collapse=True):\n    super(TargetedClassificationSpecification, self).__init__(\n        label, num_classes, collapse=collapse)\n    batch_size = tf.shape(label)[0]\n    if len(target_class.shape) == 1:\n      target_class = tf.reshape(target_class, [batch_size, 1])\n    self._num_specifications = target_class.shape[1].value\n    if self._num_specifications is None:\n      raise ValueError('Cannot retrieve the number of target classes')\n    self._target_class = target_class\n    i = tf.range(batch_size, dtype=tf.int32)\n    self._wrong_idx = tf.stack([\n        tf.tile(tf.reshape(i, [batch_size, 1]), [1, self.num_specifications]),\n        target_class\n    ], axis=2)\n\n  @property\n  def target_class(self):\n    \"\"\"Returns the target class index.\"\"\"\n    return self._target_class\n\n  @property\n  def num_specifications(self):\n    return self._num_specifications\n\n\nclass RandomClassificationSpecification(TargetedClassificationSpecification):\n  \"\"\"Creates a single random specification that targets a random class.\"\"\"\n\n  def __init__(self, label, num_classes, num_targets=1, seed=None,\n               collapse=True):\n    # Overwrite the target indices. Each session.run() call gets new target\n    # indices, the indices should remain the same across restarts.\n    batch_size = tf.shape(label)[0]\n    j = tf.random.uniform(shape=(batch_size, num_targets), minval=1,\n                          maxval=num_classes, dtype=tf.int32, seed=seed)\n    target_class = tf.mod(tf.cast(tf.expand_dims(label, -1), tf.int32) + j,\n                          num_classes)\n    super(RandomClassificationSpecification, self).__init__(\n        label, num_classes, target_class, collapse=collapse)\n\n\nclass LeastLikelyClassificationSpecification(\n    TargetedClassificationSpecification):\n  \"\"\"Creates a single specification that targets the least likely class.\"\"\"\n\n  def __init__(self, label, num_classes, logits, num_targets=1, collapse=True):\n    # Do not target the true class. If the true class is the least likely to\n    # be predicted, it is fine to target any other class as the attack will\n    # be successful anyways.\n    j = tf.nn.top_k(-logits, k=num_targets, sorted=False).indices\n    l = tf.expand_dims(label, 1)\n    target_class = tf.mod(\n        j + tf.cast(tf.equal(j, tf.cast(l, tf.int32)), tf.int32), num_classes)\n    super(LeastLikelyClassificationSpecification, self).__init__(\n        label, num_classes, target_class, collapse=collapse)\n"
  },
  {
    "path": "interval_bound_propagation/src/utils.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Helpers.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport collections\nimport re\n\nfrom absl import logging\nfrom interval_bound_propagation.src import attacks\nfrom interval_bound_propagation.src import bounds\nfrom interval_bound_propagation.src import layers\nfrom interval_bound_propagation.src import loss\nfrom interval_bound_propagation.src import specification\nimport numpy as np\nimport tensorflow.compat.v1 as tf\n\n\n# Defines a dataset sample.\"\"\"\nSample = collections.namedtuple('Sample', ['image', 'label'])\n\n\ndef build_dataset(raw_data, batch_size=50, sequential=True):\n  \"\"\"Builds a dataset from raw NumPy tensors.\"\"\"\n  images, labels = raw_data\n  # We need width, height and channel.\n  if len(images.shape) == 3:\n    images = np.expand_dims(images, -1)\n  samples = Sample(images.astype(np.float32) / 255., labels.astype(np.int64))\n  data = tf.data.Dataset.from_tensor_slices(samples)\n  if not sequential:\n    data = data.shuffle(1000)\n  return data.repeat().batch(batch_size).make_one_shot_iterator().get_next()\n\n\ndef randomize(images, init_shape, expand_shape=None, crop_shape=None,\n              vertical_flip=False):\n  \"\"\"Returns a function that randomly translates and flips images.\"\"\"\n  def random_image(image):\n    \"\"\"Randmly translates and flips images.\"\"\"\n    image = tf.reshape(image, init_shape)\n    current_shape = init_shape\n    if expand_shape is not None and expand_shape != current_shape:\n      if expand_shape[-1] != current_shape[-1]:\n        raise ValueError('Number channels is not specified correctly.')\n      image = tf.image.resize_image_with_crop_or_pad(\n          image, expand_shape[0], expand_shape[1])\n      current_shape = expand_shape\n    if crop_shape is not None and crop_shape != current_shape:\n      image = tf.random_crop(image, crop_shape)\n    if vertical_flip:\n      image = tf.image.random_flip_left_right(image)\n    return image\n  return tf.map_fn(random_image, images)\n\n\ndef linear_schedule(step, init_step, final_step, init_value, final_value):\n  \"\"\"Linear schedule.\"\"\"\n  assert final_step >= init_step\n  if init_step == final_step:\n    return final_value\n  rate = tf.cast(step - init_step, tf.float32) / float(final_step - init_step)\n  linear_value = rate * (final_value - init_value) + init_value\n  return tf.clip_by_value(linear_value, min(init_value, final_value),\n                          max(init_value, final_value))\n\n\ndef smooth_schedule(step, init_step, final_step, init_value, final_value,\n                    mid_point=.25, beta=4.):\n  \"\"\"Smooth schedule that slowly morphs into a linear schedule.\"\"\"\n  assert final_value > init_value\n  assert final_step >= init_step\n  assert beta >= 2.\n  assert mid_point >= 0. and mid_point <= 1.\n  mid_step = int((final_step - init_step) * mid_point) + init_step\n  if mid_step <= init_step:\n    alpha = 1.\n  else:\n    t = (mid_step - init_step) ** (beta - 1.)\n    alpha = (final_value - init_value) / ((final_step - mid_step) * beta * t +\n                                          (mid_step - init_step) * t)\n  mid_value = alpha * (mid_step - init_step) ** beta + init_value\n  # Tensorflow operation.\n  is_ramp = tf.cast(step > init_step, tf.float32)\n  is_linear = tf.cast(step >= mid_step, tf.float32)\n  return (is_ramp * (\n      (1. - is_linear) * (\n          init_value +\n          alpha * tf.pow(tf.cast(step - init_step, tf.float32), beta)) +\n      is_linear * linear_schedule(\n          step, mid_step, final_step, mid_value, final_value)) +\n          (1. - is_ramp) * init_value)\n\n\ndef build_loss_schedule(step, warmup_steps, rampup_steps, init, final,\n                        warmup=None):\n  \"\"\"Linear schedule builder.\n\n  Args:\n    step: Current step number.\n    warmup_steps: When step < warmup_steps, set value to warmup.\n    rampup_steps: Ramp up schedule value from init to final in rampup_step.\n    init: Initial schedule value after warmup_steps.\n    final: Final schedule value after warmup_steps + rampup_steps.\n    warmup: Schedule value before warmup_steps. When set to None, the warmup\n      period value is set to init.\n\n  Returns:\n    A schedule tensor.\n  \"\"\"\n  if warmup is None and init == final:\n    return init\n  if rampup_steps < 0:\n    if warmup is not None:\n      return tf.cond(step < warmup_steps, lambda: tf.constant(warmup),\n                     lambda: tf.constant(final))\n    return final\n  schedule = linear_schedule(\n      step, warmup_steps, warmup_steps + rampup_steps, init, final)\n  if warmup is not None:\n    # Set the value to warmup during warmup process.\n    return tf.cond(step < warmup_steps,\n                   lambda: tf.constant(warmup), lambda: schedule)\n  return schedule\n\n\ndef add_image_normalization(model, mean, std):\n  def _model(x, *args, **kwargs):\n    return model(layers.ImageNorm(mean, std)(x), *args, **kwargs)\n  return _model\n\n\ndef create_specification(label, num_classes, logits,\n                         specification_type='one_vs_all', collapse=True):\n  \"\"\"Creates a specification of the desired type.\"\"\"\n  def _num_targets(name):\n    tokens = name.rsplit('_', 1)\n    return int(tokens[1]) if len(tokens) > 1 else 1\n  if specification_type == 'one_vs_all':\n    return specification.ClassificationSpecification(label, num_classes,\n                                                     collapse=collapse)\n  elif specification_type.startswith('random'):\n    return specification.RandomClassificationSpecification(\n        label, num_classes, _num_targets(specification_type), collapse=collapse)\n  elif specification_type.startswith('least_likely'):\n    return specification.LeastLikelyClassificationSpecification(\n        label, num_classes, logits, _num_targets(specification_type),\n        collapse=collapse)\n  else:\n    raise ValueError('Unknown specification type: \"{}\"'.format(\n        specification_type))\n\n\ndef create_classification_losses(\n    global_step, inputs, label, predictor_network, epsilon, loss_weights,\n    warmup_steps=0, rampup_steps=-1, input_bounds=(0., 1.),\n    loss_builder=loss.Losses, options=None):\n  \"\"\"Create the training loss.\"\"\"\n  # Whether to elide the last linear layer with the specification.\n  elide = True\n  # Which loss to use for the IBP loss.\n  loss_type = 'xent'\n  # If the loss_type is 'hinge', which margin to use.\n  loss_margin = 10.\n  # Amount of label smoothing.\n  label_smoothing = 0.\n  # If True, batch normalization stops training after warm-up.\n  is_training_off_after = -1\n  # If True, epsilon changes more smoothly.\n  smooth_epsilon_schedule = False\n  # Either 'one_vs_all', 'random_n', 'least_likely_n' or 'none'.\n  verified_specification = 'one_vs_all'\n  # Attack options.\n  attack_specification = 'UntargetedPGDAttack_7x1x1_UnrolledAdam_.1'\n  attack_scheduled = False\n  attack_random_init = 1.\n  # Model arguments.\n  nominal_args = dict(is_training=True, test_local_stats=False, reuse=False)\n  attack_args = {\n      'intermediate': dict(is_training=False, test_local_stats=False,\n                           reuse=True),\n      'final': dict(is_training=False, test_local_stats=False, reuse=True),\n  }\n  if options is not None:\n    elide = options.get('elide_last_layer', elide)\n    loss_type = options.get('verified_loss_type', loss_type)\n    loss_margin = options.get('verified_loss_margin', loss_type)\n    label_smoothing = options.get('label_smoothing', label_smoothing)\n    is_training_off_after = options.get(\n        'is_training_off_after', is_training_off_after)\n    smooth_epsilon_schedule = options.get(\n        'smooth_epsilon_schedule', smooth_epsilon_schedule)\n    verified_specification = options.get(\n        'verified_specification', verified_specification)\n    attack_specification = options.get(\n        'attack_specification', attack_specification)\n    attack_scheduled = options.get('attack_scheduled', attack_scheduled)\n    attack_random_init = options.get('attack_random_init', attack_random_init)\n    nominal_args = dict(options.get('nominal_args', nominal_args))\n    attack_args = dict(options.get('attack_args', attack_args))\n\n  def _get_schedule(init, final, warmup=None):\n    return build_loss_schedule(global_step, warmup_steps, rampup_steps, init,\n                               final, warmup)\n  def _is_loss_active(init, final, warmup=None):\n    return init > 0. or final > 0. or (warmup is not None and warmup > 0.)\n  nominal_xent = _get_schedule(**loss_weights.get('nominal'))\n  attack_xent = _get_schedule(**loss_weights.get('attack'))\n  use_attack = _is_loss_active(**loss_weights.get('attack'))\n  verified_loss = _get_schedule(**loss_weights.get('verified'))\n  use_verification = _is_loss_active(**loss_weights.get('verified'))\n  if verified_specification == 'none':\n    use_verification = False\n  weight_mixture = loss.ScalarLosses(\n      nominal_cross_entropy=nominal_xent,\n      attack_cross_entropy=attack_xent,\n      verified_loss=verified_loss)\n\n  # Ramp-up.\n  if rampup_steps < 0:\n    train_epsilon = tf.constant(epsilon)\n  else:\n    if smooth_epsilon_schedule:\n      train_epsilon = smooth_schedule(\n          global_step, warmup_steps, warmup_steps + rampup_steps, 0., epsilon)\n    else:\n      train_epsilon = linear_schedule(\n          global_step, warmup_steps, warmup_steps + rampup_steps, 0., epsilon)\n\n  # Set is_training according to options.\n  if is_training_off_after >= 0:\n    is_training = global_step < is_training_off_after\n  else:\n    is_training = True\n  # If the build arguments want training off, we set is_training to False.\n  # Otherwise, we respect the is_training_off_after option.\n  def _update_is_training(kwargs):\n    if 'is_training' in kwargs:\n      kwargs['is_training'] &= is_training\n  _update_is_training(nominal_args)\n  _update_is_training(attack_args['intermediate'])\n  _update_is_training(attack_args['final'])\n\n  logits = predictor_network(inputs, override=True, **nominal_args)\n  num_classes = predictor_network.output_size\n  if use_verification:\n    logging.info('Verification active.')\n    input_interval_bounds = bounds.IntervalBounds(\n        tf.maximum(inputs - train_epsilon, input_bounds[0]),\n        tf.minimum(inputs + train_epsilon, input_bounds[1]))\n    predictor_network.propagate_bounds(input_interval_bounds)\n    spec = create_specification(label, num_classes, logits,\n                                verified_specification, collapse=elide)\n  else:\n    logging.info('Verification disabled.')\n    spec = None\n  if use_attack:\n    logging.info('Attack active.')\n    pgd_attack = create_attack(\n        attack_specification, predictor_network, label,\n        train_epsilon if attack_scheduled else epsilon,\n        input_bounds=input_bounds, random_init=attack_random_init,\n        predictor_kwargs=attack_args)\n\n  else:\n    logging.info('Attack disabled.')\n    pgd_attack = None\n  losses = loss_builder(predictor_network, spec, pgd_attack,\n                        interval_bounds_loss_type=loss_type,\n                        interval_bounds_hinge_margin=loss_margin,\n                        label_smoothing=label_smoothing)\n  losses(label)\n  train_loss = sum(l * w for l, w in zip(losses.scalar_losses,\n                                         weight_mixture))\n  # Add a regularization loss.\n  regularizers = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)\n  train_loss = train_loss + tf.reduce_sum(regularizers)\n  return losses, train_loss, train_epsilon\n\n\n# Additional helper code to build specific PGD attacks.\ndef get_attack_builder(logits, label, name='UntargetedPGDAttack',\n                       random_seed=None, manual_target_class=None):\n  \"\"\"Returns a callable with the same arguments as PGDAttack.\n\n  In addition to the callable, this function also returns the targeted class\n  indices as a Tensor with the same shape as label.\n\n  Usage is as follows:\n    logits = model(inputs)\n    attack_cls, specification, target_class = get_attack_builder(logits, labels)\n    # target_class is None, if attack_cls is not a targeted attack.\n    attack_instance = attack_cls(model, specification, epsilon)\n    perturbed_inputs = attack_instance(inputs, labels)\n\n  Args:\n    logits: Tensor of nominal logits of shape [batch_size, num_classes].\n    label: Tensor of labels of shape [batch_size].\n    name: Name of a PGDAttack class or any of \"RandomMoreLikelyPGDAttack\",\n      \"RandomMostLikelyPGDAttack\", \"LeastLikelyMoreLikelyPGDAttack\",\n      \"LeastLikelyMostLikelyPGDAttack\", \"ManualMoreLikelyPGDAttack\",\n      \"ManualMostLikelyPGDAttack\". Any attack name can be postfixed by\n      \"Xent\" to use the cross-entropy loss rather than margin loss.\n    random_seed: Sets the random seed for \"Random*\" attacks.\n    manual_target_class: For \"Manual*\" attacks, Tensor of target class indices\n      of shape [batch_size].\n\n  Returns:\n    A callable, a Specification and a Tensor of target label (or None if the\n    attack is not targeted).\n  \"\"\"\n  if name.endswith('Xent'):\n    use_xent = True\n    name = name[:-4]\n  else:\n    use_xent = False\n  if name.endswith('Linf'):\n    use_l2 = False\n    name = name[:-4]  # Just for syntactic sugar.\n  elif name.endswith('L2'):\n    use_l2 = True\n    name = name[:-2]\n  else:\n    use_l2 = False\n  num_classes = logits.shape[1].value\n  if num_classes is None:\n    raise ValueError('Cannot determine the number of classes from logits.')\n\n  # Special case for multi-targeted attacks.\n  m = re.match(r'((?:MemoryEfficient)?MultiTargetedPGDAttack)'\n               r'(?:(Top|Random)(\\d)*)?', name)\n  if m is not None:\n    # Request for a multi-targeted attack.\n    is_multitargeted = True\n    name = m.group(1)\n    is_random = (m.group(2) == 'Random')\n    max_specs = m.group(3)\n    max_specs = int(max_specs) if max_specs is not None else 0\n  else:\n    is_multitargeted = False\n\n  # Any of the readily available attack classes use the standard classification\n  # specification (one-vs-all) and are untargeted.\n  if hasattr(attacks, name):\n    attack_cls = getattr(attacks, name)\n    parameters = {}\n    if use_xent:\n      parameters['objective_fn'] = _maximize_cross_entropy\n    if use_l2:\n      parameters['project_perturbation'] = _get_projection(2)\n    if is_multitargeted:\n      parameters['max_specifications'] = max_specs\n      parameters['random_specifications'] = is_random\n    if parameters:\n      attack_cls = _change_parameters(attack_cls, **parameters)\n    attack_specification = specification.ClassificationSpecification(\n        label, num_classes)\n    return attack_cls, attack_specification, None\n\n  # Attacks can use an adaptive scheme.\n  if name.endswith('AdaptivePGDAttack'):\n    name = name[:-len('AdaptivePGDAttack')] + 'PGDAttack'\n    is_adaptive = True\n  else:\n    is_adaptive = False\n\n  # Attacks can be preceded by a number to indicate the number of target\n  # classes. For efficiency, this is only available for *MoreLikely attacks.\n  m = re.match(r'(\\d*)(.*MoreLikelyPGDAttack)', name)\n  if m is not None:\n    num_targets = int(m.group(1))\n    name = m.group(2)\n  else:\n    num_targets = 1\n\n  # All attacks that are not directly listed in the attacks library are\n  # targeted attacks that need to be manually constructed.\n  if name not in ('RandomMoreLikelyPGDAttack', 'RandomMostLikelyPGDAttack',\n                  'LeastLikelyMoreLikelyPGDAttack',\n                  'LeastLikelyMostLikelyPGDAttack',\n                  'ManualMoreLikelyPGDAttack', 'ManualMostLikelyPGDAttack'):\n    raise ValueError('Unknown attack \"{}\".'.format(name))\n\n  base_attack_cls = (attacks.AdaptiveUntargetedPGDAttack if is_adaptive else\n                     attacks.UntargetedPGDAttack)\n  if 'More' in name:\n    if use_xent:\n      raise ValueError('Using cross-entropy is not supported by '\n                       '\"*MoreLikelyPGDAttack\".')\n    attack_cls = base_attack_cls\n  else:\n    # We need to reverse the attack direction w.r.t. the specifications.\n    attack_cls = _change_parameters(\n        base_attack_cls,\n        objective_fn=(_minimize_cross_entropy if use_xent else\n                      _minimize_margin),\n        success_fn=_all_smaller)\n  if use_l2:\n    attack_cls = _change_parameters(\n        attack_cls, project_perturbation=_get_projection(2))\n\n  # Set attack specification and target class.\n  if name == 'RandomMoreLikelyPGDAttack':\n    # A random target class should become more likely than the true class.\n    attack_specification = specification.RandomClassificationSpecification(\n        label, num_classes, num_targets=num_targets, seed=random_seed)\n    target_class = (tf.squeeze(attack_specification.target_class, 1)\n                    if num_targets == 1 else None)\n\n  elif name == 'LeastLikelyMoreLikelyPGDAttack':\n    attack_specification = specification.LeastLikelyClassificationSpecification(\n        label, num_classes, logits, num_targets=num_targets)\n    target_class = (tf.squeeze(attack_specification.target_class, 1)\n                    if num_targets == 1 else None)\n\n  elif name == 'ManualMoreLikelyPGDAttack':\n    attack_specification = specification.TargetedClassificationSpecification(\n        label, num_classes, manual_target_class)\n    target_class = (tf.squeeze(attack_specification.target_class, 1)\n                    if num_targets == 1 else None)\n\n  elif name == 'RandomMostLikelyPGDAttack':\n    # This attack needs to make the random target the highest logits for\n    # it is be successful.\n    target_class = _get_random_class(label, num_classes, seed=random_seed)\n    attack_specification = specification.ClassificationSpecification(\n        target_class, num_classes)\n\n  elif name == 'LeastLikelyMostLikelyPGDAttack':\n    # This attack needs to make the least likely target the highest logits\n    # for it is be successful.\n    target_class = _get_least_likely_class(label, num_classes, logits)\n    attack_specification = specification.ClassificationSpecification(\n        target_class, num_classes)\n\n  else:\n    assert name == 'ManualMostLikelyPGDAttack'\n    target_class = manual_target_class\n    attack_specification = specification.ClassificationSpecification(\n        target_class, num_classes)\n\n  return attack_cls, attack_specification, target_class\n\n\ndef create_attack(attack_config, predictor, label, epsilon,\n                  input_bounds=(0., 1.), random_init=1., random_seed=None,\n                  predictor_kwargs=None, logits=None):\n  \"\"\"Creates an attack from a textual configuration.\n\n  Args:\n    attack_config: String with format \"[AttackClass]_[steps]x\n      [inner_restarts]x[outer_restarts]_[OptimizerClass]_[step_size]\". Inner\n      restarts involve tiling the input (they are more runtime efficient but\n      use more memory), while outer restarts use a tf.while_loop.\n    predictor: A VerifiableModelWrapper or StandardModelWrapper instance.\n    label: A Tensor of labels.\n    epsilon: Perturbation radius.\n    input_bounds: Tuple with minimum and maximum value allowed on inputs.\n    random_init: Probability of starting from random location rather than\n      nominal input image.\n    random_seed: Sets the random seed for \"Random*\" attacks.\n    predictor_kwargs: Dict of arguments passed to the predictor network.\n    logits: Logits corresponding to the nominal inputs. If None, it assumes that\n      predictor has a property named `logits`.\n\n  Returns:\n    An Attack instance.\n  \"\"\"\n  if attack_config:\n    name, steps_and_restarts, optimizer, step_size = re.split(\n        r'_\\s*(?![^()]*\\))', attack_config, maxsplit=3)\n    # Optimizers can specify contructor arguments using\n    # (arg1=value1;arg2=value2) syntax.\n    m = re.match(r'([^\\(]*)\\(([^\\)]*)\\)', optimizer)\n    if m is not None:\n      optimizer = m.group(1)\n      kwargs = 'dict(' + m.group(2).replace(';', ',') + ')'\n      kwargs = eval(kwargs)  # pylint: disable=eval-used\n    else:\n      kwargs = {}\n    optimizer = getattr(attacks, optimizer)\n    # Wrap optimizer if needed.\n    if kwargs:\n      optimizer = attacks.wrap_optimizer(optimizer, **kwargs)\n    num_steps, inner_restarts, outer_restarts = (\n        int(i) for i in steps_and_restarts.split('x', 3))\n    step_size = step_size.replace(':', ',')\n  else:\n    name = 'UntargetedPGDAttack'\n    num_steps = 200\n    inner_restarts = 1\n    outer_restarts = 1\n    optimizer = attacks.UnrolledAdam\n    step_size = .1\n\n  def attack_learning_rate_fn(t):\n    return parse_learning_rate(t, step_size)\n\n  if logits is None:\n    logits = predictor.logits\n  attack_cls, attack_specification, target_class = get_attack_builder(\n      logits, label, name=name, random_seed=random_seed)\n  attack_strategy = attack_cls(\n      predictor, attack_specification, epsilon, num_steps=num_steps,\n      num_restarts=inner_restarts, input_bounds=input_bounds,\n      optimizer_builder=optimizer, lr_fn=attack_learning_rate_fn,\n      random_init=random_init, predictor_kwargs=predictor_kwargs)\n  attack_strategy.target_class = target_class\n  if outer_restarts > 1:\n    attack_strategy = attacks.RestartedAttack(\n        attack_strategy, num_restarts=outer_restarts)\n  return attack_strategy\n\n\ndef parse_learning_rate(step, learning_rate):\n  \"\"\"Returns the learning rate as a tensor.\"\"\"\n  if isinstance(learning_rate, float):\n    return learning_rate\n  # Learning rate schedule of the form:\n  # initial_learning_rate[,learning@steps]*. E.g., \"1e-3\" or\n  # \"1e-3,1e-4@15000,1e-5@25000\". We use eval to allow learning specified as\n  # fractions (e.g., 2/255).\n  tokens = learning_rate.split(',')\n  first_lr = float(eval(tokens[0]))  # pylint: disable=eval-used\n  if len(tokens) == 1:\n    return tf.constant(first_lr, dtype=tf.float32)\n  # Parse steps.\n  init_values = [first_lr]\n  final_values = []\n  init_step = [0]\n  final_step = []\n  for t in tokens[1:]:\n    if '@' in t:\n      lr, boundary = t.split('@', 1)\n      is_linear = False\n    elif 'S' in t:  # Syntactic sugar to indicate a step.\n      lr, boundary = t.split('S', 1)\n      is_linear = False\n    elif 'L' in t:\n      lr, boundary = t.split('L', 1)\n      is_linear = True\n    else:\n      raise ValueError('Unknown specification.')\n    lr = float(eval(lr))  # pylint: disable=eval-used\n    init_values.append(lr)\n    if is_linear:\n      final_values.append(lr)\n    else:\n      final_values.append(init_values[-2])\n    boundary = int(boundary)\n    init_step.append(boundary)\n    final_step.append(boundary)\n  large_step = max(final_step) + 1\n  final_step.append(large_step)\n  final_values.append(lr)\n\n  # Find current index.\n  boundaries = list(final_step) + [large_step + 2]\n  boundaries = tf.convert_to_tensor(boundaries, dtype=tf.int64)\n  b = boundaries - tf.minimum(step + 1, large_step + 1)\n  large_step = tf.constant(\n      large_step, shape=boundaries.shape, dtype=step.dtype)\n  b = tf.where(b < 0, large_step, b)\n  idx = tf.minimum(tf.argmin(b), len(init_values) - 1)\n\n  init_step = tf.convert_to_tensor(init_step, dtype=tf.float32)\n  final_step = tf.convert_to_tensor(final_step, dtype=tf.float32)\n  init_values = tf.convert_to_tensor(init_values, dtype=tf.float32)\n  final_values = tf.convert_to_tensor(final_values, dtype=tf.float32)\n  x1 = tf.gather(init_step, idx)\n  x2 = tf.gather(final_step, idx)\n  y1 = tf.gather(init_values, idx)\n  y2 = tf.gather(final_values, idx)\n  return (tf.cast(step, tf.float32) - x1) / (x2 - x1) * (y2 - y1) + y1\n\n\ndef _change_parameters(attack_cls, **updated_kwargs):\n  def _build_new_attack(*args, **kwargs):\n    kwargs.update(updated_kwargs)\n    return attack_cls(*args, **kwargs)\n  return _build_new_attack\n\n\ndef _get_random_class(label, num_classes, seed=None):\n  batch_size = tf.shape(label)[0]\n  target_label = tf.random.uniform(\n      shape=(batch_size,), minval=1, maxval=num_classes, dtype=tf.int64,\n      seed=seed)\n  return tf.mod(tf.cast(label, tf.int64) + target_label, num_classes)\n\n\ndef _get_least_likely_class(label, num_classes, logits):\n  target_label = tf.argmin(logits, axis=1, output_type=tf.int64)\n  # In the off-chance that the least likely class is the true class, the target\n  # class is changed to the be the next index.\n  return tf.mod(target_label + tf.cast(\n      tf.equal(target_label, tf.cast(label, tf.int64)), tf.int64), num_classes)\n\n\ndef _maximize_cross_entropy(specification_bounds):\n  \"\"\"Used to maximize the cross entropy loss.\"\"\"\n  # Bounds has shape [num_restarts, batch_size, num_specs].\n  shape = tf.shape(specification_bounds)\n  added_shape = [shape[0], shape[1], 1]\n  v = tf.concat([\n      specification_bounds,\n      tf.zeros(added_shape, dtype=specification_bounds.dtype)], axis=2)\n  l = tf.concat([\n      tf.zeros_like(specification_bounds),\n      tf.ones(added_shape, dtype=specification_bounds.dtype)], axis=2)\n  # Minimize the cross-entropy loss w.r.t. target.\n  return tf.nn.softmax_cross_entropy_with_logits_v2(\n      labels=tf.stop_gradient(l), logits=v)\n\n\ndef _minimize_cross_entropy(specification_bounds):\n  return -_maximize_cross_entropy(specification_bounds)\n\n\ndef _maximize_margin(specification_bounds):\n  # Bounds has shape [num_restarts, batch_size, num_specs].\n  return tf.reduce_max(specification_bounds, axis=-1)\n\n\ndef _minimize_margin(specification_bounds):\n  return -_maximize_margin(specification_bounds)\n\n\ndef _all_smaller(specification_bounds):\n  specification_bounds = tf.reduce_max(specification_bounds, axis=-1)\n  return specification_bounds < 0\n\n\ndef _get_projection(p):\n  \"\"\"Returns a projection function.\"\"\"\n  if p == np.inf:\n    def _projection(perturbation, epsilon, input_image, image_bounds):\n      clipped_perturbation = tf.clip_by_value(perturbation, -epsilon, epsilon)\n      new_image = tf.clip_by_value(input_image + clipped_perturbation,\n                                   image_bounds[0], image_bounds[1])\n      return new_image - input_image\n    return _projection\n\n  elif p == 2:\n    def _projection(perturbation, epsilon, input_image, image_bounds):\n      axes = list(range(1, len(perturbation.get_shape())))\n      clipped_perturbation = tf.clip_by_norm(perturbation, epsilon, axes=axes)\n      new_image = tf.clip_by_value(input_image + clipped_perturbation,\n                                   image_bounds[0], image_bounds[1])\n      return new_image - input_image\n    return _projection\n\n  else:\n    raise ValueError('p must be np.inf or 2.')\n"
  },
  {
    "path": "interval_bound_propagation/src/verifiable_wrapper.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Wrapper around modules that provides additional facilities.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport abc\nimport types\n\nfrom absl import logging\nfrom interval_bound_propagation.src import layers\nimport six\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\n@six.add_metaclass(abc.ABCMeta)\nclass VerifiableWrapper(object):\n  \"\"\"Abstract wrapper class.\"\"\"\n\n  def __init__(self, module):\n    self._module = module\n    self._input_bounds = None\n    self._output_bounds = None\n\n  @property\n  def input_bounds(self):\n    assert self._input_bounds is not None\n    return self._input_bounds\n\n  @property\n  def output_bounds(self):\n    return self._output_bounds\n\n  @property\n  def module(self):\n    return self._module\n\n  def __str__(self):\n    if isinstance(self._module, tf.Tensor):\n      return str(self._module)\n    if isinstance(self._module, types.LambdaType):\n      return self._module.__name__\n    if isinstance(self._module, snt.AbstractModule):\n      return self._module.module_name\n    if hasattr(self._module, '__class__'):\n      return self._module.__class__.__name__\n    return str(self._module)\n\n  def propagate_bounds(self, *input_bounds):\n    \"\"\"Propagates bounds and saves input and output bounds.\"\"\"\n    output_bounds = self._propagate_through(self.module, *input_bounds)\n\n    if len(input_bounds) == 1:\n      self._input_bounds = input_bounds[0]\n    else:\n      self._input_bounds = tuple(input_bounds)\n    self._output_bounds = output_bounds\n\n    return output_bounds\n\n  @abc.abstractmethod\n  def _propagate_through(self, module, *input_bounds):\n    \"\"\"Propagates bounds through a verifiable wrapper.\n\n    Args:\n      module: This wrapped module, through which bounds are to be propagated.\n      *input_bounds: Bounds on the node's input(s).\n\n    Returns:\n      New bounds on the node's output.\n    \"\"\"\n\n\nclass ModelInputWrapper(object):\n  \"\"\"Virtual node representing the network's inputs.\"\"\"\n\n  def __init__(self, index):\n    super(ModelInputWrapper, self).__init__()\n    self._index = index\n    self._output_bounds = None\n\n  @property\n  def index(self):\n    return self._index\n\n  @property\n  def output_bounds(self):\n    return self._output_bounds\n\n  @output_bounds.setter\n  def output_bounds(self, bounds):\n    self._output_bounds = bounds\n\n  def __str__(self):\n    return 'Model input {}'.format(self.index)\n\n\nclass ConstWrapper(VerifiableWrapper):\n  \"\"\"Wraps a constant tensor.\"\"\"\n\n  def _propagate_through(self, module):\n    # Make sure that the constant value can be converted to a tensor.\n    return tf.convert_to_tensor(module)\n\n\nclass LinearFCWrapper(VerifiableWrapper):\n  \"\"\"Wraps fully-connected layers.\"\"\"\n\n  def __init__(self, module):\n    if not isinstance(module, snt.Linear):\n      raise ValueError('Cannot wrap {} with a LinearFCWrapper.'.format(module))\n    super(LinearFCWrapper, self).__init__(module)\n\n  def _propagate_through(self, module, input_bounds):\n    w = module.w\n    b = module.b if module.has_bias else None\n    return input_bounds.apply_linear(self, w, b)\n\n\nclass LinearConvWrapper(VerifiableWrapper):\n  \"\"\"Wraps convolutional layers.\"\"\"\n\n\nclass LinearConv1dWrapper(LinearConvWrapper):\n  \"\"\"Wraps 1-D convolutional layers.\"\"\"\n\n  def __init__(self, module):\n    if not isinstance(module, snt.Conv1D):\n      raise ValueError('Cannot wrap {} with a LinearConv1dWrapper.'.format(\n          module))\n    super(LinearConv1dWrapper, self).__init__(module)\n\n  def _propagate_through(self, module, input_bounds):\n    w = module.w\n    b = module.b if module.has_bias else None\n    padding = module.padding\n    stride = module.stride[1]\n    return input_bounds.apply_conv1d(self, w, b, padding, stride)\n\n\nclass LinearConv2dWrapper(LinearConvWrapper):\n  \"\"\"Wraps 2-D convolutional layers.\"\"\"\n\n  def __init__(self, module):\n    if not isinstance(module, snt.Conv2D):\n      raise ValueError('Cannot wrap {} with a LinearConv2dWrapper.'.format(\n          module))\n    super(LinearConv2dWrapper, self).__init__(module)\n\n  def _propagate_through(self, module, input_bounds):\n    w = module.w\n    b = module.b if module.has_bias else None\n    padding = module.padding\n    strides = module.stride[1:-1]\n    return input_bounds.apply_conv2d(self, w, b, padding, strides)\n\n\nclass IncreasingMonotonicWrapper(VerifiableWrapper):\n  \"\"\"Wraps monotonically increasing functions of the inputs.\"\"\"\n\n  def __init__(self, module, **parameters):\n    super(IncreasingMonotonicWrapper, self).__init__(module)\n    self._parameters = parameters\n\n  @property\n  def parameters(self):\n    return self._parameters\n\n  def _propagate_through(self, module, main_bounds, *other_input_bounds):\n    return main_bounds.apply_increasing_monotonic_fn(self, module,\n                                                     *other_input_bounds,\n                                                     **self.parameters)\n\n\nclass SoftmaxWrapper(VerifiableWrapper):\n  \"\"\"Wraps softmax layers.\"\"\"\n\n  def __init__(self):\n    super(SoftmaxWrapper, self).__init__(None)\n\n  def _propagate_through(self, module, input_bounds):\n    return input_bounds.apply_softmax(self)\n\n\nclass PiecewiseMonotonicWrapper(VerifiableWrapper):\n  \"\"\"Wraps a piecewise (not necessarily increasing) monotonic function.\"\"\"\n\n  def __init__(self, module, boundaries=()):\n    super(PiecewiseMonotonicWrapper, self).__init__(module)\n    self._boundaries = boundaries\n\n  @property\n  def boundaries(self):\n    return self._boundaries\n\n  def _propagate_through(self, module, main_bounds, *other_input_bounds):\n    return main_bounds.apply_piecewise_monotonic_fn(self, module,\n                                                    self.boundaries,\n                                                    *other_input_bounds)\n\n\nclass ImageNormWrapper(IncreasingMonotonicWrapper):\n  \"\"\"Convenience wrapper for getting track of the ImageNorm layer.\"\"\"\n\n  def __init__(self, module):\n    if not isinstance(module, layers.ImageNorm):\n      raise ValueError('Cannot wrap {} with a ImageNormWrapper.'.format(module))\n    super(ImageNormWrapper, self).__init__(module.apply)\n    self._inner_module = module\n\n  @property\n  def inner_module(self):\n    return self._inner_module\n\n\nclass BatchNormWrapper(VerifiableWrapper):\n  \"\"\"Wraps batch normalization.\"\"\"\n\n  def __init__(self, module):\n    if not isinstance(module, snt.BatchNorm):\n      raise ValueError('Cannot wrap {} with a BatchNormWrapper.'.format(\n          module))\n    super(BatchNormWrapper, self).__init__(module)\n\n  def _propagate_through(self, module, input_bounds):\n    if isinstance(module, layers.BatchNorm):\n      # This IBP-specific batch-norm implementation exposes stats recorded\n      # the most recent time the BatchNorm module was connected.\n      # These will be either the batch stats (e.g. if training) or the moving\n      # averages, depending on how the module was called.\n      mean = module.mean\n      variance = module.variance\n      epsilon = module.epsilon\n      scale = module.scale\n      bias = module.bias\n\n    else:\n      # This plain Sonnet batch-norm implementation only exposes the\n      # moving averages.\n      logging.warn('Sonnet BatchNorm module encountered: %s. '\n                   'IBP will always use its moving averages, not the local '\n                   'batch stats, even in training mode.', str(module))\n      mean = module.moving_mean\n      variance = module.moving_variance\n      epsilon = module._eps  # pylint: disable=protected-access\n      try:\n        bias = module.beta\n      except snt.Error:\n        bias = None\n      try:\n        scale = module.gamma\n      except snt.Error:\n        scale = None\n\n    return input_bounds.apply_batch_norm(self, mean, variance,\n                                         scale, bias, epsilon)\n\n\nclass BatchReshapeWrapper(VerifiableWrapper):\n  \"\"\"Wraps batch reshape.\"\"\"\n\n  def __init__(self, module, shape):\n    if not isinstance(module, snt.BatchReshape):\n      raise ValueError('Cannot wrap {} with a BatchReshapeWrapper.'.format(\n          module))\n    super(BatchReshapeWrapper, self).__init__(module)\n    self._shape = shape\n\n  @property\n  def shape(self):\n    return self._shape\n\n  def _propagate_through(self, module, input_bounds):\n    return input_bounds.apply_batch_reshape(self, self.shape)\n\n\nclass BatchFlattenWrapper(BatchReshapeWrapper):\n  \"\"\"Wraps batch flatten.\"\"\"\n\n  def __init__(self, module):\n    if not isinstance(module, snt.BatchFlatten):\n      raise ValueError('Cannot wrap {} with a BatchFlattenWrapper.'.format(\n          module))\n    super(BatchFlattenWrapper, self).__init__(module, [-1])\n"
  },
  {
    "path": "interval_bound_propagation/tests/attacks_test.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for attacks.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom absl.testing import parameterized\n\nimport interval_bound_propagation as ibp\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\nclass MockWithIsTraining(object):\n  \"\"\"Mock wrapper around the predictor network.\"\"\"\n\n  def __init__(self, module, test):\n    self._module = module\n    self._test = test\n\n  def __call__(self, z0, is_training=False):\n    # is_training should be False.\n    self._test.assertFalse(is_training)\n    return self._module(z0)\n\n\nclass MockWithoutIsTraining(object):\n  \"\"\"Mock wrapper around the predictor network.\"\"\"\n\n  def __init__(self, module, test):\n    self._module = module\n    self._test = test\n\n  def __call__(self, z0):\n    return self._module(z0)\n\n\nclass AttacksTest(parameterized.TestCase, tf.test.TestCase):\n\n  @parameterized.named_parameters(\n      ('UntargetedWithGradientDescent', MockWithIsTraining,\n       ibp.UntargetedPGDAttack, ibp.UnrolledGradientDescent, 1.),\n      ('UntargetedWithAdam', MockWithIsTraining,\n       ibp.UntargetedPGDAttack, ibp.UnrolledAdam, 1.),\n      ('MultiTargetedWithGradientDescent', MockWithIsTraining,\n       ibp.MultiTargetedPGDAttack, ibp.UnrolledGradientDescent, 1.),\n      ('MultiTargetedWithAdam', MockWithIsTraining,\n       ibp.MultiTargetedPGDAttack, ibp.UnrolledAdam, 1.),\n      ('DiverseEpsilon', MockWithIsTraining,\n       ibp.MultiTargetedPGDAttack, ibp.UnrolledAdam, [1., 1.]),\n      ('WithoutIsTraining', MockWithoutIsTraining,\n       ibp.UntargetedPGDAttack, ibp.UnrolledGradientDescent, 1.),\n      ('Restarted', MockWithIsTraining,\n       ibp.UntargetedPGDAttack, ibp.UnrolledGradientDescent, 1., True),\n      ('SPSA', MockWithIsTraining,\n       ibp.UntargetedPGDAttack, ibp.UnrolledSPSAAdam, 1.))\n  def testEndToEnd(self, predictor_cls, attack_cls, optimizer_cls, epsilon,\n                   restarted=False):\n    # l-\\infty norm of perturbation ball.\n    if isinstance(epsilon, list):\n      # We test the ability to have different epsilons across dimensions.\n      epsilon = tf.constant([epsilon], dtype=tf.float32)\n    bounds = (-.5, 2.5)\n    # Create a simple network.\n    m = snt.Linear(1, initializers={\n        'w': tf.constant_initializer(1.),\n        'b': tf.constant_initializer(1.),\n    })\n    z = tf.constant([[1, 2]], dtype=tf.float32)\n    predictor = predictor_cls(m, self)\n    # Not important for the test but needed.\n    labels = tf.constant([1], dtype=tf.int64)\n\n    # We create two attacks to maximize and then minimize the output.\n    max_spec = ibp.LinearSpecification(tf.constant([[[1.]]]))\n    max_attack = attack_cls(predictor, max_spec, epsilon, input_bounds=bounds,\n                            optimizer_builder=optimizer_cls)\n    if restarted:\n      max_attack = ibp.RestartedAttack(max_attack, num_restarts=10)\n    z_max = max_attack(z, labels)\n    min_spec = ibp.LinearSpecification(tf.constant([[[-1.]]]))\n    min_attack = attack_cls(predictor, min_spec, epsilon, input_bounds=bounds,\n                            optimizer_builder=optimizer_cls)\n    if restarted:\n      min_attack = ibp.RestartedAttack(min_attack, num_restarts=10)\n    z_min = min_attack(z, labels)\n\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      z_max_values, z_min_values = sess.run([z_max, z_min])\n      z_max_values = z_max_values[0]\n      z_min_values = z_min_values[0]\n      self.assertAlmostEqual(2., z_max_values[0])\n      self.assertAlmostEqual(2.5, z_max_values[1])\n      self.assertAlmostEqual(0., z_min_values[0])\n      self.assertAlmostEqual(1., z_min_values[1])\n\n\nif __name__ == '__main__':\n  tf.test.main()\n"
  },
  {
    "path": "interval_bound_propagation/tests/bounds_test.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for bounds.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom absl.testing import parameterized\nimport interval_bound_propagation as ibp\nimport numpy as np\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\nclass IntervalBoundsTest(parameterized.TestCase, tf.test.TestCase):\n\n  def testFCIntervalBounds(self):\n    m = snt.Linear(1, initializers={\n        'w': tf.constant_initializer(1.),\n        'b': tf.constant_initializer(2.),\n    })\n    z = tf.constant([[1, 2, 3]], dtype=tf.float32)\n    m(z)  # Connect to create weights.\n    m = ibp.LinearFCWrapper(m)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    output_bounds = m.propagate_bounds(input_bounds)\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      l, u = sess.run([output_bounds.lower, output_bounds.upper])\n      l = l.item()\n      u = u.item()\n      self.assertAlmostEqual(5., l)\n      self.assertAlmostEqual(11., u)\n\n  def testConv1dIntervalBounds(self):\n    m = snt.Conv1D(\n        output_channels=1,\n        kernel_shape=2,\n        padding='VALID',\n        stride=1,\n        use_bias=True,\n        initializers={\n            'w': tf.constant_initializer(1.),\n            'b': tf.constant_initializer(2.),\n        })\n    z = tf.constant([3, 4], dtype=tf.float32)\n    z = tf.reshape(z, [1, 2, 1])\n    m(z)  # Connect to create weights.\n    m = ibp.LinearConv1dWrapper(m)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    output_bounds = m.propagate_bounds(input_bounds)\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      l, u = sess.run([output_bounds.lower, output_bounds.upper])\n      l = l.item()\n      u = u.item()\n      self.assertAlmostEqual(7., l)\n      self.assertAlmostEqual(11., u)\n\n  def testConv2dIntervalBounds(self):\n    m = snt.Conv2D(\n        output_channels=1,\n        kernel_shape=(2, 2),\n        padding='VALID',\n        stride=1,\n        use_bias=True,\n        initializers={\n            'w': tf.constant_initializer(1.),\n            'b': tf.constant_initializer(2.),\n        })\n    z = tf.constant([1, 2, 3, 4], dtype=tf.float32)\n    z = tf.reshape(z, [1, 2, 2, 1])\n    m(z)  # Connect to create weights.\n    m = ibp.LinearConv2dWrapper(m)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    output_bounds = m.propagate_bounds(input_bounds)\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      l, u = sess.run([output_bounds.lower, output_bounds.upper])\n      l = l.item()\n      u = u.item()\n      self.assertAlmostEqual(8., l)\n      self.assertAlmostEqual(16., u)\n\n  def testReluIntervalBounds(self):\n    m = tf.nn.relu\n    z = tf.constant([[-2, 3]], dtype=tf.float32)\n    m = ibp.IncreasingMonotonicWrapper(m)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    output_bounds = m.propagate_bounds(input_bounds)\n    with self.test_session() as sess:\n      l, u = sess.run([output_bounds.lower, output_bounds.upper])\n      self.assertAlmostEqual([[0., 2.]], l.tolist())\n      self.assertAlmostEqual([[0., 4.]], u.tolist())\n\n  def testMulIntervalBounds(self):\n    m = tf.multiply\n    z = tf.constant([[-2, 3, 0]], dtype=tf.float32)\n    m = ibp.PiecewiseMonotonicWrapper(m, (0,))\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    output_bounds = m.propagate_bounds(input_bounds, input_bounds)\n    with self.test_session() as sess:\n      l, u = sess.run([output_bounds.lower, output_bounds.upper])\n      self.assertAlmostEqual([[1., 4., -1.]], l.tolist())\n      self.assertAlmostEqual([[9., 16., 1.]], u.tolist())\n\n  def testSubIntervalBounds(self):\n    m = tf.subtract\n    z = tf.constant([[-2, 3, 0]], dtype=tf.float32)\n    m = ibp.PiecewiseMonotonicWrapper(m)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    output_bounds = m.propagate_bounds(input_bounds, input_bounds)\n    with self.test_session() as sess:\n      l, u = sess.run([output_bounds.lower, output_bounds.upper])\n      self.assertAlmostEqual([[-2., -2., -2.]], l.tolist())\n      self.assertAlmostEqual([[2., 2., 2.]], u.tolist())\n\n  @parameterized.named_parameters(\n      ('DefaultAxis', -1, [[[1., 0.5, 0.5], [1., 0.5, 0.5]],\n                           [[1. / 3, 0., 0.], [1. / 3, 0., 0.]]]),\n      ('NonDefaultAxis', 0, [[[1., 1., 1.], [1., 1., 1.]],\n                             [[0., 0., 0.], [0., 0., 0.]]]))\n  def testSoftmaxIntervalBounds(self, axis, expected_outputs):\n    z = tf.constant([[1., -10., -10.], [1., -10., -10.]])\n    input_bounds = ibp.IntervalBounds(z - 1.0, z + 10.0)\n\n    softmax_fn = lambda x: tf.nn.softmax(x, axis=axis)\n    softmax_fn = ibp.VerifiableModelWrapper(softmax_fn)\n    softmax_fn(z)\n    output_bounds = softmax_fn.propagate_bounds(input_bounds)\n\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      l, u = sess.run([output_bounds.lower, output_bounds.upper])\n    self.assertTrue(np.all(np.abs(expected_outputs[0] - u) < 1e-3))\n    self.assertTrue(np.all(np.abs(expected_outputs[1] - l) < 1e-3))\n\n  def testBatchNormIntervalBounds(self):\n    z = tf.constant([[1, 2, 3]], dtype=tf.float32)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    g = tf.reshape(tf.range(-1, 2, dtype=tf.float32), [1, 3])\n    b = tf.reshape(tf.range(3, dtype=tf.float32), [1, 3])\n    batch_norm = ibp.BatchNorm(scale=True, offset=True, eps=0., initializers={\n        'gamma': lambda *args, **kwargs: g,\n        'beta': lambda *args, **kwargs: b,\n        'moving_mean': tf.constant_initializer(1.),\n        'moving_variance': tf.constant_initializer(4.),\n    })\n    batch_norm(z, is_training=False)\n    batch_norm = ibp.BatchNormWrapper(batch_norm)\n    # Test propagation.\n    output_bounds = batch_norm.propagate_bounds(input_bounds)\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      l, u = sess.run([output_bounds.lower, output_bounds.upper])\n      self.assertAlmostEqual([[-.5, 1., 2.5]], l.tolist())\n      self.assertAlmostEqual([[.5, 1., 3.5]], u.tolist())\n\n  def testCaching(self):\n    m = snt.Linear(1, initializers={\n        'w': tf.constant_initializer(1.),\n        'b': tf.constant_initializer(2.),\n    })\n    z = tf.placeholder(shape=(1, 3), dtype=tf.float32)\n    m(z)  # Connect to create weights.\n    m = ibp.LinearFCWrapper(m)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    output_bounds = m.propagate_bounds(input_bounds)\n\n    input_bounds.enable_caching()\n    output_bounds.enable_caching()\n    update_all_caches_op = tf.group([input_bounds.update_cache_op,\n                                     output_bounds.update_cache_op])\n\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n\n      # Initialise the caches based on the model inputs.\n      sess.run(update_all_caches_op, feed_dict={z: [[1., 2., 3.]]})\n\n      l, u = sess.run([output_bounds.lower, output_bounds.upper])\n      l = l.item()\n      u = u.item()\n      self.assertAlmostEqual(5., l)\n      self.assertAlmostEqual(11., u)\n\n      # Update the cache based on a different set of inputs.\n      sess.run([output_bounds.update_cache_op], feed_dict={z: [[2., 3., 7.]]})\n      # We only updated the output bounds' cache.\n      # This asserts that the computation depends on the underlying\n      # input bounds tensor, not on cached version of it.\n      # (Thus it doesn't matter what order the caches are updated.)\n\n      l, u = sess.run([output_bounds.lower, output_bounds.upper])\n      l = l.item()\n      u = u.item()\n      self.assertAlmostEqual(11., l)\n      self.assertAlmostEqual(17., u)\n\n\nif __name__ == '__main__':\n  tf.test.main()\n"
  },
  {
    "path": "interval_bound_propagation/tests/crown_test.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for CROWN bounds.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport interval_bound_propagation as ibp\nimport numpy as np\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\ndef _generate_identity_spec(modules, shape, dimension=1):\n  spec = ibp.LinearSpecification(tf.reshape(tf.eye(dimension), shape),\n                                 prune_irrelevant=False)\n  initial_bound = ibp.crown.create_initial_backward_bounds(spec, modules)\n  return initial_bound\n\n\nclass CROWNBoundsTest(tf.test.TestCase):\n\n  def testFCBackwardBounds(self):\n    m = snt.Linear(1, initializers={\n        'w': tf.constant_initializer(1.),\n        'b': tf.constant_initializer(2.),\n    })\n    z = tf.constant([[1, 2, 3]], dtype=tf.float32)\n    m(z)  # Connect to create weights.\n    m = ibp.LinearFCWrapper(m)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    m.propagate_bounds(input_bounds)  # Create IBP bounds.\n    crown_init_bounds = _generate_identity_spec([m], shape=(1, 1, 1))\n    output_bounds = m.propagate_bounds(crown_init_bounds)\n    concrete_bounds = output_bounds.concretize()\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      lw, uw, lb, ub, cl, cu = sess.run([output_bounds.lower.w,\n                                         output_bounds.upper.w,\n                                         output_bounds.lower.b,\n                                         output_bounds.upper.b,\n                                         concrete_bounds.lower,\n                                         concrete_bounds.upper])\n      self.assertTrue(np.all(lw == 1.))\n      self.assertTrue(np.all(lb == 2.))\n      self.assertTrue(np.all(uw == 1.))\n      self.assertTrue(np.all(ub == 2.))\n      cl = cl.item()\n      cu = cu.item()\n      self.assertAlmostEqual(5., cl)\n      self.assertAlmostEqual(11., cu)\n\n  def testConv2dBackwardBounds(self):\n    m = snt.Conv2D(\n        output_channels=1,\n        kernel_shape=(2, 2),\n        padding='VALID',\n        stride=1,\n        use_bias=True,\n        initializers={\n            'w': tf.constant_initializer(1.),\n            'b': tf.constant_initializer(2.),\n        })\n    z = tf.constant([1, 2, 3, 4], dtype=tf.float32)\n    z = tf.reshape(z, [1, 2, 2, 1])\n    m(z)  # Connect to create weights.\n    m = ibp.LinearConv2dWrapper(m)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    m.propagate_bounds(input_bounds)   # Create IBP bounds.\n    crown_init_bounds = _generate_identity_spec([m], shape=(1, 1, 1, 1, 1))\n    output_bounds = m.propagate_bounds(crown_init_bounds)\n    concrete_bounds = output_bounds.concretize()\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      l, u = sess.run([concrete_bounds.lower, concrete_bounds.upper])\n      l = l.item()\n      u = u.item()\n      self.assertAlmostEqual(8., l)\n      self.assertAlmostEqual(16., u)\n\n  def testReluBackwardBounds(self):\n    m = tf.nn.relu\n    z = tf.constant([[-2, 3]], dtype=tf.float32)\n    m = ibp.IncreasingMonotonicWrapper(m)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    m.propagate_bounds(input_bounds)  # Create IBP bounds.\n    crown_init_bounds = _generate_identity_spec([m], shape=(1, 2, 2),\n                                                dimension=2)\n    output_bounds = m.propagate_bounds(crown_init_bounds)\n    concrete_bounds = output_bounds.concretize()\n    with self.test_session() as sess:\n      l, u = sess.run([concrete_bounds.lower, concrete_bounds.upper])\n      self.assertAlmostEqual([[0., 2.]], l.tolist())\n      self.assertAlmostEqual([[0., 4.]], u.tolist())\n\nif __name__ == '__main__':\n  tf.test.main()\n"
  },
  {
    "path": "interval_bound_propagation/tests/fastlin_test.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for symbolic bounds.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom absl.testing import parameterized\nimport interval_bound_propagation as ibp\nimport numpy as np\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\nclass SymbolicBoundsTest(parameterized.TestCase, tf.test.TestCase):\n\n  def testConvertSymbolicBounds(self):\n    z = tf.constant([[1, 2, 3, 4]], dtype=tf.float32)\n    z = tf.reshape(z, [1, 2, 2])\n    b = ibp.SymbolicBounds.convert(z)\n    for l in (b.lower, b.upper):\n      self.assertEqual([1, 4, 2, 2], l.w.shape.as_list())\n      self.assertEqual([1, 2, 2], l.b.shape.as_list())\n      self.assertEqual([1, 4], l.lower.shape.as_list())\n      self.assertEqual([1, 4], l.upper.shape.as_list())\n\n  def testFCSymbolicBounds(self):\n    m = snt.Linear(1, initializers={\n        'w': tf.constant_initializer(1.),\n        'b': tf.constant_initializer(2.),\n    })\n    z = tf.constant([[1, 2, 3]], dtype=tf.float32)\n    m(z)  # Connect to create weights.\n    m = ibp.LinearFCWrapper(m)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    input_bounds = ibp.SymbolicBounds.convert(input_bounds)\n    output_bounds = m.propagate_bounds(input_bounds)\n    concrete_bounds = ibp.IntervalBounds.convert(output_bounds)\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      l, u, cl, cu = sess.run([output_bounds.lower, output_bounds.upper,\n                               concrete_bounds.lower, concrete_bounds.upper])\n      self.assertTrue(np.all(l.w == 1.))\n      self.assertTrue(np.all(l.b == 2.))\n      self.assertAlmostEqual([[0, 1, 2]], l.lower.tolist())\n      self.assertAlmostEqual([[2, 3, 4]], l.upper.tolist())\n      self.assertTrue(np.all(u.w == 1.))\n      self.assertTrue(np.all(u.b == 2.))\n      self.assertAlmostEqual([[0, 1, 2]], u.lower.tolist())\n      self.assertAlmostEqual([[2, 3, 4]], u.upper.tolist())\n      cl = cl.item()\n      cu = cu.item()\n      self.assertAlmostEqual(5., cl)\n      self.assertAlmostEqual(11., cu)\n\n  def testConv2dSymbolicBounds(self):\n    m = snt.Conv2D(\n        output_channels=1,\n        kernel_shape=(2, 2),\n        padding='VALID',\n        stride=1,\n        use_bias=True,\n        initializers={\n            'w': tf.constant_initializer(1.),\n            'b': tf.constant_initializer(2.),\n        })\n    z = tf.constant([1, 2, 3, 4], dtype=tf.float32)\n    z = tf.reshape(z, [1, 2, 2, 1])\n    m(z)  # Connect to create weights.\n    m = ibp.LinearConv2dWrapper(m)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    input_bounds = ibp.SymbolicBounds.convert(input_bounds)\n    output_bounds = m.propagate_bounds(input_bounds)\n    output_bounds = ibp.IntervalBounds.convert(output_bounds)\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      l, u = sess.run([output_bounds.lower, output_bounds.upper])\n      l = l.item()\n      u = u.item()\n      self.assertAlmostEqual(8., l)\n      self.assertAlmostEqual(16., u)\n\n  def testConv1dSymbolicBounds(self):\n    m = snt.Conv1D(\n        output_channels=1,\n        kernel_shape=(2),\n        padding='VALID',\n        stride=1,\n        use_bias=True,\n        initializers={\n            'w': tf.constant_initializer(1.),\n            'b': tf.constant_initializer(3.),\n        })\n    z = tf.constant([3, 4], dtype=tf.float32)\n    z = tf.reshape(z, [1, 2, 1])\n    m(z)  # Connect to create weights.\n    m = ibp.LinearConv1dWrapper(m)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    input_bounds = ibp.SymbolicBounds.convert(input_bounds)\n    output_bounds = m.propagate_bounds(input_bounds)\n    output_bounds = ibp.IntervalBounds.convert(output_bounds)\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      l, u = sess.run([output_bounds.lower, output_bounds.upper])\n      l = l.item()\n      u = u.item()\n      self.assertAlmostEqual(8., l)\n      self.assertAlmostEqual(12., u)\n\n  def testReluSymbolicBounds(self):\n    m = tf.nn.relu\n    z = tf.constant([[-2, 3]], dtype=tf.float32)\n    m = ibp.IncreasingMonotonicWrapper(m)\n    input_bounds = ibp.IntervalBounds(z - 1., z + 1.)\n    input_bounds = ibp.SymbolicBounds.convert(input_bounds)\n    output_bounds = m.propagate_bounds(input_bounds)\n    output_bounds = ibp.IntervalBounds.convert(output_bounds)\n    with self.test_session() as sess:\n      l, u = sess.run([output_bounds.lower, output_bounds.upper])\n      self.assertAlmostEqual([[0., 2.]], l.tolist())\n      self.assertAlmostEqual([[0., 4.]], u.tolist())\n\n\nif __name__ == '__main__':\n  tf.test.main()\n"
  },
  {
    "path": "interval_bound_propagation/tests/layers_test.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for layers.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport interval_bound_propagation as ibp\nimport numpy as np\nimport tensorflow.compat.v1 as tf\n\n\ndef _get_inputs(dtype=tf.float32):\n  v = np.array(range(6), dtype=dtype.as_numpy_dtype)\n  input_v = np.array([v] * 7)\n  inputs = tf.constant(input_v)\n  return v, input_v, inputs\n\n\nclass LayersTest(tf.test.TestCase):\n\n  def assertBetween(self, value, minv, maxv):\n    \"\"\"Asserts that value is between minv and maxv (inclusive).\"\"\"\n    self.assertLessEqual(minv, value)\n    self.assertGreaterEqual(maxv, value)\n\n  # Subset of the tests in sonnet/python/modules/batch_norm_test.py.\n  def testBatchNormUpdateImproveStatistics(self):\n    \"\"\"Test that updating the moving_mean improves statistics.\"\"\"\n    _, _, inputs = _get_inputs()\n    # Use small decay_rate to update faster.\n    bn = ibp.BatchNorm(offset=False, scale=False, decay_rate=0.1,\n                       update_ops_collection=tf.GraphKeys.UPDATE_OPS)\n    out1 = bn(inputs, is_training=False)\n    # Build the update ops.\n    bn(inputs, is_training=True)\n\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      out_v = sess.run(out1)\n      # Before updating the moving_mean the results are off.\n      self.assertBetween(np.max(np.abs(np.zeros([7, 6]) - out_v)), 2, 5)\n      sess.run(tuple(tf.get_collection(tf.GraphKeys.UPDATE_OPS)))\n      # After updating the moving_mean the results are better.\n      out_v = sess.run(out1)\n      self.assertBetween(np.max(np.abs(np.zeros([7, 6]) - out_v)), 1, 2)\n\n  def testImageNorm(self):\n    mean = [4, 0, -4]\n    std = [1., 2., 4.]\n    image = tf.constant(4., shape=[10, 2, 2, 3])\n    normalized_image = ibp.ImageNorm(mean, std)(image)\n\n    with self.test_session() as sess:\n      out_image = sess.run(normalized_image)\n      self.assertTrue(np.all(np.isclose(out_image[:, :, :, 0], 0.)))\n      self.assertTrue(np.all(np.isclose(out_image[:, :, :, 1], 2.)))\n      self.assertTrue(np.all(np.isclose(out_image[:, :, :, 2], 2.)))\n\n\nif __name__ == '__main__':\n  tf.test.main()\n"
  },
  {
    "path": "interval_bound_propagation/tests/loss_test.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for loss.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport interval_bound_propagation as ibp\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\nclass FixedNN(snt.AbstractModule):\n\n  def _build(self, z0, is_training=False):\n    self._m = snt.Linear(2, initializers={\n        'w': tf.constant_initializer(1.),\n        'b': lambda *unsed_args, **unused_kwargs: tf.constant([0., 1.]),\n    })\n    return self._m(z0)\n\n\nclass LossTest(tf.test.TestCase):\n\n  def testEndToEnd(self):\n    predictor = FixedNN()\n    predictor = ibp.VerifiableModelWrapper(predictor)\n    # Labels.\n    labels = tf.constant([1], dtype=tf.int64)\n    # Connect to input.\n    z = tf.constant([[1, 2, 3]], dtype=tf.float32)\n    predictor(z, is_training=True)\n    # Input bounds.\n    eps = 1.\n    input_bounds = ibp.IntervalBounds(z - eps, z + eps)\n    predictor.propagate_bounds(input_bounds)\n    # Create output specification (that forces the first logits to be greater).\n    c = tf.constant([[[1, -1]]], dtype=tf.float32)\n    d = tf.constant([[0]], dtype=tf.float32)\n    # Turn elision off for more interesting results.\n    spec = ibp.LinearSpecification(c, d, collapse=False)\n    # Create an attack.\n    attack = ibp.UntargetedPGDAttack(\n        predictor, spec, eps, num_steps=1, input_bounds=(-100., 100))\n    # Build loss.\n    losses = ibp.Losses(predictor, spec, attack,\n                        interval_bounds_loss_type='hinge',\n                        interval_bounds_hinge_margin=0.)\n    losses(labels)\n\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      # We expect the worst-case logits from IBP to be [9, 4].\n      # The adversarial attack should fail since logits are always [l, l + 1].\n      # Similarly, the nominal predictions are correct.\n      accuracy_values, loss_values = sess.run(\n          [losses.scalar_metrics, losses.scalar_losses])\n      self.assertAlmostEqual(1., accuracy_values.nominal_accuracy)\n      self.assertAlmostEqual(0., accuracy_values.verified_accuracy)\n      self.assertAlmostEqual(1., accuracy_values.attack_accuracy)\n      expected_xent = 0.31326168751822947\n      self.assertAlmostEqual(expected_xent, loss_values.nominal_cross_entropy,\n                             places=5)\n      self.assertAlmostEqual(expected_xent, loss_values.attack_cross_entropy,\n                             places=5)\n      expected_hinge = 5.\n      self.assertAlmostEqual(expected_hinge, loss_values.verified_loss)\n\n\nif __name__ == '__main__':\n  tf.test.main()\n"
  },
  {
    "path": "interval_bound_propagation/tests/model_test.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for model.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom absl.testing import parameterized\n\nimport interval_bound_propagation as ibp\nimport numpy as np\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\ndef _build_model():\n  num_classes = 3\n  layer_types = (\n      ('conv2d', (2, 2), 4, 'VALID', 1),\n      ('activation', 'relu'),\n      ('linear', 10),\n      ('activation', 'relu'))\n  return ibp.DNN(num_classes, layer_types)\n\n\nclass ModelTest(parameterized.TestCase, tf.test.TestCase):\n\n  def testDNN(self):\n    predictor = _build_model()\n    # Input.\n    z = tf.constant([1, 2, 3, 4], dtype=tf.float32)\n    z = tf.reshape(z, [1, 2, 2, 1])\n    predictor(z)\n    # Verify the variables that are created.\n    expected_shapes = {\n        'predictor/conv2d_0/w:0': (2, 2, 1, 4),\n        'predictor/conv2d_0/b:0': (4,),\n        'predictor/linear_0/w:0': (4, 10),\n        'predictor/linear_0/b:0': (10,),\n        'predictor/linear_1/w:0': (10, 3),\n        'predictor/linear_1/b:0': (3,),\n    }\n    for v in predictor.get_variables():\n      self.assertEqual(expected_shapes[v.name], v.shape)\n\n  def _propagation_test(self, wrapper, inputs, outputs):\n    input_bounds = ibp.IntervalBounds(inputs, inputs)\n    output_bounds = wrapper.propagate_bounds(input_bounds)\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      o, l, u = sess.run([outputs, output_bounds.lower, output_bounds.upper])\n      self.assertAlmostEqual(o.tolist(), l.tolist())\n      self.assertAlmostEqual(o.tolist(), u.tolist())\n\n  def testVerifiableModelWrapperDNN(self):\n    predictor = _build_model()\n    # Input.\n    z = tf.constant([1, 2, 3, 4], dtype=tf.float32)\n    z = tf.reshape(z, [1, 2, 2, 1])\n    wrapper = ibp.VerifiableModelWrapper(predictor)\n    wrapper(z)\n    # Verify basic wrapping.\n    self.assertEqual(predictor, wrapper.wrapped_network)\n    self.assertEqual(3, wrapper.output_size)\n    self.assertEqual((1, 3), tuple(wrapper.logits.shape.as_list()))\n    self.assertEqual(z, wrapper.inputs)\n    # Build another input and test reuse.\n    z2 = tf.constant([1, 2, 3, 4], dtype=tf.float32)\n    z2 = tf.reshape(z, [1, 2, 2, 1])\n    logits = wrapper(z2, reuse=True)\n    self.assertEqual(z, wrapper.inputs)\n    self.assertNotEqual(z2, wrapper.inputs)\n    # Check that the verifiable modules are constructed.\n    self.assertLen(wrapper.input_wrappers, 1)\n    self.assertLen(wrapper.modules, 6)\n    self.assertIsInstance(wrapper.modules[0].module, snt.Conv2D)\n    self.assertEqual(wrapper.modules[1].module, tf.nn.relu)\n    self.assertIsInstance(wrapper.modules[2].module, snt.BatchFlatten)\n    self.assertIsInstance(wrapper.modules[3].module, snt.Linear)\n    self.assertEqual(wrapper.modules[4].module, tf.nn.relu)\n    self.assertIsInstance(wrapper.modules[5].module, snt.Linear)\n    # It's a sequential network, so all nodes (including input) have fanout 1.\n    self.assertEqual(wrapper.fanout_of(wrapper.input_wrappers[0]), 1)\n    for module in wrapper.modules:\n      self.assertEqual(wrapper.fanout_of(module), 1)\n    # Check propagation.\n    self._propagation_test(wrapper, z2, logits)\n\n  def testVerifiableModelWrapperResnet(self):\n    def _build(z0, is_training=False):  # pylint: disable=unused-argument\n      input_size = np.prod(z0.shape[1:])\n      # We make a resnet-like structure.\n      z = snt.Linear(input_size)(z0)\n      z_left = tf.nn.relu(z)\n      z_left = snt.Linear(input_size)(z_left)\n      z = z_left + z0\n      return snt.Linear(2)(z)\n\n    z = tf.constant([[1, 2, 3, 4]], dtype=tf.float32)\n    wrapper = ibp.VerifiableModelWrapper(_build)\n    logits = wrapper(z)\n    self.assertLen(wrapper.input_wrappers, 1)\n    self.assertLen(wrapper.modules, 5)\n    # Check input has fanout 2, as it is the start of the resnet block.\n    self.assertEqual(wrapper.fanout_of(wrapper.input_wrappers[0]), 2)\n    for module in wrapper.modules:\n      self.assertEqual(wrapper.fanout_of(module), 1)\n    # Check propagation.\n    self._propagation_test(wrapper, z, logits)\n\n  def testVerifiableModelWrapperPool(self):\n    def _build(z0):\n      z = tf.reduce_mean(z0, axis=1, keep_dims=True)\n      z = tf.reduce_max(z, axis=2, keep_dims=False)\n      return snt.Linear(2)(z)\n\n    z = tf.constant([[1, 2, 3, 4]], dtype=tf.float32)\n    z = tf.reshape(z, [1, 2, 2])\n    wrapper = ibp.VerifiableModelWrapper(_build)\n    logits = wrapper(z)\n    self.assertLen(wrapper.modules, 3)\n    # Check propagation.\n    self._propagation_test(wrapper, z, logits)\n\n  def testVerifiableModelWrapperConcat(self):\n    def _build(z0):\n      z = snt.Linear(10)(z0)\n      z = tf.concat([z, z0], axis=1)\n      return snt.Linear(2)(z)\n\n    z = tf.constant([[1, 2, 3, 4]], dtype=tf.float32)\n    wrapper = ibp.VerifiableModelWrapper(_build)\n    logits = wrapper(z)\n    self.assertLen(wrapper.modules, 3)\n    # Check propagation.\n    self._propagation_test(wrapper, z, logits)\n\n  def testVerifiableModelWrapperExpandAndSqueeze(self):\n    def _build(z0):\n      z = snt.Linear(10)(z0)\n      z = tf.expand_dims(z, axis=-1)\n      z = tf.squeeze(z, axis=-1)\n      return snt.Linear(2)(z)\n\n    z = tf.constant([[1, 2, 3, 4]], dtype=tf.float32)\n    wrapper = ibp.VerifiableModelWrapper(_build)\n    logits = wrapper(z)\n    self.assertLen(wrapper.modules, 4)\n    # Check propagation.\n    self._propagation_test(wrapper, z, logits)\n\n  @parameterized.named_parameters(\n      ('Add', lambda z: z + z, 3),\n      ('Sub', lambda z: z - z, 3),\n      ('Identity', tf.identity, 3),\n      ('Mul', lambda z: z * z, 3),\n      ('Slice', lambda z: tf.slice(z, [0, 0], [-1, 5]), 3),\n      ('StridedSlice', lambda z: z[:, :5], 3),\n      ('Reshape', lambda z: tf.reshape(z, [2, 5]), 3),\n      ('Const', lambda z: z + tf.ones_like(z), 5))\n  def testVerifiableModelWrapperSimple(self, fn, expected_modules):\n    def _build(z0):\n      z = snt.Linear(10)(z0)\n      z = fn(z)\n      return snt.Linear(2)(z)\n\n    z = tf.constant([[1, 2, 3, 4]], dtype=tf.float32)\n    wrapper = ibp.VerifiableModelWrapper(_build)\n    logits = wrapper(z)\n    self.assertLen(wrapper.modules, expected_modules)\n    # Check propagation.\n    self._propagation_test(wrapper, z, logits)\n\n  def testPointlessReshape(self):\n    def _build(z0):\n      z = snt.Linear(10)(z0)\n      z = snt.BatchFlatten()(z)  # This is a no-op; no graph nodes created.\n      return snt.Linear(2)(z)\n\n    z = tf.constant([[1, 2, 3, 4]], dtype=tf.float32)\n    wrapper = ibp.VerifiableModelWrapper(_build)\n    logits = wrapper(z)\n    # Expect the batch flatten to have been skipped.\n    self.assertLen(wrapper.modules, 2)\n    self.assertIsInstance(wrapper.modules[0], ibp.LinearFCWrapper)\n    self.assertIsInstance(wrapper.modules[1], ibp.LinearFCWrapper)\n    # Check propagation.\n    self._propagation_test(wrapper, z, logits)\n\n  def testLeakyRelu(self):\n    def _build(z0):\n      z = snt.Linear(10)(z0)\n      z = tf.nn.leaky_relu(z0, alpha=0.375)\n      return snt.Linear(2)(z)\n\n    z = tf.constant([[1, 2, 3, 4]], dtype=tf.float32)\n    wrapper = ibp.VerifiableModelWrapper(_build)\n    logits = wrapper(z)\n    self.assertLen(wrapper.modules, 3)\n    self.assertEqual(wrapper.modules[1].module.__name__, 'leaky_relu')\n    self.assertEqual(wrapper.modules[1].parameters['alpha'], 0.375)\n    # Check propagation.\n    self._propagation_test(wrapper, z, logits)\n\n  def testMultipleInputs(self):\n    # Tensor to overwrite.\n    def _build(z0, z1):\n      return z0 + z1\n\n    z0 = tf.constant([[1, 2, 3, 4]], dtype=tf.float32)\n    z1 = tf.constant([[2, 2, 4, 4]], dtype=tf.float32)\n    wrapper = ibp.VerifiableModelWrapper(_build)\n    logits = wrapper(z0, z1)\n    input_bounds0 = ibp.IntervalBounds(z0 - 2, z0 + 1)\n    input_bounds1 = ibp.IntervalBounds(z1, z1 + 10)\n    output_bounds = wrapper.propagate_bounds(input_bounds0, input_bounds1)\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      o, l, u = sess.run([logits, output_bounds.lower, output_bounds.upper])\n      print(o, l, u)\n      self.assertAlmostEqual([[3., 4., 7., 8.]], o.tolist())\n      self.assertAlmostEqual([[1., 2., 5., 6.]], l.tolist())\n      self.assertAlmostEqual([[14., 15., 18., 19.]], u.tolist())\n\n\nif __name__ == '__main__':\n  tf.test.main()\n"
  },
  {
    "path": "interval_bound_propagation/tests/relative_bounds_test.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for relative_bounds.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom absl.testing import parameterized\nimport interval_bound_propagation as ibp\nfrom interval_bound_propagation import layer_utils\nimport numpy as np\nimport sonnet as snt\nimport tensorflow.compat.v1 as tf\n\n\nclass RelativeIntervalBoundsTest(tf.test.TestCase, parameterized.TestCase):\n\n  @parameterized.named_parameters(('float32', tf.float32),\n                                  ('float64', tf.float64))\n  def test_linear_bounds_shape(self, dtype):\n    batch_size = 11\n    input_size = 7\n    output_size = 5\n\n    w = tf.placeholder(dtype=dtype, shape=(input_size, output_size))\n    b = tf.placeholder(dtype=dtype, shape=(output_size,))\n    lb_rel_in = tf.placeholder(dtype=dtype, shape=(batch_size, input_size))\n    ub_rel_in = tf.placeholder(dtype=dtype, shape=(batch_size, input_size))\n    nominal = tf.placeholder(dtype=dtype, shape=(batch_size, input_size))\n\n    bounds_in = ibp.RelativeIntervalBounds(lb_rel_in, ub_rel_in, nominal)\n    bounds_out = bounds_in.apply_linear(None, w, b)\n    lb_out, ub_out = bounds_out.lower, bounds_out.upper\n\n    self.assertEqual(dtype, lb_out.dtype)\n    self.assertEqual(dtype, ub_out.dtype)\n    self.assertEqual((batch_size, output_size), lb_out.shape)\n    self.assertEqual((batch_size, output_size), ub_out.shape)\n\n  @parameterized.named_parameters(('float32', tf.float32, 1.e-6),\n                                  ('float64', tf.float64, 1.e-8))\n  def test_linear_bounds(self, dtype, tol):\n    w = tf.constant([[1.0, 2.0, 3.0], [4.0, -5.0, 6.0]], dtype=dtype)\n    b = tf.constant([0.1, 0.2, 0.3], dtype=dtype)\n    lb_in = tf.constant([[-1.0, -1.0]], dtype=dtype)\n    ub_in = tf.constant([[2.0, 2.0]], dtype=dtype)\n    nominal = tf.constant([[3.1, 4.2]], dtype=dtype)\n\n    bounds_in = ibp.RelativeIntervalBounds(lb_in - nominal,\n                                           ub_in - nominal, nominal)\n    bounds_out = bounds_in.apply_linear(None, w, b)\n    lb_out, ub_out = bounds_out.lower, bounds_out.upper\n\n    lb_out_exp = np.array([[-4.9, -11.8, -8.7]])\n    ub_out_exp = np.array([[10.1, 9.2, 18.3]])\n\n    with self.test_session() as session:\n      lb_out_act, ub_out_act = session.run((lb_out, ub_out))\n      self.assertAllClose(lb_out_exp, lb_out_act, atol=tol, rtol=tol)\n      self.assertAllClose(ub_out_exp, ub_out_act, atol=tol, rtol=tol)\n\n  @parameterized.named_parameters(('float32', tf.float32),\n                                  ('float64', tf.float64))\n  def test_conv2d_bounds_shape(self, dtype):\n    batch_size = 23\n    input_height = 17\n    input_width = 7\n    kernel_height = 3\n    kernel_width = 4\n    input_channels = 3\n    output_channels = 5\n    padding = 'VALID'\n    strides = (2, 1)\n\n    # Expected output dimensions, based on convolution settings.\n    output_height = 8\n    output_width = 4\n\n    w = tf.placeholder(dtype=dtype, shape=(\n        kernel_height, kernel_width, input_channels, output_channels))\n    b = tf.placeholder(dtype=dtype, shape=(output_channels,))\n    lb_rel_in = tf.placeholder(dtype=dtype, shape=(\n        batch_size, input_height, input_width, input_channels))\n    ub_rel_in = tf.placeholder(dtype=dtype, shape=(\n        batch_size, input_height, input_width, input_channels))\n    nominal = tf.placeholder(dtype=dtype, shape=(\n        batch_size, input_height, input_width, input_channels))\n\n    bounds_in = ibp.RelativeIntervalBounds(lb_rel_in, ub_rel_in, nominal)\n    bounds_out = bounds_in.apply_conv2d(None, w, b, padding, strides)\n    lb_out, ub_out = bounds_out.lower, bounds_out.upper\n\n    self.assertEqual(dtype, lb_out.dtype)\n    self.assertEqual(dtype, ub_out.dtype)\n    self.assertEqual((batch_size, output_height, output_width, output_channels),\n                     lb_out.shape)\n    self.assertEqual((batch_size, output_height, output_width, output_channels),\n                     ub_out.shape)\n\n  @parameterized.named_parameters(('float32', tf.float32, 1.e-5),\n                                  ('float64', tf.float64, 1.e-8))\n  def test_conv2d_bounds(self, dtype, tol):\n    batch_size = 53\n    input_height = 17\n    input_width = 7\n    kernel_height = 3\n    kernel_width = 4\n    input_channels = 3\n    output_channels = 2\n    padding = 'VALID'\n    strides = (2, 1)\n\n    w = tf.random_normal(dtype=dtype, shape=(\n        kernel_height, kernel_width, input_channels, output_channels))\n    b = tf.random_normal(dtype=dtype, shape=(output_channels,))\n    lb_in = tf.random_normal(dtype=dtype, shape=(\n        batch_size, input_height, input_width, input_channels))\n    ub_in = tf.random_normal(dtype=dtype, shape=(\n        batch_size, input_height, input_width, input_channels))\n    lb_in, ub_in = tf.minimum(lb_in, ub_in), tf.maximum(lb_in, ub_in)\n    nominal = tf.random_normal(dtype=dtype, shape=(\n        batch_size, input_height, input_width, input_channels))\n\n    bounds_in = ibp.RelativeIntervalBounds(lb_in - nominal,\n                                           ub_in - nominal, nominal)\n    bounds_out = bounds_in.apply_conv2d(None, w, b, padding, strides)\n    lb_out, ub_out = bounds_out.lower, bounds_out.upper\n\n    # Compare against equivalent linear layer.\n    bounds_out_lin = _materialised_conv_bounds(\n        w, b, padding, strides, bounds_in)\n    lb_out_lin, ub_out_lin = bounds_out_lin.lower, bounds_out_lin.upper\n\n    with self.test_session() as session:\n      (lb_out_val, ub_out_val,\n       lb_out_lin_val, ub_out_lin_val) = session.run((lb_out, ub_out,\n                                                      lb_out_lin, ub_out_lin))\n      self.assertAllClose(lb_out_val, lb_out_lin_val, atol=tol, rtol=tol)\n      self.assertAllClose(ub_out_val, ub_out_lin_val, atol=tol, rtol=tol)\n\n  @parameterized.named_parameters(('float32', tf.float32),\n                                  ('float64', tf.float64))\n  def test_conv1d_bounds_shape(self, dtype):\n    batch_size = 23\n    input_length = 13\n    kernel_length = 3\n    input_channels = 3\n    output_channels = 5\n    padding = 'VALID'\n    strides = (2,)\n\n    # Expected output dimensions, based on convolution settings.\n    output_length = 6\n\n    w = tf.placeholder(dtype=dtype, shape=(\n        kernel_length, input_channels, output_channels))\n    b = tf.placeholder(dtype=dtype, shape=(output_channels,))\n    lb_rel_in = tf.placeholder(dtype=dtype, shape=(\n        batch_size, input_length, input_channels))\n    ub_rel_in = tf.placeholder(dtype=dtype, shape=(\n        batch_size, input_length, input_channels))\n    nominal = tf.placeholder(dtype=dtype, shape=(\n        batch_size, input_length, input_channels))\n\n    bounds_in = ibp.RelativeIntervalBounds(lb_rel_in, ub_rel_in, nominal)\n    bounds_out = bounds_in.apply_conv1d(None, w, b, padding, strides[0])\n    lb_out, ub_out = bounds_out.lower, bounds_out.upper\n\n    self.assertEqual(dtype, lb_out.dtype)\n    self.assertEqual(dtype, ub_out.dtype)\n    self.assertEqual((batch_size, output_length, output_channels),\n                     lb_out.shape)\n    self.assertEqual((batch_size, output_length, output_channels),\n                     ub_out.shape)\n\n  @parameterized.named_parameters(('float32', tf.float32, 1.e-5),\n                                  ('float64', tf.float64, 1.e-8))\n  def test_conv1d_bounds(self, dtype, tol):\n    batch_size = 53\n    input_length = 13\n    kernel_length = 5\n    input_channels = 3\n    output_channels = 2\n    padding = 'VALID'\n    strides = (2,)\n\n    w = tf.random_normal(dtype=dtype, shape=(\n        kernel_length, input_channels, output_channels))\n    b = tf.random_normal(dtype=dtype, shape=(output_channels,))\n    lb_in = tf.random_normal(dtype=dtype, shape=(\n        batch_size, input_length, input_channels))\n    ub_in = tf.random_normal(dtype=dtype, shape=(\n        batch_size, input_length, input_channels))\n    lb_in, ub_in = tf.minimum(lb_in, ub_in), tf.maximum(lb_in, ub_in)\n    nominal = tf.random_normal(dtype=dtype, shape=(\n        batch_size, input_length, input_channels))\n\n    bounds_in = ibp.RelativeIntervalBounds(lb_in - nominal,\n                                           ub_in - nominal, nominal)\n    bounds_out = bounds_in.apply_conv1d(None, w, b, padding, strides[0])\n    lb_out, ub_out = bounds_out.lower, bounds_out.upper\n\n    # Compare against equivalent linear layer.\n    bounds_out_lin = _materialised_conv_bounds(\n        w, b, padding, strides, bounds_in)\n    lb_out_lin, ub_out_lin = bounds_out_lin.lower, bounds_out_lin.upper\n\n    with self.test_session() as session:\n      (lb_out_val, ub_out_val,\n       lb_out_lin_val, ub_out_lin_val) = session.run((lb_out, ub_out,\n                                                      lb_out_lin, ub_out_lin))\n      self.assertAllClose(lb_out_val, lb_out_lin_val, atol=tol, rtol=tol)\n      self.assertAllClose(ub_out_val, ub_out_lin_val, atol=tol, rtol=tol)\n\n  @parameterized.named_parameters(\n      ('float32_snt', snt.BatchNorm, tf.float32, 1.e-5, False),\n      ('float64_snt', snt.BatchNorm, tf.float64, 1.e-8, False),\n      ('float32', ibp.BatchNorm, tf.float32, 1.e-5, False),\n      ('float64', ibp.BatchNorm, tf.float64, 1.e-8, False),\n      ('float32_train', ibp.BatchNorm, tf.float32, 1.e-5, True),\n      ('float64_train', ibp.BatchNorm, tf.float64, 1.e-8, True))\n  def test_batchnorm_bounds(self, batchnorm_class, dtype, tol, is_training):\n    batch_size = 11\n    input_size = 7\n    output_size = 5\n\n    lb_in = tf.random_normal(dtype=dtype, shape=(batch_size, input_size))\n    ub_in = tf.random_normal(dtype=dtype, shape=(batch_size, input_size))\n    lb_in, ub_in = tf.minimum(lb_in, ub_in), tf.maximum(lb_in, ub_in)\n    nominal = tf.random_normal(dtype=dtype, shape=(batch_size, input_size))\n\n    # Linear layer.\n    w = tf.random_normal(dtype=dtype, shape=(input_size, output_size))\n    b = tf.random_normal(dtype=dtype, shape=(output_size,))\n\n    # Batch norm layer.\n    epsilon = 1.e-2\n    bn_initializers = {\n        'beta': tf.random_normal_initializer(),\n        'gamma': tf.random_uniform_initializer(.1, 3.),\n        'moving_mean': tf.random_normal_initializer(),\n        'moving_variance': tf.random_uniform_initializer(.1, 3.)\n    }\n    batchnorm_module = batchnorm_class(offset=True, scale=True, eps=epsilon,\n                                       initializers=bn_initializers)\n    # Connect the batchnorm module to the graph.\n    batchnorm_module(tf.random_normal(dtype=dtype,\n                                      shape=(batch_size, output_size)),\n                     is_training=is_training)\n\n    bounds_in = ibp.RelativeIntervalBounds(lb_in - nominal,\n                                           ub_in - nominal, nominal)\n    bounds_out = bounds_in.apply_linear(None, w, b)\n    bounds_out = bounds_out.apply_batch_norm(\n        batchnorm_module,\n        batchnorm_module.mean if is_training else batchnorm_module.moving_mean,\n        batchnorm_module.variance if is_training\n        else batchnorm_module.moving_variance,\n        batchnorm_module.gamma,\n        batchnorm_module.beta,\n        epsilon)\n    lb_out, ub_out = bounds_out.lower, bounds_out.upper\n\n    # Separately, calculate dual objective by adjusting the linear layer.\n    wn, bn = layer_utils.combine_with_batchnorm(w, b, batchnorm_module)\n    bounds_out_lin = bounds_in.apply_linear(None, wn, bn)\n    lb_out_lin, ub_out_lin = bounds_out_lin.lower, bounds_out_lin.upper\n\n    init_op = tf.global_variables_initializer()\n\n    with self.test_session() as session:\n      session.run(init_op)\n      (lb_out_val, ub_out_val,\n       lb_out_lin_val, ub_out_lin_val) = session.run((lb_out, ub_out,\n                                                      lb_out_lin, ub_out_lin))\n      self.assertAllClose(lb_out_val, lb_out_lin_val, atol=tol, rtol=tol)\n      self.assertAllClose(ub_out_val, ub_out_lin_val, atol=tol, rtol=tol)\n\n\ndef _materialised_conv_bounds(w, b, padding, strides, bounds_in):\n  \"\"\"Calculates bounds on output of an N-D convolution layer.\n\n  The calculation is performed by first materialising the convolution as a\n  (sparse) fully-connected linear layer. Doing so will affect performance, but\n  may be useful for investigating numerical stability issues.\n\n  Args:\n    w: (N+2)D tensor of shape (kernel_height, kernel_width, input_channels,\n      output_channels) containing weights for the convolution.\n    b: 1D tensor of shape (output_channels) containing biases for the\n      convolution, or `None` if no bias.\n    padding: `\"VALID\"` or `\"SAME\"`, the convolution's padding algorithm.\n    strides: Integer list of length N: `[vertical_stride, horizontal_stride]`.\n    bounds_in: bounds of shape (batch_size, input_height, input_width,\n      input_channels) containing bounds on the inputs to the\n      convolution layer.\n\n  Returns:\n    bounds of shape (batch_size, output_height, output_width,\n      output_channels) with bounds on the outputs of the\n      convolution layer.\n\n  Raises:\n    ValueError: if an unsupported convolution dimensionality is encountered.\n  \"\"\"\n  # Flatten the inputs, as the materialised convolution will have no\n  # spatial structure.\n  bounds_in_flat = bounds_in.apply_batch_reshape(None, [-1])\n\n  # Materialise the convolution as a (sparse) fully connected linear layer.\n  input_shape = bounds_in.shape[1:]\n  w_lin, b_lin = layer_utils.materialise_conv(w, b, input_shape,\n                                              padding=padding, strides=strides)\n  bounds_out_flat = bounds_in_flat.apply_linear(None, w_lin, b_lin)\n\n  # Unflatten the output bounds.\n  output_shape = layer_utils.conv_output_shape(input_shape, w, padding, strides)\n  return bounds_out_flat.apply_batch_reshape(None, output_shape)\n\n\nif __name__ == '__main__':\n  tf.test.main()\n"
  },
  {
    "path": "interval_bound_propagation/tests/simplex_bounds_test.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for naive_bounds.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nfrom absl.testing import parameterized\nimport interval_bound_propagation as ibp\nfrom interval_bound_propagation import layer_utils\nimport numpy as np\nimport tensorflow.compat.v1 as tf\n\n\nclass SimplexBoundsTest(tf.test.TestCase, parameterized.TestCase):\n\n  @parameterized.named_parameters(('float32', tf.float32),\n                                  ('float64', tf.float64))\n  def test_linear_simplex_bounds_shape(self, dtype):\n    vocab_size = 103\n    batch_size = 11\n    input_size = 7\n    output_size = 5\n\n    w = tf.placeholder(dtype=dtype, shape=(input_size, output_size))\n    b = tf.placeholder(dtype=dtype, shape=(output_size,))\n    embedding = tf.placeholder(dtype=dtype, shape=(vocab_size, input_size))\n    centres = tf.placeholder(dtype=dtype, shape=(batch_size, input_size))\n    r = .2\n\n    bounds_in = ibp.SimplexBounds(embedding, centres, r)\n    bounds_out = bounds_in.apply_linear(None, w, b)\n    lb_out, ub_out = bounds_out.lower, bounds_out.upper\n\n    self.assertEqual(dtype, lb_out.dtype)\n    self.assertEqual(dtype, ub_out.dtype)\n    self.assertEqual((batch_size, output_size), lb_out.shape)\n    self.assertEqual((batch_size, output_size), ub_out.shape)\n\n  @parameterized.named_parameters(('float32', tf.float32, 1.e-6),\n                                  ('float64', tf.float64, 1.e-8))\n  def test_linear_bounds_on_embedding_layer(self, dtype, tol):\n    w = tf.constant([[1.0, 2.0, 3.0], [4.0, -5.0, 6.0]], dtype=dtype)\n    b = tf.constant([0.01, -0.02, 0.03], dtype=dtype)\n    embedding = tf.constant([[0.0, 0.0], [10.0, 10.0], [0.0, -20.0]],\n                            dtype=dtype)\n    centres = tf.constant([[7.0, 6.0]], dtype=dtype)\n    r = .1\n    # Simplex vertices: [6.3, 5.4], [7.3, 6.4], and [6.3, 3.4].\n    # They map to: [27.91, -14.42, 51.33], [32.91, -17.42, 60.33],\n    # and [19.91, -4.42, 39.33].\n\n    bounds_in = ibp.SimplexBounds(embedding, centres, r)\n    bounds_out = bounds_in.apply_linear(None, w, b)\n    lb_out, ub_out = bounds_out.lower, bounds_out.upper\n\n    lb_out_exp = np.array([[19.91, -17.42, 39.33]])\n    ub_out_exp = np.array([[32.91, -4.42, 60.33]])\n\n    with self.test_session() as session:\n      lb_out_act, ub_out_act = session.run((lb_out, ub_out))\n      self.assertAllClose(lb_out_exp, lb_out_act, atol=tol, rtol=tol)\n      self.assertAllClose(ub_out_exp, ub_out_act, atol=tol, rtol=tol)\n\n  @parameterized.named_parameters(('float32', tf.float32),\n                                  ('float64', tf.float64))\n  def test_conv1d_simplex_bounds_shape(self, dtype):\n    num_vertices = 41\n    batch_size = 11\n    input_length = 13\n    kernel_length = 5\n    input_channels = 3\n    output_channels = 2\n    padding = 'VALID'\n    strides = (2,)\n\n    # Expected output dimensions, based on convolution settings.\n    output_length = 5\n\n    w = tf.placeholder(dtype=dtype, shape=(\n        kernel_length, input_channels, output_channels))\n    b = tf.placeholder(dtype=dtype, shape=(output_channels,))\n    vertices = tf.placeholder(dtype=dtype, shape=(\n        batch_size, num_vertices, input_length, input_channels))\n    centres = tf.placeholder(dtype=dtype, shape=(\n        batch_size, input_length, input_channels))\n    r = .2\n\n    bounds_in = ibp.SimplexBounds(vertices, centres, r)\n    bounds_out = bounds_in.apply_conv1d(None, w, b, padding, strides)\n    lb_out, ub_out = bounds_out.lower, bounds_out.upper\n\n    self.assertEqual(dtype, lb_out.dtype)\n    self.assertEqual(dtype, ub_out.dtype)\n    self.assertEqual((batch_size, output_length, output_channels),\n                     lb_out.shape)\n    self.assertEqual((batch_size, output_length, output_channels),\n                     ub_out.shape)\n\n  @parameterized.named_parameters(('float32', tf.float32, 2.e-6),\n                                  ('float64', tf.float64, 1.e-8))\n  def test_conv1d_simplex_bounds(self, dtype, tol):\n    num_vertices = 37\n    batch_size = 53\n    input_length = 17\n    kernel_length = 7\n    input_channels = 3\n    output_channels = 2\n    padding = 'VALID'\n    strides = (2,)\n\n    w = tf.random_normal(dtype=dtype, shape=(\n        kernel_length, input_channels, output_channels))\n    b = tf.random_normal(dtype=dtype, shape=(output_channels,))\n    vertices = tf.random_normal(dtype=dtype, shape=(\n        batch_size, num_vertices, input_length, input_channels))\n    centres = tf.random_normal(dtype=dtype, shape=(\n        batch_size, input_length, input_channels))\n    r = .2\n\n    bounds_in = ibp.SimplexBounds(vertices, centres, r)\n    bounds_out = bounds_in.apply_conv1d(None, w, b, padding, strides[0])\n    lb_out, ub_out = bounds_out.lower, bounds_out.upper\n\n    # Compare against equivalent linear layer.\n    bounds_out_lin = _materialised_conv_simplex_bounds(\n        w, b, padding, strides, bounds_in)\n    lb_out_lin, ub_out_lin = bounds_out_lin.lower, bounds_out_lin.upper\n\n    with self.test_session() as session:\n      (lb_out_val, ub_out_val,\n       lb_out_lin_val, ub_out_lin_val) = session.run((lb_out, ub_out,\n                                                      lb_out_lin, ub_out_lin))\n      self.assertAllClose(lb_out_val, lb_out_lin_val, atol=tol, rtol=tol)\n      self.assertAllClose(ub_out_val, ub_out_lin_val, atol=tol, rtol=tol)\n\n\ndef _materialised_conv_simplex_bounds(w, b, padding, strides, bounds_in):\n  \"\"\"Calculates naive bounds on output of an N-D convolution layer.\n\n  The calculation is performed by first materialising the convolution as a\n  (sparse) fully-connected linear layer. Doing so will affect performance, but\n  may be useful for investigating numerical stability issues.\n\n  The layer inputs and the vertices are assumed to be (N-D) sequences in an\n  embedding space. The input domain is taken to be the simplex of perturbations\n  of the centres (true inputs) towards the given vertices.\n\n  Specifically, the input domain is the convex hull of this set of vertices::\n    { (1-r)*centres + r*vertices[j] : j<num_vertices }\n\n  Args:\n    w: (N+2)D tensor of shape (kernel_length, input_channels, output_channels)\n      containing weights for the convolution.\n    b: 1D tensor of shape (output_channels) containing biases for the\n      convolution, or `None` if no bias.\n    padding: `\"VALID\"` or `\"SAME\"`, the convolution's padding algorithm.\n    strides: Integer list of length N: `[vertical_stride, horizontal_stride]`.\n    bounds_in: bounds of shape (batch_size, input_length, input_channels)\n      containing bounds on the inputs to the convolution layer.\n\n  Returns:\n    bounds of shape (batch_size, output_length, output_channels)\n      with bounds on the outputs of the convolution layer.\n\n  Raises:\n    ValueError: if an unsupported convolution dimensionality is encountered.\n  \"\"\"\n  # Flatten the inputs, as the materialised convolution will have no\n  # spatial structure.\n  bounds_in_flat = bounds_in.apply_batch_reshape(None, [-1])\n\n  # Materialise the convolution as a (sparse) fully connected linear layer.\n  input_shape = bounds_in.shape[1:]\n  w_lin, b_lin = layer_utils.materialise_conv(w, b, input_shape,\n                                              padding=padding, strides=strides)\n  bounds_out_flat = bounds_in_flat.apply_linear(None, w_lin, b_lin)\n\n  # Unflatten the output bounds.\n  output_shape = layer_utils.conv_output_shape(input_shape, w, padding, strides)\n  return bounds_out_flat.apply_batch_reshape(None, output_shape)\n\n\nif __name__ == '__main__':\n  tf.test.main()\n"
  },
  {
    "path": "interval_bound_propagation/tests/specification_test.py",
    "content": "# coding=utf-8\n# Copyright 2019 The Interval Bound Propagation Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Tests for specification.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport collections\n\nimport interval_bound_propagation as ibp\nimport numpy as np\nimport tensorflow.compat.v1 as tf\n\n\nMockLinearModule = collections.namedtuple('MockLinearModule', ['w', 'b'])\nMockModule = collections.namedtuple(\n    'MockModule', ['input_bounds', 'output_bounds', 'module'])\n\n\ndef _build_spec_input():\n  # Specifications expects a list of objects with output_bounds or input_bounds\n  # attributes.\n  w = np.identity(2, dtype=np.float32)\n  b = np.ones(2, dtype=np.float32)\n  snt_module = MockLinearModule(tf.constant(w), tf.constant(b))\n  z_lower = np.array([[1, 2]], dtype=np.float32)\n  z_upper = np.array([[3, 4]], dtype=np.float32)\n  input_bounds = ibp.IntervalBounds(tf.constant(z_lower), tf.constant(z_upper))\n  z_lower += b\n  z_upper += b\n  output_bounds = ibp.IntervalBounds(tf.constant(z_lower), tf.constant(z_upper))\n  return [MockModule(input_bounds, output_bounds, snt_module)]\n\n\ndef _build_classification_specification(label, num_classes, collapse):\n  \"\"\"Returns a LinearSpecification for adversarial classification.\"\"\"\n  # Pre-construct the specifications of the different classes.\n  eye = np.eye(num_classes - 1)\n  specifications = []\n  for i in range(num_classes):\n    specifications.append(np.concatenate(\n        [eye[:, :i], -np.ones((num_classes - 1, 1)), eye[:, i:]], axis=1))\n  specifications = np.array(specifications, dtype=np.float32)\n  specifications = tf.constant(specifications)\n  # We can then use gather.\n  c = tf.gather(specifications, label)\n  # By construction all specifications are relevant.\n  d = tf.zeros(shape=(tf.shape(label)[0], num_classes - 1))\n  return ibp.LinearSpecification(c, d, prune_irrelevant=False,\n                                 collapse=collapse)\n\n\nclass SpecificationTest(tf.test.TestCase):\n\n  def testLinearSpecification(self):\n    # c has shape [batch_size, num_specifications, num_outputs]\n    # d has shape [batch_size, num_specifications]\n    c = tf.constant([[[1, 2]]], dtype=tf.float32)\n    d = tf.constant([[3]], dtype=tf.float32)\n    # The above is equivalent to z_{K,1} + 2 * z_{K,2} + 3 <= 0\n    spec = ibp.LinearSpecification(c, d, collapse=False)\n    spec_collapse = ibp.LinearSpecification(c, d, collapse=True)\n    modules = _build_spec_input()\n    values = spec(modules)\n    values_collapse = spec_collapse(modules)\n    with self.test_session() as sess:\n      self.assertAlmostEqual(17., sess.run(values).item())\n      self.assertAlmostEqual(17., sess.run(values_collapse).item())\n\n  def testEquivalenceLinearClassification(self):\n    num_classes = 3\n    def _build_model():\n      layer_types = (\n          ('conv2d', (2, 2), 4, 'VALID', 1),\n          ('activation', 'relu'),\n          ('linear', 10),\n          ('activation', 'relu'))\n      return ibp.DNN(num_classes, layer_types)\n\n    # Input.\n    batch_size = 100\n    width = height = 2\n    channels = 3\n    num_restarts = 10\n    z = tf.random.uniform((batch_size, height, width, channels),\n                          minval=-1., maxval=1., dtype=tf.float32)\n    y = tf.random.uniform((batch_size,), minval=0, maxval=num_classes,\n                          dtype=tf.int64)\n    predictor = _build_model()\n    predictor = ibp.VerifiableModelWrapper(predictor)\n    logits = predictor(z)\n    random_logits1 = tf.random.uniform((num_restarts, batch_size, num_classes))\n    random_logits2 = tf.random.uniform((num_restarts, num_classes - 1,\n                                        batch_size, num_classes))\n    input_bounds = ibp.IntervalBounds(z - 2., z + 4.)\n    predictor.propagate_bounds(input_bounds)\n\n    # Specifications.\n    s1 = ibp.ClassificationSpecification(y, num_classes, collapse=False)\n    s1_collapse = ibp.ClassificationSpecification(y, num_classes, collapse=True)\n    s2 = _build_classification_specification(y, num_classes, collapse=False)\n    s2_collapse = _build_classification_specification(y, num_classes,\n                                                      collapse=True)\n    def _build_values(s, s_collapse):\n      return [\n          s(predictor.modules),\n          s_collapse(predictor.modules),\n          s.evaluate(logits),\n          s.evaluate(random_logits1),\n          s.evaluate(random_logits2)\n      ]\n    v1 = _build_values(s1, s1_collapse)\n    v2 = _build_values(s2, s2_collapse)\n\n    with self.test_session() as sess:\n      sess.run(tf.global_variables_initializer())\n      output1, output2 = sess.run([v1, v2])\n    for a, b in zip(output1, output2):\n      self.assertTrue(np.all(np.abs(a - b) < 1e-5))\n\n\nif __name__ == '__main__':\n  tf.test.main()\n"
  },
  {
    "path": "setup.py",
    "content": "# Copyright 2018 The Interval Bound Propagation Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#    http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Setup for pip package.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport unittest\nfrom setuptools import find_packages\nfrom setuptools import setup\n\n\nREQUIRED_PACKAGES = ['six', 'absl-py', 'numpy']\nEXTRA_PACKAGES = {\n    'tensorflow': ['tensorflow>=1.8.0'],\n    'tensorflow with gpu': ['tensorflow-gpu>=1.8.0'],\n    'sonnet': ['dm-sonnet>=1.26'],\n    'sonnet with gpu': ['dm-sonnet-gpu>=1.26'],\n}\n\n\ndef ibp_test_suite():\n  test_loader = unittest.TestLoader()\n  test_suite = test_loader.discover('interval_bound_propagation/tests',\n                                    pattern='*_test.py')\n  return test_suite\n\nsetup(\n    name='interval_bound_propagation',\n    version='1.1',\n    description='A library to train verifiably robust neural networks.',\n    url='https://github.com/deepmind/interval_bound_propagation',\n    author='DeepMind',\n    author_email='no-reply@google.com',\n    # Contained modules and scripts.\n    packages=find_packages(),\n    install_requires=REQUIRED_PACKAGES,\n    extras_require=EXTRA_PACKAGES,\n    platforms=['any'],\n    license='Apache 2.0',\n    test_suite='setup.ibp_test_suite',\n)\n"
  }
]