[
  {
    "path": ".gitignore",
    "content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n\n# Distribution / packaging\n.Python\n*.egg-info/\n*.egg\n.eggs/\ndist/\nsdist/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# Condensa-specific\ncompressed/\ntrained/\nresults/\ndata/\nbuild/\n# version.py is auto-generated by setup.py\nversion.py\n"
  },
  {
    "path": ".style.yapf",
    "content": "[style]\nbased_on_style = pep8\nblank_lines_around_top_level_definition = 1\ncolumn_limit = 79\nindent_width = 4\n"
  },
  {
    "path": "LICENSE",
    "content": "                                 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 2019 NVIDIA Corporation\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": "# A Programming System for Neural Network Compression\n\n**Note:** the original version of Condensa (contained in this branch) is no longer actively maintained. Please check out the [lite branch](https://github.com/NVlabs/condensa/tree/lite) for the most up-to-date version.\n\nCondensa is a framework for _programmable model compression_ in Python.\nIt comes with a set of built-in compression operators which may be used to\ncompose complex compression schemes targeting specific combinations of DNN architecture,\nhardware platform, and optimization objective.\nTo recover any accuracy lost during compression, Condensa uses a constrained\noptimization formulation of model compression and employs an Augmented Lagrangian-based\nalgorithm as the optimizer.\n\n**Status**: Condensa is under active development, and bug reports, pull requests, and other feedback are all highly appreciated. See the contributions section below for more details on how to contribute.\n\n## Supported Operators and Schemes\n\nCondensa provides the following set of pre-built compression schemes:\n\n* [Unstructured Pruning](https://nvlabs.github.io/condensa/modules/schemes.html#unstructured-pruning)\n* [Filter and Neuron Pruning](https://nvlabs.github.io/condensa/modules/schemes.html#neuron-pruning)\n* [Block Pruning](https://nvlabs.github.io/condensa/modules/schemes.html#block-pruning)\n* [Quantization](https://nvlabs.github.io/condensa/modules/schemes.html#quantization)\n* [Scheme Composition](https://nvlabs.github.io/condensa/modules/schemes.html#composition)\n\nThe schemes above are built using one or more [compression operators](https://nvlabs.github.io/condensa/modules/pi.html), which may be combined in various ways to define your own custom schemes.\n\nPlease refer to the [documentation](https://nvlabs.github.io/condensa/index.html) for a detailed description of available operators and schemes.\n\n## Prerequisites\n\nCondensa requires:\n\n* A working Linux installation (we use Ubuntu 18.04)\n* NVIDIA drivers and CUDA 10+ for GPU support\n* Python 3.5 or newer\n* PyTorch 1.0 or newer\n\n## Installation\n\nThe most straightforward way of installing Condensa is via `pip`:\n\n```bash\npip install condensa\n```\n\n### Installation from Source\n\nRetrieve the latest source code from the Condensa repository:\n\n```bash\ngit clone https://github.com/NVlabs/condensa.git\n```\n\nNavigate to the source code directory and run the following:\n\n```bash\npip install -e .\n```\n\n### Test out the Installation\n\nTo check the installation, run the unit test suite:\n\n```bash\nbash run_all_tests.sh -v\n```\n\n## Getting Started\n\nThe [AlexNet Notebook](https://github.com/NVlabs/condensa/blob/master/notebooks/AlexNet.ipynb) contains a simple step-by-step walkthrough of compressing a pre-trained model using Condensa.\nCheck out the [`examples` folder](https://github.com/NVlabs/condensa/tree/master/examples/cifar) for additional, more complex examples of using Condensa (**note**: some examples require the `torchvision` package to be installed).\n\n## Documentation\n\nDocumentation is available [here](https://nvlabs.github.io/condensa/). Please also check out the [Condensa paper](https://arxiv.org/abs/1911.02497) for a detailed\ndescription of Condensa's motivation, features, and performance results.\n\n## Contributing\n\nWe appreciate all contributions, including bug fixes, new features and documentation, and additional tutorials. You can initiate\ncontributions via Github pull requests. When making code contributions, please follow the `PEP 8` Python coding standard and provide\nunit tests for the new features. Finally, make sure to sign off your commits using the `-s` flag or adding \n`Signed-off-By: Name<Email>` in the commit message.\n\n## Citing Condensa\n\nIf you use Condensa for research, please consider citing the following paper:\n\n```\n@article{condensa2020,\n  title={A Programmable Approach to Neural Network Compression}, \n  author={V. {Joseph} and G. L. {Gopalakrishnan} and S. {Muralidharan} and M. {Garland} and A. {Garg}},\n  journal={IEEE Micro}, \n  year={2020},\n  volume={40},\n  number={5},\n  pages={17-25},\n  doi={10.1109/MM.2020.3012391}\n}\n```\n\n## Disclaimer\n\nCondensa is a research prototype and not an official NVIDIA product. Many features are still experimental and yet to be properly documented.\n"
  },
  {
    "path": "condensa/__init__.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nname = \"condensa\"\n\nfrom .version import __version__\nfrom . import opt\n\nfrom .dtypes import *\nfrom .compressor import *\nfrom .finetune import *\nfrom .pi import *\nfrom .delta import *\nfrom .util import *\n\nfrom . import schemes\nfrom . import data\n"
  },
  {
    "path": "condensa/cfg.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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__CONDENSA_RECORD_MODE__ = False\n__CONDENSA_PI_PRECHECK__ = True\n"
  },
  {
    "path": "condensa/compressor.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport torch.nn\n\nclass Compressor(object):\n    \"\"\"Condensa model compressor class.\"\"\"\n    def __init__(self,\n                 opt,\n                 scheme,\n                 model,\n                 trainloader,\n                 testloader,\n                 valloader,\n                 criterion):\n        \"\"\"\n        Creates a `Compressor` instance.\n\n        :param opt: Optimizer.\n        :type opt: `condensa.Optimizer`\n        :param scheme: Compression scheme (class).\n        :param model: PyTorch model.\n        :type model: `torch.nn.Module`\n        :param trainloader: Training dataloader.\n        :param testloader: Test dataloader.\n        :param valloader: Validation dataloader.\n        :param criterion: Loss criterion.\n        \"\"\"\n        assert isinstance(model, torch.nn.Module)\n\n        self.opt = opt\n        self.pi = scheme.pi\n        self.delta = scheme.delta\n        self.model = model\n        self.trainloader = trainloader\n        self.testloader = testloader\n        self.valloader = valloader\n        self.criterion = criterion\n\n        self._statistics = None\n\n    @property\n    def statistics(self):\n        \"\"\"\n        Retrieves compressed model statistics.\n\n        :return: Model statistics.\n        :rtype: `dict`\n        \"\"\"\n        return self._statistics\n\n    def run(self):\n        \"\"\"\n        Executes model compressor.\n\n        :return: Compressed model.\n        :rtype: `torch.nn.Module`\n        \"\"\"\n        w, statistics = self.opt.compress(self.model, self.pi, self.delta,\n                                          self.trainloader, self.testloader,\n                                          self.valloader, self.criterion)\n        self._statistics = statistics\n        return w\n"
  },
  {
    "path": "condensa/data.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport sys\nimport numpy as np\nimport torch\nimport torch.utils.data as data\nimport PIL\n\ndef fast_collate(batch):\n    \"\"\"Fast batch collation. Based on version from\n       NVIDIA Apex: https://github.com/NVIDIA/apex.\"\"\"\n    imgs = [img[0] for img in batch]\n    targets = torch.tensor([target[1] for target in batch], dtype=torch.int64)\n    w = imgs[0].size[0]\n    h = imgs[0].size[1]\n    tensor = torch.zeros((len(imgs), 3, h, w), dtype=torch.uint8)\n    for i, img in enumerate(imgs):\n        nump_array = np.asarray(img, dtype=np.uint8)\n        if (nump_array.ndim < 3):\n            nump_array = np.expand_dims(nump_array, axis=-1)\n        nump_array = np.rollaxis(nump_array, 2)\n        tensor[i] += torch.from_numpy(nump_array)\n    return tensor, targets\n\nclass GPUDataLoader(object):\n    \"\"\"Custom data loader with support for prefetching and fast collation.\n       Based on version from NVIDIA Apex: https://github.com/NVIDIA/apex.\"\"\"\n    def __init__(self,\n                 dataset,\n                 batch_size,\n                 shuffle,\n                 num_workers,\n                 sampler=None,\n                 meanstd=None):\n        if isinstance(dataset[0][0], PIL.Image.Image):\n            nc = len(dataset[0][0].getbands())\n        else:\n            raise RuntimeError(\n                '[Condensa] GPUDataLoader only supports PIL image datasets')\n\n        if not torch.cuda.is_available():\n            raise RuntimeError(\n                '[Condensa] GPUDataLoader requires PyTorch CUDA support')\n\n        if nc == 3:\n            loader = data.DataLoader(dataset,\n                                     batch_size=batch_size,\n                                     shuffle=shuffle,\n                                     pin_memory=True,\n                                     sampler=sampler,\n                                     collate_fn=fast_collate,\n                                     num_workers=num_workers)\n        else:\n            raise NotImplementedError(\n                '[Condensa] GPUDataLoader currently only supports 3-channel images'\n            )\n\n        self.base_loader = loader\n        self.loader = iter(loader)\n        self.stream = torch.cuda.Stream()\n        if meanstd is not None:\n            mean, std = meanstd\n            self.mean = torch.tensor([x * 255\n                                      for x in mean]).cuda().view(1, nc, 1, 1)\n            self.std = torch.tensor([x * 255\n                                     for x in std]).cuda().view(1, nc, 1, 1)\n        self.preload()\n\n    def __len__(self):\n        return len(self.base_loader)\n\n    def __iter__(self):\n        self.loader = iter(self.base_loader)\n        self.preload()\n        return self\n\n    def __next__(self):\n        torch.cuda.current_stream().wait_stream(self.stream)\n        input = self.next_input\n        target = self.next_target\n        if input is None and target is None:\n            raise StopIteration\n        input.record_stream(torch.cuda.current_stream())\n        target.record_stream(torch.cuda.current_stream())\n        self.preload()\n        return input, target\n\n    def preload(self):\n        try:\n            self.next_input, self.next_target = next(self.loader)\n        except StopIteration:\n            self.next_input = None\n            self.next_target = None\n            return\n        with torch.cuda.stream(self.stream):\n            self.next_input = self.next_input.cuda(non_blocking=True)\n            self.next_target = self.next_target.cuda(non_blocking=True)\n\n            self.next_input = self.next_input.float()\n            if self.mean is not None and self.std is not None:\n                self.next_input = self.next_input.sub_(self.mean).div_(\n                    self.std)\n"
  },
  {
    "path": "condensa/delta.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nfrom condensa import dtypes\n\ndef dequantize(module, dtype):\n    \"\"\"\n    De-quantizes module to given data type (inplace).\n\n    :param module: PyTorch module.\n    :type module: `torch.nn.Module`\n    :param dtype: Target data type.\n    \"\"\"\n    if dtype.as_dtype_enum == dtypes.DT_FLOAT32:\n        module.float()\n    elif dtype.as_dtype_enum == dtypes.DT_FLOAT64:\n        module.double()\n    else:\n        raise TypeError('Unknown data type specified for de-quantization')\n"
  },
  {
    "path": "condensa/dtypes.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport numpy as np\n\nfrom .type_enums import *\n\nclass DType(object):\n    \"\"\"Data type for quantization.\"\"\"\n    def __init__(self, dtype):\n        self._dtype = dtype\n\n    @property\n    def name(self):\n        return _DTYPE_TO_STRING[self._dtype]\n\n    @property\n    def as_numpy_dtype(self):\n        return _TO_NP[self._dtype]\n\n    @property\n    def as_dtype_enum(self):\n        return self._dtype\n\n    def __int__(self):\n        return self._dtype\n\n    def __str__(self):\n        return \"<dtype: %r>\" % self.name\n\nfloat16 = DType(DT_FLOAT16)\nfloat32 = DType(DT_FLOAT32)\nfloat64 = DType(DT_FLOAT64)\nint8 = DType(DT_INT8)\nuint8 = DType(DT_UINT8)\nint16 = DType(DT_INT16)\nuint16 = DType(DT_UINT16)\n\n_DTYPE_TO_STRING = {\n    DT_FLOAT16: \"float16\",\n    DT_FLOAT32: \"float32\",\n    DT_FLOAT64: \"float64\",\n    DT_INT8: \"int8\",\n    DT_UINT8: \"uint8\",\n    DT_INT16: \"int16\",\n    DT_UINT16: \"uint16\"\n}\n\n_TO_NP = {\n    DT_FLOAT16: np.float16,\n    DT_FLOAT32: np.float32,\n    DT_FLOAT64: np.float64,\n    DT_INT8: np.int8,\n    DT_UINT8: np.uint8,\n    DT_INT16: np.int16,\n    DT_UINT16: np.uint16\n}\n"
  },
  {
    "path": "condensa/finetune.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport sys\nimport logging\nimport numpy as np\nfrom copy import deepcopy\n\nimport torch\nimport torch.backends.cudnn as cudnn\nfrom tqdm import tqdm\n\nimport condensa.tensor as T\nimport condensa.util as util\n\nlogger = logging.getLogger(__name__)\n\nclass FineTuner(object):\n    \"\"\"Condensa model fine-tuner. Can be used for retraining compressed\n       models while keeping all zero-valued parameters clipped to zero.\"\"\"\n    def __init__(self, w, layer_types=None, biases=True):\n        self.w = w\n        self.layer_types = layer_types\n        self.biases = biases\n        self._compute_mask_inplace()\n\n    def _compute_mask_inplace(self):\n        with torch.no_grad():\n            for m in self.w.modules():\n                if type(m) in self.layer_types\\\n                   and not hasattr(m, 'condensa_nocompress'):\n                    if hasattr(m, 'weight'):\n                        m.mask_w = torch.gt(m.weight.data.abs(), 0.)\n                    if self.biases:\n                        if hasattr(m, 'bias') and m.bias is not None:\n                            m.mask_b = torch.gt(m.bias.data.abs(), 0.)\n\n    def _apply_mask(self):\n        with torch.no_grad():\n            for m in self.w.modules():\n                if hasattr(m, 'mask_w'):\n                    T.apply_mask_inplace(m.weight.data, m.mask_w)\n                if hasattr(m, 'mask_b'):\n                    T.apply_mask_inplace(m.bias.data, m.mask_b)\n\n    def run(self,\n            epochs,\n            lr,\n            lr_end,\n            momentum,\n            weight_decay,\n            criterion,\n            trainloader,\n            testloader,\n            valloader,\n            debugging_flags={}):\n        \"\"\"\n        Fine-tunes a compressed model. Currently only supports SGD.\n\n        :param epochs: Number of epochs\n        :type epochs: `int`\n        :param lr: Learning rate\n        :type lr: `float`\n        :param lr_end: End learning rate\n        :type lr_end: `float`\n        :param momentum: Momentum\n        :type momentum: float\n        :param weight_decay: Weight decay\n        :type weight_decay: float\n        :param criterion: Loss criterion\n        :param trainloader: Training dataloader\n        :param testloader: Test dataloader\n        :param valloader: Validation dataloader\n        :param debugging_flags: Debugging flags\n        :type debugging_flags: dict\n        \"\"\"\n        use_cuda = torch.cuda.is_available()\n\n        validate = (valloader is not None)\n        test = (testloader is not None)\n\n        if use_cuda:\n            cudnn.benchmark = True\n            self.w = self.w.cuda()\n\n        _model_stat_fn = debugging_flags['custom_model_statistics']\\\n                   if 'custom_model_statistics' in debugging_flags\\\n                   else util.empty_stat_fn\n\n        if validate:\n            val_loss, val_stats = _model_stat_fn(self.w, criterion, valloader)\n            logging.info(\n                '[Condensa:FineTuner] Original model val_loss: {:.2f}, {}'\n                .format(val_loss,\n                ', '.join(['{}:{}'.format(k, v) for k,v in val_stats.items()])))\n        if test:\n            test_loss, test_stats = _model_stat_fn(\n                self.w, criterion, testloader)\n            logging.info(\n                '[Condensa:FineTuner] Original model test_loss: {:.2f}, {} '\n                .format(test_loss,\n                ', '.join(['{}:{}'.format(k, v) for k,v in test_stats.items()])))\n\n        l_alpha = np.exp((np.log(lr_end) - np.log(lr)) / float(epochs))\n        optimizer = torch.optim.SGD(self.w.parameters(),\n                                    lr=lr,\n                                    momentum=momentum,\n                                    weight_decay=weight_decay,\n                                    nesterov=False)\n        with torch.no_grad():\n            best_model = deepcopy(self.w)\n        best_loss = sys.float_info.max\n        for epoch in range(epochs):\n            # Switch to training mode\n            self.w.train()\n            nbatches = len(trainloader)\n            if logger.isEnabledFor(logging.INFO):\n                pbar = tqdm(total=nbatches, ascii=True)\n            for input, target in trainloader:\n                if torch.cuda.is_available():\n                    if not input.is_cuda: input = input.cuda()\n                    if not target.is_cuda: target = target.cuda()\n                output = self.w(input)\n                loss = criterion(output, target)\n                optimizer.zero_grad()\n                loss.backward()\n                optimizer.step()\n                # Apply mask\n                self._apply_mask()\n                if logger.isEnabledFor(logging.INFO):\n                    pbar.update()\n            if logger.isEnabledFor(logging.INFO):\n                pbar.close()\n\n            # Switch to eval mode\n            self.w.eval()\n\n            if validate:\n                val_loss, val_stats = _model_stat_fn(\n                    self.w, criterion, valloader)\n                logging.info(\n                    '[Condensa:FineTuner] Epoch [{}], VAL loss: {:.2f}, {}'\n                    .format(epoch, val_loss,\n                    ', '.join(['{}:{}'.format(k, v) for k,v in val_stats.items()])))\n            if test:\n                test_loss, test_stats = _model_stat_fn(\n                    self.w, criterion, testloader)\n                logging.info(\n                    '[Condensa:FineTuner] Epoch [{}], TEST loss: {:.2f}, {}'\n                    .format(epoch, test_loss,\n                    ', '.join(['{}:{}'.format(k, v) for k,v in test_stats.items()])))\n\n            if validate:\n                if val_loss < best_loss:\n                    logger.info(\n                        '[Condensa:FineTuner] SAVING MODEL based on VAL')\n                    best_loss = val_loss\n                    best_model = deepcopy(self.w)\n            elif test:\n                if test_loss < best_loss:\n                    logger.info(\n                        '[Condensa:FineTuner] SAVING MODEL based on TEST')\n                    best_loss = test_loss\n                    best_model = deepcopy(self.w)\n            else:\n                logger.info(\n                    '[Condensa:FineTuner] SAVING MODEL based on most recent')\n                best_model = deepcopy(self.w)\n\n            lr *= l_alpha\n            for g in optimizer.param_groups:\n                g['lr'] = lr\n\n        return best_model\n"
  },
  {
    "path": "condensa/functional.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport torch\n\ndef l2norm(tensor, dim, keepdim):\n    \"\"\"\n    Computes the l2-norm of elements in input tensor.\n\n    :param tensor: PyTorch tensor.\n    :type tensor: `torch.nn.Module`\n    :param dim: Reduction dimension.\n    :type dim: `int`\n    :param keepdim: Whether the output has `dim` retained.\n    :type keepdim: `bool`\n    :return: l2-norm of input tensor.\n    \"\"\"\n    return torch.norm(tensor, 2, dim, keepdim)\n\ndef max(tensor, dim, keepdim):\n    \"\"\"\n    Computes the maximum value of elements in input tensor.\n\n    :param tensor: PyTorch tensor.\n    :type tensor: `torch.nn.Module`\n    :param dim: Reduction dimension.\n    :type dim: `int`\n    :param keepdim: Whether the output has `dim` retained.\n    :type keepdim: `bool`\n    :return: Max of input tensor.\n    \"\"\"\n    return torch.max(tensor, dim, keepdim)[0]\n\ndef min(tensor, dim, keepdim):\n    \"\"\"\n    Computes the minimum value of elements in input tensor.\n\n    :param tensor: PyTorch tensor.\n    :type tensor: `torch.nn.Module`\n    :param dim: Reduction dimension.\n    :type dim: `int`\n    :param keepdim: Whether the output has `dim` retained.\n    :type keepdim: `bool`\n    :return: Min of input tensor.\n    \"\"\"\n    return torch.min(tensor, dim, keepdim)[0]\n\ndef mean(tensor, dim, keepdim):\n    \"\"\"\n    Computes the mean value of elements in input tensor.\n\n    :param tensor: PyTorch tensor.\n    :type tensor: `torch.nn.Module`\n    :param dim: Reduction dimension.\n    :type dim: `int`\n    :param keepdim: Whether the output has `dim` retained.\n    :type keepdim: `bool`\n    :return: Mean value of input tensor.\n    \"\"\"\n    return torch.mean(tensor, dim, keepdim)\n\ndef sum(tensor, dim, keepdim):\n    \"\"\"\n    Computes the sum of elements in input tensor.\n\n    :param tensor: PyTorch tensor.\n    :type tensor: `torch.nn.Module`\n    :param dim: Reduction dimension.\n    :type dim: `int`\n    :param keepdim: Whether the output has `dim` retained.\n    :type keepdim: `bool`\n    :return: Sum of input tensor.\n    \"\"\"\n    return torch.sum(tensor, dim, keepdim)\n"
  },
  {
    "path": "condensa/lr.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport numpy as np\n\nclass IntervalLR(object):\n    \"\"\"Decays learning rate between two values.\"\"\"\n    def __init__(self, begin, end, n):\n        \"\"\"\n        Construct an instance of `IntervalLR`.\n\n        :param begin: Starting learning rate (LR).\n        :type begin: `float`\n        :param end: Ending LR.\n        :type end: `float`\n        :param n: Number of iterations.\n        :type n: `int`\n        \"\"\"\n        self.alpha = np.exp((np.log(end) - np.log(begin)) / n)\n        self.lr = begin\n\n    def step(self):\n        \"\"\"Signal end of iteration.\"\"\"\n        self.lr *= self.alpha\n\n    @property\n    def learning_rate(self):\n        \"\"\"Returns current learning rate.\"\"\"\n        return self.lr\n\nclass DecayedLR(object):\n    \"\"\"Decays learning rate at fixed intervals.\"\"\"\n    def __init__(self, begin, schedule, gamma=0.1):\n        \"\"\"\n        Construct an instance of `DecayedLR`.\n\n        :param begin: Starting LR.\n        :type begin: `float`\n        :param schedule: List of iterations when LR must be adjusted.\n        :type schedule: `List/Tuple`\n        :param gamma: LR multiplier.\n        :type gamma: `float`\n        \"\"\"\n        self.gamma = gamma\n        self.lr = begin\n        self.schedule = schedule\n        self.counter = 0\n\n    def step(self):\n        \"\"\"Signal end of iteration.\"\"\"\n        if self.counter in self.schedule:\n            self.lr *= self.gamma\n        self.counter += 1\n\n    @property\n    def learning_rate(self):\n        \"\"\"Returns current learning rate.\"\"\"\n        return self.lr\n\nclass ExpDecayedLR(object):\n    \"\"\"Decays learning rate exponentially.\"\"\"\n    def __init__(self, begin, gamma):\n        \"\"\"\n        Construct an instance of `ExpDecayedLR`.\n\n        :param begin: Starting LR.\n        :type begin: `float`\n        :param gamma: LR multiplier.\n        :type gamma: `float`\n        \"\"\"\n        self.gamma = gamma\n        self.lr = begin\n        self.counter = 0\n\n    def step(self):\n        \"\"\"Signal end of iteration.\"\"\"\n        self.counter += 1\n\n    @property\n    def learning_rate(self):\n        \"\"\"Returns current learning rate.\"\"\"\n        return self.lr * (self.gamma**self.counter)\n"
  },
  {
    "path": "condensa/opt/__init__.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nfrom .direct.dc import DC\nfrom .lc.lc import LC\n\nfrom . import lc\n"
  },
  {
    "path": "condensa/opt/direct/__init__.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\n"
  },
  {
    "path": "condensa/opt/direct/dc.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nfrom copy import deepcopy\nimport torch\n\nfrom condensa.util import EventTimer\n\nclass DC(object):\n    \"\"\"Condensa direct compression optimizer.\"\"\"\n    def compress(self,\n                 w,\n                 pi,\n                 delta,\n                 trainloader,\n                 testloader,\n                 valloader,\n                 criterion):\n        \"\"\"\n        Performs model compression using direct optimization.\n\n        :param w: PyTorch model.\n        :type w: `torch.nn.Module`\n        :param pi: Compression function.\n        :param delta: Decompression function.\n        :param trainloader: Training dataloader.\n        :param testloader: Test dataloader.\n        :param valloader: Validation dataloader.\n        :param criterion: Loss criterion.\n        \"\"\"\n        statistics = dict()\n        timer_dc = EventTimer()\n        with torch.no_grad():\n            compressed = deepcopy(w)\n        pi(compressed)\n        statistics['total_elapsed'] = timer_dc.elapsed_seconds\n\n        return compressed, statistics\n"
  },
  {
    "path": "condensa/opt/lc/__init__.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nfrom .sgd import SGD\nfrom .adam import Adam\n"
  },
  {
    "path": "condensa/opt/lc/adam.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport math\nfrom collections import defaultdict\n\nimport torch\n\nclass Adam(object):\n    \"\"\"Custom Adam implementation for L-C optimizer.\"\"\"\n    def __init__(self,\n                 w,\n                 lr=1e-3,\n                 betas=(0.9, 0.999),\n                 eps=1e-8,\n                 weight_decay=0,\n                 amsgrad=False):\n        if not 0.0 <= lr:\n            raise ValueError(\"Invalid learning rate: {}\".format(lr))\n        if not 0.0 <= eps:\n            raise ValueError(\"Invalid epsilon value: {}\".format(eps))\n        if not 0.0 <= betas[0] < 1.0:\n            raise ValueError(\"Invalid beta parameter at index 0: {}\".format(\n                betas[0]))\n        if not 0.0 <= betas[1] < 1.0:\n            raise ValueError(\"Invalid beta parameter at index 1: {}\".format(\n                betas[1]))\n\n        try:\n            self.w = w.module\n        except AttributeError:\n            self.w = w\n\n        self.lr = lr\n        self.betas = betas\n        self.eps = eps\n        self.weight_decay = weight_decay\n        self.amsgrad = amsgrad\n\n        self.state = defaultdict(dict)\n\n    def zero_grad(self):\n        \"\"\"Zeroes out all gradients.\"\"\"\n        for p in self.w.parameters():\n            if p.grad is not None:\n                p.grad.detach_()\n                p.grad.zero_()\n\n    def reset_state(self):\n        \"\"\"Resets optimizer state.\"\"\"\n        for p in self.w.parameters():\n            if 'state' in self.state[p]:\n                self.state[p]['step'] = 0\n            if 'exp_avg' in self.state[p]:\n                self.state[p]['exp_avg'] = torch.zeros_like(p.data)\n            if 'exp_avg_sq' in self.state[p]:\n                self.state[p]['exp_avg_sq'] = torch.zeros_like(p.data)\n            if self.amsgrad and 'max_exp_avg_sq' in self.state[p]:\n                self.state[p]['max_exp_avg_sq'] = torch.zeros_like(p.data)\n\n    def _step(self, p, condense=False, mu=None, p_theta=None, p_lm=None):\n        if p.grad is None:\n            return\n\n        grad = p.grad.data\n        if grad.is_sparse:\n            raise RuntimeError('Adam does not support sparse gradients.')\n\n        state = self.state[p]\n        # State initialization\n        if len(state) == 0:\n            state['step'] = 0\n            # Exponential moving average of gradient values\n            state['exp_avg'] = torch.zeros_like(p.data)\n            # Exponential moving average of squared gradient values\n            state['exp_avg_sq'] = torch.zeros_like(p.data)\n            if self.amsgrad:\n                # Maintains max of all exp. moving avg. of sq. grad. values\n                state['max_exp_avg_sq'] = torch.zeros_like(p.data)\n\n        exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']\n        if self.amsgrad:\n            max_exp_avg_sq = state['max_exp_avg_sq']\n        beta1, beta2 = self.betas\n\n        state['step'] += 1\n\n        if self.weight_decay != 0:\n            grad.add_(self.weight_decay, p.data)\n\n        if condense is True:\n            assert (mu is not None\n                    and p_theta is not None\n                    and p_lm is not None)\n            grad.add_(mu * (p.data - p_theta.data) - p_lm.data)\n\n        # Decay the first and second moment running average coefficient\n        exp_avg.mul_(beta1).add_(1 - beta1, grad)\n        exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)\n        if self.amsgrad:\n            # Maintains the maximum of all 2nd moment running avg. till now\n            torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)\n            # Use the max. for normalizing running avg. of gradient\n            denom = max_exp_avg_sq.sqrt().add_(self.eps)\n        else:\n            denom = exp_avg_sq.sqrt().add_(self.eps)\n\n        bias_correction1 = 1 - beta1**state['step']\n        bias_correction2 = 1 - beta2**state['step']\n        step_size = self.learning_rate * math.sqrt(\n            bias_correction2) / bias_correction1\n\n        p.data.addcdiv_(-step_size, exp_avg, denom)\n\n    def step(self, lr, mu, theta, lm, closure=None):\n        loss = None\n        if closure is not None:\n            loss = closure()\n\n        self.learning_rate = lr\n        for w_m, theta_m, lm_m in zip(self.w.modules(), theta.modules(),\n                                      lm.modules()):\n            if hasattr(theta_m, 'condense'):\n                for pname in theta_m.condense:\n                    self._step(getattr(w_m, pname), True, mu,\n                               getattr(theta_m, pname), getattr(lm_m, pname))\n                params = set([name for name, _ in theta_m.named_parameters()])\n                rparams = params - theta_m.condense\n                for pname in rparams:\n                    self._step(getattr(w_m, pname))\n            else:\n                for w_p in w_m.parameters(recurse=False):\n                    self._step(w_p)\n        return loss\n"
  },
  {
    "path": "condensa/opt/lc/lc.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport sys\nfrom copy import deepcopy\nimport logging\nfrom collections import defaultdict\n\nimport numpy as np\nimport torch\nimport torch.utils.data as data\nimport torch.backends.cudnn as cudnn\nfrom tqdm import tqdm\n\nfrom condensa.util import EventTimer\nfrom condensa import cfg\nfrom condensa import util\nfrom .sgd import SGD\nfrom condensa.lr import *\n\nlogger = logging.getLogger(__name__)\n\nclass record_mode(object):\n    def __enter__(self):\n        cfg.__CONDENSA_RECORD_MODE__ = True\n\n    def __exit__(self, *args):\n        cfg.__CONDENSA_RECORD_MODE__ = False\n        return False\n\nclass LC(object):\n    \"\"\"Condensa L-C compression engine.\"\"\"\n    def __init__(self,\n                 steps=30,\n                 l_optimizer=None,\n                 l_optimizer_params={},\n                 lr=None,\n                 lr_end=None,\n                 lr_decay=None,\n                 lr_schedule=None,\n                 lr_multiplier=None,\n                 mb_iterations_per_l=0,\n                 mb_iterations_first_l=0,\n                 mu_init=0.,\n                 mu_multiplier=1.,\n                 mu_cap=10000,\n                 distributed=False,\n                 debugging_flags={}):\n        \"\"\"\n        Constructs an `LC` class instance.\n\n        :param steps: Number of L-C iterations.\n        :type steps: float\n        :param l_optimizer: L-step optimizer to use.\n        :param l_optimizer_params: L-step optimizer hyper-parameters.\n        :type l_optimizer_params: dict\n        :param lr: Starting learning rate.\n        :type lr: float\n        :param lr_end: Ending learning rate.\n        :type lr_end: float\n        :param lr_schedule: Learning rate schedule.\n        :type lr_schedule: List\n        :param lr_multiplier: Learning rate multiplier.\n        :type lr_multiplier: float\n        :param mb_iterations_per_l: Number of mini-batch iterations per L-step.\n        :type mb_iterations_per_l: int\n        :param mb_iterations_first_l: Number of mini-batch iterations for first L-step.\n        :type mb_iterations_first_l: int\n        :param mu_init: Initial value of `mu`.\n        :type mu_init: float\n        :param mu_multiplier: Mu multiplier.\n        :type mu_multiplier: float\n        :param mu_cap: Maximum permitted value for `mu`.\n        :type mu_cap: float\n        :param distributed: Enable/disable data-parallelism in L-step.\n        :type distributed: bool\n        :param debugging_flags: Debugging flags\n        :type debugging_flags: dict\n        \"\"\"\n        self._engine_config = {\n            k: v\n            for k, v in locals().items() if k != 'self'\n        }\n        logger.info('[Condensa] LC ENGINE CONFIG [' +\n                    ', '.join('{!s}={!r}'.format(k, v)\n                              for k, v in self._engine_config.items()) + ']')\n\n        if not 0 <= steps:\n            raise ValueError(\n                'Invalid steps specified: {}'.format(steps))\n        if not isinstance(l_optimizer_params, dict):\n            raise TypeError('l_optimizer_params must be a dictionary')\n        if not 0. <= lr:\n            raise ValueError('Invalid learning rate: {}'.format(lr))\n        if lr_schedule is not None and lr_multiplier is None:\n            raise TypeError(\n                'Please specify multiplier when using fixed LR schedule')\n        if not 0 < mb_iterations_per_l:\n            raise ValueError(\n                'Invalid mb_iterations_per_l specified: {}'.format(mb_iterations_per_l))\n        if not 0 < mb_iterations_first_l:\n            raise ValueError(\n                'Invalid mb_iterations_first_l specified: {}'.format(mb_iterations_first_l))\n        if not isinstance(debugging_flags, dict):\n            raise TypeError('debugging_flags must be a dictionary')\n\n        self.use_cuda = torch.cuda.is_available()\n        self.steps = steps\n        self.l_optimizer = l_optimizer if l_optimizer else SGD\n        self.l_optimizer_params = l_optimizer_params\n        self.lr = lr\n        self.lr_end = lr_end\n        self.lr_decay = lr_decay\n        self.lr_schedule = lr_schedule\n        self.lr_multiplier = lr_multiplier\n        self.mb_iterations_per_l = mb_iterations_per_l\n        self.mb_iterations_first_l = mb_iterations_first_l\n        self.mu_init = mu_init\n        self.mu_multiplier = mu_multiplier\n        self.mu_cap = mu_cap\n        self.distributed = distributed\n        self.debugging_flags = debugging_flags\n\n    def zero_(self, model):\n        \"\"\"\n        Zeroes out model parameters.\n    \n        :param model: PyTorch model.\n        :type model: torch.nn.Module\n        \"\"\"\n        with torch.no_grad():\n            pflat = torch.nn.utils.parameters_to_vector(\n                model.parameters()).fill_(0.)\n            torch.nn.utils.vector_to_parameters(pflat, model.parameters())\n\n    def compress(self, w, pi, delta, trainloader, testloader, valloader,\n                 loss_fn):\n        \"\"\"\n        Main L-C compression method.\n    \n        :param w: Input model.\n        :type w: torch.nn.Module\n        :param pi: Compression function.\n        :param delta: Decompression function.\n        :param trainloader: Training dataloader.\n        :param testloader: Test dataloader.\n        :param valloader: Validation dataloader.\n        :param loss_fn: Loss criterion.\n        \"\"\"\n        statistics = {}\n        # Save engine configuration\n        statistics.update(self._engine_config)\n\n        _model_stat_fn = self.debugging_flags['custom_model_statistics']\\\n                   if 'custom_model_statistics' in self.debugging_flags\\\n                   else util.empty_stat_fn\n        _disable_train_stats = self.debugging_flags['disable_train_stats']\\\n                     if 'disable_train_stats' in self.debugging_flags\\\n                     else False\n        timer_lc = EventTimer()\n\n        if self.use_cuda: cudnn.benchmark = True\n        logger.debug(\"[Condensa] cuDNN VERSION: {}\".format(cudnn.version()))\n\n        validate = (valloader is not None)\n        test = (testloader is not None)\n\n        # Copy model to GPU0 memory\n        if self.use_cuda: w = w.cuda(0)\n\n        # Mark all compressible modules in w\n        with record_mode():\n            pi(w)\n\n        with torch.no_grad():\n            theta = deepcopy(w)\n        self.zero_(theta)\n\n        with torch.no_grad():\n            lm = deepcopy(w)\n        self.zero_(lm)\n\n        with torch.no_grad():\n            best_model = deepcopy(w)\n\n        # Enable data-parallelism in  L step\n        if self.use_cuda and self.distributed:\n            ngpus = torch.cuda.device_count()\n            logger.info('[Condensa] {} GPUs enabled for L-step'.format(ngpus))\n            w = torch.nn.DataParallel(w)\n\n        mu = 0.\n        learning_rate = self.lr\n\n        optimizer = self.l_optimizer(w,\n                                     lr=learning_rate,\n                                     **self.l_optimizer_params)\n        optimizer.reset_state()\n\n        if not _disable_train_stats:\n            w_train_loss, w_train_stats = _model_stat_fn(w, loss_fn, trainloader)\n            logger.info('[Condensa] w TRAIN\\tloss={:.5f}, {}'\n                .format(w_train_loss,\n                ', '.join(['{}:{}'.format(k, v) for k,v in w_train_stats.items()])))\n        if validate:\n            w_val_loss, w_val_stats = _model_stat_fn(w, loss_fn, valloader)\n            logger.info('[Condensa] w VAL\\tloss={:.5f}, {}'\n                .format(w_val_loss,\n                ', '.join(['{}:{}'.format(k, v) for k,v in w_val_stats.items()])))\n        if test:\n            w_test_loss, w_test_stats = _model_stat_fn(w, loss_fn, testloader)\n            logger.info('[Condensa] w TEST\\tloss={:.5f}, {}'\n                .format(w_test_loss,\n                ', '.join(['{}:{}'.format(k, v) for k,v in w_test_stats.items()])))\n\n        best_loss = sys.float_info.max\n        train_losses = []\n        if validate: val_losses = []\n        if test: test_losses = []\n        outer_lr_scheduler = None\n        if self.lr_decay is not None:\n            outer_lr_scheduler = ExpDecayedLR(self.lr, self.lr_decay)\n        elif self.lr_schedule is not None:\n            outer_lr_scheduler = DecayedLR(self.lr, self.lr_schedule,\n                                           self.lr_multiplier)\n        for j in range(0, self.steps):\n            n_sgd_iter = (self.mb_iterations_first_l\n                          if j == 1 else self.mb_iterations_per_l)\n\n            # Set up outer learning rate\n            learning_rate = self.lr\n            if outer_lr_scheduler is not None:\n                learning_rate = outer_lr_scheduler.learning_rate\n\n            logger.info(\n                '[Condensa] LC Iteration {}:\\tmu={:.5f}, lr={:.5f}'.format(\n                    j, mu, learning_rate))\n\n            inner_lr_scheduler = None\n            if self.lr_end is not None:\n                inner_lr_scheduler = IntervalLR(learning_rate, self.lr_end,\n                                                n_sgd_iter)\n\n            # L step\n            # Switch to training mode\n            i = 0\n            w.train()\n            iterator = iter(trainloader)\n            if logger.isEnabledFor(logging.INFO) and j>0:\n                pbar = tqdm(total=n_sgd_iter, ascii=True)\n            while True:\n                if j == 0:\n                    logger.info('[Condensa] Skipping first L-step')\n                    break\n                if j == 1 and i >= self.mb_iterations_first_l:\n                    break\n                if j > 1 and i >= self.mb_iterations_per_l:\n                    break\n\n                try:\n                    inputs, targets = next(iterator)\n                except StopIteration:\n                    iterator = iter(trainloader)\n                    inputs, targets = next(iterator)\n\n                if self.use_cuda:\n                    if not inputs.is_cuda: inputs = inputs.cuda()\n                    if not targets.is_cuda:\n                        targets = targets.cuda(non_blocking=True)\n                outputs = w(inputs)\n                loss = loss_fn(outputs, targets)\n\n                optimizer.zero_grad()\n\n                loss.backward()\n                optimizer.step(learning_rate, mu, theta, lm)\n\n                if inner_lr_scheduler is not None:\n                    inner_lr_scheduler.step()\n                    learning_rate = inner_lr_scheduler.learning_rate\n\n                if logger.isEnabledFor(logging.INFO):\n                    pbar.update()\n                i += 1\n\n            if logger.isEnabledFor(logging.INFO) and j>0:\n                pbar.close()\n            logger.info('')\n\n            if self.use_cuda: torch.cuda.synchronize()\n\n            w.eval()\n            # C step and theta update\n            try:\n                theta.load_state_dict(w.module.state_dict())\n            except AttributeError:\n                theta.load_state_dict(w.state_dict())\n            if mu > 0:\n                try:\n                    wmodules = w.module.modules()\n                except AttributeError:\n                    wmodules = w.modules()\n                with record_mode():\n                    pi(theta)\n                with torch.no_grad():\n                    for w_m, theta_m, lm_m in zip(wmodules, theta.modules(),\n                                                  lm.modules()):\n                        if hasattr(theta_m, 'condense'):\n                            for pname in theta_m.condense:\n                                getattr(theta_m, pname).data = (\n                                    getattr(w_m, pname).detach() -\n                                    getattr(lm_m, pname).data / mu)\n\n            pi(theta)\n\n            if not _disable_train_stats:\n                nested_train_loss, nested_train_stats = _model_stat_fn(theta, loss_fn, trainloader)\n                train_losses.append(nested_train_loss)\n                logger.info(\n                    '[Condensa] Nested (theta) TRAIN\\tloss={:.5f}, {}'\n                    .format(nested_train_loss,\n                    ', '.join(['{}:{}'.format(k, v) for k,v in nested_train_stats.items()])))\n            if validate:\n                nested_val_loss, nested_val_stats = _model_stat_fn(theta, loss_fn, valloader)\n                val_losses.append(nested_val_loss)\n                logger.info(\n                    '[Condensa] Nested (theta) VAL\\tloss={:.5f}, {}'\n                    .format(nested_val_loss,\n                    ', '.join(['{}:{}'.format(k, v) for k,v in nested_val_stats.items()])))\n            if test:\n                nested_test_loss, nested_test_stats = _model_stat_fn(theta, loss_fn, testloader)\n                test_losses.append(nested_test_loss)\n                logger.info(\n                    '[Condensa] Nested (theta) TEST\\tloss={:.5f}, {}'\n                    .format(nested_test_loss,\n                    ', '.join(['{}:{}'.format(k, v) for k,v in nested_test_stats.items()])))\n\n            if validate:\n                if nested_val_loss < best_loss:\n                    logger.info('[Condensa] Saving model based on VAL')\n                    best_loss = nested_val_loss\n                    # Deep-copy required here to preserve dtypes\n                    best_model = deepcopy(theta)\n            elif test:\n                if nested_test_loss < best_loss:\n                    logger.info('[Condensa] Saving model based on TEST')\n                    best_loss = nested_test_loss\n                    # Deep-copy required here to preserve dtypes\n                    best_model = deepcopy(theta)\n            else:\n                logger.info('[Condensa] Saving model based on most recent')\n                best_model = deepcopy(theta)\n\n            # theta <- delta(theta)\n            delta(theta)\n\n            # LM update\n            if mu > 0:\n                try:\n                    wmodules = w.module.modules()\n                except AttributeError:\n                    wmodules = w.modules()\n                for w_m, theta_m, lm_m in zip(wmodules, theta.modules(),\n                                              lm.modules()):\n                    if hasattr(theta_m, 'condense'):\n                        for pname in theta_m.condense:\n                            getattr(\n                                lm_m,\n                                pname).data = (getattr(lm_m, pname).data - mu *\n                                               (getattr(w_m, pname).detach() -\n                                                getattr(theta_m, pname).data))\n\n            optimizer.reset_state()\n            # Update mu\n            mu = self._update_mu(mu, self.mu_init, self.mu_multiplier,\n                                 self.mu_cap)\n            # Update LR schedule\n            if outer_lr_scheduler is not None: outer_lr_scheduler.step()\n\n        statistics['elapsed_lc'] = timer_lc.elapsed_seconds\n        statistics['train_losses'] = train_losses\n        if test: statistics['test_losses'] = test_losses\n        if validate: statistics['val_losses'] = val_losses\n        return best_model, statistics\n\n    def _update_mu(self, mu, mu_init, mu_multiplier, mu_cap):\n        if mu > mu_cap:\n            return mu\n        if mu != 0:\n            return mu * mu_multiplier\n        else:\n            return mu_init\n"
  },
  {
    "path": "condensa/opt/lc/sgd.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nfrom collections import defaultdict\nimport torch\n\nclass SGD(object):\n    \"\"\"Custom SGD implementation for L-C optimizer.\"\"\"\n    def __init__(self, w, lr=None, momentum=None, weight_decay=0):\n        \"\"\"\n        Creates instance of `SGD`.\n\n        :param w: PyTorch model.\n        :type w: torch.nn.Module\n        :param lr: Learning rate.\n        :type lr: float\n        :param momentum: SGD momentum.\n        :type momentum: float\n        :param weight_decay: Weight decay amount (L2 regularation).\n        :type weight_decay: float\n        \"\"\"\n        if lr is None or momentum is None:\n            raise ValueError('Learning rate and momentum are required')\n        if lr < 0.0:\n            raise ValueError('Invalid learning rate: {}'.format(lr))\n        if momentum < 0.0:\n            raise ValueError('Invalid momentum value: {}'.format(momentum))\n        if weight_decay < 0.0:\n            raise ValueError(\n                'Invalid weight decay value: {}'.format(weight_decay))\n\n        try:\n            self.w = w.module\n        except AttributeError:\n            self.w = w\n\n        self.lr = lr\n        self.momentum = momentum\n        self.weight_decay = weight_decay\n\n        self.state = defaultdict(dict)\n\n    def zero_grad(self):\n        \"\"\"Zeroes out all gradients.\"\"\"\n        for p in self.w.parameters():\n            if p.grad is not None:\n                p.grad.detach_()\n                p.grad.zero_()\n\n    def reset_state(self):\n        \"\"\"Resets optimizer state.\"\"\"\n        for p in self.w.parameters():\n            if 'velocity' in self.state[p]:\n                self.state[p]['velocity'] = torch.zeros_like(p.data)\n\n    def _step(self, p, condense=False, mu=None, p_theta=None, p_lm=None):\n        if p.grad is None:\n            return\n        lr = self.learning_rate\n        d_p = p.grad.data\n        if self.weight_decay != 0:\n            d_p.add_(self.weight_decay, p.data)\n        if condense is True:\n            assert (mu is not None\n                    and p_theta is not None\n                    and p_lm is not None)\n            d_p.add_(mu * (p.data - p_theta.data) - p_lm.data)\n        update = p.data - lr * (d_p)\n        if 'velocity' not in self.state[p]:\n            velocity = torch.zeros_like(p.data)\n        else:\n            velocity = self.state[p]['velocity']\n        x = self.momentum * velocity + update - p.data\n        self.state[p]['velocity'] = x\n        p.data = self.momentum * x + update\n\n    def step(self, lr, mu, theta, lm, closure=None):\n        \"\"\"\n        Takes one optimizer step.\n\n        :param lr: Current learning rate.\n        :type lr: float\n        :param mu: L-C mu hyper-parameter value.\n        :type mu: float\n        :param theta: Compressed model.\n        :type theta: torch.nn.Module\n        :param lm: Lagrange multiplier.\n        :type lm: torch.nn.Module\n        :param closure: Loss closure.\n        \"\"\"\n        loss = None\n        if closure is not None:\n            loss = closure()\n\n        self.learning_rate = lr\n        for w_m, theta_m, lm_m in zip(self.w.modules(), theta.modules(),\n                                      lm.modules()):\n            if hasattr(theta_m, 'condense'):\n                for pname in theta_m.condense:\n                    self._step(getattr(w_m, pname), True, mu,\n                               getattr(theta_m, pname), getattr(lm_m, pname))\n                params = set([name for name, _ in theta_m.named_parameters()])\n                rparams = params - theta_m.condense\n                for pname in rparams:\n                    self._step(getattr(w_m, pname))\n            else:\n                for w_p in w_m.parameters(recurse=False):\n                    self._step(w_p)\n        return loss\n"
  },
  {
    "path": "condensa/pi.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport torch\nimport torch.nn\n\nfrom condensa import dtypes\nfrom condensa import cfg\nimport condensa.tensor as T\n\ndef __precheck(module):\n    if not cfg.__CONDENSA_PI_PRECHECK__: return\n\n    if len(list(module.children())) > 0:\n        raise RuntimeError('Only leaf modules may be compressed')\n    for name, _ in module.named_parameters():\n        if name != 'weight' and name != 'bias':\n            raise NotImplementedError(\n                'Unknown parameter {} detected'.format(name))\n\ndef quantize(module, dtype):\n    \"\"\"\n    Quantizes module to given data type (inplace).\n\n    :param module: PyTorch module.\n    :type module: `torch.nn.Module`\n    :param dtype: Target data type.\n    \"\"\"\n    __precheck(module)\n\n    parameters = ['weight']\n    #parameters = [name for name, _ in module.named_parameters()]\n    if hasattr(module, 'condense'): module.condense |= set(parameters)\n    else: module.condense = set(parameters)\n\n    if not cfg.__CONDENSA_RECORD_MODE__:\n        if dtype.as_dtype_enum == dtypes.DT_FLOAT16:\n            module.half()\n        elif dtype.as_dtype_enum == dtypes.DT_FLOAT32:\n            module.float()\n        else:\n            raise TypeError('Unknown data type specified for quantization')\n\ndef prune(module, threshold, parameter='weight'):\n    \"\"\"\n    Prunes module parameters based on magnitude (inplace).\n\n    :param module: PyTorch module.\n    :type module: `torch.nn.Module`\n    :param threshold: Magnitude threshold for pruning.\n    :type threshold: `float`\n    :param parameter: Module parameter to prune (default: 'weight')\n    :type parameter: str\n    \"\"\"\n    __precheck(module)\n\n    if not hasattr(module, parameter):\n        raise ValueError('Could not find parameter \\'{}\\' in module',\n                         parameter)\n\n    if hasattr(module, 'condense'): module.condense.add(parameter)\n    else: module.condense = set([parameter])\n\n    if not cfg.__CONDENSA_RECORD_MODE__:\n        p = getattr(module, parameter)\n        pdata = p.data.view(-1)\n        mask = T.simple_mask(pdata, threshold).type(pdata.type())\n        T.apply_mask_inplace(pdata, mask)\n        p.data = pdata.view_as(p).data\n        #if cfg.__CONDENSA_SAVE_MASK__: module.mask = mask.view_as(p).data\n\ndef blockprune(module,\n               threshold,\n               block_size,\n               criteria,\n               align=None,\n               parameter='weight'):\n    \"\"\"\n    Prunes blocks of module parameters based on magnitude (inplace).\n\n    :param module: PyTorch module.\n    :type module: `torch.nn.Module`\n    :param threshold: Magnitude threshold for pruning.\n    :type threshold: `float`\n    :param block_size: Block size for pruning.\n    :type block_size: `Tuple`\n    :param criteria: Aggregation function for thresholding.\n    :type criteria: `condensa.functional`\n    :param align: Alignment of compressed parameters.\n    :type align: `int`\n    :param parameter: Module parameter to prune (default: 'weight')\n    :type parameter: str\n    \"\"\"\n    __precheck(module)\n\n    if not hasattr(module, parameter):\n        raise ValueError('Could not find parameter \\'{}\\' in module',\n                         parameter)\n\n    p = getattr(module, parameter)\n    ndim = p.dim()\n    bdim = len(block_size)\n    if ndim != bdim:\n        raise RuntimeError(\n            'Block must have same dimensions as parameter \\'{}\\''.format(\n                parameter))\n\n    if hasattr(module, 'condense'): module.condense.add(parameter)\n    else: module.condense = set([parameter])\n\n    if not cfg.__CONDENSA_RECORD_MODE__:\n        mask = T.block_mask(p.data, threshold, block_size, criteria, align)\n        T.apply_mask_inplace(p.data, mask)\n        return mask\n    return None\n\ndef neuron_prune(module, threshold, criteria, align=None, prune_bias=True):\n    \"\"\"\n    Prunes neurons based on magnitude (inplace).\n\n    :param module: PyTorch module.\n    :type module: `torch.nn.Module`\n    :param threshold: Magnitude threshold for pruning.\n    :type threshold: `float`\n    :param criteria: Aggregation function for thresholding.\n    :type criteria: `condensa.functional`\n    :param align: Alignment of compressed parameters.\n    :type align: `int`\n    :param prune_bias: Whether to prune corresponding biases.\n    :type prune_bias: `bool`\n    \"\"\"\n    __precheck(module)\n\n    parameter = 'weight'\n    if not hasattr(module, parameter):\n        raise ValueError('Could not find parameter \\'{}\\' in module',\n                         parameter)\n\n    shape = getattr(module, parameter).data.shape\n    if len(shape) != 2:\n        raise NotImplementedError(\n            'Row pruning currently only supported for 2D parameters')\n\n    if hasattr(module, 'condense'): module.condense.add(parameter)\n    else: module.condense = set([parameter])\n\n    if not cfg.__CONDENSA_RECORD_MODE__:\n        block_size = (1, shape[1])\n        mask = blockprune(module, threshold, block_size, criteria, align,\n                          parameter)\n        # Prune corresponding bias tensor\n        if module.bias is not None and prune_bias is True:\n            assert mask.ndimension() == 2\n            T.apply_mask_inplace(module.bias.data, mask[:, 0])\n\ndef filter_prune(module, threshold, criteria, align=None, prune_bias=True):\n    \"\"\"\n    Prunes 3D blocks (filters) of module parameters based on magnitude (inplace).\n\n    :param module: PyTorch module.\n    :type module: `torch.nn.Module`\n    :param threshold: Magnitude threshold for pruning.\n    :type threshold: `float`\n    :param criteria: Aggregation function for thresholding.\n    :type criteria: `condensa.functional`\n    :param align: Alignment of compressed parameters.\n    :type align: `int`\n    :param prune_bias: Whether to prune corresponding biases.\n    :type prune_bias: `bool`\n    \"\"\"\n    __precheck(module)\n\n    parameter = 'weight'\n    if not hasattr(module, parameter):\n        raise ValueError('Could not find parameter \\'{}\\' in module',\n                         parameter)\n\n    p = getattr(module, parameter)\n    ndim = p.dim()\n    if ndim != 4:\n        raise RuntimeError('Filter pruning requires a 4D parameter')\n\n    if hasattr(module, 'condense'): module.condense.add(parameter)\n    else: module.condense = set([parameter])\n\n    if not cfg.__CONDENSA_RECORD_MODE__:\n        block_size = (1, *p.data.shape[1:])\n        mask = T.block_mask(p.data, threshold, block_size, criteria, align)\n        T.apply_mask_inplace(p.data, mask)\n        # Prune corresponding bias tensor\n        if module.bias is not None and prune_bias is True:\n            assert mask.ndimension() == 4\n            T.apply_mask_inplace(module.bias.data, mask[:, 0, 0, 0])\n"
  },
  {
    "path": "condensa/schemes.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport torch\n\nimport condensa\nimport condensa.tensor as T\nimport condensa.functional as F\n\nclass Compose(object):\n    \"\"\"Composes two or more schemes together.\"\"\"\n    def __init__(self, schemes):\n        \"\"\"\n        Creates a `Compose` instance.\n\n        :param schemes: List of schemes to compose.\n        :type schemes: `list`\n        \"\"\"\n        if not isinstance(schemes, list):\n            raise TypeError('Please specify schemes to compose as a list')\n        self.schemes = schemes\n\n    def pi(self, module):\n        \"\"\"\n        Applies compression scheme to module.\n    \n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        for s in self.schemes:\n            s.pi(module)\n\n    def delta(self, module):\n        \"\"\"\n        Applies de-compression scheme to module.\n\n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        for s in reversed(self.schemes):\n            s.delta(module)\n\n    def __repr__(self):\n        return '<Compose: {}>'.format(self.schemes)\n\nclass Prune(object):\n    \"\"\"Prunes network to given density.\"\"\"\n    def __init__(self, density):\n        \"\"\"\n        Creates a `Prune` instance.\n\n        :param density: Target density.\n        :type density: `float`\n        \"\"\"\n        self.density = density\n        self.layer_types = [torch.nn.Linear, torch.nn.Conv2d]\n\n    def threshold(self, module):\n        \"\"\"\n        Computes magnitude threshold.\n\n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        vec = []\n        for m in module.modules():\n            if type(m) in self.layer_types and not hasattr(\n                    m, 'condensa_nocompress'):\n                vec.append(m.weight.data.view(-1))\n        return T.threshold(torch.cat(vec), self.density)\n\n    def pi(self, module):\n        \"\"\"\n        Applies compression scheme to module.\n    \n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        threshold = self.threshold(module)\n        for m in module.modules():\n            if type(m) in self.layer_types and not hasattr(\n                    m, 'condensa_nocompress'):\n                condensa.prune(m, threshold)\n\n    def delta(self, module):\n        \"\"\"\n        Applies de-compression scheme to module.\n\n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        pass\n\n    def __repr__(self):\n        return '<Prune: density:{}>'.format(self.density)\n\nclass Quantize(object):\n    \"\"\"Quantizes network to given data-type.\"\"\"\n    def __init__(self, dtype=condensa.float16):\n        \"\"\"\n        Creates `Quantize` class instance.\n\n        :param dtype: Target data type (default: float16).\n        \"\"\"\n        self.dtype = dtype\n        self.layer_types = [torch.nn.Linear, torch.nn.Conv2d]\n\n    def pi(self, module):\n        \"\"\"\n        Applies compression scheme to module.\n    \n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        for m in module.modules():\n            if type(m) in self.layer_types and not hasattr(\n                    m, 'condensa_nocompress'):\n                condensa.quantize(m, self.dtype)\n\n    def delta(self, module):\n        \"\"\"\n        Applies de-compression scheme to module.\n\n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        for m in module.modules():\n            if type(m) in self.layer_types and not hasattr(\n                    m, 'condensa_nocompress'):\n                condensa.dequantize(m, condensa.float32)\n\n    def __repr__(self):\n        return '<Quantize: dtype:{}>'.format(self.dtype)\n\nclass NeuronPrune(object):\n    \"\"\"Prunes neurons from fully-connected layers.\"\"\"\n    def __init__(self, density, align=None, criteria=F.l2norm,\n                 prune_bias=True):\n        \"\"\"\n        Creates an instance of `NeuronPrune`.\n\n        :param density: Target density.\n        :type density: `float`\n        :param align: Tensor alignment in compressed model.\n        :type align: `int`\n        :param criteria: Neuron aggregation criteria (default: l2norm).\n        :type criteria: `condensa.functional`\n        :param prune_bias: Whether to prune corresponding biases (default: True).\n        :type prune_bias: `bool`\n        \"\"\"\n        self.density = density\n        self.align = align\n        self.criteria = criteria\n        self.prune_bias = prune_bias\n\n    def threshold(self, module):\n        \"\"\"\n        Computes magnitude threshold.\n\n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        vec = []\n        for m in module.modules():\n            if isinstance(m, torch.nn.Linear) and not hasattr(\n                    m, 'condensa_nocompress'):\n                agg = T.aggregate_neurons(m.weight.data, self.criteria)\n                vec.append(agg.view(-1))\n        return T.threshold(torch.cat(vec), self.density)\n\n    def pi(self, module):\n        \"\"\"\n        Applies compression scheme to module.\n    \n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        threshold = self.threshold(module)\n        for m in module.modules():\n            if isinstance(m, torch.nn.Linear) and not hasattr(\n                    m, 'condensa_nocompress'):\n                condensa.neuron_prune(m,\n                                      threshold,\n                                      align=self.align,\n                                      criteria=self.criteria,\n                                      prune_bias=self.prune_bias)\n\n    def delta(self, module):\n        \"\"\"\n        Applies de-compression scheme to module.\n\n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        pass\n\n    def __repr__(self):\n        return '<NeuronPrune: density:{}, align:{}, criteria:{}, prune_bias:{}>'.format(\n            self.density, self.align, self.criteria, self.prune_bias)\n\nclass FilterPrune(object):\n    \"\"\"Prunes filters from convolutional layers.\"\"\"\n    def __init__(self, density, align=None, criteria=F.l2norm,\n                 prune_bias=True):\n        \"\"\"\n        Creates an instance of `FilterPrune`.\n\n        :param density: Target density.\n        :type density: `float`\n        :param align: Tensor alignment in compressed model.\n        :type align: `int`\n        :param criteria: Filter aggregation criteria (default: l2norm).\n        :type criteria: `condensa.functional`\n        :param prune_bias: Whether to prune corresponding biases (default: True).\n        :type prune_bias: `bool`\n        \"\"\"\n        self.density = density\n        self.align = align\n        self.criteria = criteria\n        self.prune_bias = prune_bias\n\n    def threshold(self, module):\n        \"\"\"\n        Computes magnitude threshold.\n\n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        vec = []\n        for m in module.modules():\n            if isinstance(m, torch.nn.Conv2d) and not hasattr(\n                    m, 'condensa_nocompress'):\n                agg = T.aggregate_filters(m.weight.data, self.criteria)\n                vec.append(agg.view(-1))\n        return T.threshold(torch.cat(vec), self.density)\n\n    def pi(self, module):\n        \"\"\"\n        Applies compression scheme to module.\n    \n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        threshold = self.threshold(module)\n        for m in module.modules():\n            if isinstance(m, torch.nn.Conv2d) and not hasattr(\n                    m, 'condensa_nocompress'):\n                condensa.filter_prune(m,\n                                      threshold,\n                                      align=self.align,\n                                      criteria=self.criteria,\n                                      prune_bias=self.prune_bias)\n\n    def delta(self, module):\n        \"\"\"\n        Applies de-compression scheme to module.\n\n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        pass\n\n    def __repr__(self):\n        return '<FilterPrune: density:{}, align:{}, criteria:{}, prune_bias:{}>'.format(\n            self.density, self.align, self.criteria, self.prune_bias)\n\nclass StructurePrune(object):\n    \"\"\"Combines neuron and filter pruning using a single threshold value.\"\"\"\n    def __init__(self, density, align=None, criteria=F.l2norm,\n                 prune_bias=True):\n        \"\"\"\n        Creates an instance of `StructurePrune`.\n\n        :param density: Target density.\n        :type density: `float`\n        :param align: Tensor alignment in compressed model.\n        :type align: `int`\n        :param criteria: Structure aggregation criteria (default: l2norm).\n        :type criteria: `condensa.functional`\n        :param prune_bias: Whether to prune corresponding biases (default: True).\n        :type prune_bias: `bool`\n        \"\"\"\n        self.density = density\n        self.align = align\n        self.criteria = criteria\n        self.prune_bias = prune_bias\n\n    def threshold(self, module):\n        \"\"\"\n        Computes magnitude threshold.\n\n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        vec = []\n        for m in module.modules():\n            if isinstance(m, torch.nn.Linear) and not hasattr(\n                    m, 'condensa_nocompress'):\n                agg = T.aggregate_neurons(m.weight.data, self.criteria)\n                vec.append(agg.view(-1))\n            if isinstance(m, torch.nn.Conv2d) and not hasattr(\n                    m, 'condensa_nocompress'):\n                agg = T.aggregate_filters(m.weight.data, self.criteria)\n                vec.append(agg.view(-1))\n        return T.threshold(torch.cat(vec), self.density)\n\n    def pi(self, module):\n        \"\"\"\n        Applies compression scheme to module.\n    \n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        threshold = self.threshold(module)\n        for m in module.modules():\n            if isinstance(m, torch.nn.Linear) and not hasattr(\n                    m, 'condensa_nocompress'):\n                condensa.neuron_prune(m,\n                                      threshold,\n                                      align=self.align,\n                                      criteria=self.criteria,\n                                      prune_bias=self.prune_bias)\n            if isinstance(m, torch.nn.Conv2d) and not hasattr(\n                    m, 'condensa_nocompress'):\n                condensa.filter_prune(m,\n                                      threshold,\n                                      align=self.align,\n                                      criteria=self.criteria,\n                                      prune_bias=self.prune_bias)\n\n    def delta(self, module):\n        \"\"\"\n        Applies de-compression scheme to module.\n\n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        pass\n\n    def __repr__(self):\n        return '<StructurePrune: density:{}, align:{}, criteria:{}, prune_bias:{}>'.format(\n            self.density, self.align, self.criteria, self.prune_bias)\n\nclass BlockPrune(object):\n    \"\"\"Prunes blocks in Linear layers.\"\"\"\n    def __init__(self, density, block_size, criteria=F.l2norm):\n        \"\"\"\n        Creates an instance of `BlockPrune`.\n\n        :param density: Target density.\n        :type density: `float`\n        :param block_size: Target block size.\n        :type block_size: `Tuple`\n        :param criteria: Structure aggregation criteria (default: l2norm).\n        :type criteria: `condensa.functional`\n        \"\"\"\n        self.density = density\n        self.block_size = block_size\n        self.criteria = criteria\n        self.layer_types = [torch.nn.Linear]\n\n    def threshold(self, module):\n        \"\"\"\n        Computes magnitude threshold.\n\n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        vec = []\n        for m in module.modules():\n            if type(m) in self.layer_types and not hasattr(\n                    m, 'condensa_nocompress'):\n                agg = T.aggregate(m.weight.data, self.block_size,\n                                  self.criteria)\n                vec.append(agg.view(-1))\n        return T.threshold(torch.cat(vec), self.density)\n\n    def pi(self, module):\n        \"\"\"\n        Applies compression scheme to module.\n    \n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        threshold = self.threshold(module)\n        for m in module.modules():\n            if type(m) in self.layer_types and not hasattr(\n                    m, 'condensa_nocompress'):\n                condensa.blockprune(m,\n                                    threshold,\n                                    block_size=self.block_size,\n                                    criteria=self.criteria)\n\n    def delta(self, module):\n        \"\"\"\n        Applies de-compression scheme to module.\n\n        :param module: PyTorch module.\n        :type module: `torch.nn.Module`\n        \"\"\"\n        pass\n\n    def __repr__(self):\n        return '<BlockPrune: density:{}, block_size:{}>'.format(\n            self.density, self.block_size)\n"
  },
  {
    "path": "condensa/tensor.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport numpy as np\nimport torch\n\ndef density(tensor):\n    \"\"\"\n    Computes the ratio of nonzeros to total elements in a tensor.\n\n    :param tensor: PyTorch tensor\n    :type tensor: `torch.Tensor`\n    :return: Ratio of nonzeros to total elements\n    :rtype: `float`\n    \"\"\"\n    t = tensor.view(-1)\n    return float(t.nonzero().numel()) / float(t.numel())\n\ndef sparsity(tensor):\n    \"\"\"\n    Computes the ratio of zeros to total elements in a tensor.\n\n    :param tensor: PyTorch tensor\n    :type tensor: torch.Tensor\n    :return: Ratio of zeros to total elements\n    :rtype: `float`\n    \"\"\"\n    return 1. - density(tensor)\n\ndef threshold(tensor, density):\n    \"\"\"\n    Computes a magnitude-based threshold for given tensor.\n\n    :param tensor: PyTorch tensor\n    :type tensor: `torch.Tensor`\n    :param density: Desired ratio of nonzeros to total elements\n    :type density: `float`\n    :return: Magnitude threshold\n    :rtype: `float`\n    \"\"\"\n    tf = tensor.abs().view(-1)\n    numel = int(density * tf.numel())\n    if numel == 0:\n        raise RuntimeError('Provided density value causes model to be zero.')\n\n    topk, _ = torch.topk(tf.abs(), numel, sorted=True)\n    return topk.data[-1]\n\ndef aggregate(tensor, blocksize, criteria):\n    \"\"\"\n    Aggregates tensor dimensions according to criteria.\n\n    :param tensor: PyTorch tensor\n    :type tensor: `torch.Tensor`\n    :param blocksize: Size of blocks to aggregate\n    :type blocksize: `Tuple(int)`\n    :param criteria: Aggregation criteria\n    :type criteria: `condensa.functional`\n    :return: Aggregated tensor\n    :rtype: `torch.Tensor`\n    \"\"\"\n    if tensor.dim() != len(blocksize):\n        raise RuntimeError('Tensor and block dimensions do not match')\n    ndim = tensor.dim()\n\n    blocksize_flat = np.prod(np.array(blocksize))\n    shape = np.array(tensor.shape)\n    repeats = (shape / blocksize).astype(int)\n    divcheck = (shape % blocksize).astype(int)\n\n    if not np.all(divcheck == 0):\n        raise TypeError('Block size must be divisible by tensor size')\n\n    tmpshape = np.column_stack([repeats, blocksize]).ravel()\n    order = np.arange(len(tmpshape))\n    order = np.concatenate([order[::2], order[1::2]])\n    blocks = tensor.abs().reshape(tuple(tmpshape))\n    blocks = blocks.permute(tuple(order)).reshape(-1, *blocksize)\n    agg = criteria(blocks.reshape(-1, blocksize_flat), dim=1, keepdim=True)\n    return agg\n\ndef aggregate_neurons(tensor, criteria):\n    \"\"\"\n    Aggregates neurons (rows) in given weight matrix.\n  \n    :param tensor: PyTorch tensor\n    :type tensor: `torch.Tensor`\n    :param criteria: Aggregation criteria\n    :type criteria: `condensa.functional`\n    :return: Neuron-aggregated tensor\n    :rtype: `torch.Tensor`\n    \"\"\"\n    return aggregate(tensor, (1, tensor.shape[1]), criteria)\n\ndef aggregate_filters(tensor, criteria):\n    \"\"\"\n    Aggregates 3D filters in given weight tensor.\n  \n    :param tensor: PyTorch tensor\n    :type tensor: `torch.Tensor`\n    :param criteria: Aggregation criteria\n    :type criteria: `condensa.functional`\n    :return: Filter-aggregated tensor\n    :rtype: `torch.Tensor`\n    \"\"\"\n    return aggregate(tensor, (1, *tensor.shape[1:]), criteria)\n\ndef simple_mask(tensor, threshold, align=None):\n    \"\"\"\n    Computes a simple binary mask for given magnitude threshold.\n\n    :param tensor: PyTorch tensor\n    :type tensor: `torch.Tensor`\n    :param threshold: magnitude threshold for pruning\n    :type threshold: `float`\n    :return: Mask\n    :rtype: `torch.Tensor`\n    \"\"\"\n    assert tensor.dim() == 1\n    if align is None:\n        return torch.ge(tensor.abs(), threshold)\n    else:\n        size = tensor.size(0)\n        if size < align:\n            raise RuntimeError('Tensor too small for given alignment')\n        t = tensor.abs()\n        nnz = torch.ge(t, threshold).nonzero().size(0)\n        nnz = int(nnz / align) * align\n        _, indices = torch.topk(t, nnz)\n        ones = torch.ones(nnz,\n                          dtype=tensor.dtype,\n                          layout=tensor.layout,\n                          device=tensor.device)\n        mask = torch.zeros_like(tensor).scatter_(0, indices, ones)\n        return mask\n\ndef block_mask(tensor, threshold, blocksize, criteria, align=None):\n    \"\"\"\n    Computes an n-D binary mask for given magnitude threshold.\n\n    :param tensor: PyTorch tensor\n    :type tensor: `torch.Tensor`\n    :param threshold: magnitude threshold for pruning\n    :type threshold: `float`\n    :param blocksize: desired block size (Tuple)\n    :type blocksize: `Tuple`\n    :param criteria: aggregation function for thresholding (default: max)\n    :type criteria: `condensa.functional`\n    :return: Mask\n    :rtype: `torch.Tensor`\n    \"\"\"\n    # Original implementation at: https://stackoverflow.com/questions/42297115\n    # /numpy-split-cube-into-cubes/42298440#42298440\n    if tensor.dim() != len(blocksize):\n        raise RuntimeError('Tensor and block dimensions do not match')\n    ndim = tensor.dim()\n\n    blocksize_flat = np.prod(np.array(blocksize))\n    shape = np.array(tensor.shape)\n    repeats = (shape / blocksize).astype(int)\n    divcheck = (shape % blocksize).astype(int)\n\n    if not np.all(divcheck == 0):\n        raise TypeError('Block size must be divisible by tensor size')\n\n    tmpshape = np.column_stack([repeats, blocksize]).ravel()\n    order = np.arange(len(tmpshape))\n    order = np.concatenate([order[::2], order[1::2]])\n    blocks = tensor.abs().reshape(tuple(tmpshape))\n    blocks = blocks.permute(tuple(order)).reshape(-1, *blocksize)\n    agg = criteria(blocks.reshape(-1, blocksize_flat), dim=1, keepdim=True)\n\n    mask = simple_mask(agg.view(-1), threshold, align)\n    mask = mask.view(agg.shape).expand(-1,\n                                       blocksize_flat).reshape(blocks.shape)\n\n    N, newshape = mask.shape[0], mask.shape[1:]\n    repeats = (shape / newshape).astype(int)\n    tmpshape = np.concatenate([repeats, newshape])\n    order = np.arange(len(tmpshape)).reshape(2, -1).ravel(order='F')\n    return mask.reshape(tuple(tmpshape)).permute(tuple(order)).reshape(\n        tuple(shape))\n\ndef apply_mask(tensor, mask):\n    \"\"\"\n    Computes masked version of tensor.\n\n    :param tensor: PyTorch tensor\n    :type tensor: `torch.Tensor`\n    :param mask: Binary mask\n    :type mask: `torch.Tensor`\n    :return: Masked version of `tensor`\n    :rtype: `torch.Tensor`\n    \"\"\"\n    #assert isinstance(tensor, torch.Tensor)\n    return torch.mul(tensor, mask.type(tensor.type()))\n\ndef apply_mask_inplace(tensor, mask):\n    \"\"\"\n    Applies binary mask in-place.\n\n    :param tensor: PyTorch tensor\n    :type tensor: `torch.Tensor`\n    :param mask: Binary mask\n    :type mask: `torch.Tensor`\n    \"\"\"\n    #assert isinstance(tensor, torch.Tensor)\n    tensor.mul_(mask.type(tensor.type()))\n"
  },
  {
    "path": "condensa/type_enums.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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# Supported data types\nDT_FLOAT16 = 1\nDT_FLOAT32 = 2\nDT_FLOAT64 = 3\nDT_INT8 = 4\nDT_UINT8 = 5\nDT_INT16 = 6\nDT_UINT16 = 7\n"
  },
  {
    "path": "condensa/util.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport time\nimport sys\nimport numpy as np\nimport logging\nfrom tqdm import tqdm\n\nimport torch.nn.utils\nimport torch.utils.data as data\nfrom torch.autograd import Variable\n\nimport condensa.tensor as T\n\nlogger = logging.getLogger(__name__)\n\nclass AverageMeter(object):\n    \"\"\"Computes and stores the average and current value\"\"\"\n    def __init__(self):\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\ndef to_python_float(t):\n    if hasattr(t, 'item'):\n        return t.item()\n    else:\n        return t[0]\n\ndef is_leaf_node(module):\n    \"\"\"\n    Checks if given module is a leaf module.\n\n    :param module: PyTorch module\n    :type module: `torch.nn.Module`\n    :return: Boolean value representing whether module is a leaf.\n    :rtype: `bool`\n    \"\"\"\n    return list(module.children()) == []\n\ndef magnitude_threshold(module, density):\n    \"\"\"\n    Computes a magnitude-based threshold for given module.\n\n    :param module: PyTorch module\n    :type module: `torch.nn.Module`\n    :param density: Desired ratio of nonzeros to total elements\n    :type density: `float`\n    :return: Magnitude threshold\n    :rtype: `float`\n    \"\"\"\n    params = torch.nn.utils.parameters_to_vector(module.parameters())\n    return T.threshold(params, density)\n\ndef empty_stat_fn(model, criterion, dataloader):\n    \"\"\"\n    Empty model statistics function: returns loss.\n\n    :param model: PyTorch model\n    :type model: `torch.nn.Module`\n    :param loss_fn: Loss function\n    :param dataloader: Data loader to use\n    :return: Tuple of loss, dictionary of statistics\n    :rtype: `Tuple(float, dict)`\n    \"\"\"\n    return (loss(model, criterion, dataloader), {})\n\ndef accuracy(output, target, topk=(1, )):\n    \"\"\"\n    Computes the precision@k for the specified values of k\n\n    :param output: Predicted output batch\n    :type output: `torch.Tensor`\n    :param target: Actual output batch\n    :type target: `torch.Tensor`\n    :param topk: Top-k value\n    :type topk: `Tuple`\n    :return: Model accuracy\n    :rtype: `float`\n    \"\"\"\n    maxk = max(topk)\n    batch_size = target.size(0)\n\n    _, pred = output.topk(maxk, 1, True, True)\n    pred = pred.t()\n    correct = pred.eq(target.view(1, -1).expand_as(pred))\n\n    res = []\n    for k in topk:\n        correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)\n        res.append(correct_k.mul_(100.0 / batch_size))\n    return res\n\ndef loss(model, criterion, dataloader):\n    \"\"\"\n    Computes loss on given dataset.\n  \n    :param model: PyTorch model\n    :type model: `torch.nn.Module`\n    :param loss_fn: Loss function\n    :param dataloader: Data loader to use\n    :return: Loss\n    :rtype: `float`\n    \"\"\"\n    losses = AverageMeter()\n    model.eval()\n    pzero = list(model.parameters())[0]\n    if (pzero.dtype != torch.float32 and pzero.dtype != torch.float16):\n        raise NotImplementedError('Only FP16 and FP32 weights are supported')\n    cast2fp16 = (isinstance(pzero, torch.HalfTensor)\n                 or isinstance(pzero, torch.cuda.HalfTensor))\n    loss = 0.\n    with torch.no_grad():\n        for input, target in dataloader:\n            if torch.cuda.is_available():\n                input = input.cuda(non_blocking=True)\n                target = target.cuda(non_blocking=True)\n            if cast2fp16:\n                input = input.half()\n            output = model(input)\n            loss = criterion(output, target)\n            losses.update(to_python_float(loss.data), input.size(0))\n    return losses.avg\n\ndef cnn_statistics(model, criterion, dataloader):\n    \"\"\"\n    Computes accuracy of given CNN model.\n  \n    :param model: PyTorch model\n    :type model: `torch.nn.Module`\n    :param criterion: Loss function\n    :param dataloader: Data loader to use\n    :return: Top-1 and Top-5 accuracies\n    :rtype: Tuple(top1, top5)\n    \"\"\"\n    losses = AverageMeter()\n    top1 = AverageMeter()\n    top5 = AverageMeter()\n\n    model.eval()\n    pzero = list(model.parameters())[0]\n    if (pzero.dtype != torch.float32 and pzero.dtype != torch.float16):\n        raise NotImplementedError('Only FP16 and FP32 weights are supported')\n    cast2fp16 = (isinstance(pzero, torch.HalfTensor)\n                 or isinstance(pzero, torch.cuda.HalfTensor))\n    loss = 0.\n    correct = 0.\n    with torch.no_grad():\n        for input, target in dataloader:\n            if torch.cuda.is_available():\n                input = input.cuda(non_blocking=True)\n                target = target.cuda(non_blocking=True)\n            if cast2fp16:\n                input = input.half()\n            output = model(input)\n            loss = criterion(output, target)\n            prec1, prec5 = accuracy(output.data, target, topk=(1, 5))\n            losses.update(to_python_float(loss.data), input.size(0))\n            top1.update(to_python_float(prec1), input.size(0))\n            top5.update(to_python_float(prec5), input.size(0))\n    return (losses.avg,\n            {'top1': top1.avg, 'top5': top5.avg})\n\ndef compressed_model_stats(w, wc):\n    \"\"\"\n    Retrieve various statistics for compressed model.\n\n    :param w: Original model\n    :type w: `torch.nn.Module`\n    :param wc: Compressed model\n    :type wc: `torch.nn.Module`\n    :return: Dictionary of compressed model statistics\n    :rtype: `dict`\n    \"\"\"\n    stats = dict()\n    nparams_w = dict()\n    nparams_wc = dict()\n\n    nparams_w['total_nnz'] = torch.nn.utils.parameters_to_vector(\n        w.parameters()).view(-1).nonzero().numel()\n    nparams_wc['total_nnz'] = torch.nn.utils.parameters_to_vector(\n        wc.parameters()).view(-1).nonzero().numel()\n\n    for (name_w, m_w), (name_wc, m_wc) in zip(w.named_modules(),\n                                              wc.named_modules()):\n        if type(m_w) == torch.nn.Linear or type(m_w) == torch.nn.Conv2d:\n            nparams_w[name_w] = torch.nn.utils.parameters_to_vector(\n                m_w.parameters()).view(-1).nonzero().numel()\n            nparams_wc[name_wc] = torch.nn.utils.parameters_to_vector(\n                m_wc.parameters()).view(-1).nonzero().numel()\n\n    stats['num_params'] = nparams_w\n    stats['num_params_compressed'] = nparams_wc\n    return stats\n\ndef pretrain(epochs, model, trainloader, criterion, optimizer):\n    \"\"\"\n    No-frills pre-training method.\n  \n    :param epochs: Number of epochs\n    :type epochs: `int`\n    :param model: PyTorch model\n    :type model: `torch.nn.Module`\n    :param trainloader: Training dataloader\n    :param criterion: Loss criterion\n    :param optimizer: Optimizer to use\n    \"\"\"\n    _config = {'epochs': epochs}\n    logging.info('[Condensa] PRETRAIN CONFIG [' +\n                 ', '.join('{!s}={!r}'.format(k, v)\n                           for k, v in _config.items()) + ']')\n\n    use_cuda = torch.cuda.is_available()\n    if use_cuda:\n        model = model.cuda()\n        model = torch.nn.DataParallel(model)\n    mb_iterator = iter(trainloader)\n    model.train()\n    for j in range(0, epochs):\n        if logger.isEnabledFor(logging.INFO):\n            pbar = tqdm(total=len(trainloader),\n                        ascii=True,\n                        desc='Epoch {}'.format(j))\n        for input, target in trainloader:\n            if torch.cuda.is_available():\n                input = input.cuda(non_blocking=True)\n                target = target.cuda(non_blocking=True)\n            input, target = Variable(input), Variable(target)\n            optimizer.zero_grad()\n            output = model(input)\n            loss = criterion(output, target)\n            loss.backward()\n            optimizer.step()\n            if logger.isEnabledFor(logging.INFO):\n                pbar.update()\n    if logger.isEnabledFor(logging.INFO):\n        pbar.close()\n    logging.info('')\n\nclass EventTimer(object):\n    \"\"\"Simple timer class.\"\"\"\n    def __init__(self):\n        \"\"\"Constructor. Begins timing.\"\"\"\n        self.begin = time.perf_counter()\n\n    def reset(self):\n        \"\"\"Reset timer.\"\"\"\n        self.begin = time.perf_counter()\n\n    @property\n    def elapsed_seconds(self):\n        \"\"\"Returns elapsed seconds.\"\"\"\n        return (time.perf_counter() - self.begin)\n"
  },
  {
    "path": "docs/Makefile",
    "content": "# Minimal makefile for Sphinx documentation\n#\n\n# You can set these variables from the command line.\nSPHINXOPTS    =\nSPHINXBUILD   = sphinx-build\nSPHINXPROJ    = Condensa\nSOURCEDIR     = source\nBUILDDIR      = build\n\n# Put it first so that \"make\" without argument is like \"make help\".\nhelp:\n\t@$(SPHINXBUILD) -M help \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)\n\n.PHONY: help Makefile\n\n# Catch-all target: route all unknown targets to Sphinx using the new\n# \"make mode\" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).\n%: Makefile\n\t@$(SPHINXBUILD) -M $@ \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)"
  },
  {
    "path": "docs/make.bat",
    "content": "@ECHO OFF\r\n\r\npushd %~dp0\r\n\r\nREM Command file for Sphinx documentation\r\n\r\nif \"%SPHINXBUILD%\" == \"\" (\r\n\tset SPHINXBUILD=sphinx-build\r\n)\r\nset SOURCEDIR=source\r\nset BUILDDIR=build\r\nset SPHINXPROJ=Condensa\r\n\r\nif \"%1\" == \"\" goto help\r\n\r\n%SPHINXBUILD% >NUL 2>NUL\r\nif errorlevel 9009 (\r\n\techo.\r\n\techo.The 'sphinx-build' command was not found. Make sure you have Sphinx\r\n\techo.installed, then set the SPHINXBUILD environment variable to point\r\n\techo.to the full path of the 'sphinx-build' executable. Alternatively you\r\n\techo.may add the Sphinx directory to PATH.\r\n\techo.\r\n\techo.If you don't have Sphinx installed, grab it from\r\n\techo.http://sphinx-doc.org/\r\n\texit /b 1\r\n)\r\n\r\n%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS%\r\ngoto end\r\n\r\n:help\r\n%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS%\r\n\r\n:end\r\npopd\r\n"
  },
  {
    "path": "docs/source/_static/ga_tracker.js",
    "content": "window.dataLayer = window.dataLayer || [];\nfunction gtag(){dataLayer.push(arguments);}\ngtag('js', new Date());\ngtag('config', 'UA-146596996-1');\n"
  },
  {
    "path": "docs/source/conf.py",
    "content": "# -*- coding: utf-8 -*-\n#\n# Configuration file for the Sphinx documentation builder.\n#\n# This file does only contain a selection of the most common options. For a\n# full list see the documentation:\n# http://www.sphinx-doc.org/en/master/config\n\n# -- Path setup --------------------------------------------------------------\n\n# If extensions (or modules to document with autodoc) are in another directory,\n# add these directories to sys.path here. If the directory is relative to the\n# documentation root, use os.path.abspath to make it absolute, like shown here.\n#\nimport os\nimport sys\nsys.path.insert(0, os.path.abspath('../..'))\nsys.setrecursionlimit(1500)\n\nimport condensa\n\n# -- Project information -----------------------------------------------------\n\nproject = u'Condensa'\ncopyright = u'2019, NVIDIA Corporation'\nauthor = u'Saurav Muralidharan'\n\n# The short X.Y version\nversion = '0.5'\n# The full version, including alpha/beta/rc tags\nrelease = '0.5-beta'\n\n\n# -- General configuration ---------------------------------------------------\n\n# If your documentation needs a minimal Sphinx version, state it here.\n#\n# needs_sphinx = '1.0'\n\n# Add any Sphinx extension module names here, as strings. They can be\n# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom\n# ones.\nextensions = [\n    'sphinx.ext.autodoc',\n    'sphinx.ext.mathjax',\n    'sphinx.ext.ifconfig',\n    'sphinx.ext.napoleon',\n]\n\nnapoleon_use_ivar = True\n\n# Add any paths that contain templates here, relative to this directory.\ntemplates_path = ['_templates']\n\n# The suffix(es) of source filenames.\n# You can specify multiple suffix as a list of string:\n#\n# source_suffix = ['.rst', '.md']\nsource_suffix = '.rst'\n\n# The master toctree document.\nmaster_doc = 'index'\n\n# The language for content autogenerated by Sphinx. Refer to documentation\n# for a list of supported languages.\n#\n# This is also used if you do content translation via gettext catalogs.\n# Usually you set \"language\" from the command line for these cases.\nlanguage = None\n\n# List of patterns, relative to source directory, that match files and\n# directories to ignore when looking for source files.\n# This pattern also affects html_static_path and html_extra_path .\nexclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']\n\n# The name of the Pygments (syntax highlighting) style to use.\npygments_style = 'sphinx'\n\n\n# -- Options for HTML output -------------------------------------------------\n\n# The theme to use for HTML and HTML Help pages.  See the documentation for\n# a list of builtin themes.\n#\nif os.environ.get('READTHEDOCS') != 'True':\n  try:\n    import sphinx_rtd_theme\n  except ImportError:\n    pass  # assume we have sphinx >= 1.3\n  else:\n    html_theme_path = [sphinx_rtd_theme.get_html_theme_path()]\n  html_theme = 'sphinx_rtd_theme'\n\n# Theme options are theme-specific and customize the look and feel of a theme\n# further.  For a list of options available for each theme, see the\n# documentation.\n#\n# html_theme_options = {}\n\n# Add any paths that contain custom static files (such as style sheets) here,\n# relative to this directory. They are copied after the builtin static files,\n# so a file named \"default.css\" will overwrite the builtin \"default.css\".\nhtml_static_path = ['_static']\n\n# Custom sidebar templates, must be a dictionary that maps document names\n# to template names.\n#\n# The default sidebars (for documents that don't match any pattern) are\n# defined by theme itself.  Builtin themes are using these templates by\n# default: ``['localtoc.html', 'relations.html', 'sourcelink.html',\n# 'searchbox.html']``.\n#\n# html_sidebars = {}\n\n\n# -- Options for HTMLHelp output ---------------------------------------------\n\n# Output file base name for HTML help builder.\nhtmlhelp_basename = 'Condensadoc'\n\n\n# -- Options for LaTeX output ------------------------------------------------\n\nlatex_elements = {\n    # The paper size ('letterpaper' or 'a4paper').\n    #\n    # 'papersize': 'letterpaper',\n\n    # The font size ('10pt', '11pt' or '12pt').\n    #\n    # 'pointsize': '10pt',\n\n    # Additional stuff for the LaTeX preamble.\n    #\n    # 'preamble': '',\n\n    # Latex figure (float) alignment\n    #\n    # 'figure_align': 'htbp',\n}\n\n# Grouping the document tree into LaTeX files. List of tuples\n# (source start file, target name, title,\n#  author, documentclass [howto, manual, or own class]).\nlatex_documents = [\n    (master_doc, 'Condensa.tex', u'Condensa Documentation',\n     u'Saurav Muralidharan', 'manual'),\n]\n\n\n# -- Options for manual page output ------------------------------------------\n\n# One entry per manual page. List of tuples\n# (source start file, name, description, authors, manual section).\nman_pages = [\n    (master_doc, 'condensa', u'Condensa Documentation',\n     [author], 1)\n]\n\n\n# -- Options for Texinfo output ----------------------------------------------\n\n# Grouping the document tree into Texinfo files. List of tuples\n# (source start file, target name, title, author,\n#  dir menu entry, description, category)\ntexinfo_documents = [\n    (master_doc, 'Condensa', u'Condensa Documentation',\n     author, 'Condensa', 'Programmable Model Compression',\n     'Miscellaneous'),\n]\n\n\n# -- Extension configuration -------------------------------------------------\n\n#autoclass_content = 'both'\n\ndef setup(app):\n    \"\"\"\n    Insert Google Analytics tracker.\n    Based on this Stackoverflow suggestion: https://stackoverflow.com/a/41885884\n    \"\"\"\n    app.add_javascript(\"https://www.googletagmanager.com/gtag/js?id=UA-146596996-1\")\n    app.add_javascript(\"ga_tracker.js\")\n"
  },
  {
    "path": "docs/source/guide/install.rst",
    "content": "Installation\n============\n\nPrerequisites\n-------------\n\nCondensa requires:\n\n* A working Linux installation (we use Ubuntu 18.04)\n* NVIDIA drivers and CUDA 10+ for GPU support\n* Python 3.5 or newer\n* PyTorch 1.0 or newer\n\nInstallation from Source\n------------------------\n\nRetrieve the latest source code from the Condensa repository:\n\n.. code-block:: bash\n\n   git clone https://github.com/NVlabs/condensa.git\n\nNavigate to the source code directory and run the following:\n\n.. code-block:: bash\n\n   pip install -r requirements.txt\n\nTo check the installation, run the unit test suite:\n\n.. code-block:: bash\n\n   bash run_all_tests.sh -v\n"
  },
  {
    "path": "docs/source/guide/usage.rst",
    "content": "Usage\n=====\n\nThe `notebooks`_ folder contains Jupyter notebooks with step-by-step walkthroughs\nfor various usage scenarios. In particular, check out the `AlexNet Compression`_ notebook\nfor a simple getting started guide.\n\nThe `examples`_ folder contains additional, more complex examples of using Condensa.\n\n.. _notebooks: https://github.com/NVlabs/condensa/blob/master/notebooks\n.. _AlexNet Compression: https://github.com/NVlabs/condensa/blob/master/notebooks/alexnet.ipynb\n.. _examples: https://github.com/NVlabs/condensa/blob/master/examples\n"
  },
  {
    "path": "docs/source/index.rst",
    "content": ".. Condensa documentation master file, created by\n   sphinx-quickstart on Tue Sep  4 15:17:30 2018.\n   You can adapt this file completely to your liking, but it should at least\n   contain the root `toctree` directive.\n\nCondensa Documentation\n======================\n\nCondensa is a framework for **programmable model compression** in Python.\nIt comes with a set of built-in compression operators which may be used to\ncompose complex compression schemes targeting specific combinations of DNN,\nhardware platform, and optimization objective.\nCommon programming abstractions such as conditionals, iteration, and\nrecursion are all natively supported.\nTo recover any accuracy lost during compression, Condensa uses a constrained\noptimization formulation of model compression and employs an Augmented Lagrangian-based\nalgorithm as the optimizer.\n\nCondensa is under active development, and bug reports, pull requests, and other feedback are all highly appreciated.\n\n\n.. toctree::\n   :maxdepth: 2\n   :caption: Getting Started\n\n   guide/install\n   guide/usage\n\n.. toctree::\n   :maxdepth: 2\n   :caption: Module API Reference\n\n   modules/schemes\n   modules/pi\n   modules/compressor\n   modules/opt\n   modules/finetuner\n   modules/tensor\n   modules/functional\n   modules/util\n\n.. toctree::\n   :maxdepth: 2\n   :caption: Notes\n\n   modules/lc\n\nIndices and Tables\n------------------\n\n* :ref:`genindex`\n* :ref:`modindex`\n* :ref:`search`\n"
  },
  {
    "path": "docs/source/modules/compressor.rst",
    "content": "Model Compressor\n================\n.. autoclass:: condensa.Compressor\n   :members:\n\n   .. automethod:: __init__\n"
  },
  {
    "path": "docs/source/modules/finetuner.rst",
    "content": "Model Fine-Tuner\n================\n.. autoclass:: condensa.finetune.FineTuner\n   :members:\n\n   .. automethod:: __init__\n\n"
  },
  {
    "path": "docs/source/modules/functional.rst",
    "content": "Aggregation Functions\n=====================\n.. automodule:: condensa.functional\n   :members:\n\n\n"
  },
  {
    "path": "docs/source/modules/lc.rst",
    "content": "L-C Optimizer Usage\n===================\n"
  },
  {
    "path": "docs/source/modules/opt.rst",
    "content": "Optimizers\n==========\n.. automodule:: condensa.opt\n\nDirect Compression Optimizer\n----------------------------\n.. autoclass:: condensa.opt.DC\n   :members:\n\n   .. automethod:: __init__\n\nL-C Optimizer\n-------------\n.. autoclass:: condensa.opt.LC\n   :members:\n\n   .. automethod:: __init__\n\n"
  },
  {
    "path": "docs/source/modules/pi.rst",
    "content": "Compression Operators\n=====================\n.. automodule:: condensa.pi\n   :members:\n\n.. automodule:: condensa.delta\n   :members:\n\n"
  },
  {
    "path": "docs/source/modules/schemes.rst",
    "content": "Compression Schemes\n===================\n.. automodule:: condensa.schemes\n\nComposition\n-----------\n.. autoclass:: condensa.schemes.Compose\n   :members:\n\n   .. automethod:: __init__\n\nUnstructured Pruning\n--------------------\n.. autoclass:: condensa.schemes.Prune\n   :members:\n\n   .. automethod:: __init__\n\nQuantization\n------------\n.. autoclass:: condensa.schemes.Quantize\n   :members:\n\n   .. automethod:: __init__\n\nNeuron Pruning\n--------------\n.. autoclass:: condensa.schemes.NeuronPrune\n   :members:\n\n   .. automethod:: __init__\n\nFilter Pruning\n--------------\n.. autoclass:: condensa.schemes.FilterPrune\n   :members:\n\n   .. automethod:: __init__\n\nStructured Pruning\n------------------\n.. autoclass:: condensa.schemes.StructurePrune\n   :members:\n\n   .. automethod:: __init__\n\nBlock Pruning\n-------------\n.. autoclass:: condensa.schemes.BlockPrune\n   :members:\n\n   .. automethod:: __init__\n\n"
  },
  {
    "path": "docs/source/modules/tensor.rst",
    "content": "Tensor Operators\n================\n.. automodule:: condensa.tensor\n   :members:\n\n"
  },
  {
    "path": "docs/source/modules/util.rst",
    "content": "Utilities\n=========\n.. automodule:: condensa.util\n   :members:\n\n\n"
  },
  {
    "path": "examples/cifar/compress.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport os\nimport sys\nimport argparse\nimport logging\nimport csv\n\nimport gzip\nimport pickle\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.parallel\nimport torch.nn.utils\nimport torchvision.datasets as datasets\nimport torch.utils.data as data\nimport torch.backends.cudnn as cudnn\nfrom torchvision import datasets, transforms\n\nimport condensa\nfrom condensa import schemes\n\nimport util\nimport models\n\nif __name__ == '__main__':\n    model_names = sorted(\n        name for name in models.__dict__\n        if not name.startswith(\"__\") and callable(models.__dict__[name]))\n\n    valid_schemes = ['PRUNE', 'PQ', 'FILTER']\n    parser = argparse.ArgumentParser(description='CIFAR LC Compression Script')\n    parser.add_argument('--arch',\n                        default='AlexNet',\n                        choices=model_names,\n                        help='Model architecture: ' + ' | '.join(model_names) +\n                        ' (default: alexnet)')\n    parser.add_argument('--dataset', default='cifar10', type=str)\n    parser.add_argument('--model', help='Pretrained model filename')\n    parser.add_argument('--steps', type=int, help='Number of LC iterations')\n    parser.add_argument('--scheme',\n                        choices=valid_schemes,\n                        required=True,\n                        help='Compression scheme')\n    parser.add_argument('--density',\n                        required=True,\n                        type=float,\n                        help='Density for pruning')\n    parser.add_argument('--align',\n                        type=int,\n                        default=None,\n                        help='Alignment for structured pruning')\n    parser.add_argument('--l_batch_size',\n                        type=int,\n                        default=128,\n                        help='Batch size for L step')\n    parser.add_argument('--val_batch_size',\n                        type=int,\n                        default=100,\n                        help='Validation batch size')\n    parser.add_argument('--lr',\n                        type=float,\n                        default=0.02,\n                        help='Initial learning rate')\n    parser.add_argument('--lr_end',\n                        type=float,\n                        default=None,\n                        help='Ending learning rate')\n    parser.add_argument('--lr_decay',\n                        type=float,\n                        default=None,\n                        help='Learning rate decay')\n    parser.add_argument('--lr_schedule',\n                        type=int,\n                        nargs='+',\n                        default=None,\n                        help='Decrease learning rate at these epochs.')\n    parser.add_argument('--lr_multiplier',\n                        type=float,\n                        default=None,\n                        help='Learning rate multiplier')\n    parser.add_argument('--momentum',\n                        type=float,\n                        default=0.95,\n                        help='SGD momentum')\n    parser.add_argument('--weight_decay',\n                        type=float,\n                        default=0,\n                        help='SGD momentum')\n    parser.add_argument('--mb_iterations_per_l',\n                        type=int,\n                        default=2000,\n                        help='Minibatch iterations per L step')\n    parser.add_argument('--mb_iterations_first_l',\n                        type=int,\n                        default=10000,\n                        help='Minibatch iterations for first L step')\n    parser.add_argument('--mu_init',\n                        type=float,\n                        default=0.001,\n                        help='Initial value of mu')\n    parser.add_argument('--mu_multiplier', type=float, help='mu multiplier')\n    parser.add_argument('--mu_cap', type=float, default=10000, help='mu cap')\n    parser.add_argument('--out',\n                        default='compressed_model.pth',\n                        help='Compressed output model filename')\n    parser.add_argument('--csv',\n                        default=None,\n                        help='compression statistics CSV file')\n    parser.add_argument('-v',\n                        '--verbose',\n                        help='verbose logging output',\n                        action='store_true')\n\n    args = parser.parse_args()\n\n    logging.basicConfig(\n        level=logging.INFO if args.verbose else logging.WARNING,\n        format='%(message)s')\n\n    if args.dataset == 'cifar10':\n        dataset = datasets.CIFAR10\n        num_classes = 10\n    elif args.dataset == 'cifar100':\n        dataset = datasets.CIFAR100\n        num_classes = 100\n    else:\n        raise RuntimeError('Invalid dataset: must be CIFAR-10 or CIFAR-100')\n\n    # Load model architecture\n    if args.arch.endswith('resnet'):\n        model = models.__dict__[args.arch](num_classes=num_classes)\n    else:\n        model = models.__dict__[args.arch](num_classes=num_classes)\n\n    model.load_state_dict(torch.load(args.model))\n\n    if args.scheme == 'PRUNE':\n        scheme = schemes.Prune(args.density)\n    elif args.scheme == 'PQ':\n        scheme = schemes.Compose(\n            [schemes.Prune(args.density),\n             schemes.Quantize()])\n    elif args.scheme == 'FILTER':\n        scheme = schemes.FilterPrune(args.density)\n    else:\n        raise RuntimeError('Unknown scheme: {}'.format(args.scheme))\n\n    print('SCHEME: {}'.format(scheme))\n\n    trainloader,valloader = \\\n        util.cifar_train_val_loader(dataset,\n                                    args.l_batch_size,\n                                    args.val_batch_size)\n    testloader = util.cifar_test_loader(dataset, args.val_batch_size)\n\n    # Instantiate LC optimizer\n    sgd_params = {'momentum': args.momentum, 'weight_decay': args.weight_decay}\n    lc = condensa.opt.LC(steps=args.steps,\n                         l_optimizer=condensa.opt.lc.SGD,\n                         l_optimizer_params=sgd_params,\n                         lr=args.lr,\n                         lr_end=args.lr_end,\n                         lr_decay=args.lr_decay,\n                         lr_schedule=args.lr_schedule,\n                         lr_multiplier=args.lr_multiplier,\n                         mb_iterations_per_l=args.mb_iterations_per_l,\n                         mb_iterations_first_l=args.mb_iterations_first_l,\n                         mu_init=args.mu_init,\n                         mu_multiplier=args.mu_multiplier,\n                         mu_cap=args.mu_cap,\n                         debugging_flags={'custom_model_statistics':\n                                           condensa.util.cnn_statistics})\n\n    criterion = nn.CrossEntropyLoss().cuda()\n    # Compress model using Condensa\n    compressor = condensa.Compressor(lc, scheme, model, trainloader,\n                                     testloader, valloader, criterion)\n\n    w = compressor.run()\n\n    if args.out is not None:\n        torch.save(w.state_dict(), args.out)\n        logging.info('[Condensa] Compressed model written to disk')\n\n    print('\\n==== Profiling Results ====')\n    for k, v in compressor.statistics.items():\n        print('  ' + k + ':', v)\n    print('')\n\n    if args.csv is not None:\n        with open(args.csv, 'w') as csv_file:\n            writer = csv.writer(csv_file)\n            for k, v in compressor.statistics.items():\n                row = [k]\n                if isinstance(v, list): row += [str(x) for x in v]\n                else: row.append(str(v))\n                writer.writerow(row)\n        csv_file.close()\n        logging.info('[Condensa] Compression stats written to disk')\n"
  },
  {
    "path": "examples/cifar/compress_alexnet.sh",
    "content": "#!/usr/bin/env bash\n\nif [[ $# -eq 0 ]]; then\n  echo \"Usage: compress_alexnet.sh [scheme] [density] [#iterations]\"\n  exit 1\nfi\n\nSCHEME=${1}\nDENSITY=${2}\nSTEPS=${3}\n\nPREFIX=alexnet_${SCHEME}_${DENSITY//[\\.]/_}\n\npython compress.py\\\n       --arch alexnet --dataset cifar10\\\n       --lr 0.01 --lr_end 1e-4\\\n       --weight_decay 0\\\n       --momentum 0.95\\\n       --mb_iterations_per_l 3000\\\n       --mb_iterations_first_l 30000\\\n       --mu_init 1e-3 --mu_multiplier 1.1\\\n       --l_batch_size 128\\\n       --model trained/alexnet.pth\\\n       --scheme ${SCHEME}\\\n       --density ${DENSITY}\\\n       --out compressed/${PREFIX}.pth\\\n       --csv results/${PREFIX}.csv\\\n       -v --steps ${STEPS}\n\n"
  },
  {
    "path": "examples/cifar/finetune.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport os\nimport sys\nimport argparse\nimport logging\nimport csv\n\nimport gzip\nimport pickle\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.parallel\nimport torch.nn.utils\nimport torchvision.datasets as datasets\nimport torch.utils.data as data\nimport torch.backends.cudnn as cudnn\nfrom torchvision import datasets, transforms\n\nimport condensa\n\nimport util\nimport models\n\nif __name__ == '__main__':\n    model_names = sorted(\n        name for name in models.__dict__\n        if not name.startswith(\"__\") and callable(models.__dict__[name]))\n\n    parser = argparse.ArgumentParser(description='CIFAR fine-tuning script')\n    parser.add_argument('--arch',\n                        default='AlexNet',\n                        choices=model_names,\n                        help='Model architecture: ' + ' | '.join(model_names) +\n                        ' (default: alexnet)')\n    parser.add_argument('--dataset', default='cifar10', type=str)\n    parser.add_argument('--model', help='Pretrained model filename')\n    parser.add_argument('--epochs',\n                        type=int,\n                        help='Number of fine-tuning epochs')\n    parser.add_argument('--batch_size',\n                        type=int,\n                        default=128,\n                        help='Batch size for training')\n    parser.add_argument('--val_batch_size',\n                        type=int,\n                        default=128,\n                        help='Validation batch size')\n    parser.add_argument('--lr', type=float, default=0.1, help='Learning rate')\n    parser.add_argument('--lr_end',\n                        type=float,\n                        default=0.01,\n                        help='Ending learning rate')\n    parser.add_argument('--momentum',\n                        type=float,\n                        default=0.9,\n                        help='SGD momentum')\n    parser.add_argument('--weight_decay',\n                        type=float,\n                        default=0,\n                        help='SGD weight decay')\n    parser.add_argument('--out',\n                        default='finetuned.pth',\n                        help='Fine-tuned output model filename')\n    parser.add_argument('-v',\n                        '--verbose',\n                        help='verbose logging output',\n                        action='store_true')\n\n    args = parser.parse_args()\n\n    logging.basicConfig(\n        level=logging.INFO if args.verbose else logging.WARNING,\n        format='%(message)s')\n\n    if args.dataset == 'cifar10':\n        dataset = datasets.CIFAR10\n        num_classes = 10\n    elif args.dataset == 'cifar100':\n        dataset = datasets.CIFAR100\n        num_classes = 100\n    else:\n        raise RuntimeError('Invalid dataset: must be cifar10 or cifar100')\n\n    # Load model architecture\n    if args.arch.endswith('resnet'):\n        model = models.__dict__[args.arch](num_classes=num_classes)\n    else:\n        model = models.__dict__[args.arch](num_classes=num_classes)\n\n    model.load_state_dict(torch.load(args.model))\n    # Compute #nonzeros prior to fine-tuning\n    nparams_w = torch.nn.utils.parameters_to_vector(\n        model.parameters()).view(-1).nonzero().numel()\n\n    # Only fine-tune fully-connected and convolutional layers\n    layer_types = [torch.nn.Linear, torch.nn.Conv2d]\n\n    trainloader,valloader = \\\n        util.cifar_train_val_loader(dataset,\n                                    args.batch_size,\n                                    args.val_batch_size)\n    testloader = util.cifar_test_loader(dataset, args.val_batch_size)\n    criterion = torch.nn.CrossEntropyLoss().cuda()\n    ft = condensa.FineTuner(model, layer_types)\n    w_ft = ft.run(epochs=args.epochs,\n                  lr=args.lr,\n                  lr_end=args.lr_end,\n                  momentum=args.momentum,\n                  weight_decay=args.weight_decay,\n                  criterion=criterion,\n                  trainloader=trainloader,\n                  testloader=testloader,\n                  valloader=valloader,\n                  debugging_flags={'custom_model_statistics':\n                                    condensa.util.cnn_statistics})\n    nparams_wft = torch.nn.utils.parameters_to_vector(\n        w_ft.parameters()).view(-1).nonzero().numel()\n    print('#Nonzero parameters: before [{}], after [{}]'.format(\n        nparams_w, nparams_wft))\n\n    if args.out is not None:\n        torch.save(w_ft.state_dict(), args.out)\n        logging.info('[Condensa] Fine-tuned model written to disk')\n"
  },
  {
    "path": "examples/cifar/models/__init__.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nfrom __future__ import absolute_import\n\nfrom .resnet import *\nfrom .vgg import *\nfrom .alexnet import *\n"
  },
  {
    "path": "examples/cifar/models/alexnet.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nfrom __future__ import absolute_import\n\nimport torch.nn as nn\n\n__all__ = ['alexnet']\n\nclass AlexNet(nn.Module):\n    def __init__(self, num_classes=10):\n        super(AlexNet, self).__init__()\n        self.features = nn.Sequential(\n            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=5),\n            nn.ReLU(inplace=True),\n            nn.MaxPool2d(kernel_size=2, stride=2),\n            nn.Conv2d(64, 192, kernel_size=5, padding=2),\n            nn.ReLU(inplace=True),\n            nn.MaxPool2d(kernel_size=2, stride=2),\n            nn.Conv2d(192, 384, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(384, 256, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(256, 256, kernel_size=3, padding=1),\n            nn.ReLU(inplace=True),\n            nn.MaxPool2d(kernel_size=2, stride=2),\n        )\n        self.classifier = nn.Linear(256, num_classes)\n\n    def forward(self, x):\n        x = self.features(x)\n        x = x.view(x.size(0), -1)\n        x = self.classifier(x)\n        return x\n\ndef alexnet(**kwargs):\n    model = AlexNet(**kwargs)\n    return model\n"
  },
  {
    "path": "examples/cifar/models/resnet.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nfrom __future__ import absolute_import\n\nimport torch.nn as nn\nimport math\n\n__all__ = ['resnet20', 'resnet56', 'resnet110']\n\ndef conv3x3(in_planes, out_planes, stride=1):\n    \"3x3 convolution with padding\"\n    return nn.Conv2d(in_planes,\n                     out_planes,\n                     kernel_size=3,\n                     stride=stride,\n                     padding=1,\n                     bias=False)\n\nclass BasicBlock(nn.Module):\n    expansion = 1\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None):\n        super(BasicBlock, self).__init__()\n        self.conv1 = conv3x3(inplanes, planes, stride)\n        self.bn1 = nn.BatchNorm2d(planes)\n        self.relu = nn.ReLU(inplace=True)\n        self.conv2 = conv3x3(planes, planes)\n        self.bn2 = nn.BatchNorm2d(planes)\n        self.downsample = downsample\n        self.stride = stride\n\n    def forward(self, x):\n        residual = x\n\n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n\n        out = self.conv2(out)\n        out = self.bn2(out)\n\n        if self.downsample is not None:\n            residual = self.downsample(x)\n\n        out += residual\n        out = self.relu(out)\n\n        return out\n\nclass Bottleneck(nn.Module):\n    expansion = 4\n\n    def __init__(self, inplanes, planes, stride=1, downsample=None):\n        super(Bottleneck, self).__init__()\n        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(planes)\n        self.conv2 = nn.Conv2d(planes,\n                               planes,\n                               kernel_size=3,\n                               stride=stride,\n                               padding=1,\n                               bias=False)\n        self.bn2 = nn.BatchNorm2d(planes)\n        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)\n        self.bn3 = nn.BatchNorm2d(planes * 4)\n        self.relu = nn.ReLU(inplace=True)\n        self.downsample = downsample\n        self.stride = stride\n\n    def forward(self, x):\n        residual = x\n\n        out = self.conv1(x)\n        out = self.bn1(out)\n        out = self.relu(out)\n\n        out = self.conv2(out)\n        out = self.bn2(out)\n        out = self.relu(out)\n\n        out = self.conv3(out)\n        out = self.bn3(out)\n\n        if self.downsample is not None:\n            residual = self.downsample(x)\n\n        out += residual\n        out = self.relu(out)\n\n        return out\n\nclass ResNet(nn.Module):\n    def __init__(self, block, depth, num_classes=10):\n        super(ResNet, self).__init__()\n        assert (depth - 2) % 6 == 0, 'depth should be 6n+2'\n        n = (depth - 2) // 6\n\n        self.inplanes = 16\n        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(16)\n        self.relu = nn.ReLU(inplace=True)\n        self.layer1 = self._make_layer(block, 16, n)\n        self.layer2 = self._make_layer(block, 32, n, stride=2)\n        self.layer3 = self._make_layer(block, 64, n, stride=2)\n        self.avgpool = nn.AvgPool2d(8)\n        self.fc = nn.Linear(64 * block.expansion, num_classes)\n\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n                m.weight.data.normal_(0, math.sqrt(2. / n))\n            elif isinstance(m, nn.BatchNorm2d):\n                m.weight.data.fill_(1)\n                m.bias.data.zero_()\n\n    def _make_layer(self, block, planes, blocks, stride=1):\n        downsample = None\n        if stride != 1 or self.inplanes != planes * block.expansion:\n            downsample = nn.Sequential(\n                nn.Conv2d(self.inplanes,\n                          planes * block.expansion,\n                          kernel_size=1,\n                          stride=stride,\n                          bias=False),\n                nn.BatchNorm2d(planes * block.expansion),\n            )\n            downsample[0].condensa_nocompress = True\n\n        layers = []\n        layers.append(block(self.inplanes, planes, stride, downsample))\n        self.inplanes = planes * block.expansion\n        for i in range(1, blocks):\n            layers.append(block(self.inplanes, planes))\n\n        return nn.Sequential(*layers)\n\n    def forward(self, x):\n        x = self.conv1(x)\n        x = self.bn1(x)\n        x = self.relu(x)\n\n        x = self.layer1(x)\n        x = self.layer2(x)\n        x = self.layer3(x)\n\n        x = self.avgpool(x)\n        x = x.view(x.size(0), -1)\n        x = self.fc(x)\n\n        return x\n\ndef resnet20(**kwargs):\n    return ResNet(BasicBlock, 20, **kwargs)\n\ndef resnet56(**kwargs):\n    return ResNet(Bottleneck, 56, **kwargs)\n\ndef resnet110(**kwargs):\n    return ResNet(Bottleneck, 110, **kwargs)\n"
  },
  {
    "path": "examples/cifar/models/vgg.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nfrom __future__ import absolute_import\n\nimport torch.nn as nn\nimport math\n\n__all__ = [\n    'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn',\n    'vgg16', 'vgg16_bn', 'vgg19_bn', 'vgg19'\n]\n\nclass VGG(nn.Module):\n    def __init__(self, features, num_classes=10):\n        super(VGG, self).__init__()\n        self.features = features\n        self.classifier = nn.Linear(512, num_classes)\n        self._initialize_weights()\n\n    def forward(self, x):\n        x = self.features(x)\n        x = x.view(x.size(0), -1)\n        x = self.classifier(x)\n        return x\n\n    def _initialize_weights(self):\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n                m.weight.data.normal_(0, math.sqrt(2. / n))\n                if m.bias is not None:\n                    m.bias.data.zero_()\n            elif isinstance(m, nn.BatchNorm2d):\n                m.weight.data.fill_(1)\n                m.bias.data.zero_()\n            elif isinstance(m, nn.Linear):\n                n = m.weight.size(1)\n                m.weight.data.normal_(0, 0.01)\n                m.bias.data.zero_()\n\ndef make_layers(cfg, batch_norm=False):\n    layers = []\n    in_channels = 3\n    for v in cfg:\n        if v == 'M':\n            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]\n        else:\n            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)\n            if batch_norm:\n                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]\n            else:\n                layers += [conv2d, nn.ReLU(inplace=True)]\n            in_channels = v\n    return nn.Sequential(*layers)\n\ncfg = {\n    'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],\n    'B': [64, 64, 'M', 128, 128, 'M', 256, 256,\n          'M', 512, 512, 'M', 512, 512, 'M'],\n    'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256,\n          'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],\n    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256,\n          'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],\n}\n\ndef vgg11(**kwargs):\n    model = VGG(make_layers(cfg['A']), **kwargs)\n    return model\n\ndef vgg11_bn(**kwargs):\n    model = VGG(make_layers(cfg['A'], batch_norm=True), **kwargs)\n    return model\n\ndef vgg13(**kwargs):\n    model = VGG(make_layers(cfg['B']), **kwargs)\n    return model\n\ndef vgg13_bn(**kwargs):\n    model = VGG(make_layers(cfg['B'], batch_norm=True), **kwargs)\n    return model\n\ndef vgg16(**kwargs):\n    model = VGG(make_layers(cfg['D']), **kwargs)\n    return model\n\ndef vgg16_bn(**kwargs):\n    model = VGG(make_layers(cfg['D'], batch_norm=True), **kwargs)\n    return model\n\ndef vgg19(**kwargs):\n    model = VGG(make_layers(cfg['E']), **kwargs)\n    return model\n\ndef vgg19_bn(**kwargs):\n    model = VGG(make_layers(cfg['E'], batch_norm=True), **kwargs)\n    return model\n"
  },
  {
    "path": "examples/cifar/util.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport os\n\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.data as data\n\nimport torchvision\nimport torchvision.transforms as transforms\nfrom torch.utils.data.sampler import SubsetRandomSampler\n\nimport condensa.data\n\ndef cifar_train_val_loader(dataset,\n                           train_batch_size,\n                           val_batch_size,\n                           root='./data',\n                           random_seed=42,\n                           shuffle=True):\n    \"\"\"\n    Splits the CIFAR training set into training and validation\n    sets (9:1 split) and returns the corresponding data loaders.\n    \"\"\"\n    transform_train = transforms.Compose([\n        transforms.RandomCrop(32, padding=4),\n        transforms.RandomHorizontalFlip(),\n    ])\n    trainset = dataset(root=root,\n                       train=True,\n                       download=True,\n                       transform=transform_train)\n    valset = dataset(root=root, train=True, download=True, transform=None)\n    num_train = len(trainset)\n    indices = list(range(num_train))\n    split = 5000\n\n    if shuffle:\n        np.random.seed(random_seed)\n        np.random.shuffle(indices)\n\n    train_idx, val_idx = indices[split:], indices[:split]\n    trainsampler = SubsetRandomSampler(train_idx)\n    valsampler = SubsetRandomSampler(val_idx)\n\n    meanstd = ((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))\n    trainloader = condensa.data.GPUDataLoader(trainset,\n                                              batch_size=train_batch_size,\n                                              shuffle=False,\n                                              num_workers=8,\n                                              sampler=trainsampler,\n                                              meanstd=meanstd)\n    valloader =   condensa.data.GPUDataLoader(valset,\n                                              batch_size=val_batch_size,\n                                              shuffle=False,\n                                              num_workers=8,\n                                              sampler=valsampler,\n                                              meanstd=meanstd)\n\n    return (trainloader, valloader)\n\ndef cifar_test_loader(dataset, batch_size, root='./data'):\n    \"\"\"\n    Construct a CIFAR test dataset loader.\n    \"\"\"\n    testset = dataset(root=root, train=False, download=True, transform=None)\n    meanstd = ((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))\n    testloader = condensa.data.GPUDataLoader(testset,\n                                             batch_size=batch_size,\n                                             shuffle=False,\n                                             num_workers=8,\n                                             meanstd=meanstd)\n    return testloader\n"
  },
  {
    "path": "notebooks/AlexNet.ipynb",
    "content": "{\n \"cells\": [\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"# Tutorial: Compressing AlexNet with Condensa\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"In this tutorial, we will walk through compressing the [AlexNet neural network](https://en.wikipedia.org/wiki/AlexNet) on the CIFAR-10 dataset using Condensa. We will target two different objectives: reducing total model memory footprint, and reducing the inference latency of the compressed model. \\n\",\n    \"\\n\",\n    \"We assume that Condensa is already installed and working (check out the [Installation Guide](https://nvlabs.github.io/condensa/guide/install.html) for instructions). If you'd like to follow along by executing the code in this notebook, please also make sure that [Jupyter](https://jupyter.org/) is installed on your local system.\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"## Defining the Network\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"Let's start by defining the AlexNet network architecture in PyTorch as shown below:\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"import torch\\n\",\n    \"import torch.nn as nn\\n\",\n    \"\\n\",\n    \"class AlexNet(nn.Module):\\n\",\n    \"    def __init__(self, num_classes=10):\\n\",\n    \"        super(AlexNet, self).__init__()\\n\",\n    \"        self.features = nn.Sequential(\\n\",\n    \"            nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=5),\\n\",\n    \"            nn.ReLU(inplace=True),\\n\",\n    \"            nn.MaxPool2d(kernel_size=2, stride=2),\\n\",\n    \"            nn.Conv2d(64, 192, kernel_size=5, padding=2),\\n\",\n    \"            nn.ReLU(inplace=True),\\n\",\n    \"            nn.MaxPool2d(kernel_size=2, stride=2),\\n\",\n    \"            nn.Conv2d(192, 384, kernel_size=3, padding=1),\\n\",\n    \"            nn.ReLU(inplace=True),\\n\",\n    \"            nn.Conv2d(384, 256, kernel_size=3, padding=1),\\n\",\n    \"            nn.ReLU(inplace=True),\\n\",\n    \"            nn.Conv2d(256, 256, kernel_size=3, padding=1),\\n\",\n    \"            nn.ReLU(inplace=True),\\n\",\n    \"            nn.MaxPool2d(kernel_size=2, stride=2),\\n\",\n    \"        )\\n\",\n    \"        self.classifier = nn.Linear(256, num_classes)\\n\",\n    \"\\n\",\n    \"    def forward(self, x):\\n\",\n    \"        x = self.features(x)\\n\",\n    \"        x = x.view(x.size(0), -1)\\n\",\n    \"        x = self.classifier(x)\\n\",\n    \"        return x\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"We instantiate this class into `model`:\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"model = AlexNet()\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"## Load Pre-Trained Weights\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"Now that we have defined the network architecture, let us load a pre-trained set of weights into the model from the `AlexNet.pth` file included with this notebook.\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"model.load_state_dict(torch.load('AlexNet.pth'))\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"## Preparing for Compression\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"Let's make sure CUDA is enabled in PyTorch.\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"assert torch.cuda.is_available()\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"We now create PyTorch data loaders for the training, test, and validation datasets. To save space, we wrap the data loading code into two utility functions: `cifar_train_val_loader` and `cifar_test_loader` (please refer to `util.py` in the current `notebooks` folder for the full code).\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"import util\\n\",\n    \"import torchvision.datasets as datasets\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"dataset = datasets.CIFAR10\\n\",\n    \"\\n\",\n    \"trainloader,valloader = util.cifar_train_val_loader(dataset, train_batch_size=128, val_batch_size=128)\\n\",\n    \"testloader = util.cifar_test_loader(dataset, batch_size=128)\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"The utilities above split the original training set into training and validation sets (using a 9:1 split) and perform data normalization for all datasets. They also utilize Condensa's `GPUDataLoader` to enable fast data prefetching and collation.\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"We now define our loss criterion:\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"criterion = nn.CrossEntropyLoss().cuda()\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"Finally, we set our logging level to `INFO` so that Condensa prints out intermediate updates.\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"import logging\\n\",\n    \"logging.basicConfig(level=logging.INFO, format='%(message)s')\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"## Two Different Compression Strategies\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"In this tutorial, we will explore two different ways of compressing the AlexNet network: one targeted at reducing the total model memory footprint (named `MEM`) and the other at reducing inference runtime latency (named `FLOP`).\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"### MEM Scheme\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"The `MEM` scheme aims to reduce the total model memory footprint (number of bytes required to store the non-zero elements of the compressed model). To this end, we perform a combination of _pruning_ (clipping model parameters to zero) and _quantization_ (using 16-bit floating point representation to store model weights instead of 32-bit). Expressing this scheme in Condensa is fairly straightforward using the built-in [`Compose`](https://nvlabs.github.io/condensa/modules/schemes.html#composition) scheme as shown below:\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"import condensa\\n\",\n    \"from condensa.schemes import Compose, Prune, Quantize\\n\",\n    \"\\n\",\n    \"MEM = Compose([Prune(0.02), Quantize(condensa.float16)])\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"Here, the [`Compose`](https://nvlabs.github.io/condensa/modules/schemes.html#composition) operator successively applies pruning followed by quantization to the model. The pruning density, or the ratio of non-zero parameters in the compressed model to the original one, is specified as 0.02 (2%). Condensa includes a number of other common schemes, including structured and block pruning, among others. For a list of available schemes, please refer to [this page](https://nvlabs.github.io/condensa/modules/schemes.html) in the API documentation. Users may also define their own custom schemes as Python functions that invoke the compression and decompression operators available in Condensa (see [`schemes.py`](https://github.com/NVlabs/condensa/blob/master/condensa/schemes.py) for examples of how to define custom schemes).\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"### FLOP Scheme\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"While the `MEM` scheme is effective at reducing the number of non-zero elements in a model, this may not directly translate into improvements in actual inference runtime. Most modern CPUs and GPUs are unable to detect individual zero elements and bypass computations on them in hardware. Instead, to realize speedups on such architectures, we perform filter pruning, which removes entire filters (3D blocks) at once from convolutional layers. This enables the weight tensors to be physically reshaped in the compressed model. We call this the `FLOP` scheme in this tutorial, and use the [`FilterPrune`](https://nvlabs.github.io/condensa/modules/schemes.html#filter-pruning) scheme in Condensa to define it.\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"from condensa.schemes import FilterPrune\\n\",\n    \"FLOP = condensa.schemes.FilterPrune(0.5)\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"## Setting up the Optimizer\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"To recover any accuracy lost due to compression, Condensa comes with a set of _optimizers_. Each optimizer takes a pre-trained model, applies the compression scheme, and tries to recover the original accuracy either directly or iteratively. In this tutorial, we'll be using Condensa's [L-C optimizer](https://nvlabs.github.io/condensa/modules/opt.html#l-c-optimizer). We instantiate it as follows:\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"lc = condensa.opt.LC(steps=35,                             # L-C iterations\\n\",\n    \"                     l_optimizer=condensa.opt.lc.SGD,      # L-step sub-optimizer\\n\",\n    \"                     l_optimizer_params={'momentum':0.95}, # L-step sub-optimizer parameters\\n\",\n    \"                     lr=0.01,                              # Initial learning rate\\n\",\n    \"                     lr_end=1e-4,                          # Final learning rate\\n\",\n    \"                     mb_iterations_per_l=3000,             # Mini-batch iterations per L-step\\n\",\n    \"                     mb_iterations_first_l=30000,          # Mini-batch iterations for first L-step\\n\",\n    \"                     mu_init=1e-3,                         # Initial value of `mu`\\n\",\n    \"                     mu_multiplier=1.1,                    # Multiplier for `mu`\\n\",\n    \"                     mu_cap=10000,                         # Maximum value of `mu`\\n\",\n    \"                     debugging_flags={'custom_model_statistics':\\n\",\n    \"                                      condensa.util.cnn_statistics})\\n\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"Each optimizer in Condensa has its own set of hyper-parameters which must be specified manually by the user. A full description of hyper-parameter tuning is beyond the scope of this tutorial, but for additional information on what each hyper-parameter represents and tips on finding its optimal value, we refer you to the [Condensa paper](https://arxiv.org/abs/1911.02497). In this notebook, we run the L-C algorithm for 35 iterations using the hyper-parameter values shown above. L-C hyper-parameter values for a number of common convolutional neural networks are also included in the [`examples`](https://github.com/NVlabs/condensa/blob/master/examples/) folder.\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"## Compressing the Model\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"Once the optimizer is instantiated, we can go ahead and perform the actual compression using the [`Compressor`](https://nvlabs.github.io/condensa/modules/compressor.html#model-compressor) class and its [`run`](https://nvlabs.github.io/condensa/modules/compressor.html#condensa.compressor.Compressor.run) method. **Note:** the next two lines may take a while to execute!\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"compressor_MEM  = condensa.Compressor(lc,\\n\",\n    \"                                      MEM,\\n\",\n    \"                                      model,\\n\",\n    \"                                      trainloader,\\n\",\n    \"                                      testloader,\\n\",\n    \"                                      valloader,\\n\",\n    \"                                      criterion)\\n\",\n    \"w_MEM  = compressor_MEM.run()\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"compressor_FLOP = condensa.Compressor(lc,\\n\",\n    \"                                      FLOP,\\n\",\n    \"                                      model,\\n\",\n    \"                                      trainloader,\\n\",\n    \"                                      testloader,\\n\",\n    \"                                      valloader,\\n\",\n    \"                                      criterion)\\n\",\n    \"\\n\",\n    \"w_FLOP = compressor_FLOP.run()\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"We specify the optimizer, scheme, input model, training, test, and validation sets, and the loss criterion to create an instance of the [`Compressor`](https://nvlabs.github.io/condensa/modules/compressor.html#model-compressor) class. Since the optimizer is specified as a parameter, we are able to easily experiment with alternative optimizers in Condensa.\\n\",\n    \"\\n\",\n    \"In the above snippets, `w_MEM` and `w_FLOP` contain the models compressed using the `MEM` and `FLOP` schemes, respectively. We can now save these to disk:\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"torch.save(w_MEM.state_dict(), 'AlexNet_MEM.pth')\\n\",\n    \"torch.save(w_FLOP.state_dict(), 'AlexNet_FLOP.pth')\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"Condensa also records various statistics about the compression process. These can be retrieved using the `statistics` member of the compressor object as follows:\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"for k,v in compressor_MEM.statistics.items():\\n\",\n    \"    print('{}: {}'.format(k, v))\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"metadata\": {},\n   \"outputs\": [],\n   \"source\": [\n    \"for k,v in compressor_FLOP.statistics.items():\\n\",\n    \"    print('{}: {}'.format(k, v))\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"## Results\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"We notice that Condensa achieves top-1 test accuracies of **77.49%** and **76.81%** for the `MEM` and `FLOP` schemes, respectively (compared to the baseline accuracy of **77.07%** for AlexNet). For more complex models, it is possible to further improve accuracies via [model fine-tuning](https://nvlabs.github.io/condensa/modules/finetuner.html).\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"### Memory and Runtime Reductions\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"Using the `MEM` scheme, we reduce the model memory footprint by **97.83x**. Additionally, we achieve a **55.6%** reduction in FLOPs using the `FLOP` scheme.\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"## More Info\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"metadata\": {},\n   \"source\": [\n    \"We provide additional real-world compression examples targeting complex networks such as ResNet50 and VGG-19 in Condensa's [examples folder](https://github.com/NVlabs/condensa/tree/master/examples). Be sure to check them out!\\n\",\n    \"\\n\",\n    \"For more details on the design and implementation of Condensa, and its performance on real-world networks, please refer to the [Condensa paper](https://arxiv.org/abs/1911.02497).\"\n   ]\n  }\n ],\n \"metadata\": {\n  \"kernelspec\": {\n   \"display_name\": \"Python 3\",\n   \"language\": \"python\",\n   \"name\": \"python3\"\n  },\n  \"language_info\": {\n   \"codemirror_mode\": {\n    \"name\": \"ipython\",\n    \"version\": 3\n   },\n   \"file_extension\": \".py\",\n   \"mimetype\": \"text/x-python\",\n   \"name\": \"python\",\n   \"nbconvert_exporter\": \"python\",\n   \"pygments_lexer\": \"ipython3\",\n   \"version\": \"3.7.3\"\n  }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 2\n}\n"
  },
  {
    "path": "notebooks/util.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport os\n\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nimport torch.utils.data as data\n\nimport torchvision\nimport torchvision.transforms as transforms\nfrom torch.utils.data.sampler import SubsetRandomSampler\n\nimport condensa.data\n\ndef cifar_train_val_loader(dataset,\n                           train_batch_size,\n                           val_batch_size,\n                           root='./data',\n                           random_seed=42,\n                           shuffle=True):\n    \"\"\"\n    Splits the CIFAR training set into training and validation\n    sets (9:1 split) and returns the corresponding data loaders.\n    \"\"\"\n    transform_train = transforms.Compose([\n        transforms.RandomCrop(32, padding=4),\n        transforms.RandomHorizontalFlip(),\n    ])\n    trainset = dataset(root=root,\n                       train=True,\n                       download=True,\n                       transform=transform_train)\n    valset = dataset(root=root, train=True, download=True, transform=None)\n    num_train = len(trainset)\n    indices = list(range(num_train))\n    split = 5000\n\n    if shuffle:\n        np.random.seed(random_seed)\n        np.random.shuffle(indices)\n\n    train_idx, val_idx = indices[split:], indices[:split]\n    trainsampler = SubsetRandomSampler(train_idx)\n    valsampler = SubsetRandomSampler(val_idx)\n\n    meanstd = ((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))\n    trainloader = condensa.data.GPUDataLoader(trainset,\n                                              batch_size=train_batch_size,\n                                              shuffle=False,\n                                              num_workers=8,\n                                              sampler=trainsampler,\n                                              meanstd=meanstd)\n    valloader =   condensa.data.GPUDataLoader(valset,\n                                              batch_size=val_batch_size,\n                                              shuffle=False,\n                                              num_workers=8,\n                                              sampler=valsampler,\n                                              meanstd=meanstd)\n\n    return (trainloader, valloader)\n\ndef cifar_test_loader(dataset, batch_size, root='./data'):\n    \"\"\"\n    Construct a CIFAR test dataset loader.\n    \"\"\"\n    testset = dataset(root=root, train=False, download=True, transform=None)\n    meanstd = ((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))\n    testloader = condensa.data.GPUDataLoader(testset,\n                                             batch_size=batch_size,\n                                             shuffle=False,\n                                             num_workers=8,\n                                             meanstd=meanstd)\n    return testloader\n"
  },
  {
    "path": "run_all_tests.sh",
    "content": "#!/bin/bash\n\nVERBOSE=0\nif [[ $1 == \"-v\" ]] || [[ $1 == \"--verbose\" ]]; then\n  VERBOSE=1\nfi\n\nfor f in $(find test -name '*.py'); do\n  if [[ $VERBOSE -eq 1 ]]; then\n    echo \"[Condensa Test] $f\"\n  fi\n  python3 $f\ndone\n"
  },
  {
    "path": "setup.cfg",
    "content": "[metadata]\ndescription-file = README.md\n"
  },
  {
    "path": "setup.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport os\nfrom setuptools import setup\nfrom setuptools import find_packages\n\ncwd = os.path.dirname(os.path.abspath(__file__))\nversion = '0.5.0-beta'\n\ndef build_deps():\n  version_path = os.path.join(cwd, 'condensa', 'version.py')\n  with open(version_path, 'w') as f:\n    f.write(\"__version__ = '{}'\\n\".format(version))\n\nbuild_deps()\n\nwith open(os.path.join(cwd, 'README.md'), encoding='utf-8') as f:\n    long_description = f.read()\n\ninstall_requires = ['numpy',\n                    'torch>=1.0.0',\n                    'tqdm']\n\nsetup(name='condensa',\n      version=version,\n      description='Condensa Programmable Model Compression Framework',\n      long_description=long_description,\n      long_description_content_type='text/markdown',\n      url='https://github.com/NVLabs/condensa',\n      author='Saurav Muralidharan',\n      author_email='sauravm@nvidia.com',\n      license='Apache License 2.0',\n      keywords=['compression', 'quantization', 'pruning'],\n      install_requires=install_requires,\n      packages=find_packages(),\n      classifiers=[\n        'Development Status :: 4 - Beta',\n        'Intended Audience :: Developers',\n        'Intended Audience :: Science/Research',\n        'Topic :: Software Development :: Build Tools',\n        'License :: OSI Approved :: Apache Software License',\n        'Programming Language :: Python :: 3',\n        'Topic :: Software Development :: Libraries',\n        'Topic :: Scientific/Engineering :: Artificial Intelligence'\n      ],\n      )\n"
  },
  {
    "path": "test/schemes/test_prune.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport torch\n\nimport condensa\nimport condensa.schemes as schemes\nimport condensa.tensor as T\nimport condensa.functional as F\n\ndef test_prune(device):\n    fc = torch.nn.Linear(100, 10, bias=True).to(device)\n    scheme = schemes.Prune(0.5)\n    threshold = scheme.threshold(fc)\n    scheme.pi(fc)\n\n    t = fc.weight.data.abs().view(-1)\n    nzs = torch.index_select(t, 0, t.nonzero().view(-1))\n    assert (nzs >= threshold).all()\n\ndef test_filter_prune(device):\n    conv = torch.nn.Conv2d(3,\n                           64,\n                           kernel_size=11,\n                           stride=4,\n                           padding=5,\n                           bias=True).to(device)\n\n    criteria = F.l2norm\n    scheme = schemes.FilterPrune(0.5, criteria=criteria, prune_bias=True)\n    threshold = scheme.threshold(conv)\n    scheme.pi(conv)\n\n    # Check against threshold\n    agg = T.aggregate_filters(conv.weight.data, criteria).view(-1)\n    nzs = torch.index_select(agg, 0, agg.nonzero().view(-1))\n    assert (nzs >= threshold).all()\n\n    # Check biases: all zero filters must have corresponding zero biases\n    zero_indices = (agg == 0).nonzero().view(-1)\n    z = torch.index_select(conv.bias.data, 0, zero_indices)\n    assert (z == 0.).all()\n\ndef test_neuron_prune(device):\n    fc = torch.nn.Linear(100, 10, bias=True).to(device)\n\n    criteria = F.l2norm\n    scheme = schemes.NeuronPrune(0.5, criteria=criteria, prune_bias=True)\n    threshold = scheme.threshold(fc)\n    scheme.pi(fc)\n\n    # Check against threshold\n    agg = T.aggregate_neurons(fc.weight.data, criteria).view(-1)\n    nzs = torch.index_select(agg, 0, agg.nonzero().view(-1))\n    assert (nzs >= threshold).all()\n\n    # Check biases: all zero neurons must have corresponding zero biases\n    zero_indices = (agg == 0).nonzero().view(-1)\n    z = torch.index_select(fc.bias.data, 0, zero_indices)\n    assert (z == 0.).all()\n\ndef test_block_prune(device, blocksize=(10,10)):\n    fc = torch.nn.Linear(100, 100, bias=False).to(device)\n\n    criteria = F.l2norm\n    scheme = schemes.BlockPrune(0.5, criteria=criteria, block_size=blocksize)\n    threshold = scheme.threshold(fc)\n    scheme.pi(fc)\n\n    # Check against threshold\n    agg = T.aggregate(fc.weight.data, blocksize, criteria).view(-1)\n    nzs = torch.index_select(agg, 0, agg.nonzero().view(-1))\n    assert (nzs >= threshold).all()\n\nif __name__ == '__main__':\n    test_prune('cpu')\n    test_filter_prune('cpu')\n    test_neuron_prune('cpu')\n    test_block_prune('cpu')\n\n    if torch.cuda.is_available():\n        test_prune('cuda:0')\n        test_filter_prune('cuda:0')\n        test_neuron_prune('cuda:0')\n        test_block_prune('cuda:0')\n"
  },
  {
    "path": "test/schemes/test_qz.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport torch\n\nimport condensa\nfrom condensa import schemes\n\ndef test_float16(device):\n    scheme = schemes.Quantize(condensa.float16)\n    fc = torch.nn.Linear(100, 10).float().to(device)\n\n    scheme.pi(fc)\n    assert fc.weight.dtype == torch.float16\n    scheme.delta(fc)\n    assert fc.weight.dtype == torch.float32\n\nif __name__ == '__main__':\n    test_float16('cpu')\n    if torch.cuda.is_available():\n        test_float16('cpu')\n"
  },
  {
    "path": "test/tensor/test_mask_apply.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport numpy as np\nimport torch\n\nimport condensa.tensor as T\n\ndef test_apply_mask(device):\n    a = torch.randn(20).to(device)\n    threshold = T.threshold(a, 0.3)\n    mask = T.simple_mask(a, threshold)\n    T.apply_mask_inplace(a, mask)\n\n    for i in range(len(a)):\n        assert a[i] == 0. or abs(a[i]) >= threshold\n\nif __name__ == '__main__':\n    test_apply_mask('cpu')\n    if torch.cuda.is_available():\n        test_apply_mask('cuda:0')\n"
  },
  {
    "path": "test/tensor/test_maskgen.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport numpy as np\nimport torch\n\nimport condensa.tensor as T\n\ndef test_simple_mask(device):\n    a = torch.randn(20).to(device)\n    threshold = T.threshold(a, 0.3)\n    mask = T.simple_mask(a, threshold)\n\n    for i in range(len(a)):\n        if abs(a[i]) >= threshold: assert mask[i] == 1\n        else: assert mask[i] == 0\n\ndef test_block_mask(device):\n    pass\n\nif __name__ == '__main__':\n    test_simple_mask('cpu')\n    test_block_mask('cpu')\n    if torch.cuda.is_available():\n        test_simple_mask('cuda:0')\n        test_block_mask('cuda:0')\n"
  },
  {
    "path": "test/tensor/test_util.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport numpy as np\nimport torch\n\nimport condensa.tensor as T\n\ndef test_density(device):\n    zeros = torch.zeros(10).to(device)\n    ones = torch.ones(30).to(device)\n    assert T.density(zeros) == 0.\n    assert T.density(ones) == 1.\n    assert T.density(torch.cat((zeros, ones))) == 0.75\n\ndef test_sparsity(device):\n    zeros = torch.zeros(10).to(device)\n    ones = torch.ones(30).to(device)\n    assert T.sparsity(zeros) == 1.\n    assert T.sparsity(ones) == 0.\n    assert T.sparsity(torch.cat((zeros, ones))) == 0.25\n\ndef test_threshold(device):\n    a = torch.IntTensor(np.arange(0, 30)).to(device)\n    threshold2 = T.threshold(a, 0.2)\n    threshold3 = T.threshold(a, 0.3)\n    threshold5 = T.threshold(a, 0.5)\n\n    assert threshold2.item() == 24\n    assert threshold3.item() == 21\n    assert threshold5.item() == 15\n\nif __name__ == '__main__':\n    test_density('cpu')\n    test_sparsity('cpu')\n    test_threshold('cpu')\n\n    if torch.cuda.is_available():\n        test_density('cpu')\n        test_sparsity('cpu')\n        test_threshold('cpu')\n"
  },
  {
    "path": "test/test_lr.py",
    "content": "# Copyright 2019 NVIDIA Corporation\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\nimport numpy as np\n\nimport condensa\nimport condensa.lr as lr\n\ndef test_interval_lr():\n    schedule = lr.IntervalLR(1., 1e-6, 100)\n    assert schedule.learning_rate == 1.\n    for i in range(0, 100):\n        schedule.step()\n    assert np.isclose(schedule.learning_rate, 1e-6)\n\ndef test_decayed_lr():\n    schedule = lr.DecayedLR(100.0, [10, 20], gamma=0.1)\n    for i in range(0, 30):\n        schedule.step()\n        if i == 10: assert schedule.learning_rate == 10.0\n        elif i == 20: assert schedule.learning_rate == 1.0\n\ndef test_exp_decayed_lr():\n    schedule = lr.ExpDecayedLR(1.0, 0.1)\n    for i in range(0, 100):\n        schedule.step()\n    assert schedule.learning_rate == 1.0 * (0.1**100)\n\nif __name__ == '__main__':\n    test_interval_lr()\n    test_decayed_lr()\n    test_exp_decayed_lr()\n"
  }
]