├── .github ├── ISSUE_TEMPLATE │ ├── bug_report.md │ ├── feature_request.md │ └── regression_bug_report.md ├── labeler.yml └── workflows │ └── labeler.yml ├── BUILD ├── CONTRIBUTING.md ├── CONTRIBUTING_TECHNIQUE.md ├── LICENSE ├── MAINTAINING.md ├── OFFICIAL_MODELS.md ├── README.md ├── RELEASE.md ├── SECURITY.md ├── WORKSPACE ├── bazel.rc ├── ci └── kokoro │ ├── README.md │ ├── gcp_ubuntu │ ├── Dockerfile │ ├── build.sh │ ├── common.cfg │ ├── continuous.cfg │ └── presubmit.cfg │ ├── run_bazel_unittests.sh │ └── run_install_import_test.sh ├── community └── rfcs │ ├── README.md │ └── techniques │ ├── design │ └── 20200507-weight-clustering-design.md │ └── proposals │ └── 20200507-weight-clustering-proposal.md ├── pip_pkg.sh ├── requirements.txt ├── setup.py └── tensorflow_model_optimization ├── BUILD ├── __init__.py ├── g3doc ├── _book.yaml ├── _index.yaml ├── guide │ ├── clustering │ │ ├── clustering_comprehensive_guide.ipynb │ │ ├── clustering_example.ipynb │ │ └── index.md │ ├── combine │ │ ├── collaborative_optimization.md │ │ ├── cqat_example.ipynb │ │ ├── images │ │ │ ├── collaborative_optimization.png │ │ │ └── collaborative_optimization_dist.png │ │ ├── pcqat_example.ipynb │ │ ├── pqat_example.ipynb │ │ └── sparse_clustering_example.ipynb │ ├── get_started.md │ ├── index.md │ ├── install.md │ ├── optimize_further.md │ ├── pruning │ │ ├── README.md │ │ ├── comprehensive_guide.ipynb │ │ ├── index.md │ │ ├── pruning_for_on_device_inference.ipynb │ │ ├── pruning_with_keras.ipynb │ │ └── pruning_with_sparsity_2_by_4.ipynb │ ├── quantization │ │ ├── post_training.md │ │ ├── training.md │ │ ├── training_comprehensive_guide.ipynb │ │ └── training_example.ipynb │ └── roadmap.md └── tools │ └── build_docs.py ├── python ├── BUILD ├── __init__.py ├── core │ ├── BUILD │ ├── README.md │ ├── __init__.py │ ├── api │ │ ├── BUILD │ │ ├── __init__.py │ │ ├── clustering │ │ │ ├── __init__.py │ │ │ └── keras │ │ │ │ ├── __init__.py │ │ │ │ └── experimental │ │ │ │ └── __init__.py │ │ ├── experimental │ │ │ ├── __init__.py │ │ │ └── combine │ │ │ │ └── __init__.py │ │ ├── quantization │ │ │ ├── __init__.py │ │ │ └── keras │ │ │ │ ├── __init__.py │ │ │ │ ├── collab_opts │ │ │ │ └── __init__.py │ │ │ │ ├── default_8bit │ │ │ │ ├── __init__.py │ │ │ │ └── default_8bit_transforms │ │ │ │ │ └── __init__.py │ │ │ │ ├── experimental │ │ │ │ ├── __init__.py │ │ │ │ └── default_n_bit │ │ │ │ │ ├── __init__.py │ │ │ │ │ └── default_n_bit_transforms │ │ │ │ │ └── __init__.py │ │ │ │ ├── graph_transformations │ │ │ │ ├── __init__.py │ │ │ │ ├── model_transformer │ │ │ │ │ └── __init__.py │ │ │ │ └── transforms │ │ │ │ │ └── __init__.py │ │ │ │ └── quantizers │ │ │ │ └── __init__.py │ │ └── sparsity │ │ │ ├── __init__.py │ │ │ └── keras │ │ │ └── __init__.py │ ├── clustering │ │ ├── BUILD │ │ ├── __init__.py │ │ └── keras │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── cluster.py │ │ │ ├── cluster_config.py │ │ │ ├── cluster_distributed_test.py │ │ │ ├── cluster_integration_test.py │ │ │ ├── cluster_test.py │ │ │ ├── cluster_wrapper.py │ │ │ ├── cluster_wrapper_test.py │ │ │ ├── clusterable_layer.py │ │ │ ├── clustering_algorithm.py │ │ │ ├── clustering_callbacks.py │ │ │ ├── clustering_centroids.py │ │ │ ├── clustering_centroids_test.py │ │ │ ├── clustering_registry.py │ │ │ ├── clustering_registry_test.py │ │ │ ├── experimental │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ └── cluster.py │ │ │ ├── mnist_clusterable_layer_test.py │ │ │ └── mnist_clustering_test.py │ ├── common │ │ └── keras │ │ │ └── compression │ │ │ ├── BUILD │ │ │ ├── README.md │ │ │ ├── algorithm.py │ │ │ ├── algorithms │ │ │ ├── BUILD │ │ │ ├── bias_only.py │ │ │ ├── bias_only_test.py │ │ │ ├── different_training_and_inference.py │ │ │ ├── different_training_and_inference_test.py │ │ │ ├── epr.py │ │ │ ├── epr_test.py │ │ │ ├── periodical_update_and_scheduling.py │ │ │ ├── periodical_update_and_scheduling_test.py │ │ │ ├── same_training_and_inference.py │ │ │ ├── same_training_and_inference_test.py │ │ │ ├── weight_clustering.py │ │ │ └── weight_clustering_test.py │ │ │ ├── internal │ │ │ ├── BUILD │ │ │ └── optimize.py │ │ │ ├── schedules.py │ │ │ └── schedules_test.py │ ├── internal │ │ ├── BUILD │ │ ├── README.md │ │ ├── __init__.py │ │ └── tensor_encoding │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── core │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── core_encoder.py │ │ │ ├── core_encoder_test.py │ │ │ ├── encoding_stage.py │ │ │ ├── encoding_stage_test.py │ │ │ ├── gather_encoder.py │ │ │ ├── gather_encoder_test.py │ │ │ ├── simple_encoder.py │ │ │ └── simple_encoder_test.py │ │ │ ├── encoders │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── common_encoders.py │ │ │ └── common_encoders_test.py │ │ │ ├── stages │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── research │ │ │ │ ├── BUILD │ │ │ │ ├── __init__.py │ │ │ │ ├── clipping.py │ │ │ │ ├── clipping_test.py │ │ │ │ ├── kashin.py │ │ │ │ ├── kashin_test.py │ │ │ │ ├── misc.py │ │ │ │ ├── misc_test.py │ │ │ │ ├── quantization.py │ │ │ │ └── quantization_test.py │ │ │ ├── stages_impl.py │ │ │ └── stages_impl_test.py │ │ │ ├── testing │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── test_utils.py │ │ │ └── test_utils_test.py │ │ │ └── utils │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── py_utils.py │ │ │ ├── py_utils_test.py │ │ │ ├── tf_utils.py │ │ │ └── tf_utils_test.py │ ├── keras │ │ ├── BUILD │ │ ├── __init__.py │ │ ├── compat.py │ │ ├── metrics.py │ │ ├── metrics_test.py │ │ ├── test_utils.py │ │ ├── testing │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ └── test_utils_mnist.py │ │ └── utils.py │ ├── quantization │ │ ├── BUILD │ │ ├── __init__.py │ │ └── keras │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── collab_opts │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── cluster_preserve │ │ │ │ ├── BUILD │ │ │ │ ├── __init__.py │ │ │ │ ├── cluster_preserve_integration_test.py │ │ │ │ ├── cluster_preserve_quantize_registry.py │ │ │ │ ├── cluster_preserve_quantize_registry_test.py │ │ │ │ ├── cluster_utils.py │ │ │ │ ├── default_8bit_cluster_preserve_quantize_scheme.py │ │ │ │ └── mnist_prune_cluster_preserve_qat_test.py │ │ │ └── prune_preserve │ │ │ │ ├── BUILD │ │ │ │ ├── __init__.py │ │ │ │ ├── default_8bit_prune_preserve_quantize_scheme.py │ │ │ │ ├── prune_preserve_quantize_registry.py │ │ │ │ └── prune_preserve_quantize_registry_test.py │ │ │ ├── default_8bit │ │ │ ├── BUILD │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ ├── default_8bit_quantize_configs.py │ │ │ ├── default_8bit_quantize_layout_transform.py │ │ │ ├── default_8bit_quantize_registry.py │ │ │ ├── default_8bit_quantize_registry_test.py │ │ │ ├── default_8bit_quantize_scheme.py │ │ │ ├── default_8bit_quantizers.py │ │ │ ├── default_8bit_quantizers_test.py │ │ │ ├── default_8bit_transforms.py │ │ │ ├── default_8bit_transforms_test.py │ │ │ └── quantize_numerical_test.py │ │ │ ├── experimental │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ └── default_n_bit │ │ │ │ ├── BUILD │ │ │ │ ├── README.md │ │ │ │ ├── __init__.py │ │ │ │ ├── default_n_bit_quantize_configs.py │ │ │ │ ├── default_n_bit_quantize_layout_transform.py │ │ │ │ ├── default_n_bit_quantize_registry.py │ │ │ │ ├── default_n_bit_quantize_registry_test.py │ │ │ │ ├── default_n_bit_quantize_scheme.py │ │ │ │ ├── default_n_bit_quantizers.py │ │ │ │ ├── default_n_bit_quantizers_test.py │ │ │ │ ├── default_n_bit_transforms.py │ │ │ │ └── default_n_bit_transforms_test.py │ │ │ ├── graph_transformations │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── model_transformer.py │ │ │ ├── model_transformer_test.py │ │ │ ├── transforms.py │ │ │ └── transforms_test.py │ │ │ ├── layers │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── conv_batchnorm_test_utils.py │ │ │ └── dense_batchnorm_test_utils.py │ │ │ ├── quant_ops.py │ │ │ ├── quant_ops_test.py │ │ │ ├── quantize.py │ │ │ ├── quantize_annotate.py │ │ │ ├── quantize_annotate_test.py │ │ │ ├── quantize_aware_activation.py │ │ │ ├── quantize_aware_activation_test.py │ │ │ ├── quantize_config.py │ │ │ ├── quantize_functional_test.py │ │ │ ├── quantize_integration_test.py │ │ │ ├── quantize_layer.py │ │ │ ├── quantize_layer_test.py │ │ │ ├── quantize_layout_transform.py │ │ │ ├── quantize_models_test.py │ │ │ ├── quantize_registry.py │ │ │ ├── quantize_scheme.py │ │ │ ├── quantize_test.py │ │ │ ├── quantize_wrapper.py │ │ │ ├── quantize_wrapper_test.py │ │ │ ├── quantizers.py │ │ │ ├── quantizers_test.py │ │ │ └── utils.py │ ├── sparsity │ │ ├── BUILD │ │ ├── __init__.py │ │ └── keras │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── estimator_utils.py │ │ │ ├── prunable_layer.py │ │ │ ├── prune.py │ │ │ ├── prune_distributed_test.py │ │ │ ├── prune_integration_test.py │ │ │ ├── prune_registry.py │ │ │ ├── prune_registry_test.py │ │ │ ├── prune_test.py │ │ │ ├── pruning_callbacks.py │ │ │ ├── pruning_callbacks_test.py │ │ │ ├── pruning_impl.py │ │ │ ├── pruning_impl_test.py │ │ │ ├── pruning_policy.py │ │ │ ├── pruning_policy_test.py │ │ │ ├── pruning_schedule.py │ │ │ ├── pruning_schedule_test.py │ │ │ ├── pruning_utils.py │ │ │ ├── pruning_utils_test.py │ │ │ ├── pruning_wrapper.py │ │ │ ├── pruning_wrapper_test.py │ │ │ ├── test_utils.py │ │ │ └── tools │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── check_sparsity_m_by_n.py │ │ │ ├── evaluate_pruning.py │ │ │ ├── sparsity_tooling.py │ │ │ └── sparsity_tooling_test.py │ └── version.py └── examples │ ├── __init__.py │ ├── cluster_preserve_qat │ ├── __init__.py │ └── keras │ │ ├── BUILD │ │ ├── __init__.py │ │ └── mnist_cnn.py │ ├── clustering │ └── keras │ │ ├── imdb │ │ ├── BUILD │ │ ├── imdb_gru.py │ │ ├── imdb_lstm.py │ │ ├── imdb_multiple_cells.py │ │ ├── imdb_rnn.py │ │ └── imdb_utils.py │ │ └── mnist │ │ ├── BUILD │ │ ├── __init__.py │ │ ├── mnist_cnn.py │ │ └── mnist_mha.py │ ├── quantization │ ├── __init__.py │ └── keras │ │ ├── BUILD │ │ ├── __init__.py │ │ ├── mnist_cnn.py │ │ └── mnist_cnn_cont_quant.py │ ├── quantization_with_sparsity │ ├── __init__.py │ └── keras │ │ ├── BUILD │ │ ├── __init__.py │ │ └── mnist_cnn.py │ └── sparsity │ └── keras │ ├── imdb │ ├── BUILD │ └── imdb_lstm.py │ └── mnist │ ├── BUILD │ ├── __init__.py │ ├── mnist_cnn.py │ ├── mnist_e2e.py │ ├── mnist_e2e_sparsity2x4.py │ └── mnist_mha.py └── tensorflow_model_optimization.bzl /.github/ISSUE_TEMPLATE/bug_report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Bug report 3 | about: Create a report to help us improve 4 | title: '' 5 | labels: bug 6 | assignees: '' 7 | 8 | --- 9 | Prior to filing: check that this should be a bug instead of a feature request. Everything supported, including the compatible versions of TensorFlow, is listed in the overview page of each technique. For example, the overview page of quantization-aware training is [here](https://www.tensorflow.org/model_optimization/guide/quantization/training). An issue for anything not supported should be a feature request. 10 | 11 | **Describe the bug** 12 | A clear and concise description of what the bug is. 13 | 14 | **System information** 15 | 16 | TensorFlow version (installed from source or binary): 17 | 18 | TensorFlow Model Optimization version (installed from source or binary): 19 | 20 | Python version: 21 | 22 | **Describe the expected behavior** 23 | 24 | **Describe the current behavior** 25 | 26 | **Code to reproduce the issue** 27 | Provide a reproducible code that is the bare minimum necessary to generate the 28 | problem. 29 | 30 | **Screenshots** 31 | If applicable, add screenshots to help explain your problem. 32 | 33 | **Additional context** 34 | Add any other context about the problem here. 35 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Feature request 3 | about: Use this template for raising a feature request 4 | title: '' 5 | labels: feature request 6 | assignees: '' 7 | 8 | --- 9 | 10 | **System information** 11 | 12 | - TensorFlow version (you are using): 13 | - Are you willing to contribute it (Yes/No): 14 | 15 | **Motivation** 16 | 17 | What is the use case and how does it broadly benefits users? Prioritization and whether a feature is added is based on how it 18 | helps the community and the feature's maintenance costs. 19 | 20 | As examples: 21 | 22 | 1. Instead of, "Enable the technique for my model," "Enable this technique to work better with standard object detection models, including R-CNN (link) and SSD (link)" is stronger. 23 | 24 | 2. Instead of, "Try something more customized with the technique," "Implement a variant of the algorithm described in equations X and Y of this paper" is clearer. 25 | 26 | **Describe the feature** 27 | 28 | **Describe how the feature helps achieve the use case** 29 | 30 | **Describe how existing APIs don't satisfy your use case (optional if obvious)** 31 | 32 | As examples: 33 | 34 | 1. You tried using APIs X and Y and were able to do Z. However, that was not sufficient because of ... 35 | 36 | 2. You achieved your use case with the code snippet W. However, this was more difficult than it should be because of ... (e.g. ran into issue X or had 37 | to do Y). 38 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/regression_bug_report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Regression bug report 3 | about: Create a report for regressions 4 | title: '' 5 | labels: regression bug 6 | assignees: '' 7 | 8 | --- 9 | Prior to filing: check that this should be a bug instead of a feature request. Everything supported, including the compatible versions of TensorFlow, is listed in the overview page of each technique. For example, the overview page of quantization-aware training is [here](https://www.tensorflow.org/model_optimization/guide/quantization/training). An issue for anything not supported should be a feature request. 10 | 11 | **Describe the bug** 12 | A clear and concise description of what the bug is. 13 | 14 | **System information** 15 | 16 | TensorFlow version (installed from source or binary): 17 | 18 | TensorFlow Model Optimization version (installed from source or binary): 19 | 20 | Python version: 21 | 22 | The behavior was correct previously when using versions: 23 | - If this behavior was never correct, please file a report for a regular 24 | bug, not a regression. 25 | - (Optional) Commit/PR that introduced the regression: 26 | 27 | **Describe the expected behavior** 28 | 29 | **Describe the current behavior** 30 | 31 | **Code to reproduce the issue** 32 | Provide a reproducible code that is the bare minimum necessary to generate the 33 | problem. 34 | 35 | **Screenshots** 36 | If applicable, add screenshots to help explain your problem. 37 | 38 | **Additional context** 39 | Add any other context about the problem here. 40 | -------------------------------------------------------------------------------- /.github/labeler.yml: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | technique:clustering: 16 | - tensorflow_model_optimization/python/core/clustering/**/* 17 | - tensorflow_model_optimization/python/core/api/clustering/**/* 18 | - tensorflow_model_optimization/python/examples/clustering/**/* 19 | - tensorflow_model_optimization/g3doc/guide/clustering/**/* 20 | 21 | technique:pruning: 22 | - tensorflow_model_optimization/python/core/sparsity/**/* 23 | - tensorflow_model_optimization/python/core/api/sparsity/**/* 24 | - tensorflow_model_optimization/python/examples/sparsity/**/* 25 | - tensorflow_model_optimization/g3doc/guide/pruning/**/* 26 | 27 | technique:qat: 28 | - tensorflow_model_optimization/python/core/quantization/**/* 29 | - tensorflow_model_optimization/python/core/api/quantization/**/* 30 | - tensorflow_model_optimization/python/examples/quantization/**/* 31 | - tensorflow_model_optimization/g3doc/guide/quantization/**/* 32 | 33 | comp:weight compression: 34 | - tensorflow_model_optimization/python/core/common/keras/compression/**/* 35 | 36 | # Note: this does not capture all PRs that need a API review. 37 | api-review: 38 | - tensorflow_model_optimization/python/core/api/**/* 39 | -------------------------------------------------------------------------------- /.github/workflows/labeler.yml: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | name: "Pull Request Labeler" 16 | on: 17 | - pull_request_target 18 | 19 | jobs: 20 | triage: 21 | runs-on: ubuntu-latest 22 | steps: 23 | - uses: actions/labeler@v2 24 | with: 25 | repo-token: "${{ secrets.GITHUB_TOKEN }}" 26 | -------------------------------------------------------------------------------- /BUILD: -------------------------------------------------------------------------------- 1 | sh_binary( 2 | name = "pip_pkg", 3 | srcs = ["pip_pkg.sh"], 4 | data = [ 5 | "setup.py", 6 | "//tensorflow_model_optimization", 7 | ], 8 | ) 9 | -------------------------------------------------------------------------------- /MAINTAINING.md: -------------------------------------------------------------------------------- 1 | This document is for maintainers. 2 | 3 | It is a work in progress and suggestions are welcome. 4 | 5 | # Triaging 6 | A issue or PR should only be assigned to someone if they would be actively working on 7 | it or would be working on it shortly. This avoids the scenario when we think 8 | everything is being addressed when that is not the case, particularly when people are 9 | busy, which may cause issues to invisibly languish. 10 | 11 | The [maintainer for a particular 12 | technique](https://github.com/tensorflow/model-optimization#maintainers) is 13 | responsible for that technique. 14 | 15 | Github PRs 16 | - The maintainers should watch out for PRs, filtered by the corresponding 17 | technique label. For example, the label for quantization-aware training is 18 | [here](https://github.com/tensorflow/model-optimization/pulls?q=is%3Apr+is%3Aopen+label%3Atechnique%3Aqat+). 19 | This labeling is automated 20 | [here](https://github.com/tensorflow/model-optimization/blob/master/.github/workflows/labeler.yml). 21 | 22 | Github Issues 23 | - The maintainers should skim through issues and attach the appropriate technique label. This has not been automated yet. 24 | - Use other labels as needed, according to [CONTRIBUTING.md](https://github.com/tensorflow/model-optimization/blob/master/CONTRIBUTING.md). 25 | 26 | # Code Reviews 27 | Ensure that contributions adhere to the guidelines suggested in 28 | [CONTRIBUTING.md](https://github.com/tensorflow/model-optimization/blob/master/CONTRIBUTING.md). 29 | Comments that are frequently made in PRs should be added there. 30 | 31 | If there are public API changes, please attach the ["api-review" 32 | label](https://github.com/tensorflow/model-optimization/pulls?q=is%3Apr+is%3Aopen+label%3Aapi-review) 33 | for a closer look from the TFMOT team. 34 | 35 | Once you think a PR is ready to merge, apply the ["ready to pull" 36 | label](https://github.com/tensorflow/model-optimization/pulls?q=is%3Apr+is%3Aopen+label%3A%22ready+to+pull%22) 37 | and someone will merge it. 38 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # TensorFlow Model Optimization Toolkit 2 | 3 | The **TensorFlow Model Optimization Toolkit** is a suite of tools that users, 4 | both novice and advanced, can use to optimize machine learning models for 5 | deployment and execution. 6 | 7 | Supported techniques include quantization and pruning for sparse weights. 8 | There are APIs built specifically for Keras. 9 | 10 | For an overview of this project and individual tools, the optimization gains, 11 | and our roadmap refer to 12 | [tensorflow.org/model_optimization](https://www.tensorflow.org/model_optimization). 13 | The website also provides various tutorials and API docs. 14 | 15 | The toolkit provides stable Python APIs. 16 | 17 | ## Installation 18 | For installation instructions, see 19 | [tensorflow.org/model_optimization/guide/install](https://www.tensorflow.org/model_optimization/guide/install). 20 | 21 | ## Contribution guidelines 22 | 23 | **If you want to contribute to TensorFlow Model Optimization, be sure to review 24 | the [contribution guidelines](CONTRIBUTING.md). This project adheres to 25 | TensorFlow's 26 | [code of conduct](https://github.com/tensorflow/tensorflow/blob/master/CODE_OF_CONDUCT.md). 27 | By participating, you are expected to uphold this code.** 28 | 29 | **We use 30 | [GitHub issues](https://github.com/tensorflow/model-optimization/issues) for 31 | tracking requests and bugs.** 32 | 33 | ## Maintainers 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 |
SubpackageMaintainers
tfmot.clusteringArm ML Tooling
tfmot.quantizationTensorFlow Model Optimization
tfmot.sparsityTensorFlow Model Optimization
53 | 54 | ## Community 55 | 56 | As part of TensorFlow, we're committed to fostering an open and welcoming 57 | environment. 58 | 59 | * [TensorFlow Blog](https://blog.tensorflow.org): Stay up to date on content 60 | from the TensorFlow team and best articles from the community. 61 | -------------------------------------------------------------------------------- /SECURITY.md: -------------------------------------------------------------------------------- 1 | # TensorFlow Model Optimization Toolkit Security 2 | 3 | Please refer to [TensorFlow’s security model and guidelines][tf-security]. 4 | 5 | [tf-security]: https://github.com/tensorflow/tensorflow/blob/master/SECURITY.md 6 | -------------------------------------------------------------------------------- /WORKSPACE: -------------------------------------------------------------------------------- 1 | workspace(name = "tensorflow_model_optimization") 2 | -------------------------------------------------------------------------------- /bazel.rc: -------------------------------------------------------------------------------- 1 | # Import from tools location for backward compat with old bazel versions. 2 | import %workspace%/tools/bazel.rc 3 | -------------------------------------------------------------------------------- /ci/kokoro/README.md: -------------------------------------------------------------------------------- 1 | The entrypoint build script is gcp_ubuntu/build.sh. 2 | 3 | 4 | To manually trigger the unit test run as the CI system does: 5 | 6 | ``` 7 | export GIT_REPO_DIR= 8 | export KOKORO_ARTIFACTS_DIR="${GIT_REPO_DIR}" 9 | ci/kokoro/gcp_ubuntu/build.sh 10 | ``` 11 | 12 | That is, if the current directory is the root of the Git repository: 13 | ``` 14 | GIT_REPO_DIR="`pwd`" KOKORO_ARTIFACTS_DIR="`pwd`" ci/kokoro/gcp_ubuntu/build.sh 15 | ``` 16 | -------------------------------------------------------------------------------- /ci/kokoro/gcp_ubuntu/Dockerfile: -------------------------------------------------------------------------------- 1 | # Copyright 2020 Google LLC 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # https://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | 15 | # An image for building and testing TensorFlow Model Optimization on Ubuntu. 16 | # 17 | # Usage, assuming that the current directory is the root of the GitHub repos: 18 | # 19 | # Build: 20 | # docker build --tag={TAG} ci/kokoro/gcp_ubuntu/ 21 | # 22 | # Run interactively: 23 | # docker run -it --volume `pwd`:/tfmot --workdir /tfmot {TAG} 24 | 25 | # TODO(b/185727356): generalize to different versions of TensorFlow to 26 | # run CI against. 27 | 28 | # Latest Ubuntu LTS (Focal), at the moment. 29 | FROM ubuntu:22.04 30 | 31 | ARG BAZEL_VERSION=7.0.2 32 | ARG TENSORFLOW_VERSION=2.15.0 33 | 34 | 35 | RUN apt-get update -y 36 | 37 | # Install Python3 and set it as default 38 | RUN DEBIAN_FRONTEND=noninteractive apt-get install -y \ 39 | python3 \ 40 | python3-pip \ 41 | python3-virtualenv \ 42 | && update-alternatives --install /usr/bin/python python /usr/bin/python3 10 43 | 44 | # Install Bazel. 45 | RUN DEBIAN_FRONTEND=noninteractive apt-get install -y unzip zip wget \ 46 | && wget -O bazel-installer.sh "https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION?}/bazel-${BAZEL_VERSION?}-installer-linux-x86_64.sh" \ 47 | && chmod +x "bazel-installer.sh" \ 48 | && "./bazel-installer.sh" \ 49 | && rm "bazel-installer.sh" 50 | 51 | # Install TensorFlow 52 | RUN pip install "tensorflow==${TENSORFLOW_VERSION}" 53 | -------------------------------------------------------------------------------- /ci/kokoro/gcp_ubuntu/common.cfg: -------------------------------------------------------------------------------- 1 | # Format: //devtools/kokoro/config/proto/build.proto 2 | 3 | # Copyright 2021 Google LLC 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # https://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | 17 | build_file: "tensorflow_model_optimization/ci/kokoro/gcp_ubuntu/build.sh" 18 | 19 | action { 20 | define_artifacts { 21 | regex: "**/*sponge_log.xml" 22 | regex: "**/*sponge_log.log" 23 | } 24 | } -------------------------------------------------------------------------------- /ci/kokoro/gcp_ubuntu/continuous.cfg: -------------------------------------------------------------------------------- 1 | # Format: //devtools/kokoro/config/proto/build.proto 2 | 3 | # Copyright 2020 Google LLC 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # https://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | 17 | build_file: "tensorflow_model_optimization/ci/kokoro/gcp_ubuntu/build.sh" 18 | -------------------------------------------------------------------------------- /ci/kokoro/gcp_ubuntu/presubmit.cfg: -------------------------------------------------------------------------------- 1 | # Format: //devtools/kokoro/config/proto/build.proto 2 | 3 | # Copyright 2020 Google LLC 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # https://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | 17 | build_file: "tensorflow_model_optimization/ci/kokoro/gcp_ubuntu/build.sh" 18 | -------------------------------------------------------------------------------- /ci/kokoro/run_bazel_unittests.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Copyright 2021 Google LLC 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # https://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | 17 | # 18 | # Runs the unit tests with Bazel. 19 | # 20 | 21 | # Make Bash more strict, for easier debugging. 22 | set -e # Exit on the first error. 23 | set -u # Treat unset variables as error. 24 | set -o pipefail # Treat the failure of a command in a pipeline as error. 25 | 26 | # Display commands being run. 27 | # WARNING: please only enable 'set -x' if necessary for debugging, and be very 28 | # careful if you handle credentials (e.g. from Keystore) with 'set -x': 29 | # statements like "export VAR=$(cat /tmp/keystore/credentials)" will result in 30 | # the credentials being printed in build logs. 31 | # Additionally, recursive invocation with credentials as command-line 32 | # parameters, will print the full command, with credentials, in the build logs. 33 | # set -x 34 | 35 | pip install --requirement "requirements.txt" 36 | pip install tensorflow-compression>=2.11.0 37 | 38 | # Run the tests. 39 | # Some tests requiring more RAM than the CI machine provides are disabled. 40 | bazel test --test_size_filters="-enormous" --test_output=errors \ 41 | //tensorflow_model_optimization/python/core/... 42 | -------------------------------------------------------------------------------- /ci/kokoro/run_install_import_test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Copyright 2021 Google LLC 4 | # 5 | # Licensed under the Apache License, Version 2.0 (the "License"); 6 | # you may not use this file except in compliance with the License. 7 | # You may obtain a copy of the License at 8 | # 9 | # https://www.apache.org/licenses/LICENSE-2.0 10 | # 11 | # Unless required by applicable law or agreed to in writing, software 12 | # distributed under the License is distributed on an "AS IS" BASIS, 13 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | # See the License for the specific language governing permissions and 15 | # limitations under the License. 16 | 17 | # 18 | # A smoke test that installs TFMOT with `setup.py install` and then tries to 19 | # import it. 20 | # 21 | 22 | # Make Bash more strict, for easier debugging. 23 | set -e # Exit on the first error. 24 | set -u # Treat unset variables as error. 25 | set -o pipefail # Treat the failure of a command in a pipeline as error. 26 | 27 | # Display commands being run. 28 | # WARNING: please only enable 'set -x' if necessary for debugging, and be very 29 | # careful if you handle credentials (e.g. from Keystore) with 'set -x': 30 | # statements like "export VAR=$(cat /tmp/keystore/credentials)" will result in 31 | # the credentials being printed in build logs. 32 | # Additionally, recursive invocation with credentials as command-line 33 | # parameters, will print the full command, with credentials, in the build logs. 34 | # set -x 35 | 36 | 37 | # Create an isolated Python environment. 38 | mkdir -p venv 39 | virtualenv ./venv 40 | source ./venv/bin/activate 41 | 42 | # Install Tensorflow, then TFMOT from source. 43 | pip install tensorflow 44 | pip install --requirement "requirements.txt" 45 | python setup.py install 46 | 47 | # Go to another directory and try to import TFMOT. 48 | TEMP_PROJECT_DIR=$(mktemp --directory) 49 | pushd "${TEMP_PROJECT_DIR}" 50 | python -c "import tensorflow_model_optimization" 51 | -------------------------------------------------------------------------------- /community/rfcs/README.md: -------------------------------------------------------------------------------- 1 | # TensorFlow Model Optimization RFCs 2 | 3 | This directory stores approved RFCs for new techniques. 4 | 5 | ## Process 6 | 7 | Please carefully read our [new technique contribution guidelines](https://github.com/tensorflow/model-optimization/blob/master/CONTRIBUTING_TECHNIQUE.md). 8 | -------------------------------------------------------------------------------- /pip_pkg.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 3 | # 4 | # Licensed under the Apache License, Version 2.0 (the "License"); 5 | # you may not use this file except in compliance with the License. 6 | # You may obtain a copy of the License at 7 | # 8 | # http://www.apache.org/licenses/LICENSE-2.0 9 | # 10 | # Unless required by applicable law or agreed to in writing, software 11 | # distributed under the License is distributed on an "AS IS" BASIS, 12 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | # See the License for the specific language governing permissions and 14 | # limitations under the License. 15 | # ============================================================================== 16 | set -e 17 | set -x 18 | 19 | PLATFORM="$(uname -s | tr 'A-Z' 'a-z')" 20 | 21 | if [[ $# -lt 1 ]] ; then 22 | echo "Usage:" 23 | echo "pip_pkg /path/to/destination/directory" 24 | echo "all additional arguments (e.g. --flag1=v1 --flag2=v2) are passed to setup.py" 25 | exit 1 26 | fi 27 | 28 | # Create the destination directory, then do dirname on a non-existent file 29 | # inside it to give us a path with tilde characters resolved (readlink -f is 30 | # another way of doing this but is not available on a fresh macOS install). 31 | # Finally, use cd and pwd to get an absolute path, in case a relative one was 32 | # given. 33 | mkdir -p "$1" 34 | DEST=$(dirname "${1}/does_not_exist") 35 | DEST=$(cd "$DEST" && pwd) 36 | 37 | cd bazel-bin/pip_pkg.runfiles/_main 38 | 39 | # Pass through remaining arguments (following the first argument, which 40 | # specifies the output dir) to setup.py, e.g., 41 | # ./pip_pkg /tmp/tensorflow_model_optimization_pkg --release 42 | # passes `--release` to setup.py. 43 | python3 setup.py bdist_wheel --universal ${@:2} --dist-dir="$DEST" # >/dev/null 44 | 45 | set +x 46 | echo -e "\nBuild complete. Wheel files are in $DEST" 47 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | absl-py>=1.2.0 2 | numpy~=1.23.0 3 | six~=1.14 4 | scipy 5 | enum34~=1.1 6 | mock 7 | dm-tree~=0.1.1 8 | tf-keras>=2.14.1 9 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/BUILD: -------------------------------------------------------------------------------- 1 | # Description: 2 | # 3 | # TensorFlow Optimization is a repository for the training-time 4 | # portion of the Tensorflow Model Optimization Toolkit, 5 | # used to optimize machine learning models for deployment and execution. 6 | # 7 | # https://github.com/tensorflow/model-optimization 8 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_binary", "py_strict_library") 9 | 10 | package(default_visibility = ["//visibility:public"]) 11 | 12 | licenses(["notice"]) 13 | 14 | exports_files(["LICENSE"]) 15 | 16 | py_strict_library( 17 | name = "tensorflow_model_optimization", 18 | srcs = ["__init__.py"], 19 | deps = [ 20 | # distutils dep1, 21 | "//tensorflow_model_optimization/python", # buildcleaner: keep 22 | "//tensorflow_model_optimization/python/core:version", 23 | "//tensorflow_model_optimization/python/core/api", 24 | ], 25 | ) 26 | 27 | py_strict_binary( 28 | name = "build_docs", 29 | srcs = ["build_docs.py"], 30 | deps = [ 31 | # absl:app dep1, 32 | # absl/flags dep1, 33 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 34 | # tensorflow_docs/api_generator:generate_lib dep1, 35 | "//third_party/tensorflow_model_optimization", 36 | ], 37 | ) 38 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/g3doc/_book.yaml: -------------------------------------------------------------------------------- 1 | upper_tabs: 2 | # Tabs left of dropdown menu 3 | - include: /_upper_tabs_left.yaml 4 | - include: /api_docs/_upper_tabs_api.yaml 5 | # Dropdown menu 6 | - name: Resources 7 | path: /resources 8 | is_default: true 9 | menu: 10 | - include: /resources/_menu_toc.yaml 11 | lower_tabs: 12 | # Subsite tabs 13 | other: 14 | - name: Guide 15 | contents: 16 | - title: Model optimization guide 17 | path: /model_optimization/guide 18 | - title: Get started 19 | path: /model_optimization/guide/get_started 20 | - title: Optimize further 21 | path: /model_optimization/guide/optimize_further 22 | - title: Install 23 | path: /model_optimization/guide/install 24 | - title: Roadmap 25 | path: /model_optimization/guide/roadmap 26 | 27 | - heading: Weight pruning 28 | - title: Overview 29 | path: /model_optimization/guide/pruning 30 | - title: Pruning with Keras 31 | path: /model_optimization/guide/pruning/pruning_with_keras 32 | - title: Pruning comprehensive guide 33 | path: /model_optimization/guide/pruning/comprehensive_guide 34 | - title: Pruning for on-device inference with XNNPACK 35 | path: /model_optimization/guide/pruning/pruning_for_on_device_inference 36 | - title: Sparse weights using structural pruning 37 | path: /model_optimization/guide/pruning/pruning_with_sparsity_2_by_4 38 | 39 | - heading: Quantization 40 | - title: Quantization aware training overview 41 | path: /model_optimization/guide/quantization/training 42 | - title: Quantization aware training example 43 | path: /model_optimization/guide/quantization/training_example 44 | - title: Quantization aware training comprehensive guide 45 | path: /model_optimization/guide/quantization/training_comprehensive_guide 46 | - title: Post-training quantization 47 | path: /model_optimization/guide/quantization/post_training 48 | 49 | - heading: Weight clustering 50 | - title: Overview 51 | path: /model_optimization/guide/clustering 52 | - title: Weight clustering example 53 | path: /model_optimization/guide/clustering/clustering_example 54 | - title: Weight clustering comprehensive guide 55 | path: /model_optimization/guide/clustering/clustering_comprehensive_guide 56 | 57 | - heading: Collaborative optimization 58 | - title: Overview 59 | path: /model_optimization/guide/combine/collaborative_optimization 60 | - title: Cluster preserving quantization 61 | path: /model_optimization/guide/combine/cqat_example 62 | - title: Sparsity preserving quantization 63 | path: /model_optimization/guide/combine/pqat_example 64 | - title: Sparsity and cluster preserving quantization 65 | path: /model_optimization/guide/combine/pcqat_example 66 | - title: Sparsity preserving clustering 67 | path: /model_optimization/guide/combine/sparse_clustering_example 68 | 69 | - name: API 70 | skip_translation: true 71 | contents: 72 | - include: /model_optimization/api_docs/python/tfmot/_toc.yaml 73 | 74 | - include: /_upper_tabs_right.yaml 75 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/g3doc/guide/combine/images/collaborative_optimization.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/model-optimization/acd0717cebfde7dd0c707ddbf4915669c97e4b6a/tensorflow_model_optimization/g3doc/guide/combine/images/collaborative_optimization.png -------------------------------------------------------------------------------- /tensorflow_model_optimization/g3doc/guide/combine/images/collaborative_optimization_dist.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/model-optimization/acd0717cebfde7dd0c707ddbf4915669c97e4b6a/tensorflow_model_optimization/g3doc/guide/combine/images/collaborative_optimization_dist.png -------------------------------------------------------------------------------- /tensorflow_model_optimization/g3doc/guide/get_started.md: -------------------------------------------------------------------------------- 1 | # Get started with TensorFlow model optimization 2 | 3 | ## 1. Choose the best model for the task 4 | 5 | Depending on the task, you will need to make a tradeoff between model complexity 6 | and size. If your task requires high accuracy, then you may need a large and 7 | complex model. For tasks that require less precision, it is better to use a 8 | smaller model because they not only use less disk space and memory, but they are 9 | also generally faster and more energy efficient. 10 | 11 | ## 2. Pre-optimized models 12 | 13 | See if any existing 14 | [TensorFlow Lite pre-optimized models](https://www.tensorflow.org/lite/models) 15 | provide the efficiency required by your application. 16 | 17 | ## 3. Post-training tooling 18 | 19 | If you cannot use a pre-trained model for your application, try using 20 | [TensorFlow Lite post-training quantization tools](./quantization/post_training) 21 | during [TensorFlow Lite conversion](https://www.tensorflow.org/lite/convert), 22 | which can optimize your already-trained TensorFlow model. 23 | 24 | See the 25 | [post-training quantization tutorial](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/performance/post_training_quant.ipynb) 26 | to learn more. 27 | 28 | ## Next steps: Training-time tooling 29 | 30 | If the above simple solutions don't satisfy your needs, you may need to involve 31 | training-time optimization techniques. 32 | [Optimize further](optimize_further.md) with our training-time tools and dig 33 | deeper. 34 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/g3doc/guide/index.md: -------------------------------------------------------------------------------- 1 | # TensorFlow model optimization 2 | 3 | The *TensorFlow Model Optimization Toolkit* minimizes the complexity of 4 | optimizing machine learning inference. 5 | 6 | Inference efficiency is a critical concern when deploying machine learning 7 | models because of latency, memory utilization, and in many cases power 8 | consumption. Particularly on edge devices, such as mobile and Internet of Things 9 | (IoT), resources are further constrained, and model size and efficiency of 10 | computation become a major concern. 11 | 12 | Computational demand for *training* grows with the number of models trained on 13 | different architectures, whereas the computational demand for *inference* grows 14 | in proportion to the number of users. 15 | 16 | ## Use cases 17 | 18 | Model optimization is useful, among other things, for: 19 | 20 | * Reducing latency and cost for inference for both cloud and edge devices 21 | (e.g. mobile, IoT). 22 | * Deploying models on edge devices with restrictions on processing, memory 23 | and/or power-consumption. 24 | * Reducing payload size for over-the-air model updates. 25 | * Enabling execution on hardware restricted-to or optimized-for fixed-point 26 | operations. 27 | * Optimizing models for special purpose hardware accelerators. 28 | 29 | ## Optimization techniques 30 | 31 | The area of model optimization can involve various techniques: 32 | 33 | * Reduce parameter count with pruning and structured pruning. 34 | * Reduce representational precision with quantization. 35 | * Update the original model topology to a more efficient one with reduced 36 | parameters or faster execution. For example, tensor decomposition methods 37 | and distillation 38 | 39 | Our toolkit supports 40 | [post-training quantization](./quantization/post_training.md), 41 | [quantization aware training](./quantization/training.md), 42 | [pruning](./pruning/index.md), and [clustering](./clustering/index.md). The 43 | toolkit also provides experimental support for 44 | [collaborative optimization](./combine/collaborative_optimization.md) to combine 45 | various techniques. 46 | 47 | ### Quantization 48 | 49 | Quantized models are those where we represent the models with lower precision, 50 | such as 8-bit integers as opposed to 32-bit float. Lower precision is a 51 | requirement to leverage certain hardware. 52 | 53 | ### Sparsity and pruning 54 | 55 | Sparse models are those where connections in between operators (i.e. neural 56 | network layers) have been pruned, introducing zeros to the parameter tensors. 57 | 58 | ### Clustering 59 | 60 | Clustered models are those where the original model's parameters are replaced 61 | with a smaller number of unique values. 62 | 63 | ### Collaborative optimizaiton 64 | 65 | The toolkit provides experimental support for collaborative optimization. This 66 | enables you to benefit from combining several model compression techniques and 67 | simultaneously achieve improved accuracy through quantization aware training. 68 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/g3doc/guide/install.md: -------------------------------------------------------------------------------- 1 | # Install TensorFlow Model Optimization 2 | 3 | It is recommended to create a Python virtual environment before proceeding to 4 | the installation. Please see the TensorFlow installation 5 | [guide](https://www.tensorflow.org/install/pip#2.-create-a-virtual-environment-recommended) 6 | for more information. 7 | 8 | ### Stable Builds 9 | 10 | To install the latest version, run the following: 11 | 12 | ```shell 13 | # Installing with the `--upgrade` flag ensures you'll get the latest version. 14 | pip install --user --upgrade tensorflow-model-optimization 15 | ``` 16 | 17 | For release details, see our 18 | [release notes](https://github.com/tensorflow/model-optimization/releases). 19 | 20 | For the required version of TensorFlow and other compatibility information, see 21 | the API Compatibility Matrix section of the Overview page for the technique you 22 | intend to use. For instance, for pruning, the Overview page is 23 | [here](https://www.tensorflow.org/model_optimization/guide/pruning). 24 | 25 | Since TensorFlow is *not* included as a dependency of the TensorFlow Model 26 | Optimization package (in `setup.py`), you must explicitly install the TensorFlow 27 | package (`tf-nightly` or `tf-nightly-gpu`). This allows us to maintain one 28 | package instead of separate packages for CPU and GPU-enabled TensorFlow. 29 | 30 | ### Installing from Source 31 | 32 | You can also install from source. This requires the 33 | [Bazel](https://bazel.build/) build system. 34 | 35 | ```shell 36 | # To install dependencies on Ubuntu: 37 | # sudo apt-get install bazel git python-pip 38 | # For other platforms, see Bazel docs above. 39 | git clone https://github.com/tensorflow/model-optimization.git 40 | cd model-optimization 41 | bazel build --copt=-O3 --copt=-march=native :pip_pkg 42 | PKGDIR=$(mktemp -d) 43 | ./bazel-bin/pip_pkg $PKGDIR 44 | pip install --user --upgrade $PKGDIR/*.whl 45 | ``` 46 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/g3doc/guide/optimize_further.md: -------------------------------------------------------------------------------- 1 | # Optimize further 2 | 3 | When pre-optimized models and post-training tools do not satisfy your use case, 4 | the next step is to try the different training-time tools. 5 | 6 | Training time tools piggyback on the model's loss function over the training 7 | data such that the model can "adapt" to the changes brought by the optimization 8 | technique. 9 | 10 | The starting point to use our training APIs is a Keras training script, which 11 | can be optionally initialized from a pre-trained Keras model to further fine 12 | tune. 13 | 14 | 15 | Training time tools available for you to try: 16 | 17 | * [Weight pruning](./pruning/) 18 | * [Quantization](./quantization/training) 19 | * [Weight clustering](./clustering/) 20 | * [Collaborative optimization](./combine/collaborative_optimization) 21 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/g3doc/guide/pruning/README.md: -------------------------------------------------------------------------------- 1 | See https://tensorboard.dev/ for uploading new experiments for non-Colab users. 2 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/g3doc/guide/quantization/post_training.md: -------------------------------------------------------------------------------- 1 | # Post-training quantization 2 | 3 | Post-training quantization includes general techniques to reduce CPU and 4 | hardware accelerator latency, processing, power, and model size with little 5 | degradation in model accuracy. These techniques can be performed on an 6 | already-trained float TensorFlow model and applied during TensorFlow Lite 7 | conversion. These techniques are enabled as options in the 8 | [TensorFlow Lite converter](https://www.tensorflow.org/lite/convert/). 9 | 10 | To jump right into end-to-end examples, see the following tutorials: 11 | 12 | - [Post-training dynamic range 13 | quantization](https://www.tensorflow.org/lite/performance/post_training_quant) 14 | - [Post-training full integer 15 | quantization](https://www.tensorflow.org/lite/performance/post_training_integer_quant) 16 | - [Post-training float16 17 | quantization](https://www.tensorflow.org/lite/performance/post_training_float16_quant) 18 | 19 | 20 | ## Quantizing weights 21 | 22 | Weights can be converted to types with reduced precision, such as 16 bit floats 23 | or 8 bit integers. We generally recommend 16-bit floats for GPU acceleration and 24 | 8-bit integer for CPU execution. 25 | 26 | For example, here is how to specify 8 bit integer weight quantization: 27 | 28 | ``` 29 | import tensorflow as tf 30 | converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir) 31 | converter.optimizations = [tf.lite.Optimize.DEFAULT] 32 | tflite_quant_model = converter.convert() 33 | ``` 34 | 35 | At inference, the most critically intensive parts are computed with 8 bits 36 | instead of floating point. There is some inference-time performance overhead, 37 | relative to quantizing both weights and activations below. 38 | 39 | For more information, see the TensorFlow Lite 40 | [post-training quantization](https://www.tensorflow.org/lite/performance/post_training_quantization) 41 | guide. 42 | 43 | ## Full integer quantization of weights and activations 44 | 45 | Improve latency, processing, and power usage, and get access to integer-only 46 | hardware accelerators by making sure both weights and activations are quantized. 47 | This requires a small representative data set. 48 | 49 | ``` 50 | import tensorflow as tf 51 | 52 | def representative_dataset_gen(): 53 | for _ in range(num_calibration_steps): 54 | # Get sample input data as a numpy array in a method of your choosing. 55 | yield [input] 56 | 57 | converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir) 58 | converter.optimizations = [tf.lite.Optimize.DEFAULT] 59 | converter.representative_dataset = representative_dataset_gen 60 | tflite_quant_model = converter.convert() 61 | ``` 62 | 63 | The resulting model will still take float input and output for convenience. 64 | 65 | For more information, see the TensorFlow Lite 66 | [post-training quantization](https://www.tensorflow.org/lite/performance/post_training_quantization#full_integer_quantization_of_weights_and_activations) 67 | guide. 68 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/g3doc/guide/roadmap.md: -------------------------------------------------------------------------------- 1 | **Updated: June, 2021** 2 | 3 | TensorFlow’s Model Optimization Toolkit (MOT) has been used widely for 4 | converting/optimizing TensorFlow models to TensorFlow Lite models with smaller 5 | size, better performance and acceptable accuracy to run them on mobile and IoT 6 | devices. We are now working to extend MOT techniques and tooling beyond 7 | TensorFlow Lite to support TensorFlow SavedModel as well. 8 | 9 | The following represents a high level overview of our roadmap. You should be 10 | aware that this roadmap may change at any time and the order below does not 11 | reflect any type of priority. We strongly encourage you to comment on our 12 | roadmap and provide us feedback in the 13 | [discussion group](https://groups.google.com/a/tensorflow.org/g/tflite). 14 | 15 | ## Quantization 16 | 17 | #### TensorFlow Lite 18 | 19 | * Selective post-training quantization to exclude certain layers from 20 | quantization. 21 | * Quantization debugger to inspect quantization error losses per layer. 22 | * Applying quantization-aware training on more model coverage e.g. TensorFlow 23 | Model Garden. 24 | * Quality and performance improvements for post-training dynamic-range. 25 | quantization. 26 | 27 | #### TensorFlow 28 | 29 | * Post Training Quantization (bf16 * int8 dynamic range). 30 | * Quantization Aware Training ((bf16 * int8 weight-only with fake quant). 31 | * Selective post-training quantization to exclude certain layers from 32 | quantization. 33 | * Quantization debugger to inspect quantization error losses per layer. 34 | 35 | ## Sparsity 36 | 37 | #### TensorFlow Lite 38 | 39 | * Sparse model execution support for more models. 40 | * Target aware authoring for Sparsity. 41 | * Extend sparse op set with performant x86 kernels. 42 | 43 | #### TensorFlow 44 | 45 | * Sparity support in TensorFlow. 46 | 47 | ## Cascading compression techniques 48 | 49 | * Quantization + Tensor Compression + Sparsity: demonstrate all 50 | 3 techniques working together. 51 | 52 | ## Compression 53 | 54 | * Tensor compression API to help compression algorithm developers implement 55 | their own model compression algorithm (e.g. Weight Clustering) including 56 | providing a standard way to test/benchmark. 57 | 58 | 59 | 60 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | 3 | package(default_visibility = ["//visibility:public"]) 4 | 5 | licenses(["notice"]) 6 | 7 | py_strict_library( 8 | name = "python", 9 | srcs = ["__init__.py"], 10 | deps = [ 11 | "//tensorflow_model_optimization/python/core", # buildcleaner: keep 12 | ], 13 | ) 14 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | 3 | package(default_visibility = ["//visibility:public"]) 4 | 5 | licenses(["notice"]) 6 | 7 | py_strict_library( 8 | name = "core", 9 | srcs = ["__init__.py"], 10 | deps = [ 11 | "//tensorflow_model_optimization/python/core/api", # buildcleaner: keep 12 | "//tensorflow_model_optimization/python/core/clustering", # buildcleaner: keep 13 | "//tensorflow_model_optimization/python/core/internal", # buildcleaner: keep 14 | "//tensorflow_model_optimization/python/core/keras", # buildcleaner: keep 15 | "//tensorflow_model_optimization/python/core/quantization", # buildcleaner: keep 16 | "//tensorflow_model_optimization/python/core/sparsity", # buildcleaner: keep 17 | ], 18 | ) 19 | 20 | py_strict_library( 21 | name = "version", 22 | srcs = ["version.py"], 23 | ) 24 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/README.md: -------------------------------------------------------------------------------- 1 | Directories: 2 | - Keras: shared code between implementations of different techniques in Keras, 3 | meant to be independent of the techniques themselves. Some of this code may 4 | move to internal/keras eventually. 5 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Import API modules for Tensorflow Model Optimization.""" 16 | from tensorflow_model_optimization.python.core.api import clustering 17 | from tensorflow_model_optimization.python.core.api import experimental 18 | from tensorflow_model_optimization.python.core.api import quantization 19 | from tensorflow_model_optimization.python.core.api import sparsity 20 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/clustering/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing code for clustering.""" 16 | from tensorflow_model_optimization.python.core.api.clustering import keras 17 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/clustering/keras/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing clustering code built on Keras abstractions.""" 16 | # pylint: disable=g-bad-import-order 17 | from tensorflow_model_optimization.python.core.clustering.keras import experimental 18 | 19 | from tensorflow_model_optimization.python.core.clustering.keras.cluster import cluster_scope 20 | from tensorflow_model_optimization.python.core.clustering.keras.cluster import cluster_weights 21 | from tensorflow_model_optimization.python.core.clustering.keras.cluster import strip_clustering 22 | from tensorflow_model_optimization.python.core.clustering.keras import clusterable_layer 23 | 24 | from tensorflow_model_optimization.python.core.clustering.keras.cluster_config import CentroidInitialization 25 | from tensorflow_model_optimization.python.core.clustering.keras.clustering_algorithm import ClusteringAlgorithm 26 | from tensorflow_model_optimization.python.core.clustering.keras.clustering_callbacks import ClusteringSummaries 27 | from tensorflow_model_optimization.python.core.clustering.keras.clusterable_layer import ClusterableLayer 28 | # pylint: enable=g-bad-import-order 29 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/clustering/keras/experimental/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing experimental clustering code built on Keras abstractions.""" 16 | from tensorflow_model_optimization.python.core.clustering.keras.experimental.cluster import cluster_weights 17 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/experimental/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing code for experimental features.""" 16 | from tensorflow_model_optimization.python.core.api.experimental import combine 17 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/experimental/combine/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing collaborative optimization code.""" 16 | 17 | from tensorflow_model_optimization.python.core.quantization.keras.collab_opts.cluster_preserve.cluster_utils import ( 18 | strip_clustering_cqat,) 19 | 20 | from tensorflow_model_optimization.python.core.quantization.keras.collab_opts.cluster_preserve.default_8bit_cluster_preserve_quantize_scheme import ( 21 | Default8BitClusterPreserveQuantizeScheme,) 22 | 23 | # Deprecated import. 24 | # Please import from tfmot.quantization.keras.collab_opts 25 | from tensorflow_model_optimization.python.core.quantization.keras.collab_opts.prune_preserve.default_8bit_prune_preserve_quantize_scheme import ( 26 | Default8BitPrunePreserveQuantizeScheme,) 27 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/quantization/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing code for quantization.""" 16 | from tensorflow_model_optimization.python.core.api.quantization import keras 17 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/quantization/keras/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing quantization code built on Keras abstractions.""" 16 | # pylint: disable=g-bad-import-order 17 | 18 | # submodules 19 | from tensorflow_model_optimization.python.core.api.quantization.keras import quantizers 20 | from tensorflow_model_optimization.python.core.api.quantization.keras import default_8bit 21 | from tensorflow_model_optimization.python.core.api.quantization.keras import graph_transformations 22 | from tensorflow_model_optimization.python.core.api.quantization.keras import collab_opts 23 | from tensorflow_model_optimization.python.core.api.quantization.keras import experimental 24 | 25 | # quantize all layers with default quantization implementation. 26 | from tensorflow_model_optimization.python.core.quantization.keras.quantize import quantize_model 27 | 28 | # quantize some layers with default or custom quantization implementation. 29 | from tensorflow_model_optimization.python.core.quantization.keras.quantize import fix_input_output_range 30 | from tensorflow_model_optimization.python.core.quantization.keras.quantize import quantize_annotate_layer 31 | from tensorflow_model_optimization.python.core.quantization.keras.quantize import quantize_annotate_model 32 | from tensorflow_model_optimization.python.core.quantization.keras.quantize import quantize_apply 33 | from tensorflow_model_optimization.python.core.quantization.keras.quantize import remove_input_range 34 | 35 | # quantize with custom quantization parameterization or implementation, or 36 | # handle custom Keras layers. 37 | from tensorflow_model_optimization.python.core.quantization.keras.quantize_config import QuantizeConfig 38 | from tensorflow_model_optimization.python.core.quantization.keras.quantize_wrapper import QuantizeWrapper 39 | from tensorflow_model_optimization.python.core.quantization.keras.quantize_wrapper import QuantizeWrapperV2 40 | # Deserialize quantized model for Keras h5 format. 41 | from tensorflow_model_optimization.python.core.quantization.keras.quantize import quantize_scope 42 | 43 | # Quantization Scheme classes. 44 | from tensorflow_model_optimization.python.core.quantization.keras.quantize_scheme import QuantizeScheme 45 | from tensorflow_model_optimization.python.core.quantization.keras.quantize_layout_transform import QuantizeLayoutTransform 46 | from tensorflow_model_optimization.python.core.quantization.keras.quantize_registry import QuantizeRegistry 47 | 48 | # pylint: enable=g-bad-import-order 49 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/quantization/keras/collab_opts/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing collaborative optimization code.""" 16 | 17 | from tensorflow_model_optimization.python.core.quantization.keras.collab_opts.prune_preserve.default_8bit_prune_preserve_quantize_scheme import ( 18 | Default8BitPrunePreserveQuantizeScheme,) 19 | 20 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/quantization/keras/default_8bit/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing 8bit default quantization scheme.""" 16 | # pylint: disable=g-bad-import-order 17 | 18 | # submodules 19 | from tensorflow_model_optimization.python.core.api.quantization.keras.default_8bit import default_8bit_transforms 20 | 21 | # The 8bit default quantization scheme classes. 22 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_quantize_scheme import Default8BitQuantizeScheme 23 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_quantize_layout_transform import Default8BitQuantizeLayoutTransform 24 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_quantize_registry import Default8BitQuantizeRegistry 25 | 26 | # pylint: enable=g-bad-import-order 27 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/quantization/keras/default_8bit/default_8bit_transforms/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing 8bit default transforms.""" 16 | 17 | # The 8bit default transform classes. 18 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import ConcatTransform 19 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import ConcatTransform3Inputs 20 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import ConcatTransform4Inputs 21 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import ConcatTransform5Inputs 22 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import ConcatTransform6Inputs 23 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import Conv2DBatchNormActivationQuantize 24 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import Conv2DBatchNormQuantize 25 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import Conv2DBatchNormReLUQuantize 26 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import Conv2DReshapeBatchNormActivationQuantize 27 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import Conv2DReshapeBatchNormQuantize 28 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import Conv2DReshapeBatchNormReLUQuantize 29 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import InputLayerQuantize 30 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import LayerReluActivationQuantize 31 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import LayerReLUQuantize 32 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import SeparableConv1DQuantize 33 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit.default_8bit_transforms import SeparableConvQuantize 34 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/quantization/keras/experimental/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing experimental quantization features.""" 16 | # pylint: disable=g-bad-import-order 17 | 18 | # submodules 19 | from tensorflow_model_optimization.python.core.api.quantization.keras.experimental import default_n_bit 20 | 21 | # pylint: enable=g-bad-import-order 22 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/quantization/keras/experimental/default_n_bit/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing N-bit default quantization scheme.""" 16 | # pylint: disable=g-bad-import-order 17 | 18 | # submodules 19 | from tensorflow_model_optimization.python.core.api.quantization.keras.experimental.default_n_bit import default_n_bit_transforms 20 | 21 | # The N-bit default quantization scheme classes. 22 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_quantize_scheme import DefaultNBitQuantizeScheme 23 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_quantize_layout_transform import DefaultNBitQuantizeLayoutTransform 24 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_quantize_registry import DefaultNBitQuantizeRegistry 25 | 26 | # pylint: enable=g-bad-import-order 27 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/quantization/keras/experimental/default_n_bit/default_n_bit_transforms/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing N-bit default transforms.""" 16 | 17 | # The 8bit default transform classes. 18 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import ConcatTransform 19 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import ConcatTransform3Inputs 20 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import ConcatTransform4Inputs 21 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import ConcatTransform5Inputs 22 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import ConcatTransform6Inputs 23 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import Conv2DBatchNormActivationQuantize 24 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import Conv2DBatchNormQuantize 25 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import Conv2DBatchNormReLUQuantize 26 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import Conv2DReshapeBatchNormActivationQuantize 27 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import Conv2DReshapeBatchNormQuantize 28 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import Conv2DReshapeBatchNormReLUQuantize 29 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import InputLayerQuantize 30 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import LayerReluActivationQuantize 31 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import LayerReLUQuantize 32 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import SeparableConv1DQuantize 33 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit.default_n_bit_transforms import SeparableConvQuantize 34 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/quantization/keras/graph_transformations/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing code for graph transformation.""" 16 | 17 | # submodules 18 | from tensorflow_model_optimization.python.core.api.quantization.keras.graph_transformations import model_transformer 19 | from tensorflow_model_optimization.python.core.api.quantization.keras.graph_transformations import transforms 20 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/quantization/keras/graph_transformations/model_transformer/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing classes for model transformer.""" 16 | 17 | from tensorflow_model_optimization.python.core.quantization.keras.graph_transformations.model_transformer import ModelTransformer 18 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/quantization/keras/graph_transformations/transforms/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing classes for transform.""" 16 | 17 | from tensorflow_model_optimization.python.core.quantization.keras.graph_transformations.transforms import LayerNode 18 | from tensorflow_model_optimization.python.core.quantization.keras.graph_transformations.transforms import LayerPattern 19 | from tensorflow_model_optimization.python.core.quantization.keras.graph_transformations.transforms import Transform 20 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/quantization/keras/quantizers/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing Quantization abstraction and quantizers.""" 16 | 17 | # quantize with custom quantization parameterization or implementation, or 18 | # handle custom Keras layers. 19 | from tensorflow_model_optimization.python.core.quantization.keras.quantizers import AllValuesQuantizer 20 | from tensorflow_model_optimization.python.core.quantization.keras.quantizers import FixedQuantizer 21 | from tensorflow_model_optimization.python.core.quantization.keras.quantizers import LastValueQuantizer 22 | from tensorflow_model_optimization.python.core.quantization.keras.quantizers import MovingAverageQuantizer 23 | from tensorflow_model_optimization.python.core.quantization.keras.quantizers import Quantizer 24 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/sparsity/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing code for sparsity.""" 16 | from tensorflow_model_optimization.python.core.api.sparsity import keras 17 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/api/sparsity/keras/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Module containing sparsity code built on Keras abstractions.""" 16 | # pylint: disable=g-bad-import-order 17 | 18 | from tensorflow_model_optimization.python.core.sparsity.keras.prune import prune_low_magnitude 19 | from tensorflow_model_optimization.python.core.sparsity.keras.prune import prune_scope 20 | from tensorflow_model_optimization.python.core.sparsity.keras.prune import strip_pruning 21 | 22 | from tensorflow_model_optimization.python.core.sparsity.keras.pruning_callbacks import UpdatePruningStep 23 | from tensorflow_model_optimization.python.core.sparsity.keras.pruning_callbacks import PruningSummaries 24 | 25 | from tensorflow_model_optimization.python.core.sparsity.keras.pruning_schedule import PruningSchedule 26 | from tensorflow_model_optimization.python.core.sparsity.keras.pruning_schedule import ConstantSparsity 27 | from tensorflow_model_optimization.python.core.sparsity.keras.pruning_schedule import PolynomialDecay 28 | 29 | from tensorflow_model_optimization.python.core.sparsity.keras.prunable_layer import PrunableLayer 30 | 31 | from tensorflow_model_optimization.python.core.sparsity.keras.pruning_policy import PruningPolicy 32 | from tensorflow_model_optimization.python.core.sparsity.keras.pruning_policy import PruneForLatencyOnXNNPack 33 | 34 | # pylint: enable=g-bad-import-order 35 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/clustering/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | 3 | package(default_visibility = ["//visibility:public"]) 4 | 5 | licenses(["notice"]) 6 | 7 | py_strict_library( 8 | name = "clustering", 9 | srcs = ["__init__.py"], 10 | deps = [ 11 | "//tensorflow_model_optimization/python/core/clustering/keras", # buildcleaner: keep 12 | ], 13 | ) 14 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/clustering/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/model-optimization/acd0717cebfde7dd0c707ddbf4915669c97e4b6a/tensorflow_model_optimization/python/core/clustering/__init__.py -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/clustering/keras/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/model-optimization/acd0717cebfde7dd0c707ddbf4915669c97e4b6a/tensorflow_model_optimization/python/core/clustering/keras/__init__.py -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/clustering/keras/cluster_config.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Configuration classes for clustering.""" 16 | 17 | import enum 18 | 19 | 20 | class CentroidInitialization(str, enum.Enum): 21 | """Specifies how the cluster centroids should be initialized. 22 | 23 | * `LINEAR`: Cluster centroids are evenly spaced between the minimum and 24 | maximum values of a given weight tensor. 25 | * `RANDOM`: Centroids are sampled using the uniform distribution between the 26 | minimum and maximum weight values in a given layer. 27 | * `DENSITY_BASED`: Density-based sampling obtained as follows: first a 28 | cumulative distribution function is built for the weights, then the Y 29 | axis is evenly spaced into as many regions as many clusters we want to 30 | have. After this the corresponding X values are obtained and used to 31 | initialize the clusters centroids. 32 | * `KMEANS_PLUS_PLUS`: cluster centroids using the kmeans++ algorithm 33 | """ 34 | LINEAR = "CentroidInitialization.LINEAR" 35 | RANDOM = "CentroidInitialization.RANDOM" 36 | DENSITY_BASED = "CentroidInitialization.DENSITY_BASED" 37 | KMEANS_PLUS_PLUS = "CentroidInitialization.KMEANS_PLUS_PLUS" 38 | 39 | 40 | class GradientAggregation(str, enum.Enum): 41 | """Specifies how the cluster gradient should be aggregated. 42 | 43 | * `SUM`: The gradient of each cluster centroid is the sum of their 44 | respective child’s weight gradient. 45 | * `AVG`: The gradient of each cluster centroid is the averaged sum of 46 | their respective child’s weight gradient. 47 | """ 48 | SUM = "GradientAggregation.SUM" 49 | AVG = "GradientAggregation.AVG" 50 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/clustering/keras/clusterable_layer.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Clusterable layer API class for Keras models.""" 16 | 17 | import abc 18 | import six 19 | 20 | 21 | @six.add_metaclass(abc.ABCMeta) 22 | class ClusterableLayer: 23 | """Abstract Base Class for making your own keras layer clusterable. 24 | 25 | Your layer could be derived from a keras built-in layer or 26 | it could be a keras custom layer. 27 | 28 | The function get_clusterable_weights should be provided in both cases. 29 | 30 | The function get_clusterable_algorithm is provided, when weights for 31 | clustering is added in the keras layer. 32 | 33 | """ 34 | 35 | @abc.abstractmethod 36 | def get_clusterable_weights(self): 37 | """Returns list of clusterable weight tensors. 38 | 39 | All the weight tensors which the layer wants to be clustered during 40 | training must be returned by this method. 41 | 42 | Returns: List of weight tensors/kernels in the keras layer which must be 43 | clustered during training. Each element in the list is a (name, kernel) 44 | 2-tuple that consists of the name of the clusterable kernel and the 45 | kernel object itself. 46 | """ 47 | raise NotImplementedError('Must be implemented in subclasses.') 48 | 49 | def get_clusterable_algorithm(self, weight_name): # pylint: disable=unused-argument 50 | """Returns class with the clustering algorithm for the given weight_name. 51 | 52 | This function needs to be implemented for the customerable layers. 53 | If the layer is derived from the built-in keras layer, the clustering 54 | algorithm for the base built-in keras layer is used. 55 | 56 | The returned class should be derived from ClusteringAlgorithm and 57 | implements the function get_pulling_indices. 58 | This function is used to provide a special lookup function for the custom 59 | weights. 60 | It reshapes and tile centroids the same way as the weights. This allows us 61 | to find pulling indices efficiently. 62 | 63 | Args: 64 | weight_name ([string]): The name of the weight variable. 65 | """ 66 | return None 67 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/clustering/keras/experimental/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | 3 | package(default_visibility = [ 4 | "//tensorflow_model_optimization:__subpackages__", 5 | ]) 6 | 7 | licenses(["notice"]) 8 | 9 | py_strict_library( 10 | name = "experimental", 11 | srcs = [ 12 | "__init__.py", 13 | ], 14 | deps = [ 15 | ":cluster", # buildcleaner: keep 16 | ], 17 | ) 18 | 19 | py_strict_library( 20 | name = "cluster", 21 | srcs = ["cluster.py"], 22 | visibility = ["//visibility:public"], 23 | deps = [ 24 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster", 25 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster_config", 26 | ], 27 | ) 28 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/clustering/keras/experimental/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/common/keras/compression/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "pytype_strict_library", "pytype_strict_test") 2 | 3 | licenses(["notice"]) 4 | 5 | pytype_strict_library( 6 | name = "algorithm", 7 | srcs = ["algorithm.py"], 8 | visibility = ["//visibility:public"], 9 | deps = [ 10 | # tensorflow dep1, 11 | "//tensorflow_model_optimization/python/core/common/keras/compression/internal:optimize", 12 | "//tensorflow_model_optimization/python/core/keras:compat", 13 | ], 14 | ) 15 | 16 | pytype_strict_library( 17 | name = "schedules", 18 | srcs = ["schedules.py"], 19 | deps = [ 20 | # tensorflow dep1, 21 | ], 22 | ) 23 | 24 | pytype_strict_test( 25 | name = "schedules_test", 26 | srcs = [ 27 | "schedules_test.py", 28 | ], 29 | deps = [ 30 | ":schedules", 31 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 32 | # tensorflow dep1, 33 | ], 34 | ) 35 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/common/keras/compression/README.md: -------------------------------------------------------------------------------- 1 | # Compression API 2 | 3 | These modules implement the TFMOT's Compression APIs. But they are experimental 4 | and are not yet officially supported. 5 | 6 | ## Modules 7 | 8 | algorithm.py : The API that algorithm developers can inherit to implement any 9 | custom model compression algorithms. 10 | 11 | algorithms/ : Repository of example algorithms that can be used to Compression 12 | API, and corresponding end-to-end tests to demonstrate the usage. 13 | 14 | internal/ : Libraries for internal usage. 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/common/keras/compression/internal/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "pytype_strict_library") 2 | 3 | package(default_visibility = [ 4 | "//tensorflow_model_optimization/python/core/common/keras/compression:__subpackages__", 5 | ]) 6 | 7 | licenses(["notice"]) 8 | 9 | pytype_strict_library( 10 | name = "optimize", 11 | srcs = ["optimize.py"], 12 | deps = [ 13 | # tensorflow dep1, 14 | "//tensorflow_model_optimization/python/core/keras:compat", 15 | ], 16 | ) 17 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | 3 | package(default_visibility = ["//visibility:public"]) 4 | 5 | licenses(["notice"]) 6 | 7 | py_strict_library( 8 | name = "internal", 9 | srcs = ["__init__.py"], 10 | deps = [ 11 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding", 12 | ], 13 | ) 14 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/README.md: -------------------------------------------------------------------------------- 1 | # Internal 2 | 3 | The internal package contains a collection of modules that are documented, 4 | fully working, and there exists a plan to be graduated to the main `tf.mot` 5 | namespace -- either a literal swap of `tf.mot.internal.` for 6 | `tf.mot.`, or creation of separate public API that calls the internal 7 | API under the hood. 8 | 9 | The `tf.mot.internal` modules can change in backwards incompatible manner. Thus, 10 | depending on these modules is not recommended. 11 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/tensor_encoding/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | 3 | package(default_visibility = ["//visibility:public"]) 4 | 5 | licenses(["notice"]) 6 | 7 | py_strict_library( 8 | name = "tensor_encoding", 9 | srcs = ["__init__.py"], 10 | deps = [ 11 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core", 12 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/encoders", 13 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/stages", 14 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/testing", 15 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/utils", 16 | ], 17 | ) 18 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/tensor_encoding/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019, The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """The Tensor Encoding package.""" 15 | 16 | from __future__ import absolute_import 17 | from __future__ import division 18 | from __future__ import print_function 19 | 20 | from tensorflow_model_optimization.python.core.internal.tensor_encoding import core 21 | from tensorflow_model_optimization.python.core.internal.tensor_encoding import encoders 22 | from tensorflow_model_optimization.python.core.internal.tensor_encoding import stages 23 | from tensorflow_model_optimization.python.core.internal.tensor_encoding import testing 24 | from tensorflow_model_optimization.python.core.internal.tensor_encoding import utils 25 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/tensor_encoding/core/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019, The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Core parts of the `tensor_encoding` package.""" 15 | 16 | from __future__ import absolute_import 17 | from __future__ import division 18 | from __future__ import print_function 19 | 20 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.core.core_encoder import Encoder 21 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.core.core_encoder import EncoderComposer 22 | 23 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.core.encoding_stage import AdaptiveEncodingStageInterface 24 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.core.encoding_stage import EncodingStageInterface 25 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.core.encoding_stage import StateAggregationMode 26 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.core.encoding_stage import tf_style_adaptive_encoding_stage 27 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.core.encoding_stage import tf_style_encoding_stage 28 | 29 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.core.gather_encoder import GatherEncoder 30 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.core.simple_encoder import SimpleEncoder 31 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/tensor_encoding/encoders/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library", "py_strict_test") 2 | 3 | package(default_visibility = [ 4 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding:__subpackages__", 5 | ]) 6 | 7 | licenses(["notice"]) 8 | 9 | py_strict_library( 10 | name = "encoders", 11 | srcs = ["__init__.py"], 12 | visibility = ["//visibility:public"], 13 | deps = [ 14 | ":common_encoders", 15 | ], 16 | ) 17 | 18 | py_strict_library( 19 | name = "common_encoders", 20 | srcs = ["common_encoders.py"], 21 | deps = [ 22 | # tensorflow dep1, 23 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:core_encoder", 24 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:gather_encoder", 25 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:simple_encoder", 26 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/stages:stages_impl", 27 | ], 28 | ) 29 | 30 | py_strict_test( 31 | name = "common_encoders_test", 32 | size = "small", 33 | srcs = ["common_encoders_test.py"], 34 | deps = [ 35 | ":common_encoders", 36 | # absl/testing:parameterized dep1, 37 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 38 | # tensorflow dep1, 39 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:core_encoder", 40 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:gather_encoder", 41 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:simple_encoder", 42 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/utils:py_utils", 43 | ], 44 | ) 45 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/tensor_encoding/encoders/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019, The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """A collection of common encoders.""" 15 | 16 | from __future__ import absolute_import 17 | from __future__ import division 18 | from __future__ import print_function 19 | 20 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.encoders.common_encoders import as_gather_encoder 21 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.encoders.common_encoders import as_simple_encoder 22 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.encoders.common_encoders import hadamard_quantization 23 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.encoders.common_encoders import identity 24 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.encoders.common_encoders import uniform_quantization 25 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/tensor_encoding/stages/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library", "py_strict_test") 2 | 3 | package(default_visibility = [ 4 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding:__subpackages__", 5 | ]) 6 | 7 | licenses(["notice"]) 8 | 9 | py_strict_library( 10 | name = "stages", 11 | srcs = ["__init__.py"], 12 | visibility = ["//visibility:public"], 13 | deps = [ 14 | ":stages_impl", 15 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/stages/research", 16 | ], 17 | ) 18 | 19 | py_strict_library( 20 | name = "stages_impl", 21 | srcs = ["stages_impl.py"], 22 | deps = [ 23 | # numpy dep1, 24 | # tensorflow dep1, 25 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:encoding_stage", 26 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/utils:tf_utils", 27 | ], 28 | ) 29 | 30 | py_strict_test( 31 | name = "stages_impl_test", 32 | size = "medium", 33 | srcs = ["stages_impl_test.py"], 34 | deps = [ 35 | ":stages_impl", 36 | # absl/testing:parameterized dep1, 37 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 38 | # numpy dep1, 39 | # tensorflow dep1, 40 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/testing:test_utils", 41 | ], 42 | ) 43 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/tensor_encoding/stages/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019, The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Implementations of the (Adaptive)EncodingStageInterface.""" 15 | 16 | from __future__ import absolute_import 17 | from __future__ import division 18 | from __future__ import print_function 19 | 20 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages import research 21 | 22 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages.stages_impl import BitpackingEncodingStage 23 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages.stages_impl import FlattenEncodingStage 24 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages.stages_impl import HadamardEncodingStage 25 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages.stages_impl import IdentityEncodingStage 26 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages.stages_impl import UniformQuantizationEncodingStage 27 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/tensor_encoding/stages/research/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019, The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Experimental implementations of encoding stages. 15 | 16 | These encoding stages can possibly change without guarantees on backward 17 | compatibility. 18 | """ 19 | 20 | from __future__ import absolute_import 21 | from __future__ import division 22 | from __future__ import print_function 23 | 24 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages.research.clipping import ClipByNormEncodingStage 25 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages.research.clipping import ClipByValueEncodingStage 26 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages.research.kashin import KashinHadamardEncodingStage 27 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages.research.misc import DifferenceBetweenIntegersEncodingStage 28 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages.research.misc import SplitBySmallValueEncodingStage 29 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages.research.quantization import PerChannelPRNGUniformQuantizationEncodingStage 30 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages.research.quantization import PerChannelUniformQuantizationEncodingStage 31 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages.research.quantization import PRNGUniformQuantizationEncodingStage 32 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/tensor_encoding/testing/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library", "py_strict_test") 2 | 3 | package(default_visibility = [ 4 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding:__subpackages__", 5 | ]) 6 | 7 | licenses(["notice"]) 8 | 9 | py_strict_library( 10 | name = "testing", 11 | srcs = ["__init__.py"], 12 | visibility = ["//visibility:public"], 13 | deps = [ 14 | ":test_utils", 15 | ], 16 | ) 17 | 18 | py_strict_library( 19 | name = "test_utils", 20 | srcs = ["test_utils.py"], 21 | deps = [ 22 | # absl/testing:parameterized dep1, 23 | # numpy dep1, 24 | # six dep1, 25 | # tensorflow dep1, 26 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:encoding_stage", 27 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/utils:py_utils", 28 | ], 29 | ) 30 | 31 | py_strict_test( 32 | name = "test_utils_test", 33 | size = "medium", 34 | srcs = ["test_utils_test.py"], 35 | deps = [ 36 | ":test_utils", 37 | # absl/testing:parameterized dep1, 38 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 39 | # mock dep1, 40 | # numpy dep1, 41 | # six dep1, 42 | # tensorflow dep1, 43 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding/core:encoding_stage", 44 | ], 45 | ) 46 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/tensor_encoding/testing/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019, The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Testing utilities for the `tensor_encoding` package.""" 15 | 16 | from __future__ import absolute_import 17 | from __future__ import division 18 | from __future__ import print_function 19 | 20 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import AdaptiveNormalizeEncodingStage 21 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import aggregate_state_update_tensors 22 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import BaseEncodingStageTest 23 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import get_tensor_with_random_shape 24 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import is_adaptive_stage 25 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import PlusOneEncodingStage 26 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import PlusOneOverNEncodingStage 27 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import PlusRandomNumEncodingStage 28 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import RandomAddSubtractOneEncodingStage 29 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import ReduceMeanEncodingStage 30 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import SignIntFloatEncodingStage 31 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import SimpleLinearEncodingStage 32 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import TestData 33 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing.test_utils import TimesTwoEncodingStage 34 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/tensor_encoding/utils/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library", "py_strict_test") 2 | 3 | package(default_visibility = [ 4 | "//tensorflow_model_optimization/python/core/internal/tensor_encoding:__subpackages__", 5 | ]) 6 | 7 | licenses(["notice"]) 8 | 9 | py_strict_library( 10 | name = "utils", 11 | srcs = ["__init__.py"], 12 | visibility = ["//visibility:public"], 13 | deps = [ 14 | ":py_utils", 15 | ":tf_utils", 16 | ], 17 | ) 18 | 19 | py_strict_library( 20 | name = "py_utils", 21 | srcs = ["py_utils.py"], 22 | deps = [ 23 | # numpy dep1, 24 | # six dep1, 25 | # tensorflow dep1, 26 | # tree dep1, 27 | ], 28 | ) 29 | 30 | py_strict_test( 31 | name = "py_utils_test", 32 | size = "small", 33 | srcs = ["py_utils_test.py"], 34 | deps = [ 35 | ":py_utils", 36 | # absl/testing:parameterized dep1, 37 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 38 | # numpy dep1, 39 | # tensorflow dep1, 40 | ], 41 | ) 42 | 43 | py_strict_library( 44 | name = "tf_utils", 45 | srcs = ["tf_utils.py"], 46 | deps = [ 47 | # numpy dep1, 48 | # tensorflow dep1, 49 | ], 50 | ) 51 | 52 | py_strict_test( 53 | name = "tf_utils_test", 54 | size = "medium", 55 | srcs = ["tf_utils_test.py"], 56 | deps = [ 57 | ":tf_utils", 58 | # absl/testing:parameterized dep1, 59 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 60 | # numpy dep1, 61 | # scipy dep1, 62 | # tensorflow dep1, 63 | ], 64 | ) 65 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/internal/tensor_encoding/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019, The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """Utilities for the `tensor_encoding` package.""" 15 | 16 | from __future__ import absolute_import 17 | from __future__ import division 18 | from __future__ import print_function 19 | 20 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.utils.py_utils import assert_compatible 21 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.utils.py_utils import merge_dicts 22 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.utils.py_utils import OrderedEnum 23 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.utils.py_utils import split_dict_py_tf 24 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.utils.py_utils import static_or_dynamic_shape 25 | 26 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.utils.tf_utils import fast_walsh_hadamard_transform 27 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.utils.tf_utils import pack_into_int 28 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.utils.tf_utils import random_floats 29 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.utils.tf_utils import random_floats_cmwc 30 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.utils.tf_utils import random_signs 31 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.utils.tf_utils import random_signs_cmwc 32 | from tensorflow_model_optimization.python.core.internal.tensor_encoding.utils.tf_utils import unpack_from_int 33 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/keras/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library", "py_strict_test") 2 | 3 | package(default_visibility = [ 4 | "//tensorflow_model_optimization:__subpackages__", 5 | ]) 6 | 7 | licenses(["notice"]) 8 | 9 | py_strict_library( 10 | name = "keras", 11 | srcs = ["__init__.py"], 12 | deps = [ 13 | ":compat", 14 | ":utils", 15 | ], 16 | ) 17 | 18 | # TODO(alanchiao): split into test_utils and utils. 19 | py_strict_library( 20 | name = "test_utils", 21 | srcs = ["test_utils.py"], 22 | deps = [ 23 | ":compat", 24 | # numpy dep1, 25 | # tensorflow dep1, 26 | ], 27 | ) 28 | 29 | py_strict_library( 30 | name = "compat", 31 | srcs = ["compat.py"], 32 | deps = [ 33 | # tensorflow dep1, 34 | ], 35 | ) 36 | 37 | py_strict_library( 38 | name = "utils", 39 | srcs = ["utils.py"], 40 | visibility = ["//visibility:public"], 41 | deps = [ 42 | # tensorflow dep1, 43 | # python/framework:smart_cond tensorflow dep2, 44 | # python/ops:variables tensorflow dep2, 45 | ], 46 | ) 47 | 48 | py_strict_library( 49 | name = "metrics", 50 | srcs = ["metrics.py"], 51 | deps = [ 52 | # python/eager:monitoring tensorflow dep2, 53 | ], 54 | ) 55 | 56 | py_strict_test( 57 | name = "metrics_test", 58 | srcs = ["metrics_test.py"], 59 | deps = [ 60 | ":compat", 61 | ":metrics", 62 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 63 | # mock dep1, 64 | # tensorflow dep1, 65 | # python/eager:monitoring tensorflow dep2, 66 | ], 67 | ) 68 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/keras/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/model-optimization/acd0717cebfde7dd0c707ddbf4915669c97e4b6a/tensorflow_model_optimization/python/core/keras/__init__.py -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/keras/metrics_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Tests for Metrics.""" 16 | 17 | import mock 18 | import tensorflow as tf 19 | 20 | from tensorflow.python.eager import monitoring 21 | from tensorflow_model_optimization.python.core.keras import metrics 22 | from tensorflow_model_optimization.python.core.keras.compat import keras 23 | 24 | 25 | class MetricsTest(tf.test.TestCase): 26 | 27 | gauge = monitoring.BoolGauge('/tfmot/metrics/testing', 'testing', 'labels') 28 | 29 | def setUp(self): 30 | super(MetricsTest, self).setUp() 31 | self.test_label = keras.layers.Conv2D(1, 1).__class__.__name__ 32 | for label in [ 33 | self.test_label, metrics.MonitorBoolGauge._SUCCESS_LABEL, 34 | metrics.MonitorBoolGauge._FAILURE_LABEL 35 | ]: 36 | MetricsTest.gauge.get_cell(label).set(False) 37 | 38 | with mock.patch.object(metrics.MonitorBoolGauge, 'get_usage_gauge', 39 | return_value=MetricsTest.gauge): 40 | self.monitor = metrics.MonitorBoolGauge('testing') 41 | 42 | def test_DecoratorTest(self): 43 | @self.monitor 44 | def func(x): 45 | return x + 1 46 | 47 | self.assertEqual(func(1), 2) 48 | self.assertTrue(MetricsTest.gauge.get_cell( 49 | metrics.MonitorBoolGauge._SUCCESS_LABEL).value()) 50 | 51 | def test_DecoratorFailureTest(self): 52 | @self.monitor 53 | def func(x): 54 | raise ValueError() 55 | 56 | with self.assertRaises(ValueError): 57 | func(1) 58 | self.assertTrue(MetricsTest.gauge.get_cell( 59 | metrics.MonitorBoolGauge._FAILURE_LABEL).value()) 60 | 61 | def test_UndecoratedTest(self): 62 | with self.assertRaises(ValueError): 63 | @metrics.MonitorBoolGauge('unknown') 64 | def func(x): 65 | return x+1 66 | func(1) 67 | 68 | def test_SetTest(self): 69 | self.monitor.set(self.test_label) 70 | self.assertTrue(MetricsTest.gauge.get_cell(self.test_label).value()) 71 | 72 | 73 | if __name__ == '__main__': 74 | tf.test.main() 75 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/keras/testing/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | 3 | package(default_visibility = [ 4 | "//tensorflow_model_optimization:__subpackages__", 5 | ]) 6 | 7 | licenses(["notice"]) 8 | 9 | py_strict_library( 10 | name = "test_utils_mnist", 11 | testonly = 1, 12 | srcs = ["test_utils_mnist.py"], 13 | visibility = ["//visibility:public"], 14 | deps = [ 15 | # numpy dep1, 16 | # tensorflow dep1, 17 | "//tensorflow_model_optimization/python/core/keras:compat", 18 | ], 19 | ) 20 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/keras/testing/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/model-optimization/acd0717cebfde7dd0c707ddbf4915669c97e4b6a/tensorflow_model_optimization/python/core/keras/testing/__init__.py -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/keras/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Utility functions common to MOT techniques. 16 | 17 | `smart_cond` is not exposed as a public TF function with stabilty guarantees, 18 | and hence changes to it break our code. So, we make a copy of it here. 19 | """ 20 | 21 | from __future__ import absolute_import 22 | from __future__ import division 23 | from __future__ import print_function 24 | 25 | import tensorflow as tf 26 | 27 | # TODO(b/151772467): Move away from depending on private APIs. 28 | from tensorflow.python.framework import smart_cond as smart_module 29 | from tensorflow.python.ops import variables 30 | 31 | 32 | def smart_cond(pred, true_fn=None, false_fn=None, name=None): # pylint: disable=invalid-name 33 | """Return either `true_fn()` if predicate `pred` is true else `false_fn()`. 34 | 35 | If `pred` is a bool or has a constant value, we return either `true_fn()` 36 | or `false_fn()`, otherwise we use `tf.cond` to dynamically route to both. 37 | 38 | Arguments: 39 | pred: A scalar determining whether to return the result of `true_fn` or 40 | `false_fn`. 41 | true_fn: The callable to be performed if pred is true. 42 | false_fn: The callable to be performed if pred is false. 43 | name: Optional name prefix when using `tf.cond`. 44 | 45 | Returns: 46 | Tensors returned by the call to either `true_fn` or `false_fn`. 47 | 48 | Raises: 49 | TypeError: If `true_fn` or `false_fn` is not callable. 50 | """ 51 | if isinstance(pred, variables.Variable): 52 | return tf.cond( 53 | pred, true_fn=true_fn, false_fn=false_fn, name=name) 54 | return smart_module.smart_cond( 55 | pred, true_fn=true_fn, false_fn=false_fn, name=name) 56 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | 3 | package(default_visibility = ["//tensorflow_model_optimization:__subpackages__"]) 4 | 5 | licenses(["notice"]) 6 | 7 | py_strict_library( 8 | name = "quantization", 9 | srcs = ["__init__.py"], 10 | deps = [ 11 | "//tensorflow_model_optimization/python/core/quantization/keras", # buildcleaner: keep 12 | ], 13 | ) 14 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/collab_opts/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | 3 | package(default_visibility = ["//tensorflow_model_optimization:__subpackages__"]) 4 | 5 | licenses(["notice"]) 6 | 7 | py_strict_library( 8 | name = "collab_opts", 9 | srcs = ["__init__.py"], 10 | deps = [ 11 | "//tensorflow_model_optimization/python/core/quantization/keras/collab_opts/cluster_preserve", # buildcleaner: keep 12 | "//tensorflow_model_optimization/python/core/quantization/keras/collab_opts/prune_preserve", # buildcleaner: keep 13 | ], 14 | ) 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/collab_opts/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/collab_opts/cluster_preserve/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/collab_opts/cluster_preserve/default_8bit_cluster_preserve_quantize_scheme.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | 16 | """Default 8 bit Cluster Preserve Quantization scheme.""" 17 | 18 | from tensorflow_model_optimization.python.core.quantization.keras.collab_opts.cluster_preserve import ( 19 | cluster_preserve_quantize_registry,) 20 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantize_scheme 21 | 22 | 23 | class Default8BitClusterPreserveQuantizeScheme( 24 | default_8bit_quantize_scheme.Default8BitQuantizeScheme): 25 | """Default 8 bit Cluster Preserve Quantization Scheme.""" 26 | 27 | def __init__(self, preserve_sparsity=True): 28 | """Same as Default8BitQuantizeScheme but preserves clustering and sparsity. 29 | 30 | Args: 31 | preserve_sparsity: the flag to enable prune-cluster-preserving QAT. 32 | """ 33 | self.preserve_sparsity = preserve_sparsity 34 | 35 | def get_quantize_registry(self): 36 | return (cluster_preserve_quantize_registry. 37 | Default8bitClusterPreserveQuantizeRegistry(self.preserve_sparsity)) 38 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/collab_opts/prune_preserve/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | # Placeholder: load py_test 3 | 4 | package(default_visibility = [ 5 | "//tensorflow_model_optimization:__subpackages__", 6 | ]) 7 | 8 | licenses(["notice"]) 9 | 10 | py_strict_library( 11 | name = "prune_preserve", 12 | srcs = [ 13 | "__init__.py", 14 | ], 15 | deps = [ 16 | ":default_8bit_prune_preserve_quantize_scheme", # buildcleaner: keep 17 | ], 18 | ) 19 | 20 | py_strict_library( 21 | name = "prune_preserve_quantize_registry", 22 | srcs = [ 23 | "prune_preserve_quantize_registry.py", 24 | ], 25 | deps = [ 26 | # tensorflow dep1, 27 | "//tensorflow_model_optimization/python/core/keras:compat", 28 | "//tensorflow_model_optimization/python/core/quantization/keras:quant_ops", 29 | "//tensorflow_model_optimization/python/core/quantization/keras:quantizers", 30 | "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_registry", 31 | "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantizers", 32 | ], 33 | ) 34 | 35 | py_test( 36 | name = "prune_preserve_quantize_registry_test", 37 | srcs = [ 38 | "prune_preserve_quantize_registry_test.py", 39 | ], 40 | visibility = ["//visibility:private"], 41 | deps = [ 42 | ":prune_preserve_quantize_registry", 43 | # absl/testing:parameterized dep1, 44 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 45 | # tensorflow dep1, 46 | "//tensorflow_model_optimization/python/core/keras:compat", 47 | "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_registry", 48 | "//tensorflow_model_optimization/python/core/sparsity/keras:prune_registry", 49 | ], 50 | ) 51 | 52 | py_strict_library( 53 | name = "default_8bit_prune_preserve_quantize_scheme", 54 | srcs = [ 55 | "default_8bit_prune_preserve_quantize_scheme.py", 56 | ], 57 | visibility = ["//visibility:public"], 58 | deps = [ 59 | ":prune_preserve_quantize_registry", 60 | "//tensorflow_model_optimization/python/core/quantization/keras/default_8bit:default_8bit_quantize_scheme", 61 | ], 62 | ) 63 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/collab_opts/prune_preserve/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/collab_opts/prune_preserve/default_8bit_prune_preserve_quantize_scheme.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Default 8 bit Prune Preserve Quantization scheme which specifies how quantization should be applied.""" 16 | 17 | from tensorflow_model_optimization.python.core.quantization.keras.collab_opts.prune_preserve import ( 18 | prune_preserve_quantize_registry) 19 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import ( 20 | default_8bit_quantize_scheme,) 21 | 22 | 23 | class Default8BitPrunePreserveQuantizeScheme( 24 | default_8bit_quantize_scheme.Default8BitQuantizeScheme): 25 | """Default 8 bit Prune Preserve Quantization Scheme.""" 26 | 27 | def get_quantize_registry(self): 28 | return (prune_preserve_quantize_registry 29 | .Default8bitPrunePreserveQuantizeRegistry()) 30 | 31 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/default_8bit/README.md: -------------------------------------------------------------------------------- 1 | This directory contains code specific to the default 8-bit behavior of the QAT 2 | API including the quantization scheme and implementation. 3 | 4 | This quantization scheme is implemented by multiple backends in the ecosystem 5 | (e.g. NNAPI, TFLite). 6 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/default_8bit/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/default_8bit/default_8bit_quantize_configs.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Default 8-bit QuantizeConfigs.""" 16 | 17 | from tensorflow_model_optimization.python.core.quantization.keras import quantize_config 18 | from tensorflow_model_optimization.python.core.quantization.keras import quantizers 19 | 20 | 21 | class Default8BitOutputQuantizeConfig(quantize_config.QuantizeConfig): 22 | """QuantizeConfig which only quantizes the output from a layer.""" 23 | 24 | def __init__(self, quantize_output: bool = True) -> None: 25 | self.quantize_output = quantize_output 26 | 27 | def get_weights_and_quantizers(self, layer): 28 | return [] 29 | 30 | def get_activations_and_quantizers(self, layer): 31 | return [] 32 | 33 | def set_quantize_weights(self, layer, quantize_weights): 34 | pass 35 | 36 | def set_quantize_activations(self, layer, quantize_activations): 37 | pass 38 | 39 | def get_output_quantizers(self, layer): 40 | if self.quantize_output: 41 | return [quantizers.MovingAverageQuantizer( 42 | num_bits=8, per_axis=False, symmetric=False, narrow_range=False)] 43 | return [] 44 | 45 | def get_config(self): 46 | return {'quantize_output': self.quantize_output} 47 | 48 | 49 | class NoOpQuantizeConfig(quantize_config.QuantizeConfig): 50 | """QuantizeConfig which does not quantize any part of the layer.""" 51 | 52 | def get_weights_and_quantizers(self, layer): 53 | return [] 54 | 55 | def get_activations_and_quantizers(self, layer): 56 | return [] 57 | 58 | def set_quantize_weights(self, layer, quantize_weights): 59 | pass 60 | 61 | def set_quantize_activations(self, layer, quantize_activations): 62 | pass 63 | 64 | def get_output_quantizers(self, layer): 65 | return [] 66 | 67 | def get_config(self): 68 | return {} 69 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/default_8bit/default_8bit_quantize_scheme.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Quantization scheme which specifies how quantization should be applied.""" 16 | 17 | from tensorflow_model_optimization.python.core.quantization.keras import quantize_scheme 18 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantize_layout_transform 19 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantize_registry 20 | 21 | 22 | class Default8BitQuantizeScheme(quantize_scheme.QuantizeScheme): 23 | """Default 8Bit Scheme supported by TFLite.""" 24 | 25 | def __init__(self, disable_per_axis=False): 26 | self._disable_per_axis = disable_per_axis 27 | 28 | def get_layout_transformer(self): 29 | return default_8bit_quantize_layout_transform.\ 30 | Default8BitQuantizeLayoutTransform() 31 | 32 | def get_quantize_registry(self): 33 | return ( 34 | default_8bit_quantize_registry.Default8BitQuantizeRegistry( 35 | disable_per_axis=self._disable_per_axis)) 36 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/default_8bit/default_8bit_quantizers.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Quantizers specific to default 8-bit behavior.""" 16 | 17 | import tensorflow as tf 18 | 19 | from tensorflow_model_optimization.python.core.keras.compat import keras 20 | from tensorflow_model_optimization.python.core.quantization.keras import quantizers 21 | 22 | 23 | class Default8BitConvWeightsQuantizer(quantizers.LastValueQuantizer): 24 | """Quantizer for handling weights in Conv2D/DepthwiseConv2D layers.""" 25 | 26 | def __init__(self): 27 | """Construct LastValueQuantizer with params specific for TFLite Convs.""" 28 | 29 | super(Default8BitConvWeightsQuantizer, self).__init__( 30 | num_bits=8, per_axis=True, symmetric=True, narrow_range=True) 31 | 32 | def build(self, tensor_shape, name, layer): 33 | min_weight = layer.add_weight( 34 | name + '_min', 35 | shape=(tensor_shape[-1],), 36 | initializer=keras.initializers.Constant(-6.0), 37 | trainable=False, 38 | ) 39 | max_weight = layer.add_weight( 40 | name + '_max', 41 | shape=(tensor_shape[-1],), 42 | initializer=keras.initializers.Constant(6.0), 43 | trainable=False, 44 | ) 45 | 46 | return {'min_var': min_weight, 'max_var': max_weight} 47 | 48 | 49 | class Default8BitConvTransposeWeightsQuantizer(quantizers.LastValueQuantizer): 50 | """Quantizer for handling weights in Conv2DTranspose layers.""" 51 | 52 | def __init__(self): 53 | """Construct LastValueQuantizer with params specific for TFLite Conv2DTranpose.""" 54 | 55 | super(Default8BitConvTransposeWeightsQuantizer, self).__init__( 56 | num_bits=8, per_axis=False, symmetric=True, narrow_range=True) 57 | 58 | def __call__(self, inputs, training, weights, **kwargs): 59 | outputs = tf.transpose(inputs, (0, 1, 3, 2)) 60 | outputs = super(Default8BitConvTransposeWeightsQuantizer, 61 | self).__call__(outputs, training, weights, **kwargs) 62 | outputs = tf.transpose(outputs, (0, 1, 3, 2)) 63 | return outputs 64 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/default_8bit/default_8bit_quantizers_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Tests for default Quantizers.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | from absl.testing import parameterized 22 | import tensorflow as tf 23 | 24 | from tensorflow_model_optimization.python.core.keras.compat import keras 25 | from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantizers 26 | 27 | 28 | Default8BitConvWeightsQuantizer = ( 29 | default_8bit_quantizers.Default8BitConvWeightsQuantizer 30 | ) 31 | 32 | 33 | class Default8BitConvWeightsQuantizerTest(tf.test.TestCase, 34 | parameterized.TestCase): 35 | 36 | @parameterized.parameters( 37 | (keras.layers.Conv2D, { 38 | 'filters': 5, 39 | 'kernel_size': (2, 2) 40 | }), 41 | (keras.layers.DepthwiseConv2D, { 42 | 'kernel_size': (2, 2), 43 | 'depth_multiplier': 5, 44 | }) 45 | ) 46 | def testConstructsMinMaxVarsCorrectShape(self, layer_type, kwargs): 47 | quantizer = Default8BitConvWeightsQuantizer() 48 | 49 | model = keras.Sequential([ 50 | layer_type(input_shape=(5, 2, 3), **kwargs)]) 51 | layer = model.layers[0] 52 | 53 | min_max_vars = quantizer.build( 54 | layer.weights[0].shape, 'kernel', layer) 55 | # TODO(pulkitb): Add value test to ensure per-axis quantization is 56 | # happening properly. Probably to quant_ops_test.py 57 | quantized_weight = quantizer(layer.weights[0], True, # pylint: disable=unused-variable 58 | weights=min_max_vars) 59 | 60 | min_var = min_max_vars['min_var'] 61 | max_var = min_max_vars['max_var'] 62 | self.assertEqual(5, min_var.shape) 63 | self.assertEqual(5, max_var.shape) 64 | 65 | 66 | if __name__ == '__main__': 67 | tf.test.main() 68 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/experimental/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | 3 | package(default_visibility = [ 4 | "//tensorflow_model_optimization:__subpackages__", 5 | ]) 6 | 7 | licenses(["notice"]) 8 | 9 | py_strict_library( 10 | name = "experimental", 11 | srcs = [ 12 | "__init__.py", 13 | ], 14 | deps = [ 15 | ":quantization", # buildcleaner: keep 16 | ], 17 | ) 18 | 19 | py_strict_library( 20 | name = "quantization", 21 | srcs = ["__init__.py"], 22 | deps = [ 23 | "//tensorflow_model_optimization/python/core/quantization/keras", # buildcleaner: keep 24 | ], 25 | ) 26 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/experimental/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/README.md: -------------------------------------------------------------------------------- 1 | This directory is modified based on default_8bit, which allows you to manually 2 | change the number of bits of weight and activation in QAT. 3 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/default_n_bit_quantize_configs.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Default N-Bit QuantizeConfigs.""" 16 | 17 | from typing import Any, Dict 18 | from tensorflow_model_optimization.python.core.quantization.keras import quantize_config 19 | from tensorflow_model_optimization.python.core.quantization.keras import quantizers 20 | 21 | 22 | class DefaultNBitOutputQuantizeConfig(quantize_config.QuantizeConfig): 23 | """QuantizeConfig which only quantizes the output from a layer.""" 24 | 25 | def __init__(self, num_bits_weight: int = 8, num_bits_activation: int = 8): 26 | self._num_bits_weight = num_bits_weight 27 | self._num_bits_activation = num_bits_activation 28 | 29 | def get_weights_and_quantizers(self, layer): 30 | return [] 31 | 32 | def get_activations_and_quantizers(self, layer): 33 | return [] 34 | 35 | def set_quantize_weights(self, layer, quantize_weights): 36 | pass 37 | 38 | def set_quantize_activations(self, layer, quantize_activations): 39 | pass 40 | 41 | def get_output_quantizers(self, layer): 42 | return [quantizers.MovingAverageQuantizer( 43 | num_bits=self._num_bits_activation, per_axis=False, 44 | symmetric=False, narrow_range=False)] # activation/output 45 | 46 | def get_config(self) -> Dict[str, Any]: 47 | return { 48 | 'num_bits_weight': self._num_bits_weight, 49 | 'num_bits_activation': self._num_bits_activation, 50 | } 51 | 52 | 53 | class NoOpQuantizeConfig(quantize_config.QuantizeConfig): 54 | """QuantizeConfig which does not quantize any part of the layer.""" 55 | 56 | def get_weights_and_quantizers(self, layer): 57 | return [] 58 | 59 | def get_activations_and_quantizers(self, layer): 60 | return [] 61 | 62 | def set_quantize_weights(self, layer, quantize_weights): 63 | pass 64 | 65 | def set_quantize_activations(self, layer, quantize_activations): 66 | pass 67 | 68 | def get_output_quantizers(self, layer): 69 | return [] 70 | 71 | def get_config(self): 72 | return {} 73 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/default_n_bit_quantize_scheme.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Quantization scheme which specifies how quantization should be applied.""" 16 | 17 | from tensorflow_model_optimization.python.core.quantization.keras import quantize_scheme 18 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit import default_n_bit_quantize_layout_transform 19 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit import default_n_bit_quantize_registry 20 | 21 | 22 | class DefaultNBitQuantizeScheme(quantize_scheme.QuantizeScheme): 23 | """Default N-Bit Scheme supported by TFLite.""" 24 | 25 | def __init__(self, disable_per_axis=False, 26 | num_bits_weight=8, num_bits_activation=8): 27 | self._disable_per_axis = disable_per_axis 28 | self._num_bits_weight = num_bits_weight 29 | self._num_bits_activation = num_bits_activation 30 | 31 | def get_layout_transformer(self): 32 | return default_n_bit_quantize_layout_transform.DefaultNBitQuantizeLayoutTransform( 33 | num_bits_weight=self._num_bits_weight, 34 | num_bits_activation=self._num_bits_activation) 35 | 36 | def get_quantize_registry(self): 37 | return ( 38 | default_n_bit_quantize_registry.DefaultNBitQuantizeRegistry( 39 | disable_per_axis=self._disable_per_axis, 40 | num_bits_weight=self._num_bits_weight, 41 | num_bits_activation=self._num_bits_activation)) 42 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/default_n_bit_quantizers.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Quantizers specific to default 8-bit behavior.""" 16 | 17 | import tensorflow as tf 18 | 19 | from tensorflow_model_optimization.python.core.keras.compat import keras 20 | from tensorflow_model_optimization.python.core.quantization.keras import quantizers 21 | 22 | 23 | class DefaultNBitConvWeightsQuantizer(quantizers.LastValueQuantizer): 24 | """Quantizer for handling weights in Conv2D/DepthwiseConv2D layers.""" 25 | 26 | def __init__(self, num_bits_weight: int = 8, num_bits_activation: int = 8): 27 | """Construct LastValueQuantizer with params specific for TFLite Convs.""" 28 | 29 | super(DefaultNBitConvWeightsQuantizer, self).__init__( 30 | num_bits=num_bits_weight, 31 | per_axis=True, 32 | symmetric=True, 33 | narrow_range=True) # weight 34 | self._num_bits_weight = num_bits_weight 35 | self._num_bits_activation = num_bits_activation 36 | 37 | def build(self, tensor_shape, name, layer): 38 | min_weight = layer.add_weight( 39 | name + '_min', 40 | shape=(tensor_shape[-1],), 41 | initializer=keras.initializers.Constant(-6.0), 42 | trainable=False, 43 | ) 44 | max_weight = layer.add_weight( 45 | name + '_max', 46 | shape=(tensor_shape[-1],), 47 | initializer=keras.initializers.Constant(6.0), 48 | trainable=False, 49 | ) 50 | 51 | return {'min_var': min_weight, 'max_var': max_weight} 52 | 53 | 54 | class DefaultNBitConvTransposeWeightsQuantizer(quantizers.LastValueQuantizer): 55 | """Quantizer for handling weights in Conv2DTranspose layers.""" 56 | 57 | def __init__(self, num_bits_weight: int = 8, num_bits_activation: int = 8): 58 | """Construct LastValueQuantizer with params specific for TFLite Conv2DTranpose.""" 59 | super(DefaultNBitConvTransposeWeightsQuantizer, self).__init__( 60 | num_bits=num_bits_weight, 61 | per_axis=False, 62 | symmetric=True, 63 | narrow_range=True) # weight 64 | self._num_bits_weight = num_bits_weight 65 | self._num_bits_activation = num_bits_activation 66 | 67 | def __call__(self, inputs, training, weights, **kwargs): 68 | outputs = tf.transpose(inputs, (0, 1, 3, 2)) 69 | outputs = super(DefaultNBitConvTransposeWeightsQuantizer, 70 | self).__call__(outputs, training, weights, **kwargs) 71 | outputs = tf.transpose(outputs, (0, 1, 3, 2)) 72 | return outputs 73 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/default_n_bit_quantizers_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Tests for default Quantizers.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | from absl.testing import parameterized 22 | import tensorflow as tf 23 | 24 | from tensorflow_model_optimization.python.core.keras.compat import keras 25 | from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit import default_n_bit_quantizers 26 | 27 | 28 | DefaultNBitConvWeightsQuantizer = ( 29 | default_n_bit_quantizers.DefaultNBitConvWeightsQuantizer 30 | ) 31 | 32 | 33 | class DefaultNBitConvWeightsQuantizerTest(tf.test.TestCase, 34 | parameterized.TestCase): 35 | 36 | @parameterized.parameters( 37 | (keras.layers.Conv2D, { 38 | 'filters': 5, 39 | 'kernel_size': (2, 2) 40 | }), 41 | (keras.layers.DepthwiseConv2D, { 42 | 'kernel_size': (2, 2), 43 | 'depth_multiplier': 5, 44 | }) 45 | ) 46 | def testConstructsMinMaxVarsCorrectShape(self, layer_type, kwargs): 47 | quantizer = DefaultNBitConvWeightsQuantizer() 48 | 49 | model = keras.Sequential([ 50 | layer_type(input_shape=(5, 2, 3), **kwargs)]) 51 | layer = model.layers[0] 52 | 53 | min_max_vars = quantizer.build( 54 | layer.weights[0].shape, 'kernel', layer) 55 | # TODO(pulkitb): Add value test to ensure per-axis quantization is 56 | # happening properly. Probably to quant_ops_test.py 57 | quantized_weight = quantizer(layer.weights[0], True, # pylint: disable=unused-variable 58 | weights=min_max_vars) 59 | 60 | min_var = min_max_vars['min_var'] 61 | max_var = min_max_vars['max_var'] 62 | self.assertEqual(5, min_var.shape) 63 | self.assertEqual(5, max_var.shape) 64 | 65 | 66 | if __name__ == '__main__': 67 | tf.test.main() 68 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library", "py_strict_test") 2 | 3 | package(default_visibility = [ 4 | "//tensorflow_model_optimization:__subpackages__", 5 | ]) 6 | 7 | licenses(["notice"]) 8 | 9 | py_strict_library( 10 | name = "graph_transformations", 11 | srcs = [ 12 | "__init__.py", 13 | ], 14 | deps = [], 15 | ) 16 | 17 | py_strict_library( 18 | name = "transforms", 19 | srcs = [ 20 | "transforms.py", 21 | ], 22 | visibility = ["//visibility:public"], 23 | deps = [ 24 | # six dep1, 25 | ], 26 | ) 27 | 28 | py_strict_test( 29 | name = "transforms_test", 30 | srcs = [ 31 | "transforms_test.py", 32 | ], 33 | visibility = ["//visibility:public"], 34 | deps = [ 35 | ":transforms", 36 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 37 | # tensorflow dep1, 38 | ], 39 | ) 40 | 41 | py_strict_library( 42 | name = "model_transformer", 43 | srcs = [ 44 | "model_transformer.py", 45 | ], 46 | visibility = ["//visibility:public"], 47 | deps = [ 48 | ":transforms", 49 | # tensorflow dep1, 50 | "//tensorflow_model_optimization/python/core/keras:compat", 51 | ], 52 | ) 53 | 54 | py_strict_test( 55 | name = "model_transformer_test", 56 | srcs = [ 57 | "model_transformer_test.py", 58 | ], 59 | visibility = ["//visibility:public"], 60 | deps = [ 61 | ":model_transformer", 62 | ":transforms", 63 | # absl/testing:parameterized dep1, 64 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 65 | # numpy dep1, 66 | # tensorflow dep1, 67 | "//tensorflow_model_optimization/python/core/keras:compat", 68 | "//tensorflow_model_optimization/python/core/quantization/keras:utils", 69 | ], 70 | ) 71 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/transforms_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Tests for transforms.py API code.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import copy 22 | 23 | import tensorflow as tf 24 | 25 | from tensorflow_model_optimization.python.core.quantization.keras.graph_transformations import transforms 26 | 27 | LayerNode = transforms.LayerNode 28 | 29 | 30 | class LayerNodeTest(tf.test.TestCase): 31 | 32 | def testEqualityLayerNode(self): 33 | conv_layer = { 34 | 'name': 'conv2d', 35 | 'class_name': 'Conv2D', 36 | 'config': { 37 | 'name': 'conv2d', 38 | } 39 | } 40 | dense_layer = { 41 | 'name': 'dense', 42 | 'class_name': 'Dense', 43 | 'config': { 44 | 'name': 'dense', 45 | } 46 | } 47 | 48 | self.assertNotEqual(LayerNode(conv_layer), LayerNode(dense_layer)) 49 | 50 | self.assertEqual(LayerNode(conv_layer), LayerNode(conv_layer)) 51 | self.assertEqual( 52 | LayerNode(conv_layer), LayerNode(copy.deepcopy(conv_layer))) 53 | 54 | self.assertNotEqual( 55 | LayerNode(conv_layer, 56 | input_layers=[LayerNode(conv_layer), LayerNode(dense_layer)]), 57 | LayerNode(conv_layer, 58 | input_layers=[LayerNode(conv_layer), LayerNode(conv_layer)])) 59 | 60 | self.assertEqual( 61 | LayerNode(conv_layer, 62 | input_layers=[LayerNode(conv_layer), LayerNode(dense_layer)]), 63 | LayerNode(conv_layer, 64 | input_layers=[LayerNode(conv_layer), LayerNode(dense_layer)])) 65 | 66 | 67 | if __name__ == '__main__': 68 | tf.test.main() 69 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/layers/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | 3 | package(default_visibility = [ 4 | "//tensorflow_model_optimization:__subpackages__", 5 | ]) 6 | 7 | licenses(["notice"]) 8 | 9 | py_strict_library( 10 | name = "layers", 11 | srcs = [ 12 | "__init__.py", 13 | ], 14 | deps = [], 15 | ) 16 | 17 | py_strict_library( 18 | name = "conv_batchnorm_test_utils", 19 | srcs = ["conv_batchnorm_test_utils.py"], 20 | deps = [ 21 | # tensorflow dep1, 22 | "//tensorflow_model_optimization/python/core/keras:compat", 23 | ], 24 | ) 25 | 26 | py_strict_library( 27 | name = "dense_batchnorm_test_utils", 28 | srcs = ["dense_batchnorm_test_utils.py"], 29 | deps = [ 30 | # tensorflow dep1, 31 | "//tensorflow_model_optimization/python/core/keras:compat", 32 | ], 33 | ) 34 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/layers/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/layers/dense_batchnorm_test_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Test utils for dense batchnorm folding.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import tensorflow as tf 22 | 23 | from tensorflow_model_optimization.python.core.keras.compat import keras 24 | 25 | 26 | class DenseModel(object): 27 | """Construct and access Dense + BatchNorm + activation models.""" 28 | 29 | params = { 30 | 'units': 32, 31 | 'input_shape': (32,), 32 | 'batch_size': 1, 33 | } 34 | 35 | @classmethod 36 | def get_batched_input_shape(cls): 37 | """Return input shape with batch size.""" 38 | shape = [cls.params['batch_size']] 39 | shape.extend(cls.params['input_shape']) 40 | return shape 41 | 42 | @classmethod 43 | def get_nonfolded_batchnorm_model(cls, 44 | post_bn_activation=None, 45 | normalization_type='BatchNormalization'): 46 | """Return nonfolded Dense + BN + optional activation model.""" 47 | if normalization_type == 'BatchNormalization': 48 | normalization = keras.layers.BatchNormalization 49 | elif normalization_type == 'SyncBatchNormalization': 50 | normalization = keras.layers.experimental.SyncBatchNormalization 51 | 52 | inp = keras.layers.Input(cls.params['input_shape'], 53 | cls.params['batch_size']) 54 | x = keras.layers.Dense(cls.params['units'])(inp) 55 | out = normalization(axis=-1)(x) 56 | if post_bn_activation is not None: 57 | out = post_bn_activation(out) 58 | return keras.Model(inp, out) 59 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/quantize_layer_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Tests for QuantizeWrapper.""" 16 | 17 | from __future__ import absolute_import 18 | from __future__ import division 19 | from __future__ import print_function 20 | 21 | import numpy as np 22 | import tensorflow as tf 23 | 24 | from tensorflow_model_optimization.python.core.keras.compat import keras 25 | from tensorflow_model_optimization.python.core.quantization.keras import quantize_layer 26 | from tensorflow_model_optimization.python.core.quantization.keras import quantizers 27 | 28 | 29 | QuantizeLayer = quantize_layer.QuantizeLayer 30 | deserialize_layer = keras.layers.deserialize 31 | serialize_layer = keras.layers.serialize 32 | 33 | 34 | class QuantizeLayerTest(tf.test.TestCase): 35 | 36 | def setUp(self): 37 | super(QuantizeLayerTest, self).setUp() 38 | self.quant_params = { 39 | 'num_bits': 8, 40 | 'narrow_range': False 41 | } 42 | self.quantizer = quantizers.LastValueQuantizer( 43 | per_axis=False, symmetric=True, **self.quant_params) 44 | 45 | def testQuantizesTensors(self): 46 | model = keras.Sequential( 47 | [QuantizeLayer(quantizer=self.quantizer, input_shape=(4,))] 48 | ) 49 | 50 | x = np.random.rand(1, 4) 51 | quant_x = tf.quantization.fake_quant_with_min_max_vars( 52 | x, -6.0, 6.0, **self.quant_params) 53 | 54 | self.assertAllClose(self.evaluate(quant_x), model.predict(x)) 55 | 56 | def testSerializationQuantizeLayer(self): 57 | layer = QuantizeLayer( 58 | quantizer=self.quantizer, 59 | input_shape=(4,)) 60 | 61 | custom_objects = { 62 | 'QuantizeLayer': QuantizeLayer, 63 | 'LastValueQuantizer': quantizers.LastValueQuantizer 64 | } 65 | 66 | serialized_layer = serialize_layer(layer) 67 | with keras.utils.custom_object_scope(custom_objects): 68 | layer_from_config = deserialize_layer(serialized_layer) 69 | 70 | self.assertEqual(layer_from_config.get_config(), layer.get_config()) 71 | 72 | def testNoQuantizeLayer(self): 73 | layer = QuantizeLayer(quantizer=None, input_shape=(4,)) 74 | model = keras.Sequential([layer]) 75 | x = np.random.rand(1, 4) 76 | self.assertAllClose(x, model.predict(x)) 77 | 78 | custom_objects = { 79 | 'QuantizeLayer': QuantizeLayer, 80 | } 81 | 82 | serialized_layer = serialize_layer(layer) 83 | with keras.utils.custom_object_scope(custom_objects): 84 | layer_from_config = deserialize_layer(serialized_layer) 85 | 86 | self.assertEqual(layer_from_config.get_config(), layer.get_config()) 87 | 88 | 89 | if __name__ == '__main__': 90 | tf.test.main() 91 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/quantize_layout_transform.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Abstract Base Class for quantization transformations to a keras model. 16 | 17 | Keras models need certain transformations for quantization to exactly match the 18 | behavior of the backend they will be implemented on. This is important for 19 | improving model performance. 20 | 21 | This interface abstracts that behavior. Different backends can implement their 22 | own version. 23 | 24 | Module: tfmot.quantization.keras 25 | """ 26 | 27 | from __future__ import absolute_import 28 | from __future__ import division 29 | from __future__ import print_function 30 | 31 | import abc 32 | import six 33 | 34 | 35 | @six.add_metaclass(abc.ABCMeta) 36 | class QuantizeLayoutTransform(object): 37 | """Apply transformations to the model. 38 | 39 | Transforms the original model to perform better while quantized 40 | and also match the layout of the target backend. 41 | """ 42 | 43 | @abc.abstractmethod 44 | def apply(self, model, layer_quantize_map): 45 | """Transform model to a quantization friendly model. 46 | 47 | Args: 48 | model: Keras model to be quantized. 49 | layer_quantize_map: Map containing list of layers to be quantized and 50 | associated metadata. Keys are layer names which need to be quantized, 51 | and values are dicts containing relevant metadata. For example, 52 | any custom `QuantizeConfig` passed with a layer is present. 53 | 54 | Returns: 55 | New keras model based on `model` which has been 56 | transformed to match the layout of the target backend. 57 | """ 58 | raise NotImplementedError('Must be implemented in subclasses.') 59 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/quantize_registry.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Quantization registry which specifies how layers should be quantized. 16 | 17 | Module: tfmot.quantization.keras 18 | """ 19 | 20 | import abc 21 | import six 22 | 23 | 24 | @six.add_metaclass(abc.ABCMeta) 25 | class QuantizeRegistry(object): 26 | """ABC interface which specifies how layers should be quantized. 27 | 28 | The Registry is designed to function as a repository of `QuantizeConfig`s 29 | linked to layers. The idea is that while applying quantization to the various 30 | layers within a Keras model, the registry can be used to query which 31 | `QuantizeConfig` can be used to quantize a specific `layer`. The 32 | `QuantizeConfig` itself contains information to quantize that specific 33 | layer. 34 | 35 | We provide a default registry for built-in Keras layers, but implementing this 36 | interface allows users the ability to write their own custom registries 37 | specific to their needs. It can also be extended to be used for any Keras 38 | layer, such as custom Keras layers. 39 | """ 40 | 41 | @abc.abstractmethod 42 | def get_quantize_config(self, layer): 43 | """Returns the quantization config for the given layer. 44 | 45 | Args: 46 | layer: input layer to return quantize config for. 47 | 48 | Returns: 49 | Returns the QuantizeConfig for the given layer. 50 | """ 51 | raise NotImplementedError('Must be implemented in subclasses.') 52 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/quantization/keras/quantize_scheme.py: -------------------------------------------------------------------------------- 1 | # Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Quantization scheme which specifies how quantization should be applied. 16 | 17 | Module: tfmot.quantization.keras 18 | """ 19 | 20 | import abc 21 | import six 22 | 23 | 24 | @six.add_metaclass(abc.ABCMeta) 25 | class QuantizeScheme(object): 26 | """ABC interface which specifies transformer and quantization registry.""" 27 | 28 | @abc.abstractmethod 29 | def get_layout_transformer(self): 30 | """Returns the layout transforms for this scheme. 31 | 32 | Returns: 33 | Returns the QuantizeLayoutTransform for this quantization scheme. 34 | """ 35 | raise NotImplementedError('Must be implemented in subclasses.') 36 | 37 | @abc.abstractmethod 38 | def get_quantize_registry(self): 39 | """Returns the quantization registry for this scheme. 40 | 41 | Returns: 42 | Returns the QuantizeRegistry for this quantization scheme. 43 | """ 44 | raise NotImplementedError('Must be implemented in subclasses.') 45 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/sparsity/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | 3 | package(default_visibility = ["//visibility:public"]) 4 | 5 | licenses(["notice"]) 6 | 7 | py_strict_library( 8 | name = "sparsity", 9 | srcs = ["__init__.py"], 10 | deps = [ 11 | "//tensorflow_model_optimization/python/core/sparsity/keras", # buildcleaner: keep 12 | ], 13 | ) 14 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/sparsity/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/sparsity/keras/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/sparsity/keras/prunable_layer.py: -------------------------------------------------------------------------------- 1 | # Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Abstract Base Class for making a custom keras layer prunable.""" 16 | 17 | import abc 18 | import six 19 | 20 | 21 | @six.add_metaclass(abc.ABCMeta) 22 | class PrunableLayer(object): 23 | """Abstract Base Class for making your own keras layer prunable. 24 | 25 | Custom keras layers which want to add pruning should implement this class. 26 | 27 | """ 28 | 29 | @abc.abstractmethod 30 | def get_prunable_weights(self): 31 | """Returns list of prunable weight tensors. 32 | 33 | All the weight tensors which the layer wants to be pruned during 34 | training must be returned by this method. 35 | 36 | Returns: List of weight tensors/kernels in the keras layer which must be 37 | pruned during training. 38 | """ 39 | raise NotImplementedError('Must be implemented in subclasses.') 40 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/sparsity/keras/tools/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | # Placeholder: load py_binary 3 | # Placeholder: load py_test 4 | 5 | package(default_visibility = [ 6 | "//tensorflow_model_optimization:__subpackages__", 7 | ]) 8 | 9 | licenses(["notice"]) 10 | 11 | py_strict_library( 12 | name = "sparsity_tooling", 13 | srcs = ["sparsity_tooling.py"], 14 | visibility = ["//visibility:public"], 15 | deps = [ 16 | # tensorflow dep1, 17 | "//tensorflow_model_optimization/python/core/keras:compat", 18 | "//tensorflow_model_optimization/python/core/keras:metrics", 19 | "//tensorflow_model_optimization/python/core/sparsity/keras:prune", 20 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_schedule", 21 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_wrapper", 22 | ], 23 | ) 24 | 25 | py_test( 26 | name = "sparsity_tooling_test", 27 | size = "medium", 28 | srcs = ["sparsity_tooling_test.py"], 29 | visibility = ["//visibility:public"], 30 | deps = [ 31 | ":sparsity_tooling", 32 | # absl/testing:parameterized dep1, 33 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 34 | # tensorflow dep1, 35 | "//tensorflow_model_optimization/python/core/keras:compat", 36 | "//tensorflow_model_optimization/python/core/sparsity/keras:test_utils", 37 | ], 38 | ) 39 | 40 | py_binary( 41 | name = "evaluate_pruning", 42 | srcs = ["evaluate_pruning.py"], 43 | deps = [ 44 | ":sparsity_tooling", 45 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 46 | # tensorflow dep1, 47 | "//tensorflow_model_optimization/python/core/keras:compat", 48 | "//tensorflow_model_optimization/python/core/sparsity/keras:prune", 49 | ], 50 | ) 51 | 52 | py_strict_library( 53 | name = "check_sparsity_m_by_n", 54 | srcs = ["check_sparsity_m_by_n.py"], 55 | deps = [ 56 | # absl:app dep1, 57 | # absl/flags dep1, 58 | # numpy dep1, 59 | # tensorflow dep1, 60 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_utils", 61 | ], 62 | ) 63 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/sparsity/keras/tools/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/sparsity/keras/tools/sparsity_tooling_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Tests for keras pruning tools in sparsity_tooling.py.""" 16 | 17 | import tensorflow as tf 18 | 19 | from tensorflow_model_optimization.python.core.keras.compat import keras 20 | from tensorflow_model_optimization.python.core.sparsity.keras import pruning_wrapper 21 | from tensorflow_model_optimization.python.core.sparsity.keras import test_utils 22 | from tensorflow_model_optimization.python.core.sparsity.keras.tools import sparsity_tooling 23 | 24 | 25 | test = tf.test 26 | 27 | 28 | class SparsityToolingTest(test.TestCase): 29 | 30 | def test_prune_model(self): 31 | model = keras.Sequential([ 32 | keras.layers.Dense(10, input_shape=(10,)), 33 | keras.layers.Dense(2), 34 | ]) 35 | pruned_model = sparsity_tooling.prune_for_benchmark( 36 | model, target_sparsity=.8, block_size=(1, 1)) 37 | 38 | for layer in pruned_model.layers: 39 | self.assertEqual((1, 1), layer.block_size) 40 | test_utils.assert_model_sparsity(self, 0.8, pruned_model) 41 | 42 | def test_prune_model_recursively(self): 43 | """Checks that models are recursively pruned.""" 44 | 45 | # Setup a model with one layer being a keras.Model. 46 | internal_model = keras.Sequential([ 47 | keras.layers.Dense(10, input_shape=(10,)), 48 | ]) 49 | model = keras.Sequential([ 50 | internal_model, 51 | keras.layers.Dense(20), 52 | ]) 53 | pruned_model = sparsity_tooling.prune_for_benchmark( 54 | model, target_sparsity=.8, block_size=(1, 1)) 55 | 56 | test_utils.assert_model_sparsity(self, 0.8, pruned_model) 57 | 58 | # Check the block size of the prunned layers 59 | prunned_dense_layers = [ 60 | layer 61 | for layer in pruned_model.submodules 62 | if isinstance(layer, pruning_wrapper.PruneLowMagnitude) 63 | ] 64 | self.assertEqual(2, len(prunned_dense_layers)) 65 | for layer in prunned_dense_layers: 66 | self.assertEqual((1, 1), layer.block_size) 67 | 68 | 69 | if __name__ == '__main__': 70 | test.main() 71 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/core/version.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Define TensorFlow Model Optimization version information.""" 16 | 17 | # We follow Semantic Versioning (https://semver.org/) 18 | _MAJOR_VERSION = '0' 19 | _MINOR_VERSION = '8' 20 | _PATCH_VERSION = '0' 21 | 22 | # When building releases, we can update this value on the release branch to 23 | # reflect the current release candidate ('rc0', 'rc1') or, finally, the official 24 | # stable release (indicated by `_VERSION_SUFFIX = ''`). Outside the context of a 25 | # release branch, the current version is by default assumed to be a 26 | # 'development' version, labeled 'dev'. 27 | _VERSION_SUFFIX = '' 28 | 29 | # Example, '0.4.0-dev' 30 | __version__ = '.'.join([ 31 | _MAJOR_VERSION, 32 | _MINOR_VERSION, 33 | _PATCH_VERSION, 34 | ]) 35 | if _VERSION_SUFFIX: 36 | __version__ = '{}-{}'.format(__version__, _VERSION_SUFFIX) 37 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/model-optimization/acd0717cebfde7dd0c707ddbf4915669c97e4b6a/tensorflow_model_optimization/python/examples/__init__.py -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/cluster_preserve_qat/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/model-optimization/acd0717cebfde7dd0c707ddbf4915669c97e4b6a/tensorflow_model_optimization/python/examples/cluster_preserve_qat/__init__.py -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/cluster_preserve_qat/keras/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_binary") 2 | 3 | licenses(["notice"]) 4 | 5 | py_strict_binary( 6 | name = "mnist_cnn", 7 | srcs = [ 8 | "mnist_cnn.py", 9 | ], 10 | deps = [ 11 | # absl:app dep1, 12 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 13 | # numpy dep1, 14 | # tensorflow dep1, 15 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster", 16 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster_config", 17 | "//tensorflow_model_optimization/python/core/keras:compat", 18 | "//tensorflow_model_optimization/python/core/quantization/keras:quantize", 19 | "//tensorflow_model_optimization/python/core/quantization/keras/collab_opts/cluster_preserve:cluster_utils", 20 | "//tensorflow_model_optimization/python/core/quantization/keras/collab_opts/cluster_preserve:default_8bit_cluster_preserve_quantize_scheme", 21 | ], 22 | ) 23 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/cluster_preserve_qat/keras/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/model-optimization/acd0717cebfde7dd0c707ddbf4915669c97e4b6a/tensorflow_model_optimization/python/examples/cluster_preserve_qat/keras/__init__.py -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/clustering/keras/imdb/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_library") 2 | # Placeholder: load py_binary 3 | 4 | package( 5 | default_visibility = ["//visibility:public"], 6 | ) 7 | 8 | licenses(["notice"]) 9 | 10 | filegroup( 11 | name = "all_files", 12 | srcs = glob(["**"]), 13 | ) 14 | 15 | py_binary( 16 | name = "imdb_lstm", 17 | srcs = [ 18 | "imdb_lstm.py", 19 | ], 20 | deps = [ 21 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 22 | # numpy dep1, 23 | # tensorflow dep1, 24 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster", 25 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster_config", 26 | ], 27 | ) 28 | 29 | py_binary( 30 | name = "imdb_rnn", 31 | srcs = [ 32 | "imdb_rnn.py", 33 | ], 34 | deps = [ 35 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 36 | # numpy dep1, 37 | # tensorflow dep1, 38 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster", 39 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster_config", 40 | ], 41 | ) 42 | 43 | py_binary( 44 | name = "imdb_gru", 45 | srcs = [ 46 | "imdb_gru.py", 47 | ], 48 | deps = [ 49 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 50 | # numpy dep1, 51 | # tensorflow dep1, 52 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster", 53 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster_config", 54 | ], 55 | ) 56 | 57 | py_strict_library( 58 | name = "imdb_utils", 59 | srcs = [ 60 | "imdb_utils.py", 61 | ], 62 | deps = [ 63 | # tensorflow dep1, 64 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster", 65 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster_config", 66 | "//tensorflow_model_optimization/python/core/keras:compat", 67 | ], 68 | ) 69 | 70 | py_binary( 71 | name = "imdb_multiple_cells", 72 | srcs = [ 73 | "imdb_multiple_cells.py", 74 | ], 75 | deps = [ 76 | ":imdb_utils", 77 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 78 | # tensorflow dep1, 79 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster", 80 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster_config", 81 | "//tensorflow_model_optimization/python/core/keras:compat", 82 | ], 83 | ) 84 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/clustering/keras/imdb/imdb_gru.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Train a GRU on the IMDB sentiment classification task. 16 | 17 | The dataset is actually too small for LSTM to be of any advantage 18 | compared to simpler, much faster methods such as TF-IDF+LogReg. 19 | """ 20 | 21 | from __future__ import print_function 22 | 23 | import tensorflow.keras as keras 24 | import tensorflow.keras.preprocessing.sequence as sequence 25 | from tensorflow_model_optimization.python.core.clustering.keras import cluster 26 | from tensorflow_model_optimization.python.core.clustering.keras import cluster_config 27 | 28 | 29 | max_features = 20000 30 | maxlen = 100 # cut texts after this number of words 31 | batch_size = 32 32 | 33 | print("Loading data...") 34 | (x_train, 35 | y_train), (x_test, 36 | y_test) = keras.datasets.imdb.load_data(num_words=max_features) 37 | print(len(x_train), "train sequences") 38 | print(len(x_test), "test sequences") 39 | 40 | print("Pad sequences (samples x time)") 41 | x_train = sequence.pad_sequences(x_train, maxlen=maxlen) 42 | x_test = sequence.pad_sequences(x_test, maxlen=maxlen) 43 | print("x_train shape:", x_train.shape) 44 | print("x_test shape:", x_test.shape) 45 | 46 | print("Build model...") 47 | model = keras.models.Sequential() 48 | 49 | model.add(keras.layers.Embedding(max_features, 128, input_length=maxlen)) 50 | model.add(keras.layers.GRU(128)) 51 | model.add(keras.layers.Dropout(0.5)) 52 | model.add(keras.layers.Dense(1)) 53 | model.add(keras.layers.Activation("sigmoid")) 54 | 55 | model = cluster.cluster_weights( 56 | model, 57 | number_of_clusters=16, 58 | cluster_centroids_init=cluster_config.CentroidInitialization 59 | .KMEANS_PLUS_PLUS, 60 | ) 61 | 62 | model.compile(loss="binary_crossentropy", 63 | optimizer="adam", 64 | metrics=["accuracy"]) 65 | 66 | 67 | print("Train...") 68 | model.fit(x_train, y_train, batch_size=batch_size, epochs=3, 69 | validation_data=(x_test, y_test)) 70 | score, acc = model.evaluate(x_test, y_test, 71 | batch_size=batch_size) 72 | 73 | print("Test score:", score) 74 | print("Test accuracy:", acc) 75 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/clustering/keras/imdb/imdb_lstm.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Train a LSTM on the IMDB sentiment classification task. 16 | 17 | The dataset is actually too small for LSTM to be of any advantage 18 | compared to simpler, much faster methods such as TF-IDF+LogReg. 19 | """ 20 | 21 | from __future__ import print_function 22 | 23 | import tensorflow.keras as keras 24 | import tensorflow.keras.preprocessing.sequence as sequence 25 | 26 | from tensorflow_model_optimization.python.core.clustering.keras import cluster 27 | from tensorflow_model_optimization.python.core.clustering.keras import cluster_config 28 | 29 | 30 | max_features = 20000 31 | maxlen = 100 # cut texts after this number of words 32 | batch_size = 32 33 | 34 | print("Loading data...") 35 | (x_train, 36 | y_train), (x_test, 37 | y_test) = keras.datasets.imdb.load_data(num_words=max_features) 38 | print(len(x_train), "train sequences") 39 | print(len(x_test), "test sequences") 40 | 41 | print("Pad sequences (samples x time)") 42 | x_train = sequence.pad_sequences(x_train, maxlen=maxlen) 43 | x_test = sequence.pad_sequences(x_test, maxlen=maxlen) 44 | print("x_train shape:", x_train.shape) 45 | print("x_test shape:", x_test.shape) 46 | 47 | print("Build model...") 48 | model = keras.models.Sequential() 49 | 50 | model.add(keras.layers.Embedding(max_features, 128, input_length=maxlen)) 51 | model.add(keras.layers.LSTM(128)) 52 | model.add(keras.layers.Dropout(0.5)) 53 | model.add(keras.layers.Dense(1)) 54 | model.add(keras.layers.Activation("sigmoid")) 55 | 56 | model = cluster.cluster_weights( 57 | model, 58 | number_of_clusters=16, 59 | cluster_centroids_init=cluster_config.CentroidInitialization 60 | .KMEANS_PLUS_PLUS, 61 | ) 62 | 63 | model.compile(loss="binary_crossentropy", 64 | optimizer="adam", 65 | metrics=["accuracy"]) 66 | 67 | print("Train...") 68 | model.fit(x_train, y_train, batch_size=batch_size, epochs=3, 69 | validation_data=(x_test, y_test)) 70 | score, acc = model.evaluate(x_test, y_test, 71 | batch_size=batch_size) 72 | 73 | print("Test score:", score) 74 | print("Test accuracy:", acc) 75 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/clustering/keras/imdb/imdb_multiple_cells.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """End-to-end tests for StackedRNNCells. 16 | 17 | The dataset is actually too small for LSTM to be of any advantage 18 | compared to simpler, much faster methods such as TF-IDF+LogReg. 19 | """ 20 | 21 | from __future__ import print_function 22 | 23 | import tensorflow as tf 24 | 25 | from tensorflow_model_optimization.python.core.keras.compat import keras 26 | from tensorflow_model_optimization.python.examples.clustering.keras.imdb.imdb_utils import cluster_train_eval_strip 27 | from tensorflow_model_optimization.python.examples.clustering.keras.imdb.imdb_utils import prepare_dataset 28 | 29 | 30 | max_features = 20000 31 | maxlen = 100 # cut texts after this number of words 32 | batch_size = 32 33 | 34 | x_train, y_train, x_test, y_test = prepare_dataset() 35 | 36 | print("Build a model with the StackedRNNCells with LSTMCell...") 37 | model = keras.models.Sequential() 38 | 39 | model.add(keras.layers.Embedding(max_features, 128, input_length=maxlen)) 40 | model.add( 41 | keras.layers.RNN( 42 | keras.layers.StackedRNNCells( 43 | [keras.layers.LSTMCell(128) for _ in range(2)] 44 | ) 45 | ) 46 | ) 47 | model.add(keras.layers.Dropout(0.5)) 48 | model.add(keras.layers.Dense(1)) 49 | model.add(keras.layers.Activation("sigmoid")) 50 | 51 | test_case = "StackedRNNCells_LSTMCell" 52 | cluster_train_eval_strip( 53 | model, x_train, y_train, x_test, y_test, batch_size, test_case) 54 | 55 | print("Build a model with the Bidirectional wrapper with LSTM layer...") 56 | model = keras.models.Sequential() 57 | 58 | model.add(keras.layers.Embedding(max_features, 128, input_length=maxlen)) 59 | model.add(keras.layers.Bidirectional(keras.layers.LSTM(128))) 60 | model.add(keras.layers.Dropout(0.5)) 61 | model.add(keras.layers.Dense(1)) 62 | model.add(keras.layers.Activation("sigmoid")) 63 | 64 | test_case = "Bidirectional_LSTM" 65 | cluster_train_eval_strip( 66 | model, x_train, y_train, x_test, y_test, batch_size, test_case) 67 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/clustering/keras/imdb/imdb_rnn.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Train a SimpleRNN on the IMDB sentiment classification task. 16 | 17 | The dataset is actually too small for LSTM to be of any advantage 18 | compared to simpler, much faster methods such as TF-IDF+LogReg. 19 | """ 20 | 21 | from __future__ import print_function 22 | 23 | import tensorflow.keras as keras 24 | import tensorflow.keras.preprocessing.sequence as sequence 25 | from tensorflow_model_optimization.python.core.clustering.keras import cluster 26 | from tensorflow_model_optimization.python.core.clustering.keras import cluster_config 27 | 28 | 29 | max_features = 20000 30 | maxlen = 100 # cut texts after this number of words 31 | batch_size = 32 32 | 33 | print("Loading data...") 34 | (x_train, 35 | y_train), (x_test, 36 | y_test) = keras.datasets.imdb.load_data(num_words=max_features) 37 | print(len(x_train), "train sequences") 38 | print(len(x_test), "test sequences") 39 | 40 | print("Pad sequences (samples x time)") 41 | x_train = sequence.pad_sequences(x_train, maxlen=maxlen) 42 | x_test = sequence.pad_sequences(x_test, maxlen=maxlen) 43 | print("x_train shape:", x_train.shape) 44 | print("x_test shape:", x_test.shape) 45 | 46 | print("Build model...") 47 | model = keras.models.Sequential() 48 | 49 | model.add(keras.layers.Embedding(max_features, 128, input_length=maxlen)) 50 | model.add(keras.layers.SimpleRNN(128)) 51 | model.add(keras.layers.Dropout(0.5)) 52 | model.add(keras.layers.Dense(1)) 53 | model.add(keras.layers.Activation("sigmoid")) 54 | 55 | model = cluster.cluster_weights( 56 | model, 57 | number_of_clusters=16, 58 | cluster_centroids_init=cluster_config.CentroidInitialization 59 | .KMEANS_PLUS_PLUS, 60 | ) 61 | 62 | model.compile(loss="binary_crossentropy", 63 | optimizer="adam", 64 | metrics=["accuracy"]) 65 | 66 | 67 | print("Train...") 68 | model.fit(x_train, y_train, batch_size=batch_size, epochs=3, 69 | validation_data=(x_test, y_test)) 70 | score, acc = model.evaluate(x_test, y_test, 71 | batch_size=batch_size) 72 | 73 | print("Test score:", score) 74 | print("Test accuracy:", acc) 75 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/clustering/keras/imdb/imdb_utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Common utils for testing RNN e2e tests. 16 | 17 | The dataset is actually too small for LSTM to be of any advantage 18 | compared to simpler, much faster methods such as TF-IDF+LogReg. 19 | """ 20 | from __future__ import print_function 21 | 22 | import tensorflow as tf 23 | 24 | from tensorflow_model_optimization.python.core.clustering.keras import cluster 25 | from tensorflow_model_optimization.python.core.clustering.keras import cluster_config 26 | from tensorflow_model_optimization.python.core.keras.compat import keras 27 | 28 | 29 | sequence = keras.preprocessing.sequence 30 | 31 | 32 | def prepare_dataset(): 33 | """Prepare the dataset.""" 34 | 35 | max_features = 20000 36 | maxlen = 100 # cut texts after this number of words 37 | 38 | print("Loading data...") 39 | (x_train, y_train), (x_test, y_test) = keras.datasets.imdb.load_data( 40 | num_words=max_features 41 | ) 42 | print(len(x_train), "train sequences") 43 | print(len(x_test), "test sequences") 44 | 45 | print("Pad sequences (samples x time)") 46 | x_train = sequence.pad_sequences(x_train, maxlen=maxlen) 47 | x_test = sequence.pad_sequences(x_test, maxlen=maxlen) 48 | 49 | return x_train, y_train, x_test, y_test 50 | 51 | 52 | def cluster_train_eval_strip( 53 | model, x_train, y_train, x_test, y_test, batch_size, test_case): 54 | """Train, evaluate and strip clustering.""" 55 | model = cluster.cluster_weights( 56 | model, 57 | number_of_clusters=16, 58 | cluster_centroids_init=cluster_config.CentroidInitialization 59 | .KMEANS_PLUS_PLUS,) 60 | 61 | model.compile(loss="binary_crossentropy", 62 | optimizer="adam", 63 | metrics=["accuracy"]) 64 | 65 | print("Train...") 66 | model.fit(x_train, y_train, batch_size=batch_size, epochs=1, 67 | validation_data=(x_test, y_test), verbose=2) 68 | score, acc = model.evaluate(x_test, y_test, 69 | batch_size=batch_size) 70 | 71 | print("Test score:", score) 72 | print("Test accuracy:", acc) 73 | 74 | print("Strip clustering wrapper...") 75 | model = cluster.strip_clustering(model) 76 | if "Bidirectional" in test_case: 77 | layer_weight = getattr(model.layers[1].forward_layer.cell, "kernel") 78 | elif "StackedRNNCells" in test_case: 79 | layer_weight = getattr(model.layers[1].cell.cells[0], "kernel") 80 | else: 81 | raise ValueError("Only Bidirectional and StackedRNNCells are tested now.") 82 | print("Number of clusters:", len(set(layer_weight.numpy().flatten()))) 83 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/clustering/keras/mnist/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_binary") 2 | 3 | package( 4 | default_visibility = ["//visibility:public"], 5 | ) 6 | 7 | licenses(["notice"]) 8 | 9 | filegroup( 10 | name = "all_files", 11 | srcs = glob(["**"]), 12 | ) 13 | 14 | py_strict_binary( 15 | name = "mnist_cnn", 16 | srcs = [ 17 | "mnist_cnn.py", 18 | ], 19 | deps = [ 20 | # absl:app dep1, 21 | # absl/flags dep1, 22 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 23 | # tensorflow dep1, 24 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster", 25 | "//tensorflow_model_optimization/python/core/clustering/keras:cluster_config", 26 | "//tensorflow_model_optimization/python/core/clustering/keras:clustering_callbacks", 27 | "//tensorflow_model_optimization/python/core/keras:compat", 28 | ], 29 | ) 30 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/clustering/keras/mnist/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/quantization/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/model-optimization/acd0717cebfde7dd0c707ddbf4915669c97e4b6a/tensorflow_model_optimization/python/examples/quantization/__init__.py -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/quantization/keras/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_binary") 2 | 3 | licenses(["notice"]) 4 | 5 | py_strict_binary( 6 | name = "mnist_cnn", 7 | srcs = [ 8 | "mnist_cnn.py", 9 | ], 10 | deps = [ 11 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 12 | # tensorflow dep1, 13 | "//tensorflow_model_optimization/python/core/keras:compat", 14 | "//tensorflow_model_optimization/python/core/quantization/keras:quantize", 15 | ], 16 | ) 17 | 18 | py_strict_binary( 19 | name = "mnist_cnn_cont_quant", 20 | srcs = [ 21 | "mnist_cnn_cont_quant.py", 22 | ], 23 | deps = [ 24 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 25 | # numpy dep1, 26 | # tensorflow dep1, 27 | "//tensorflow_model_optimization/python/core/keras:compat", 28 | "//tensorflow_model_optimization/python/core/quantization/keras:quantize", 29 | ], 30 | ) 31 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/quantization/keras/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/model-optimization/acd0717cebfde7dd0c707ddbf4915669c97e4b6a/tensorflow_model_optimization/python/examples/quantization/keras/__init__.py -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/quantization_with_sparsity/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/model-optimization/acd0717cebfde7dd0c707ddbf4915669c97e4b6a/tensorflow_model_optimization/python/examples/quantization_with_sparsity/__init__.py -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/quantization_with_sparsity/keras/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_binary") 2 | 3 | licenses(["notice"]) 4 | 5 | py_strict_binary( 6 | name = "mnist_cnn", 7 | srcs = [ 8 | "mnist_cnn.py", 9 | ], 10 | deps = [ 11 | # absl:app dep1, 12 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 13 | # numpy dep1, 14 | # tensorflow dep1, 15 | "//tensorflow_model_optimization/python/core/keras:compat", 16 | "//tensorflow_model_optimization/python/core/quantization/keras:quantize", 17 | "//tensorflow_model_optimization/python/core/quantization/keras/collab_opts/prune_preserve:default_8bit_prune_preserve_quantize_scheme", 18 | "//tensorflow_model_optimization/python/core/sparsity/keras:prune", 19 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_callbacks", 20 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_schedule", 21 | ], 22 | ) 23 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/quantization_with_sparsity/keras/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tensorflow/model-optimization/acd0717cebfde7dd0c707ddbf4915669c97e4b6a/tensorflow_model_optimization/python/examples/quantization_with_sparsity/keras/__init__.py -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/sparsity/keras/imdb/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_binary") 2 | 3 | package( 4 | default_visibility = ["//visibility:public"], 5 | ) 6 | 7 | licenses(["notice"]) 8 | 9 | filegroup( 10 | name = "all_files", 11 | srcs = glob(["**"]), 12 | ) 13 | 14 | py_strict_binary( 15 | name = "imdb_lstm", 16 | srcs = [ 17 | "imdb_lstm.py", 18 | ], 19 | deps = [ 20 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 21 | # numpy dep1, 22 | # tensorflow dep1, 23 | "//tensorflow_model_optimization/python/core/keras:compat", 24 | "//tensorflow_model_optimization/python/core/sparsity/keras:prune", 25 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_callbacks", 26 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_schedule", 27 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_wrapper", 28 | ], 29 | ) 30 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/sparsity/keras/mnist/BUILD: -------------------------------------------------------------------------------- 1 | load("//tensorflow_model_optimization:tensorflow_model_optimization.bzl", "py_strict_binary") 2 | 3 | package( 4 | default_visibility = ["//visibility:public"], 5 | ) 6 | 7 | licenses(["notice"]) 8 | 9 | filegroup( 10 | name = "all_files", 11 | srcs = glob(["**"]), 12 | ) 13 | 14 | py_strict_binary( 15 | name = "mnist_cnn", 16 | srcs = [ 17 | "mnist_cnn.py", 18 | ], 19 | deps = [ 20 | # absl:app dep1, 21 | # absl/flags dep1, 22 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 23 | # tensorflow dep1, 24 | "//tensorflow_model_optimization/python/core/keras:compat", 25 | "//tensorflow_model_optimization/python/core/sparsity/keras:prune", 26 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_callbacks", 27 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_schedule", 28 | ], 29 | ) 30 | 31 | py_strict_binary( 32 | name = "mnist_e2e", 33 | srcs = [ 34 | "mnist_e2e.py", 35 | ], 36 | deps = [ 37 | # absl:app dep1, 38 | # absl/flags dep1, 39 | # google/protobuf:use_fast_cpp_protos dep1, # Automatically added 40 | # tensorflow dep1, 41 | "//tensorflow_model_optimization/python/core/keras:compat", 42 | "//tensorflow_model_optimization/python/core/keras:test_utils", 43 | "//tensorflow_model_optimization/python/core/sparsity/keras:prune", 44 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_callbacks", 45 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_schedule", 46 | ], 47 | ) 48 | 49 | py_strict_binary( 50 | name = "mnist_e2e_sparsity2x4", 51 | srcs = ["mnist_e2e_sparsity2x4.py"], 52 | deps = [ 53 | # absl:app dep1, 54 | # tensorflow dep1, 55 | "//tensorflow_model_optimization/python/core/keras:compat", 56 | "//tensorflow_model_optimization/python/core/keras:test_utils", 57 | "//tensorflow_model_optimization/python/core/sparsity/keras:prune", 58 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_callbacks", 59 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_schedule", 60 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_utils", 61 | "//tensorflow_model_optimization/python/core/sparsity/keras:pruning_wrapper", 62 | ], 63 | ) 64 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/python/examples/sparsity/keras/mnist/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /tensorflow_model_optimization/tensorflow_model_optimization.bzl: -------------------------------------------------------------------------------- 1 | # Copyright 2021 The TensorFlow Authors. All Rights Reserved. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Definitions for custom rules.""" 16 | 17 | # Placeholder to use until bazel supports py_strict_binary. 18 | def py_strict_binary(name, **kwargs): 19 | native.py_binary(name = name, **kwargs) 20 | 21 | # Placeholder to use until bazel supports py_strict_library. 22 | def py_strict_library(name, **kwargs): 23 | native.py_library(name = name, **kwargs) 24 | 25 | # Placeholder to use until bazel supports pytype_strict_binary. 26 | def pytype_strict_binary(name, **kwargs): 27 | native.py_binary(name = name, **kwargs) 28 | 29 | # Placeholder to use until bazel supports pytype_strict_library. 30 | def pytype_strict_library(name, **kwargs): 31 | native.py_library(name = name, **kwargs) 32 | 33 | # Placeholder to use until bazel supports pytype_library. 34 | def pytype_library(name, **kwargs): 35 | native.py_library(name = name, **kwargs) 36 | 37 | # Placeholder to use until bazel supports py_strict_test. 38 | def py_strict_test(name, **kwargs): 39 | native.py_test(name = name, **kwargs) 40 | 41 | # Placeholder to use until bazel supports pytype_strict_test. 42 | def pytype_strict_test(name, **kwargs): 43 | native.py_test(name = name, **kwargs) 44 | --------------------------------------------------------------------------------