├── .github └── ISSUE_TEMPLATE │ ├── bug_report.md │ ├── documentation-issue.md │ ├── feature_request.md │ └── other-issues.md ├── .gitignore ├── .pylintrc ├── CONTRIBUTING.md ├── CONTRIBUTING_en.md ├── LICENSE.TXT ├── README.md ├── automl ├── __init__.py ├── automl.yaml ├── automl_hparam.conf ├── cvt_hparam_file.py ├── parse_results.py └── seven.yaml ├── datasets ├── __init__.py ├── abstract_dataset.py ├── cifar10_dataset.py ├── ilsvrc12_dataset.py └── pascalvoc_dataset.py ├── docs ├── README.md ├── doc-requirements.txt ├── docs │ ├── .markdownlint.json │ ├── MathJax.js │ ├── automl_based_methods.md │ ├── cp_learner.md │ ├── cpr_learner.md │ ├── dcp_learner.md │ ├── distillation.md │ ├── faq.md │ ├── index.md │ ├── installation.md │ ├── multi_gpu_training.md │ ├── nuq_learner.md │ ├── performance.md │ ├── pics │ │ ├── dcp_learner.png │ │ ├── deep_compression_algor.png │ │ ├── framework_design.png │ │ ├── rl_workflow.png │ │ ├── train_n_inference.png │ │ └── wsl_pr_schedule.png │ ├── pre_trained_models.md │ ├── reference.md │ ├── reinforcement_learning.md │ ├── self_defined_models.md │ ├── test_cases.md │ ├── tutorial.md │ ├── uq_learner.md │ └── ws_learner.md ├── mkdocs.yml └── qr_code.jpg ├── examples ├── convnet_at_fmnist.py ├── convnet_at_fmnist_run.py └── fmnist_dataset.py ├── learners ├── __init__.py ├── abstract_learner.py ├── channel_pruning │ ├── __init__.py │ ├── channel_pruner.py │ ├── learner.py │ └── model_wrapper.py ├── channel_pruning_gpu │ ├── __init__.py │ └── learner.py ├── channel_pruning_rmt │ ├── __init__.py │ └── learner.py ├── discr_channel_pruning │ ├── __init__.py │ └── learner.py ├── distillation_helper.py ├── full_precision │ ├── __init__.py │ └── learner.py ├── learner_utils.py ├── nonuniform_quantization │ ├── bit_optimizer.py │ ├── learner.py │ ├── rl_helper.py │ └── utils.py ├── uniform_quantization │ ├── bit_optimizer.py │ ├── learner.py │ ├── rl_helper.py │ └── utils.py ├── uniform_quantization_tf │ ├── __init__.py │ ├── learner.py │ └── utils.py └── weight_sparsification │ ├── __init__.py │ ├── learner.py │ ├── pr_optimizer.py │ ├── rl_helper.py │ └── utils.py ├── main.sh ├── nets ├── __init__.py ├── abstract_model_helper.py ├── faster_rcnn_at_pascalvoc.py ├── faster_rcnn_at_pascalvoc_run.py ├── lenet_at_cifar10.py ├── lenet_at_cifar10_run.py ├── mobilenet_at_ilsvrc12.py ├── mobilenet_at_ilsvrc12_run.py ├── resnet_at_cifar10.py ├── resnet_at_cifar10_run.py ├── resnet_at_ilsvrc12.py ├── resnet_at_ilsvrc12_run.py ├── vgg_at_pascalvoc.py └── vgg_at_pascalvoc_run.py ├── path.conf.template ├── requirement.txt ├── rl_agents ├── __init__.py ├── ddpg │ ├── __init__.py │ ├── actor_critic.py │ ├── agent.py │ ├── noise.py │ ├── replay_buffer.py │ └── running_mean_std.py └── unit_tests │ ├── __init__.py │ ├── move_to_target.py │ └── pendulum_v0.py ├── run.sh ├── run_pylint.sh ├── scripts ├── create_minimal.sh ├── download_minimal.sh ├── run_docker.sh ├── run_local.sh └── run_seven.sh ├── seven.yaml ├── tools ├── benchmark │ └── calc_inference_time.py ├── conversion │ ├── convert_data_format.py │ ├── export_chn_pruned_tflite_model.py │ ├── export_pb_tflite_models.py │ └── export_quant_tflite_model.py └── graph_tools │ └── add_to_collection.py └── utils ├── __init__.py ├── external ├── __init__.py ├── conv_blocks.py ├── faster_rcnn_tensorflow │ ├── configs │ │ └── cfgs.py │ ├── net │ │ ├── mobilenet_v2_faster_rcnn.py │ │ └── resnet_faster_rcnn.py │ ├── preprocessing │ │ └── faster_rcnn_preprocessing.py │ └── utility │ │ ├── __init__.py │ │ ├── anchor_target_layer_without_boxweight.py │ │ ├── anchor_utils.py │ │ ├── boxes_utils.py │ │ ├── coco_dict.py │ │ ├── draw_box_in_img.py │ │ ├── encode_and_decode.py │ │ ├── label_dict.py │ │ ├── loss_utils.py │ │ ├── proposal_opr.py │ │ ├── proposal_target_layer.py │ │ ├── remote_sensing_dict.py │ │ ├── show_box_in_tensor.py │ │ └── tf_ops.py ├── imagenet_preprocessing.py ├── mobilenet.py ├── mobilenet_v1.py ├── mobilenet_v2.py ├── resnet_model.py └── ssd_tensorflow │ ├── LICENSE │ ├── README.md │ ├── dataset │ ├── convert_tfrecords.py │ ├── dataset_common.py │ └── dataset_inspect.py │ ├── demo │ ├── demo1.jpg │ ├── demo2.jpg │ └── demo3.jpg │ ├── eval_ssd.py │ ├── net │ └── ssd_net.py │ ├── preprocessing │ ├── preprocessing_unittest.py │ └── ssd_preprocessing.py │ ├── simple_ssd_demo.py │ ├── train_ssd.py │ ├── utility │ ├── anchor_manipulator.py │ ├── anchor_manipulator_unittest.py │ ├── checkpint_inspect.py │ ├── draw_toolbox.py │ └── scaffolds.py │ └── voc_eval.py ├── get_data_dir.py ├── get_idle_gpus.py ├── get_path_args.py ├── lrn_rate_utils.py ├── misc_utils.py └── multi_gpu_wrapper.py /.github/ISSUE_TEMPLATE/bug_report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Bug report 3 | about: Create a report to help us improve 4 | 5 | --- 6 | 7 | **Describe the bug** 8 | A clear and concise description of what the bug is. 9 | 10 | **To Reproduce** 11 | Steps to reproduce the behavior: 12 | 1. Go to '...' 13 | 2. Click on '....' 14 | 3. Scroll down to '....' 15 | 4. See error 16 | 17 | **Expected behavior** 18 | A clear and concise description of what you expected to happen. 19 | 20 | **Screenshots** 21 | If applicable, add screenshots to help explain your problem. 22 | 23 | **Desktop (please complete the following information):** 24 | - OS: [e.g. iOS] 25 | - Browser [e.g. chrome, safari] 26 | - Version [e.g. 22] 27 | 28 | **Smartphone (please complete the following information):** 29 | - Device: [e.g. iPhone6] 30 | - OS: [e.g. iOS8.1] 31 | - Browser [e.g. stock browser, safari] 32 | - Version [e.g. 22] 33 | 34 | **Additional context** 35 | Add any other context about the problem here. 36 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/documentation-issue.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Documentation Issue 3 | about: Use this template for documentation related issues 4 | 5 | --- 6 | 7 | Please make sure that this is a documentation issue. 8 | 9 | 10 | **System information** 11 | - PocketFlow version: 12 | - Doc Link: 13 | 14 | 15 | **Describe the documentation issue** 16 | 17 | **We welcome contributions by users. Will you be able to update submit a PR to fix the doc Issue?** 18 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/feature_request.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Feature request 3 | about: Suggest an idea for this project 4 | 5 | --- 6 | 7 | **Is your feature request related to a problem? Please describe.** 8 | A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] 9 | 10 | **Describe the solution you'd like** 11 | A clear and concise description of what you want to happen. 12 | 13 | **Describe alternatives you've considered** 14 | A clear and concise description of any alternative solutions or features you've considered. 15 | 16 | **Additional context** 17 | Add any other context or screenshots about the feature request here. 18 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/other-issues.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Other Issues 3 | about: Use this template for any other non-support related issues 4 | 5 | --- 6 | 7 | This template is for miscellaneous issues not covered by the other issue categories. 8 | 9 | For questions on how to work with PocketFlow, or support for problems that are not verified bugs in PocketFlow, please go to [StackOverflow](https://stackoverflow.com). 10 | 11 | 12 | For high-level discussions about TensorFlow, please post to [discuss group](https://groups.google.com/forum/#!forum/pocketflow). 13 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *.swp 3 | main.py 4 | logs 5 | models 6 | models_* 7 | automl_output_* 8 | .pylint_results 9 | start_multi.sh 10 | ratio.list 11 | path.conf 12 | nvidia-smi-dump 13 | ssd_outputs 14 | dump 15 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # 为PocketFlow做出贡献 2 | 3 | [腾讯开源激励计划](https://opensource.tencent.com/contribution)鼓励所有开发者的参与和贡献，我们期待你的加入。你可以报告 [Issues](https://github.com/Tencent/PocketFlow/issues) 或者提交 [Pull Requests](https://github.com/Tencent/PocketFlow/pulls)。在贡献代码之前，请阅读以下指引。 4 | 5 | ## 问题管理 6 | 7 | 我们使用 Github Issues 以收集问题和功能需求。 8 | 9 | ### 查找已有的 Issues 10 | 11 | 在创建新 Issue 之前，请先搜索是否存在已有的或者类似的 Issue ，以避免重复。 12 | 13 | ### 创建新 Issue 14 | 15 | 当创建新 Issue 时，请提供详细的描述、截屏以及/或者短视频来帮助我们定位和复现问题。 16 | 17 | ## 分支管理 18 | 19 | 目前，为简便起见，我们仅有一个分支： 20 | 21 | 1. `master` 分支： 22 | 1. 这是稳定分支。高度稳定的版本将会被标注特定的版本号。 23 | 2. 请向该分支提交包含问题修复或者新功能的 Pull Requests。 24 | 25 | ## Pull Requests 26 | 27 | 我们欢迎所有人向 PocketFlow 贡献代码。我们的代码团队会监控 Pull Requests, 进行相应的代码测试与检查，通过测试的 PR 将会被合并至 `master` 分支。 28 | 29 | 在提交 PR 之前，请确认: 30 | 31 | 1. 从主项目中 fork 代码 32 | 2. 与主项目保持同步 33 | 3. 在代码变动后，对应地修改注释与文档 34 | 4. 在新文件中加入协议与版权声明 35 | 5. 确保一致的代码风格（可使用 `run_pylint.sh`） 36 | 6. 充分测试你的代码 37 | 7. 向 `master` 分支发起 PR 请求 38 | 39 | ## 协议 40 | 41 | [BSD 3-Clause License](https://github.com/Tencent/PocketFlow/blob/master/LICENSE.TXT)是PocketFlow的开源协议，你贡献的代码也会受此协议保护。 42 | -------------------------------------------------------------------------------- /CONTRIBUTING_en.md: -------------------------------------------------------------------------------- 1 | # Contributing to PocketFlow 2 | 3 | [Tencent Open Source Incentive Program](https://opensource.tencent.com/contribution) encourages all developers' participation and contribution and we are looking forward to you joining us. You are welcomed to report [issues](https://github.com/Tencent/PocketFlow/issues) or submit [pull requests](https://github.com/Tencent/PocketFlow/pulls). Before contributing, please read the following guideline. 4 | 5 | ## Issue Management 6 | 7 | We use Github Issues to track public bugs and feature requests. 8 | 9 | ### Search for Existing Issues First 10 | 11 | Please search for existing or similar issues before opening a new one, in order to avoid duplicated issues. 12 | 13 | ### Creating a New Issue 14 | 15 | When creating a new issue, please provide detailed descriptions, screenshots, and/or short videos to help us locate and reproduce the problem(s). 16 | 17 | ### Branch Management 18 | 19 | For the moment, we have only one branch for simplicity: 20 | 21 | 1. `master` branch: 22 | 1. This is the stable branch. Highly-stable versions will be tagged with certain version numbers. 23 | 2. Please submit hotfixs or new features as PR to this branch. 24 | 25 | ### Pull Requests (PR) 26 | 27 | We welcome everyone to contribute to PocketFlow. Our development team will monitor pull requests and perform related tests and code reviews. If passed, the PR will be accepted and merged to the master branch. 28 | 29 | Before submitting a PR, please confirm the following: 30 | 31 | 1. Fork the main repo 32 | 2. Keep updated with the main repo 33 | 3. Update comments and documentation after code changes 34 | 4. Add licence and copyright notes to new files 35 | 5. Keep the code style consistent (use `run_pylint.sh`) 36 | 6. Extensively test your code 37 | 7. Make a pull request to master branch 38 | 39 | ## License 40 | 41 | [BSD 3-Clause License](https://github.com/Tencent/PocketFlow/blob/master/LICENSE.TXT) is PocketFlow's open source license. Your contributed code will also be protected by this license. 42 | -------------------------------------------------------------------------------- /automl/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/automl/__init__.py -------------------------------------------------------------------------------- /automl/automl.yaml: -------------------------------------------------------------------------------- 1 | common: 2 | job_name: pocket-flow-automl-ws-uniform-0-8 3 | description: PocketFlow + AutoML for Weight Sparsifiation Learner with Uniform Pruning Ratios 4 | trial_num: 100 5 | parallel_num: 6 6 | algo_type: GP 7 | epoch_num: 1 8 | config_file: 9 | - automl/automl_hparam.conf 10 | extra_result: 11 | - loss 12 | - prune_ratio 13 | 14 | param_conf: 15 | - name: ws_prune_ratio_exp 16 | type: FLOAT 17 | min: 1.0 18 | max: 5.0 19 | - name: ws_iter_ratio_beg 20 | type: FLOAT 21 | min: 0.0 22 | max: 0.9 23 | - name: ws_iter_ratio_end 24 | type: FLOAT 25 | min: 0.0 26 | max: 1.0 27 | - name: ws_update_mask_step 28 | type: LOG10 29 | min: 50 30 | max: 2000 31 | 32 | notification: 33 | cycle_interval_min: 60 34 | if_alarm_notification: true 35 | if_cycle_notification: true 36 | if_email_notification: true 37 | -------------------------------------------------------------------------------- /automl/automl_hparam.conf: -------------------------------------------------------------------------------- 1 | ws_prune_ratio_exp = ##ws_prune_ratio_exp## 2 | ws_iter_ratio_beg = ##ws_iter_ratio_beg## 3 | ws_iter_ratio_end = ##ws_iter_ratio_end## 4 | ws_update_mask_step = ##ws_update_mask_step## 5 | -------------------------------------------------------------------------------- /automl/cvt_hparam_file.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Convert AutoML-generated hyper-parameter file to PocketFlow-compatible format.""" 18 | 19 | import sys 20 | 21 | # file paths 22 | assert len(sys.argv) == 2, '[HELP] python cvt_hparam_file.py ' 23 | file_path = sys.argv[1] 24 | 25 | # read hyper-parameters' values from file 26 | with open(file_path, 'r') as i_file: 27 | # obtain raw hyper-parameters' values 28 | for i_line in i_file: 29 | sub_strs = i_line.split() 30 | name, val = sub_strs[0], float(sub_strs[2]) 31 | if name == 'ws_prune_ratio_exp': 32 | ws_prune_ratio_exp = val 33 | elif name == 'ws_iter_ratio_beg': 34 | ws_iter_ratio_beg = val 35 | elif name == 'ws_iter_ratio_end': 36 | ws_iter_ratio_end = val 37 | elif name == 'ws_update_mask_step': 38 | ws_update_mask_step = val 39 | 40 | # make sure is smaller than 41 | ws_iter_ratio_end = ws_iter_ratio_beg + ws_iter_ratio_end * (1.0 - ws_iter_ratio_beg) 42 | 43 | # write hyper-parameters' values to file 44 | output_str = '' 45 | output_str += ' --ws_prune_ratio_exp %.4f' % ws_prune_ratio_exp 46 | output_str += ' --ws_iter_ratio_beg %.4f' % ws_iter_ratio_beg 47 | output_str += ' --ws_iter_ratio_end %.4f' % ws_iter_ratio_end 48 | output_str += ' --ws_update_mask_step %d' % int(ws_update_mask_step) 49 | print(output_str) 50 | -------------------------------------------------------------------------------- /automl/parse_results.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Convert PocketFlow-generated result file to AutoML-compatible format.""" 18 | 19 | import sys 20 | 21 | assert len(sys.argv) == 2, '[HELP] python parse_results.py ' 22 | file_path = sys.argv[1] 23 | 24 | with open(file_path, 'r') as i_file: 25 | for i_line in i_file: 26 | if 'INFO:tensorflow:accuracy:' in i_line: 27 | accuracy = float(i_line.split()[-1]) 28 | elif 'INFO:tensorflow:pruning ratio:' in i_line: 29 | prune_ratio = float(i_line.split()[-1]) 30 | elif 'INFO:tensorflow:loss:' in i_line: 31 | loss = float(i_line.split()[-1]) 32 | 33 | print('object_value=%f' % accuracy) 34 | print('prune_ratio=%f' % prune_ratio) 35 | print('loss=%f' % loss) 36 | -------------------------------------------------------------------------------- /automl/seven.yaml: -------------------------------------------------------------------------------- 1 | version: 1 2 | kind: standalone 3 | jobname: pocket-flow 4 | container: 5 | image: 6 | docker.oa.com/g_tfplus/tfplus:tensorflow1.8-python3.6-cuda9.0-cudnn7.0.4.31-ubuntu16.04-tfplus-v2 7 | resources: 8 | nvidia.com/gpu: 1 9 | env: 10 | - name: SEVEN_HTTP_FORWARD_PORT 11 | - name: NB_GPUS 12 | value: 1 13 | -------------------------------------------------------------------------------- /datasets/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/datasets/__init__.py -------------------------------------------------------------------------------- /datasets/abstract_dataset.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Abstract class for datasets.""" 18 | 19 | from abc import ABC 20 | import tensorflow as tf 21 | 22 | from utils.multi_gpu_wrapper import MultiGpuWrapper as mgw 23 | 24 | FLAGS = tf.app.flags.FLAGS 25 | 26 | tf.app.flags.DEFINE_string('data_disk', 'local', 'data disk\'s location (\'local\' | \'hdfs\')') 27 | tf.app.flags.DEFINE_string('data_hdfs_host', None, 'HDFS host for data files') 28 | tf.app.flags.DEFINE_string('data_dir_local', None, 'data directory - local') 29 | tf.app.flags.DEFINE_string('data_dir_hdfs', None, 'data directory - HDFS') 30 | tf.app.flags.DEFINE_integer('cycle_length', 4, '# of datasets to interleave from in parallel') 31 | tf.app.flags.DEFINE_integer('nb_threads', 8, '# of threads for preprocessing the dataset') 32 | tf.app.flags.DEFINE_integer('buffer_size', 1024, '# of elements to be buffered when prefetching') 33 | tf.app.flags.DEFINE_integer('prefetch_size', 8, '# of mini-batches to be buffered when prefetching') 34 | 35 | class AbstractDataset(ABC): 36 | '''Abstract class for datasets.''' 37 | 38 | def __init__(self, is_train): 39 | """Constructor function. 40 | 41 | Args: 42 | * is_train: whether to construct the training subset 43 | """ 44 | 45 | # following attributes must be initialized by each sub-class 46 | self.file_pattern = None 47 | self.dataset_fn = None 48 | self.parse_fn = None 49 | self.batch_size = None 50 | 51 | # determine whether data sharding is enabled 52 | self.is_train = is_train 53 | self.enbl_shard = (is_train and FLAGS.enbl_multi_gpu) # shard files for multi-GPU training 54 | 55 | def build(self, enbl_trn_val_split=False): 56 | '''Build iterator(s) for tf.data.Dataset() object. 57 | 58 | Args: 59 | * enbl_trn_val_split: whether to split into training & validation subsets 60 | 61 | Returns: 62 | * iterator_trn: iterator for the training subset 63 | * iterator_val: iterator for the validation subset 64 | OR 65 | * iterator: iterator for the chosen subset (training OR testing) 66 | 67 | Example: 68 | # build iterator(s) 69 | dataset = xxxxDataset(is_train=True) # TF operations are not created 70 | iterator = dataset.build() # TF operations are created 71 | OR 72 | iterator_trn, iterator_val = dataset.build(enbl_trn_val_split=True) # for dataset-train only 73 | 74 | # use the iterator to obtain a mini-batch of images & labels 75 | images, labels = iterator.get_next() 76 | ''' 77 | 78 | # obtain list of data files' names 79 | filenames = tf.data.Dataset.list_files(self.file_pattern, shuffle=True) 80 | if self.enbl_shard: 81 | filenames = filenames.shard(mgw.size(), mgw.rank()) 82 | 83 | # create a tf.data.Dataset from list of files 84 | dataset = filenames.apply( 85 | tf.contrib.data.parallel_interleave(self.dataset_fn, cycle_length=FLAGS.cycle_length)) 86 | dataset = dataset.map(self.parse_fn, num_parallel_calls=FLAGS.nb_threads) 87 | 88 | # create iterators for training & validation subsets separately 89 | if self.is_train and enbl_trn_val_split: 90 | iterator_val = self.__make_iterator(dataset.take(FLAGS.nb_smpls_val)) 91 | iterator_trn = self.__make_iterator(dataset.skip(FLAGS.nb_smpls_val)) 92 | return iterator_trn, iterator_val 93 | 94 | return self.__make_iterator(dataset) 95 | 96 | def __make_iterator(self, dataset): 97 | """Make an iterator from tf.data.Dataset. 98 | 99 | Args: 100 | * dataset: tf.data.Dataset object 101 | 102 | Returns: 103 | * iterator: iterator for the dataset 104 | """ 105 | 106 | dataset = dataset.apply(tf.contrib.data.shuffle_and_repeat(buffer_size=FLAGS.buffer_size)) 107 | dataset = dataset.batch(self.batch_size) 108 | dataset = dataset.prefetch(FLAGS.prefetch_size) 109 | iterator = dataset.make_one_shot_iterator() 110 | 111 | return iterator 112 | -------------------------------------------------------------------------------- /datasets/cifar10_dataset.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """CIFAR-10 dataset.""" 18 | 19 | import os 20 | import tensorflow as tf 21 | 22 | from datasets.abstract_dataset import AbstractDataset 23 | 24 | FLAGS = tf.app.flags.FLAGS 25 | 26 | tf.app.flags.DEFINE_integer('nb_classes', 10, '# of classes') 27 | tf.app.flags.DEFINE_integer('nb_smpls_train', 50000, '# of samples for training') 28 | tf.app.flags.DEFINE_integer('nb_smpls_val', 5000, '# of samples for validation') 29 | tf.app.flags.DEFINE_integer('nb_smpls_eval', 10000, '# of samples for evaluation') 30 | tf.app.flags.DEFINE_integer('batch_size', 128, 'batch size per GPU for training') 31 | tf.app.flags.DEFINE_integer('batch_size_eval', 100, 'batch size for evaluation') 32 | 33 | # CIFAR-10 specifications 34 | LABEL_BYTES = 1 35 | IMAGE_HEI = 32 36 | IMAGE_WID = 32 37 | IMAGE_CHN = 3 38 | IMAGE_BYTES = IMAGE_CHN * IMAGE_HEI * IMAGE_WID 39 | RECORD_BYTES = LABEL_BYTES + IMAGE_BYTES 40 | IMAGE_AVE = tf.constant([[[125.3, 123.0, 113.9]]], dtype=tf.float32) 41 | IMAGE_STD = tf.constant([[[63.0, 62.1, 66.7]]], dtype=tf.float32) 42 | 43 | def parse_fn(example_serialized, is_train): 44 | """Parse image & labels from the serialized data. 45 | 46 | Args: 47 | * example_serialized: serialized example data 48 | * is_train: whether data augmentation should be applied 49 | 50 | Returns: 51 | * image: image tensor 52 | * label: one-hot label tensor 53 | """ 54 | 55 | # data parsing 56 | record = tf.decode_raw(example_serialized, tf.uint8) 57 | label = tf.slice(record, [0], [LABEL_BYTES]) 58 | label = tf.one_hot(tf.reshape(label, []), FLAGS.nb_classes) 59 | image = tf.slice(record, [LABEL_BYTES], [IMAGE_BYTES]) 60 | image = tf.reshape(image, [IMAGE_CHN, IMAGE_HEI, IMAGE_WID]) 61 | image = tf.cast(tf.transpose(image, [1, 2, 0]), tf.float32) 62 | image = (image - IMAGE_AVE) / IMAGE_STD 63 | 64 | # data augmentation 65 | if is_train: 66 | image = tf.image.resize_image_with_crop_or_pad(image, IMAGE_HEI + 8, IMAGE_WID + 8) 67 | image = tf.random_crop(image, [IMAGE_HEI, IMAGE_WID, IMAGE_CHN]) 68 | image = tf.image.random_flip_left_right(image) 69 | 70 | return image, label 71 | 72 | class Cifar10Dataset(AbstractDataset): 73 | '''CIFAR-10 dataset.''' 74 | 75 | def __init__(self, is_train): 76 | """Constructor function. 77 | 78 | Args: 79 | * is_train: whether to construct the training subset 80 | """ 81 | 82 | # initialize the base class 83 | super(Cifar10Dataset, self).__init__(is_train) 84 | 85 | # choose local files or HDFS files w.r.t. FLAGS.data_disk 86 | if FLAGS.data_disk == 'local': 87 | assert FLAGS.data_dir_local is not None, ' must not be None' 88 | data_dir = FLAGS.data_dir_local 89 | elif FLAGS.data_disk == 'hdfs': 90 | assert FLAGS.data_hdfs_host is not None and FLAGS.data_dir_hdfs is not None, \ 91 | 'both and must not be None' 92 | data_dir = FLAGS.data_hdfs_host + FLAGS.data_dir_hdfs 93 | else: 94 | raise ValueError('unrecognized data disk: ' + FLAGS.data_disk) 95 | 96 | # configure file patterns & function handlers 97 | if is_train: 98 | self.file_pattern = os.path.join(data_dir, 'data_batch_*.bin') 99 | self.batch_size = FLAGS.batch_size 100 | else: 101 | self.file_pattern = os.path.join(data_dir, 'test_batch.bin') 102 | self.batch_size = FLAGS.batch_size_eval 103 | self.dataset_fn = lambda x: tf.data.FixedLengthRecordDataset(x, RECORD_BYTES) 104 | self.parse_fn = lambda x: parse_fn(x, is_train=is_train) 105 | -------------------------------------------------------------------------------- /datasets/ilsvrc12_dataset.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """ILSVRC-12 dataset.""" 18 | 19 | import os 20 | import tensorflow as tf 21 | 22 | from datasets.abstract_dataset import AbstractDataset 23 | from utils.external.imagenet_preprocessing import preprocess_image 24 | 25 | FLAGS = tf.app.flags.FLAGS 26 | 27 | tf.app.flags.DEFINE_integer('nb_classes', 1001, '# of classes') 28 | tf.app.flags.DEFINE_integer('nb_smpls_train', 1281167, '# of samples for training') 29 | tf.app.flags.DEFINE_integer('nb_smpls_val', 10000, '# of samples for validation') 30 | tf.app.flags.DEFINE_integer('nb_smpls_eval', 50000, '# of samples for evaluation') 31 | tf.app.flags.DEFINE_integer('batch_size', 64, 'batch size per GPU for training') 32 | tf.app.flags.DEFINE_integer('batch_size_eval', 100, 'batch size for evaluation') 33 | 34 | # ILSVRC-12 specifications 35 | IMAGE_HEI = 224 36 | IMAGE_WID = 224 37 | IMAGE_CHN = 3 38 | 39 | def parse_example_proto(example_serialized): 40 | """Parse image buffer, label, and bounding box from the serialized data. 41 | 42 | Args: 43 | * example_serialized: serialized example data 44 | 45 | Returns: 46 | * image_buffer: image buffer label 47 | * label: label tensor (not one-hot) 48 | * bbox: bounding box tensor 49 | """ 50 | 51 | # parse features from the serialized data 52 | feature_map = { 53 | 'image/encoded': tf.FixedLenFeature([], dtype=tf.string, default_value=''), 54 | 'image/class/label': tf.FixedLenFeature([1], dtype=tf.int64, default_value=-1), 55 | 'image/class/text': tf.FixedLenFeature([], dtype=tf.string, default_value=''), 56 | } 57 | bbox_keys = ['image/object/bbox/' + x for x in ['xmin', 'ymin', 'xmax', 'ymax']] 58 | feature_map.update({key: tf.VarLenFeature(dtype=tf.float32) for key in bbox_keys}) 59 | features = tf.parse_single_example(example_serialized, feature_map) 60 | 61 | # obtain the label and bounding boxes 62 | label = tf.cast(features['image/class/label'], dtype=tf.int32) 63 | xmin = tf.expand_dims(features['image/object/bbox/xmin'].values, 0) 64 | ymin = tf.expand_dims(features['image/object/bbox/ymin'].values, 0) 65 | xmax = tf.expand_dims(features['image/object/bbox/xmax'].values, 0) 66 | ymax = tf.expand_dims(features['image/object/bbox/ymax'].values, 0) 67 | 68 | # Note that we impose an ordering of (y, x) just to make life difficult. 69 | bbox = tf.concat(axis=0, values=[ymin, xmin, ymax, xmax]) 70 | bbox = tf.expand_dims(bbox, 0) 71 | bbox = tf.transpose(bbox, [0, 2, 1]) 72 | 73 | return features['image/encoded'], label, bbox 74 | 75 | def parse_fn(example_serialized, is_train): 76 | """Parse image & labels from the serialized data. 77 | 78 | Args: 79 | * example_serialized: serialized example data 80 | * is_train: whether data augmentation should be applied 81 | 82 | Returns: 83 | * image: image tensor 84 | * label: one-hot label tensor 85 | """ 86 | 87 | image_buffer, label, bbox = parse_example_proto(example_serialized) 88 | image = preprocess_image( 89 | image_buffer=image_buffer, bbox=bbox, output_height=IMAGE_HEI, 90 | output_width=IMAGE_WID, num_channels=IMAGE_CHN, is_training=is_train) 91 | label = tf.one_hot(tf.reshape(label, []), FLAGS.nb_classes) 92 | 93 | return image, label 94 | 95 | class Ilsvrc12Dataset(AbstractDataset): 96 | '''ILSVRC-12 dataset.''' 97 | 98 | def __init__(self, is_train): 99 | """Constructor function. 100 | 101 | Args: 102 | * is_train: whether to construct the training subset 103 | """ 104 | 105 | # initialize the base class 106 | super(Ilsvrc12Dataset, self).__init__(is_train) 107 | 108 | # choose local files or HDFS files w.r.t. FLAGS.data_disk 109 | if FLAGS.data_disk == 'local': 110 | assert FLAGS.data_dir_local is not None, ' must not be None' 111 | data_dir = FLAGS.data_dir_local 112 | elif FLAGS.data_disk == 'hdfs': 113 | assert FLAGS.data_hdfs_host is not None and FLAGS.data_dir_hdfs is not None, \ 114 | 'both and must not be None' 115 | data_dir = FLAGS.data_hdfs_host + FLAGS.data_dir_hdfs 116 | else: 117 | raise ValueError('unrecognized data disk: ' + FLAGS.data_disk) 118 | 119 | # configure file patterns & function handlers 120 | if is_train: 121 | self.file_pattern = os.path.join(data_dir, 'train-*-of-*') 122 | self.batch_size = FLAGS.batch_size 123 | else: 124 | self.file_pattern = os.path.join(data_dir, 'validation-*-of-*') 125 | self.batch_size = FLAGS.batch_size_eval 126 | self.dataset_fn = tf.data.TFRecordDataset 127 | self.parse_fn = lambda x: parse_fn(x, is_train=is_train) 128 | -------------------------------------------------------------------------------- /docs/README.md: -------------------------------------------------------------------------------- 1 | # How to Build the Documentation Site 2 | 3 | 1. Install mkdocs and other dependencies: 4 | 5 | ``` bash 6 | $ pip install -r doc-requirements.txt 7 | ``` 8 | 9 | 2. Build a directory named "site", which contains all files needed for the documentation website: 10 | 11 | ``` bash 12 | $ mkdocs build --clean 13 | ``` 14 | -------------------------------------------------------------------------------- /docs/doc-requirements.txt: -------------------------------------------------------------------------------- 1 | mkdocs 2 | markdown 3 | python-markdown-math 4 | pymdown-extensions 5 | -------------------------------------------------------------------------------- /docs/docs/.markdownlint.json: -------------------------------------------------------------------------------- 1 | { 2 | "MD013": false, 3 | "MD014": false, 4 | "MD024": {"allow_different_nesting": true} 5 | } 6 | -------------------------------------------------------------------------------- /docs/docs/MathJax.js: -------------------------------------------------------------------------------- 1 | (function () { 2 | var newMathJax = 'https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.1/MathJax.js'; 3 | var oldMathJax = 'cdn.mathjax.org/mathjax/latest/MathJax.js'; 4 | 5 | var replaceScript = function (script, src) { 6 | // 7 | // Make redirected script 8 | // 9 | var newScript = document.createElement('script'); 10 | newScript.src = newMathJax + src.replace(/.*?(\?|$)/, '$1'); 11 | // 12 | // Move onload and onerror handlers to new script 13 | // 14 | newScript.onload = script.onload; 15 | newScript.onerror = script.onerror; 16 | script.onload = script.onerror = null; 17 | // 18 | // Move any content (old-style configuration scripts) 19 | // 20 | while (script.firstChild) newScript.appendChild(script.firstChild); 21 | // 22 | // Copy script id 23 | // 24 | if (script.id != null) newScript.id = script.id; 25 | // 26 | // Replace original script with new one 27 | // 28 | script.parentNode.replaceChild(newScript, script); 29 | // 30 | // Issue a console warning 31 | // 32 | console.warn('WARNING: cdn.mathjax.org has been retired. Check https://www.mathjax.org/cdn-shutting-down/ for migration tips.') 33 | } 34 | 35 | if (document.currentScript) { 36 | var script = document.currentScript; 37 | replaceScript(script, script.src); 38 | } else { 39 | // 40 | // Look for current script by searching for one with the right source 41 | // 42 | var n = oldMathJax.length; 43 | var scripts = document.getElementsByTagName('script'); 44 | for (var i = 0; i < scripts.length; i++) { 45 | var script = scripts[i]; 46 | var src = (script.src || '').replace(/.*?:\/\//,''); 47 | if (src.substr(0, n) === oldMathJax) { 48 | replaceScript(script, src); 49 | break; 50 | } 51 | } 52 | } 53 | })(); -------------------------------------------------------------------------------- /docs/docs/automl_based_methods.md: -------------------------------------------------------------------------------- 1 | # AutoML-based Methods 2 | 3 | Under construction ... 4 | -------------------------------------------------------------------------------- /docs/docs/distillation.md: -------------------------------------------------------------------------------- 1 | # Distillation 2 | 3 | Distillation (Hinton et al., 2015) is a kind of model compression approaches in which a pre-trained large model teaches a smaller model to achieve the similar prediction performance. 4 | It is often named as the "teacher-student" training, where the large model is the teacher and the smaller model is the student. 5 | 6 | With distillation, knowledge can be transferred from the teacher model to the student by minimizing a loss function to recover the distribution of class probabilities predicted by the teacher model. 7 | In most situations, the probability of the correct class predicted by the teacher model is very high, and probabilities of other classes are close to 0, which may not be able to provide extra information beyond ground-truth labels. 8 | To overcome this issue, a commonly-used solution is to raise the temperature of the final softmax function until the cumbersome model produces a suitably soft set of targets. The soften probability $q_i$ of class $i$ is calculated from the logit $z_i$: 9 | 10 | $$ 11 | q_i = \frac{\exp \left( z_i / T \right)}{\sum_j{\exp \left( z_j / T \right)}} 12 | $$ 13 | 14 | where $T$ is the temperature. 15 | As $T$ grows, the probability distribution is more smooth, providing more information as to which classes the cumbersome model more similar to the predicted class. 16 | It is better to include the standard loss ($T = 1$) between the predicted class probabilities and ground-truth labels. 17 | The overall loss function is given by: 18 | 19 | $$ 20 | L \left( x; W \right) = H \left( y, \sigma \left( z_s; T = 1 \right) \right) + \alpha \cdot H \left( \sigma \left( z_t; T = \tau \right), \sigma \left( z_s, T = \tau \right) \right) 21 | $$ 22 | 23 | where $x$ is the input, $W$ are parameters of the distilled small model and $y$ is ground-truth labels, $\sigma$ is the softmax parameterized by temperature $T$, $H$ is the cross-entropy loss, and $\alpha$ is the coefficient of distillation loss. 24 | The coefficient $\alpha$ can be set by `--loss_w_dst` and the temperature $T$ can be set by `--tempr_dst`. 25 | 26 | ## Combination with Other Model Compression Approaches 27 | 28 | Other model model compression techniques, such as channel pruning, weight pruning, and quantization, can be augmented with distillation. To enable the distillation loss, simply append the `--enbl_dst` argument when starting the program. 29 | -------------------------------------------------------------------------------- /docs/docs/faq.md: -------------------------------------------------------------------------------- 1 | # Frequently Asked Questions 2 | 3 | **Q: Under construction ...** 4 | 5 | A: Under construction ... -------------------------------------------------------------------------------- /docs/docs/installation.md: -------------------------------------------------------------------------------- 1 | # Installation 2 | 3 | PocketFlow is developed and tested on Linux, using Python 3.6 and TensorFlow 1.10.0. We support the following three execution modes for PocketFlow: 4 | 5 | * Local mode: run PocketFlow on the local machine. 6 | * Docker mode: run PocketFlow within a docker image. 7 | * Seven mode: run PocketFlow on the seven cluster (only available within Tencent). 8 | 9 | ## Clone PocketFlow 10 | 11 | To make a local copy of the PocketFlow repository, use: 12 | 13 | ``` bash 14 | $ git clone https://github.com/Tencent/PocketFlow.git 15 | ``` 16 | 17 | ## Create a Path Configuration File 18 | 19 | PocketFlow requires a path configuration file, named `path.conf`, to setup directory paths to data sets and pre-trained models under different execution modes, as well as HDFS / HTTP connection parameters. 20 | 21 | We have provided a template file to help you create your own path configuration file. You can find it in the PocketFlow repository, named `path.conf.template`, which contains more detailed descriptions on how to customize path configurations. For instance, if you want to use CIFAR-10 and ImageNet data sets stored on the local machine, then the path configuration file should look like this: 22 | 23 | ``` bash 24 | # data files 25 | data_hdfs_host = None 26 | data_dir_local_cifar10 = /home/user_name/datasets/cifar-10-batches-bin # this line has been edited! 27 | data_dir_hdfs_cifar10 = None 28 | data_dir_seven_cifar10 = None 29 | data_dir_docker_cifar10 = /opt/ml/data # DO NOT EDIT 30 | data_dir_local_ilsvrc12 = /home/user_name/datasets/imagenet_tfrecord # this line has been edited! 31 | data_dir_hdfs_ilsvrc12 = None 32 | data_dir_seven_ilsvrc12 = None 33 | data_dir_docker_ilsvrc12 = /opt/ml/data # DO NOT EDIT 34 | 35 | # model files 36 | model_http_url = https://api.ai.tencent.com/pocketflow 37 | ``` 38 | 39 | In short, you need to replace "None" in the template file with the actual path (or HDFS / HTTP connection parameters) if available, or leave it unchanged otherwise. 40 | 41 | ## Prepare for the Local Mode 42 | 43 | We recommend to use Anaconda as the Python environment, which has many essential packages built-in. The Anaconda installer can be downloaded from [here](https://www.anaconda.com/download/#linux). To install, use the following command: 44 | 45 | ``` bash 46 | # install Anaconda; replace the installer's file name if needed 47 | $ bash Anaconda3-5.2.0-Linux-x86_64.sh 48 | 49 | # activate Anaconda's Python path 50 | $ source ~/.bashrc 51 | ``` 52 | 53 | For Anaconda 5.3.0 or later, the default Python version is 3.7, which does not support installing TensorFlow with pip directly. Therefore, you need to manually switch to Python 3.6 once Anaconda is installed: 54 | 55 | ``` bash 56 | # install Python 3.6 57 | $ conda install python=3.6 58 | ``` 59 | 60 | To install TensorFlow, you may refer to TensorFlow's official [documentation](https://www.tensorflow.org/install/pip) for detailed instructions. Specially, if GPU-based training is required, then you need to follow the [GPU support guide](https://www.tensorflow.org/install/gpu) to set up a CUDA-enabled GPU card in prior to installation. After that, install TensorFlow with: 61 | 62 | ``` bash 63 | # TensorFlow with GPU support; use if GPU is not available 64 | $ pip install tensorflow-gpu 65 | 66 | # verify the install 67 | $ python -c "import tensorflow as tf; print(tf.__version__)" 68 | ``` 69 | 70 | To run PocketFlow in the local mode, *e.g.* to train a full-precision ResNet-20 model for the CIFAR-10 classification task, use the following command: 71 | 72 | ``` bash 73 | $ ./scripts/run_local.sh nets/resnet_at_cifar10_run.py 74 | ``` 75 | 76 | ## Prepare for the Docker Mode 77 | 78 | Docker offers an alternative way to run PocketFlow within an isolated container, so that your local Python environment remains untouched. We recommend you to use the [horovod](https://github.com/uber/horovod) docker provided by Uber, which enables multi-GPU distributed training for TensorFlow with only a few lines modification. Once docker is installed, the docker image can be obtained via: 79 | 80 | ``` bash 81 | # obtain the docker image 82 | $ docker pull uber/horovod 83 | ``` 84 | 85 | To run PocketFlow in the docker mode, *e.g.* to train a full-precision ResNet-20 model for the CIFAR-10 classification task, use the following command: 86 | 87 | ``` bash 88 | $ ./scripts/run_docker.sh nets/resnet_at_cifar10_run.py 89 | ``` 90 | 91 | ## Prepare for the Seven Mode 92 | 93 | Seven is a distributed learning platform built for both CPU and GPU clusters. Users can submit tasks to the seven cluster, using built-in data sets and docker images seamlessly. 94 | 95 | To run PocketFlow in the seven mode, *e.g.* to train a full-precision ResNet-20 model for the CIFAR-10 classification task, use the following command: 96 | 97 | ``` bash 98 | $ ./scripts/run_seven.sh nets/resnet_at_cifar10_run.py 99 | ``` 100 | -------------------------------------------------------------------------------- /docs/docs/multi_gpu_training.md: -------------------------------------------------------------------------------- 1 | # Multi-GPU Training 2 | 3 | Due to the high computational complexity, it often takes hours or even days to fully train deep learning models using a single GPU. 4 | In PocketFlow, we adopt multi-GPU training to speed-up this time-consuming training process. 5 | Our implementation is compatible with: 6 | 7 | * [Horovod](https://github.com/uber/horovod): a distributed training framework for TensorFlow, Keras, and PyTorch. 8 | * TF-Plus: an optimized framework for TensorFlow-based distributed training (only available within Tencent). 9 | 10 | We have provide a wrapper class, `MultiGpuWrapper`, to seamlessly switch between the above two frameworks. 11 | It will sequentially check whether Horovod and TF-Plus can be used, and use the first available one as the underlying framework for multi-GPU training. 12 | 13 | The main reason that using Horovod or TF-Plus instead TensorFlow's original distributed training routine is that these frameworks provide many easy-to-use APIs and require far less code changes to change from single-GPU to multi-GPU training, as we shall see later. 14 | 15 | ## From Single-GPU to Multi-GPU 16 | 17 | To extend a single-GPU based training script to the multi-GPU scenario, at most 7 steps are needed: 18 | 19 | * Import the Horovod or TF-Plus module. 20 | 21 | ``` Python 22 | from utils.multi_gpu_wrapper import MultiGpuWrapper as mgw 23 | ``` 24 | 25 | * Initialize the multi-GPU training framework, as early as possible. 26 | 27 | ``` Python 28 | mgw.init() 29 | ``` 30 | 31 | * For each worker, create a session with a distinct GPU device. 32 | 33 | ``` Python 34 | config = tf.ConfigProto() 35 | config.gpu_options.visible_device_list = str(mgw.local_rank()) 36 | sess = tf.Session(config=config) 37 | ``` 38 | 39 | * (Optional) Let each worker use a distinct subset of training data. 40 | 41 | ``` Python 42 | filenames = tf.data.Dataset.list_files(file_pattern, shuffle=True) 43 | filenames = filenames.shard(mgw.size(), mgw.rank()) 44 | ``` 45 | 46 | * Wrapper the optimizer for distributed gradient communication. 47 | 48 | ``` Python 49 | optimizer = tf.train.AdamOptimizer(learning_rate=lrn_rate) 50 | optimizer = mgw.DistributedOptimizer(optimizer) 51 | train_op = optimizer.minimize(loss) 52 | ``` 53 | 54 | * Synchronize master's parameters to all the other workers. 55 | 56 | ``` Python 57 | bcast_op = mgw.broadcast_global_variables(0) 58 | sess.run(tf.global_variables_initializer()) 59 | sess.run(bcast_op) 60 | ``` 61 | 62 | * (Optional) Save checkpoint files at the master node periodically. 63 | 64 | ``` Python 65 | if mgw.rank() == 0: 66 | saver.save(sess, save_path, global_step) 67 | ``` 68 | 69 | ## Usage Example 70 | 71 | Here, we provide a code snippet to demonstrate how to use multi-GPU training to speed-up training. 72 | Please note that many implementation details are omitted for clarity. 73 | 74 | ``` Python 75 | import tensorflow as tf 76 | from utils.multi_gpu_wrapper import MultiGpuWrapper as mgw 77 | 78 | # initialization 79 | mgw.init() 80 | 81 | # create the training graph 82 | with tf.Graph().as_default(): 83 | # create a TensorFlow session 84 | config = tf.ConfigProto() 85 | config.gpu_options.visible_device_list = str(mgw.local_rank()) 86 | sess = tf.Session(config=config) 87 | 88 | # use tf.data.Dataset() to traverse images and labels 89 | filenames = tf.data.Dataset.list_files(file_pattern, shuffle=True) 90 | filenames = filenames.shard(mgw.size(), mgw.rank()) 91 | images, labels = get_images_n_labels(filenames) 92 | 93 | # define the network and its loss function 94 | logits = forward_pass(images) 95 | loss = calc_loss(labels, logits) 96 | 97 | # create an optimizer and setup training-related operations 98 | global_step = tf.train.get_or_create_global_step() 99 | optimizer = tf.train.AdamOptimizer(learning_rate=lrn_rate) 100 | optimizer = mgw.DistributedOptimizer(optimizer) 101 | train_op = optimizer.minimize(loss, global_step=global_step) 102 | bcast_op = mgw.broadcast_global_variables(0) 103 | 104 | # multi-GPU training 105 | sess.run(tf.global_variables_initializer()) 106 | sess.run(bcast_op) 107 | for idx_iter in range(nb_iters): 108 | sess.run(train_op) 109 | if mgw.rank() == 0 and (idx_iter + 1) % save_step == 0: 110 | saver.save(sess, save_path, global_step) 111 | ``` 112 | -------------------------------------------------------------------------------- /docs/docs/performance.md: -------------------------------------------------------------------------------- 1 | # Performance 2 | 3 | In this documentation, we present evaluation results for applying various model compression methods for ResNet and MobileNet models on the ImageNet classification task, including channel pruning, weight sparsification, and uniform quantization. 4 | 5 | We adopt `ChannelPrunedLearner` to shrink the number of channels for convolutional layers to reduce the computation complexity. 6 | Instead of using the same pruning ratio for all layers, we utilize the DDPG algorithm as the RL agent to iteratively search for the optimal pruning ratio of each layer. 7 | After obtaining the optimal pruning ratios, group fine-tuning is adopted to further improve the compressed model's accuracy, as demonstrated below: 8 | 9 | | Model | Pruning Ratio | Uniform | RL-based | RL-based + Group Fine-tuning | 10 | |:------------:|:-------------:|:-------:|:-------------:|:----------------------------:| 11 | | MobileNet-v1 | 50% | 66.5% | 67.8% (+1.3%) | 67.9% (+1.4%) | 12 | | MobileNet-v1 | 60% | 66.2% | 66.9% (+0.7%) | 67.0% (+0.8%) | 13 | | MobileNet-v1 | 70% | 64.4% | 64.5% (+0.1%) | 64.8% (+0.4%) | 14 | | Mobilenet-v1 | 80% | 61.4% | 61.4% (+0.0%) | 62.2% (+0.8%) | 15 | 16 | **Note:** The original uncompressed MobileNet-v1's top-1 accuracy is 70.89%. 17 | 18 | We adopt `WeightSparseLearner` to introduce the sparsity constraint so that a large portion of model weights can be removed, which leads to smaller model and lower FLOPs for inference. 19 | Comparing with the original algorithm proposed in (Zhu & Gupta, 2017), we also incorporate network distillation and reinforcement learning algorithms to further improve the compressed model's accuracy, as shown in the table below: 20 | 21 | | Model | Sparsity | (Zhu & Gupta, 2017) | RL-based | 22 | |:------------:|:--------:|:-------------------:|:-------------:| 23 | | MobileNet-v1 | 50% | 69.5% | 70.5% (+1.0%) | 24 | | MobileNet-v1 | 75% | 67.7% | 68.5% (+0.8%) | 25 | | MobileNet-v1 | 90% | 61.8% | 63.4% (+1.6%) | 26 | | MobileNet-v1 | 95% | 53.6% | 56.8% (+3.2%) | 27 | 28 | **Note:** The original uncompressed MobileNet-v1's top-1 accuracy is 70.89%. 29 | 30 | We adopt `UniformQuantTFLearner` to uniformly quantize model weights from 32-bit floating-point numbers to 8-bit fixed-point numbers. 31 | The resulting model can be converted into the TensorFlow Lite format for deployment on mobile devices. 32 | In the following two tables, we show that 8-bit quantized models can be as accurate as (or even better than) the original 32-bit ones, and the inference time can be significantly reduced after quantization. 33 | 34 | | Model | Top-1 Acc. (32-bit) | Top-5 Acc. (32-bit) | Top-1 Acc. (8-bit) | Top-5 Acc. (8-bit) | 35 | |:------------:|:-------------------:|:-------------------:|:------------------:|:------------------:| 36 | | ResNet-18 | 70.28% | 89.38% | 70.31% (+0.03%) | 89.40% (+0.02%) | 37 | | ResNet-50 | 75.97% | 92.88% | 76.01% (+0.04%) | 92.87% (-0.01%) | 38 | | MobileNet-v1 | 70.89% | 89.56% | 71.29% (+0.40%) | 89.79% (+0.23%) | 39 | | MobileNet-v2 | 71.84% | 90.60% | 72.26% (+0.42%) | 90.77% (+0.17%) | 40 | 41 | | Model | Hardware | CPU | Time (32-bit) | Time (8-bit) | Speed-up | 42 | |:------------:|:-----------:|:--------------:|:-------------:|:------------:|:------------:| 43 | | MobileNet-v1 | XiaoMi 8 SE | Snapdragon 710 | 156.33 | 62.60 | 2.50$\times$ | 44 | | MobileNet-v1 | XiaoMI 8 | Snapdragon 845 | 124.53 | 56.12 | 2.22$\times$ | 45 | | MobileNet-v1 | Huawei P20 | Kirin 970 | 152.54 | 68.43 | 2.23$\times$ | 46 | | MobileNet-v2 | XiaoMi 8 SE | Snapdragon 710 | 153.18 | 57.55 | 2.66$\times$ | 47 | | MobileNet-v2 | XiaoMi 8 | Snapdragon 845 | 120.59 | 49.04 | 2.46$\times$ | 48 | | MobileNet-v2 | Huawei P20 | Kirin 970 | 226.61 | 61.38 | 3.69$\times$ | 49 | 50 | * All the reported time are in milliseconds. 51 | -------------------------------------------------------------------------------- /docs/docs/pics/dcp_learner.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/docs/docs/pics/dcp_learner.png -------------------------------------------------------------------------------- /docs/docs/pics/deep_compression_algor.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/docs/docs/pics/deep_compression_algor.png -------------------------------------------------------------------------------- /docs/docs/pics/framework_design.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/docs/docs/pics/framework_design.png -------------------------------------------------------------------------------- /docs/docs/pics/rl_workflow.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/docs/docs/pics/rl_workflow.png -------------------------------------------------------------------------------- /docs/docs/pics/train_n_inference.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/docs/docs/pics/train_n_inference.png -------------------------------------------------------------------------------- /docs/docs/pics/wsl_pr_schedule.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/docs/docs/pics/wsl_pr_schedule.png -------------------------------------------------------------------------------- /docs/docs/pre_trained_models.md: -------------------------------------------------------------------------------- 1 | # Pre-trained Models 2 | 3 | We maintain a list of pre-trained uncompressed models, so that the training process of model compression does not need to start from scratch. 4 | 5 | For the CIFAR-10 data set, we provide following pre-trained models: 6 | 7 | | Model name | Accuracy | URL | 8 | |:----------:|:--------:|:---------------------------------------------------------------------------------:| 9 | | LeNet | 81.79% | [Link](https://api.ai.tencent.com/pocketflow/models_lenet_at_cifar_10.tar.gz) | 10 | | ResNet-20 | 91.93% | [Link](https://api.ai.tencent.com/pocketflow/models_resnet_20_at_cifar_10.tar.gz) | 11 | | ResNet-32 | 92.59% | [Link](https://api.ai.tencent.com/pocketflow/models_resnet_32_at_cifar_10.tar.gz) | 12 | | ResNet-44 | 92.76% | [Link](https://api.ai.tencent.com/pocketflow/models_resnet_44_at_cifar_10.tar.gz) | 13 | | ResNet-56 | 93.23% | [Link](https://api.ai.tencent.com/pocketflow/models_resnet_56_at_cifar_10.tar.gz) | 14 | 15 | For the ImageNet (ILSVRC-12) data set, we provide following pre-trained models: 16 | 17 | | Model name | Top-1 Acc. | Top-5 Acc. | URL | 18 | |:------------:|:----------:|:----------:|:-------------------------------------------------------------------------------------:| 19 | | ResNet-18 | 70.28% | 89.38% | [Link](https://api.ai.tencent.com/pocketflow/models_resnet_18_at_ilsvrc_12.tar.gz) | 20 | | ResNet-34 | 73.41% | 91.27% | [Link](https://api.ai.tencent.com/pocketflow/models_resnet_34_at_ilsvrc_12.tar.gz) | 21 | | ResNet-50 | 75.97% | 92.88% | [Link](https://api.ai.tencent.com/pocketflow/models_resnet_50_at_ilsvrc_12.tar.gz) | 22 | | MobileNet-v1 | 70.89% | 89.56% | [Link](https://api.ai.tencent.com/pocketflow/models_mobilenet_v1_at_ilsvrc_12.tar.gz) | 23 | | MobileNet-v2 | 71.84% | 90.60% | [Link](https://api.ai.tencent.com/pocketflow/models_mobilenet_v2_at_ilsvrc_12.tar.gz) | 24 | -------------------------------------------------------------------------------- /docs/docs/reference.md: -------------------------------------------------------------------------------- 1 | # Reference 2 | 3 | * [**Bengio et al., 2015**] Yoshua Bengio, Nicholas Leonard, and Aaron Courville. *Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation*. CoRR, abs/1308.3432, 2013. 4 | * [**Bergstra et al., 2013**] J. Bergstra, D. Yamins, and D. D. Cox. *Making a Science of Model Search: Hyperparameter Optimization in Hundreds of Dimensions for Vision Architectures*. In International Conference on Machine Learning (ICML), pages 115-123, Jun 2013. 5 | * [**Han et al., 2016**] Song Han, Huizi Mao, and William J. Dally. *Deep Compression: Compressing Deep Neural Network with Pruning, Trained Quantization and Huffman Coding*. In International Conference on Learning Representations (ICLR), 2016. 6 | * [**He et al., 2017**] Yihui He, Xiangyu Zhang, and Jian Sun. *Channel Pruning for Accelerating Very Deep Neural Networks*. In IEEE International Conference on Computer Vision (ICCV), pages 1389-1397, 2017. 7 | * [**He et al., 2018**] Yihui He, Ji Lin, Zhijian Liu, Hanrui Wang, Li-Jia Li, and Song Han. *AMC: AutoML for Model Compression and Acceleration on Mobile Devices*. In European Conference on Computer Vision (ECCV), pages 784-800, 2018. 8 | * [**Hinton et al., 2015**] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. *Distilling the Knowledge in a Neural Network*. CoRR, abs/1503.02531, 2015. 9 | * [**Jacob et al., 2018**] Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, and Dmitry Kalenichenko. *Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference*. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pages 2704-2713, 2018. 10 | * [**Lillicrap et al., 2016**] Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. *Continuous Control with Deep Reinforcement Learning*. In International Conference on Learning Representations (ICLR), 2016. 11 | * [**Mockus, 1975**] J. Mockus. *On Bayesian Methods for Seeking the Extremum*. In Optimization Techniques IFIP Technical Conference, pages 400-404, 1975. 12 | * [**Zhu & Gupta, 2017**] Michael Zhu and Suyog Gupta. *To Prune, or Not to Prune: Exploring the Efficacy of Pruning for Model Compression*. CoRR, abs/1710.01878, 2017. 13 | * [**Zhuang et al., 2018**] Zhuangwei Zhuang, Mingkui Tan, Bohan Zhuang, Jing Liu, Jiezhang Cao, Qingyao Wu, Junzhou Huang, and Jinhui Zhu. *Discrimination-aware Channel Pruning for Deep Neural Networks*. In Annual Conference on Neural Information Processing Systems (NIPS), 2018. 14 | -------------------------------------------------------------------------------- /docs/docs/reinforcement_learning.md: -------------------------------------------------------------------------------- 1 | # Reinforcement Learning 2 | 3 | For most deep learning models, the parameter redundancy differs from one layer to another. 4 | Some layers may be more robust to model compression algorithms due to larger redundancy, while others may be more sensitive. 5 | Therefore, it is often sub-optimal to use a unified pruning ratio or number of quantization bits for all layers, which completely omits the redundancy difference. 6 | However, it is also time-consuming or even impractical to manually setup the optimal value of such hyper-parameter for each layer, especially for deep networks with tens or hundreds of layers. 7 | 8 | To overcome this dilemma, in PocketFlow, we adopt reinforcement learning to automatically determine the optimal pruning ratio or number of quantization bits for each layer. 9 | Our approach is innovated from (He et al., 2018), which automatically determines each layer's optimal pruning ratio, and generalize it to hyper-parameter optimization for more model compression methods. 10 | 11 | In this documentation, we take `UniformQuantLearner` as an example to explain how the reinforcement learning method is used to iteratively optimize the number of quantization bits for each layer. 12 | It is worthy mentioning that this feature is also available for `ChannelPrunedLearner`, `WeightSparseLearner`, and `NonUniformQuantLearner`. 13 | 14 | ## Algorithm Description 15 | 16 | Here, we assume the original model to be compressed consists of $T$ layers, and denote the $t$-th layer's weight tensor as $\mathbf{W}_{t}$ and its quantization bit-width as $b_{t}$. 17 | In order to maximally exploit the parameter redundancy of each layer, we need to find the optimal combination of layer-wise quantization bit-width that achieves the highest accuracy after compression while satisfying: 18 | 19 | $$ 20 | \sum_{t = 1}^{T} b_{t} \left| \mathbf{W}_{t} \right| \le b \cdot \sum_{t = 1}^{T} \left| \mathbf{W}_{t} \right| 21 | $$ 22 | 23 | where $\left| \mathbf{W}_{t} \right|$ denotes the number of parameters in the weight tensor $\mathbf{W}_{t}$ and $b$ is the whole network's target quantization bit-width. 24 | 25 | Below, we present the overall workflow of adopting reinforcement learning, or more specifically, the DDPG algorithm (Lillicrap et al., 2016) to search for the optimal combination of layer-wise quantization bit-width: 26 | 27 | ![RL Workflow](pics/rl_workflow.png) 28 | 29 | To start with, we initialize an DDPG agent and set the best reward $r_{best}$ to negative infinity to track the optimal combination of layer-wise quantization bit-width. 30 | The search process consists of multiple roll-outs. 31 | In each roll-out, we sequentially traverse each layer in the network to determine its quantization bit-width. 32 | For the $t$-th layer, we construct its state vector with following information: 33 | 34 | * one-hot embedding of layer index 35 | * shape of weight tensor 36 | * number of parameters in the weight tensor 37 | * number of quantization bits used by previous layers 38 | * budget of quantization bits for remaining layers 39 | 40 | Afterwards, we feed this state vector into the DDPG agent to choose an action, which is then converted into the quantization bit-width under certain constraints. 41 | A commonly-used constraint is that with the selected quantization bit-width, the budget of quantization bits for remaining layers should be sufficient, *e.g.* ensuring the minimal quantization bit-width can be satisfied. 42 | 43 | After obtaining all layer's quantization bit-width, we quantize each layer's weights with the corresponding quantization bit-width, and fine-tune the quantized network for a few iteration (as supported by each learner's "Fast Fine-tuning" mode). 44 | We then evaluate the fine-tuned network' accuracy and use it as the reward signal $r_{n}$. 45 | The reward signal is compared against the best reward discovered so far, and the optimal combination of layer-wise quantization bit-width is updated if the current reward is larger. 46 | 47 | Finally, we generate a list of transitions from all the $\left( \mathbf{s}_{t}, a_{t}, r_{t}, \mathbf{s}_{t + 1} \right)$ tuples in the roll-out, and store them in the DDPG agent's replay buffer. 48 | The DDPG agent is then trained with one or more mini-batches of sampled transitions, so that it can choose better actions in the following roll-outs. 49 | 50 | After obtaining the optimal combination of layer-wise quantization bit-width, we can optionally use `UniformQuantLearner`'s "Re-training with Full Data" mode (also supported by others learners) for a complete quantization-aware training to further reduce the accuracy loss. 51 | -------------------------------------------------------------------------------- /docs/docs/test_cases.md: -------------------------------------------------------------------------------- 1 | # Test Cases 2 | 3 | This document contains various test cases to cover different combinations of learners and hyper-parameter settings. Any merge request to the master branch should be able to pass all the test cases to be approved. 4 | 5 | ## Full-Precision 6 | 7 | ``` bash 8 | # local mode 9 | $ ./scripts/run_local.sh nets/resnet_at_cifar10_run.py 10 | $ ./scripts/run_local.sh nets/resnet_at_cifar10_run.py \ 11 | --enbl_dst 12 | $ ./scripts/run_local.sh nets/resnet_at_cifar10_run.py \ 13 | --data_disk hdfs 14 | $ ./scripts/run_local.sh nets/resnet_at_cifar10_run.py \ 15 | --data_disk hdfs \ 16 | --enbl_dst 17 | 18 | # seven mode 19 | $ ./scripts/run_seven.sh nets/resnet_at_ilsvrc12_run.py 20 | $ ./scripts/run_seven.sh nets/resnet_at_ilsvrc12_run.py \ 21 | --enbl_dst 22 | $ ./scripts/run_seven.sh nets/resnet_at_ilsvrc12_run.py \ 23 | --data_disk hdfs 24 | $ ./scripts/run_seven.sh nets/resnet_at_ilsvrc12_run.py \ 25 | --data_disk hdfs \ 26 | --enbl_dst 27 | 28 | # docker mode 29 | $ ./scripts/run_docker.sh nets/lenet_at_cifar10_run.py 30 | $ ./scripts/run_docker.sh nets/lenet_at_cifar10_run.py \ 31 | --enbl_dst 32 | $ ./scripts/run_docker.sh nets/resnet_at_cifar10_run.py 33 | $ ./scripts/run_docker.sh nets/resnet_at_cifar10_run.py \ 34 | --enbl_dst 35 | ``` 36 | 37 | ## Channel Pruning 38 | 39 | ``` bash 40 | # uniform preserve ratios for all layers 41 | $ ./scripts/run_seven.sh nets/resnet_at_cifar10_run.py \ 42 | --learner channel \ 43 | --cp_prune_option uniform \ 44 | --cp_uniform_preserve_ratio 0.5 45 | 46 | # auto-tuned preserve ratios for each layer 47 | $ ./scripts/run_seven.sh nets/resnet_at_cifar10_run.py \ 48 | --cp_learner channel \ 49 | --cp_prune_option auto \ 50 | --cp_preserve_ratio 0.3 51 | ``` 52 | 53 | ## Discrimination-aware Channel Pruning 54 | 55 | ``` bash 56 | # no network distillation 57 | $ ./scripts/run_local.sh nets/resnet_at_cifar10_run.py \ 58 | --learner dis-chn-pruned \ 59 | --dcp_nb_stages 3 \ 60 | --data_disk hdfs 61 | 62 | # network distillation 63 | $ ./scripts/run_seven.sh nets/mobilenet_at_ilsvrc12_run.py \ 64 | --learner dis-chn-pruned \ 65 | --enbl_dst \ 66 | --dcp_nb_stages 4 67 | ``` 68 | 69 | ## Weight Sparsification 70 | 71 | ``` bash 72 | # uniform pruning ratios for all layers 73 | $ ./scripts/run_local.sh nets/resnet_at_cifar10_run.py \ 74 | --learner weight-sparse \ 75 | --ws_prune_ratio_prtl uniform \ 76 | --data_disk hdfs 77 | 78 | # optimal pruning ratios for each layer 79 | $ ./scripts/run_local.sh nets/resnet_at_cifar10_run.py \ 80 | --learner weight-sparse \ 81 | --ws_prune_ratio_prtl optimal \ 82 | --data_disk hdfs 83 | 84 | # heurist pruning ratios for each layer 85 | $ ./scripts/run_seven.sh nets/resnet_at_ilsvrc12_run.py \ 86 | --learner weight-sparse \ 87 | --ws_prune_ratio_prtl heurist 88 | 89 | # optimal pruning ratios for each layer 90 | $ ./scripts/run_seven.sh nets/resnet_at_ilsvrc12_run.py \ 91 | --learner weight-sparse \ 92 | --ws_prune_ratio_prtl optimal 93 | ``` 94 | 95 | ## Uniform Quantization 96 | 97 | ``` bash 98 | # channel-based bucketing 99 | $ ./scripts/run_local.sh nets/resnet_at_cifar10_run.py \ 100 | --learner uniform \ 101 | --uql_use_buckets \ 102 | --uql_bucket_type channel \ 103 | --data_disk hdfs 104 | 105 | # split-based bucketing 106 | $ ./scripts/run_local.sh nets/resnet_at_cifar10_run.py \ 107 | --learner uniform \ 108 | --uql_use_buckets \ 109 | --uql_bucket_type split \ 110 | --data_disk hdfs 111 | 112 | # channel-based bucketing + RL 113 | $ ./scripts/run_seven.sh nets/mobilenet_at_ilsvrc12_run.py -n=2 \ 114 | --learner uniform \ 115 | --uql_enbl_rl_agent \ 116 | --uql_use_buckets \ 117 | --uql_bucket_type channel 118 | 119 | # split-based bucketing + RL 120 | $ ./scripts/run_seven.sh nets/resnet_at_ilsvrc12_run.py -n=2 \ 121 | --learner uniform \ 122 | --uql_enbl_rl_agent \ 123 | --uql_use_buckets \ 124 | --uql_bucket_type split 125 | ``` 126 | 127 | ## Non-uniform Quantization 128 | 129 | ``` bash 130 | # channel-based bucketing + RL + optimize clusters 131 | $ ./scripts/run_local.sh nets/resnet_at_cifar10_run.py \ 132 | --learner non-uniform \ 133 | --nuql_enbl_rl_agent \ 134 | --nuql_use_buckets \ 135 | --nuql_bucket_type channel \ 136 | --nuql_opt_mode clusters \ 137 | --data_disk hdfs 138 | 139 | # split-based bucketing + RL + optimize weights 140 | $ ./scripts/run_local.sh nets/resnet_at_cifar10_run.py \ 141 | --learner non-uniform \ 142 | --nuql_enbl_rl_agent \ 143 | --nuql_use_buckets \ 144 | --nuql_bucket_type split \ 145 | --nuql_opt_mode weights \ 146 | --data_disk hdfs 147 | 148 | # channel-based bucketing + RL + optimize weights 149 | $ ./scripts/run_seven.sh nets/mobilenet_at_ilsvrc12_run.py -n=2 \ 150 | --learner non-uniform \ 151 | --nuql_enbl_rl_agent \ 152 | --nuql_use_buckets \ 153 | --nuql_bucket_type channel \ 154 | --nuql_opt_mode weights 155 | 156 | # split-based bucketing + RL + optimize clusters 157 | $ ./scripts/run_seven.sh nets/resnet_at_ilsvrc12_run.py -n=2 \ 158 | --learner non-uniform \ 159 | --nuql_enbl_rl_agent \ 160 | --nuql_use_buckets \ 161 | --nuql_bucket_type split \ 162 | --nuql_opt_mode clusters 163 | ``` 164 | -------------------------------------------------------------------------------- /docs/mkdocs.yml: -------------------------------------------------------------------------------- 1 | site_name: PocketFlow Docs 2 | nav: 3 | - Home: index.md 4 | - Installation: installation.md 5 | - Tutorial: tutorial.md 6 | - Learners - Algorithms: 7 | - Channel Pruning: cp_learner.md 8 | - Channel Pruning - Remastered: cpr_learner.md 9 | - Discrimination-aware Channel Pruning: dcp_learner.md 10 | - Weight Sparsification: ws_learner.md 11 | - Uniform Quantization: uq_learner.md 12 | - Non-uniform Quantization: nuq_learner.md 13 | - Learners - Misc.: 14 | - Distillation: distillation.md 15 | - Multi-GPU Training: multi_gpu_training.md 16 | - Hyper-parameter Optimizers: 17 | - Reinforcement Learning: reinforcement_learning.md 18 | - AutoML-based Methods: automl_based_methods.md 19 | - Self-defined Models: self_defined_models.md 20 | - Performance: performance.md 21 | - Frequently Asked Questions: faq.md 22 | - Appendix: 23 | - Pre-trained Models: pre_trained_models.md 24 | - Test Cases: test_cases.md 25 | - Reference: reference.md 26 | theme: readthedocs 27 | 28 | markdown_extensions: 29 | - pymdownx.arithmatex 30 | extra_javascript: 31 | - mathjax-config.js 32 | - MathJax.js?config=TeX-AMS-MML_HTMLorMML 33 | -------------------------------------------------------------------------------- /docs/qr_code.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/docs/qr_code.jpg -------------------------------------------------------------------------------- /examples/convnet_at_fmnist.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Model helper for creating a ConvNet model for the Fashion-MNIST dataset.""" 18 | 19 | import tensorflow as tf 20 | 21 | from nets.abstract_model_helper import AbstractModelHelper 22 | from datasets.fmnist_dataset import FMnistDataset 23 | from utils.lrn_rate_utils import setup_lrn_rate_piecewise_constant 24 | from utils.multi_gpu_wrapper import MultiGpuWrapper as mgw 25 | 26 | FLAGS = tf.app.flags.FLAGS 27 | 28 | tf.app.flags.DEFINE_float('nb_epochs_rat', 1.0, '# of training epochs\'s ratio') 29 | tf.app.flags.DEFINE_float('lrn_rate_init', 1e-1, 'initial learning rate') 30 | tf.app.flags.DEFINE_float('batch_size_norm', 128, 'normalization factor of batch size') 31 | tf.app.flags.DEFINE_float('momentum', 0.9, 'momentum coefficient') 32 | tf.app.flags.DEFINE_float('loss_w_dcy', 3e-4, 'weight decaying loss\'s coefficient') 33 | 34 | def forward_fn(inputs, data_format): 35 | """Forward pass function. 36 | 37 | Args: 38 | * inputs: inputs to the network's forward pass 39 | * data_format: data format ('channels_last' OR 'channels_first') 40 | 41 | Returns: 42 | * inputs: outputs from the network's forward pass 43 | """ 44 | 45 | # tranpose the image tensor if needed 46 | if data_format == 'channel_first': 47 | inputs = tf.transpose(inputs, [0, 3, 1, 2]) 48 | 49 | # conv1 50 | inputs = tf.layers.conv2d(inputs, 32, [5, 5], padding='same', 51 | data_format=data_format, activation=tf.nn.relu, name='conv1') 52 | inputs = tf.layers.max_pooling2d(inputs, [2, 2], 2, data_format=data_format, name='pool1') 53 | 54 | # conv2 55 | inputs = tf.layers.conv2d(inputs, 64, [5, 5], padding='same', 56 | data_format=data_format, activation=tf.nn.relu, name='conv2') 57 | inputs = tf.layers.max_pooling2d(inputs, [2, 2], 2, data_format=data_format, name='pool2') 58 | 59 | # fc3 60 | inputs = tf.layers.flatten(inputs, name='flatten') 61 | inputs = tf.layers.dense(inputs, 1024, activation=tf.nn.relu, name='fc3') 62 | 63 | # fc4 64 | inputs = tf.layers.dense(inputs, FLAGS.nb_classes, name='fc4') 65 | inputs = tf.nn.softmax(inputs, name='softmax') 66 | 67 | return inputs 68 | 69 | class ModelHelper(AbstractModelHelper): 70 | """Model helper for creating a ConvNet model for the Fashion-MNIST dataset.""" 71 | 72 | def __init__(self, data_format='channels_last'): 73 | """Constructor function.""" 74 | 75 | # class-independent initialization 76 | super(ModelHelper, self).__init__(data_format) 77 | 78 | # initialize training & evaluation subsets 79 | self.dataset_train = FMnistDataset(is_train=True) 80 | self.dataset_eval = FMnistDataset(is_train=False) 81 | 82 | def build_dataset_train(self, enbl_trn_val_split=False): 83 | """Build the data subset for training, usually with data augmentation.""" 84 | 85 | return self.dataset_train.build(enbl_trn_val_split) 86 | 87 | def build_dataset_eval(self): 88 | """Build the data subset for evaluation, usually without data augmentation.""" 89 | 90 | return self.dataset_eval.build() 91 | 92 | def forward_train(self, inputs): 93 | """Forward computation at training.""" 94 | 95 | return forward_fn(inputs, self.data_format) 96 | 97 | def forward_eval(self, inputs): 98 | """Forward computation at evaluation.""" 99 | 100 | return forward_fn(inputs, self.data_format) 101 | 102 | def calc_loss(self, labels, outputs, trainable_vars): 103 | """Calculate loss (and some extra evaluation metrics).""" 104 | 105 | loss = tf.losses.softmax_cross_entropy(labels, outputs) 106 | loss += FLAGS.loss_w_dcy * tf.add_n([tf.nn.l2_loss(var) for var in trainable_vars]) 107 | accuracy = tf.reduce_mean( 108 | tf.cast(tf.equal(tf.argmax(labels, axis=1), tf.argmax(outputs, axis=1)), tf.float32)) 109 | metrics = {'accuracy': accuracy} 110 | 111 | return loss, metrics 112 | 113 | def setup_lrn_rate(self, global_step): 114 | """Setup the learning rate (and number of training iterations).""" 115 | 116 | nb_epochs = 160 117 | idxs_epoch = [40, 80, 120] 118 | decay_rates = [1.0, 0.1, 0.01, 0.001] 119 | batch_size = FLAGS.batch_size * (1 if not FLAGS.enbl_multi_gpu else mgw.size()) 120 | lrn_rate = setup_lrn_rate_piecewise_constant(global_step, batch_size, idxs_epoch, decay_rates) 121 | nb_iters = int(FLAGS.nb_smpls_train * nb_epochs * FLAGS.nb_epochs_rat / batch_size) 122 | 123 | return lrn_rate, nb_iters 124 | 125 | @property 126 | def model_name(self): 127 | """Model's name.""" 128 | 129 | return 'convnet' 130 | 131 | @property 132 | def dataset_name(self): 133 | """Dataset's name.""" 134 | 135 | return 'fmnist' 136 | -------------------------------------------------------------------------------- /examples/convnet_at_fmnist_run.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Execution script for ConvNet models on the Fashion-MNIST dataset.""" 18 | 19 | import traceback 20 | import tensorflow as tf 21 | 22 | from nets.convnet_at_fmnist import ModelHelper 23 | from learners.learner_utils import create_learner 24 | 25 | FLAGS = tf.app.flags.FLAGS 26 | 27 | tf.app.flags.DEFINE_string('log_dir', './logs', 'logging directory') 28 | tf.app.flags.DEFINE_boolean('enbl_multi_gpu', False, 'enable multi-GPU training') 29 | tf.app.flags.DEFINE_string('learner', 'full-prec', 'learner\'s name') 30 | tf.app.flags.DEFINE_string('exec_mode', 'train', 'execution mode: train / eval') 31 | tf.app.flags.DEFINE_boolean('debug', False, 'debugging information') 32 | 33 | def main(unused_argv): 34 | """Main entry.""" 35 | 36 | try: 37 | # setup the TF logging routine 38 | if FLAGS.debug: 39 | tf.logging.set_verbosity(tf.logging.DEBUG) 40 | else: 41 | tf.logging.set_verbosity(tf.logging.INFO) 42 | sm_writer = tf.summary.FileWriter(FLAGS.log_dir) 43 | 44 | # display FLAGS's values 45 | tf.logging.info('FLAGS:') 46 | for key, value in FLAGS.flag_values_dict().items(): 47 | tf.logging.info('{}: {}'.format(key, value)) 48 | 49 | # build the model helper & learner 50 | model_helper = ModelHelper() 51 | learner = create_learner(sm_writer, model_helper) 52 | 53 | # execute the learner 54 | if FLAGS.exec_mode == 'train': 55 | learner.train() 56 | elif FLAGS.exec_mode == 'eval': 57 | learner.download_model() 58 | learner.evaluate() 59 | else: 60 | raise ValueError('unrecognized execution mode: ' + FLAGS.exec_mode) 61 | 62 | # exit normally 63 | return 0 64 | except ValueError: 65 | traceback.print_exc() 66 | return 1 # exit with errors 67 | 68 | if __name__ == '__main__': 69 | tf.app.run() 70 | -------------------------------------------------------------------------------- /learners/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/learners/__init__.py -------------------------------------------------------------------------------- /learners/channel_pruning/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/learners/channel_pruning/__init__.py -------------------------------------------------------------------------------- /learners/channel_pruning_gpu/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/learners/channel_pruning_gpu/__init__.py -------------------------------------------------------------------------------- /learners/channel_pruning_rmt/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/learners/channel_pruning_rmt/__init__.py -------------------------------------------------------------------------------- /learners/discr_channel_pruning/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/learners/discr_channel_pruning/__init__.py -------------------------------------------------------------------------------- /learners/full_precision/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/learners/full_precision/__init__.py -------------------------------------------------------------------------------- /learners/learner_utils.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Utility function for creating the specified learner.""" 18 | 19 | import tensorflow as tf 20 | 21 | from learners.full_precision.learner import FullPrecLearner 22 | from learners.weight_sparsification.learner import WeightSparseLearner 23 | from learners.channel_pruning.learner import ChannelPrunedLearner 24 | from learners.channel_pruning_gpu.learner import ChannelPrunedGpuLearner 25 | from learners.channel_pruning_rmt.learner import ChannelPrunedRmtLearner 26 | from learners.discr_channel_pruning.learner import DisChnPrunedLearner 27 | from learners.uniform_quantization.learner import UniformQuantLearner 28 | from learners.uniform_quantization_tf.learner import UniformQuantTFLearner 29 | from learners.nonuniform_quantization.learner import NonUniformQuantLearner 30 | 31 | FLAGS = tf.app.flags.FLAGS 32 | 33 | def create_learner(sm_writer, model_helper): 34 | """Create the learner as specified by FLAGS.learner. 35 | 36 | Args: 37 | * sm_writer: TensorFlow's summary writer 38 | * model_helper: model helper with definitions of model & dataset 39 | 40 | Returns: 41 | * learner: the specified learner 42 | """ 43 | 44 | learner = None 45 | if FLAGS.learner == 'full-prec': 46 | learner = FullPrecLearner(sm_writer, model_helper) 47 | elif FLAGS.learner == 'weight-sparse': 48 | learner = WeightSparseLearner(sm_writer, model_helper) 49 | elif FLAGS.learner == 'channel': 50 | learner = ChannelPrunedLearner(sm_writer, model_helper) 51 | elif FLAGS.learner == 'chn-pruned-gpu': 52 | learner = ChannelPrunedGpuLearner(sm_writer, model_helper) 53 | elif FLAGS.learner == 'chn-pruned-rmt': 54 | learner = ChannelPrunedRmtLearner(sm_writer, model_helper) 55 | elif FLAGS.learner == 'dis-chn-pruned': 56 | learner = DisChnPrunedLearner(sm_writer, model_helper) 57 | elif FLAGS.learner == 'uniform': 58 | learner = UniformQuantLearner(sm_writer, model_helper) 59 | elif FLAGS.learner == 'uniform-tf': 60 | learner = UniformQuantTFLearner(sm_writer, model_helper) 61 | elif FLAGS.learner == 'non-uniform': 62 | learner = NonUniformQuantLearner(sm_writer, model_helper) 63 | else: 64 | raise ValueError('unrecognized learner\'s name: ' + FLAGS.learner) 65 | 66 | return learner 67 | -------------------------------------------------------------------------------- /learners/nonuniform_quantization/rl_helper.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """ RL Helper Class for Non-Uniform Quantization """ 18 | 19 | import random 20 | import numpy as np 21 | import tensorflow as tf 22 | 23 | FLAGS = tf.app.flags.FLAGS 24 | 25 | class RLHelper(object): 26 | # pylint: disable=too-many-instance-attributes 27 | """ The Helper Class for the DDPG algorithm. Making sure the states and actions 28 | satisfy the condition of constraints given the total number bits. 29 | """ 30 | 31 | def __init__(self, sess, total_bits, num_weights, vars_list, random_layers=False): 32 | """ 33 | Args: 34 | * sess: session to run the shape 35 | * total_bits: A scalar, total num of bits as budget 36 | * num_weights: A list of number of weights in each layers 37 | * var_list: List of Tensors. Weights to quantize 38 | * random_layers: True if shuffle the layers 39 | """ 40 | 41 | self.nb_vars = len(num_weights) 42 | self.num_weights = num_weights 43 | self.total_num_weights = sum(num_weights) 44 | 45 | # TODO: the design of states is experimental 46 | self.s_dims = self.nb_vars + 6 # one hot encoding for layer ids 47 | self.total_bits = total_bits 48 | self.w_bits_used = 0 49 | self.random_layers = random_layers 50 | self.layer_idxs = list(range(self.nb_vars)) 51 | self.num_weights_to_quantize = self.total_num_weights 52 | self.quantized_layers = 0 53 | 54 | self.var_shapes = [] 55 | self.states = np.zeros((self.nb_vars, self.s_dims)) 56 | 57 | for idx, var in enumerate(vars_list): 58 | var_shape = sess.run(tf.shape(var)) 59 | shape_len = var.shape.__len__() 60 | assert shape_len in [2, 4], \ 61 | "Unknown weight shape. Must be a 2 (fc) or 4 (conv) dimensional." 62 | if shape_len == 2: 63 | var_shape = np.hstack((np.ones(2), var_shape)) 64 | self.var_shapes += [var_shape] 65 | 66 | for idx in range(self.nb_vars): 67 | state = np.zeros(self.s_dims) 68 | state[idx] = 1.0 69 | state[self.nb_vars : self.nb_vars + 4] = self.var_shapes[idx] 70 | state[self.nb_vars + 4] = self.num_weights[idx] / np.max(self.num_weights) 71 | state[self.nb_vars + 5] = np.sum(self.num_weights[idx + 1 : ]) / self.total_num_weights 72 | self.states[idx] = state 73 | 74 | def calc_state(self, idx): 75 | """ return the rl state """ 76 | state = np.copy(self.states[idx]) 77 | return state[None, :] 78 | 79 | def calc_reward(self, accuracy): 80 | """ return the rl reward via reshaping """ 81 | return accuracy * np.ones((1, 1)) 82 | 83 | def reset(self): 84 | """ reset the helper params for each rollout """ 85 | self.w_bits_used = 0 86 | self.quantized_layers = 0 87 | if self.random_layers: 88 | random.shuffle(self.layer_idxs) 89 | self.num_weights_to_quantize = self.total_num_weights 90 | 91 | def __calc_w_duty(self, idx): 92 | """ Compute the maximum bits used for layer idx """ 93 | duty = self.total_bits - self.w_bits_used - self.num_weights_to_quantize*FLAGS.nuql_w_bit_min 94 | assert duty >= 0, "No enough budget for layer {}".format(idx) 95 | return duty 96 | 97 | def calc_w(self, action, idx): 98 | """ 99 | Args: 100 | * action: An ndarray with shape (1,1), the output of actor network 101 | * idx: A scalar, the id of layer 102 | * num_weights: A list of scalars, the number of weights in each layer 103 | 104 | Return: 105 | * An ndarray with the same shape with 'action' 106 | """ 107 | 108 | duty = self.__calc_w_duty(idx) 109 | if self.quantized_layers != self.nb_vars - 1: 110 | action = np.round(action) + FLAGS.nuql_w_bit_min 111 | action = np.minimum(action, FLAGS.nuql_w_bit_min + np.floor(duty*1.0/self.num_weights[idx])) 112 | else: 113 | action = np.floor((self.total_bits - self.w_bits_used)/self.num_weights[idx]) * np.ones((1, 1)) 114 | action = np.minimum(action, FLAGS.nuql_w_bit_max) 115 | self.w_bits_used += action[0][0] * self.num_weights[idx] 116 | self.num_weights_to_quantize -= self.num_weights[idx] 117 | self.quantized_layers += 1 118 | return action 119 | -------------------------------------------------------------------------------- /learners/uniform_quantization/rl_helper.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """ RL Helper Class for Uniform Quantization """ 18 | 19 | import random 20 | import numpy as np 21 | import tensorflow as tf 22 | # import pdb 23 | 24 | FLAGS = tf.app.flags.FLAGS 25 | 26 | class RLHelper(object): 27 | # pylint: disable=too-many-instance-attributes 28 | """ The Helper Class for the DDPG algorithm. Making sure the states and actions 29 | satisfy the condition of constraints given the total number bits. 30 | """ 31 | 32 | def __init__(self, sess, total_bits, num_weights, vars_list, random_layers=False): 33 | """ 34 | Args: 35 | * sess: session to run the shape; 36 | * total_bits: a scalar, total num of bits as budget; 37 | * num_weights: a list of number of weights in each layers; 38 | * var_list: list of Tensors, weights to quantize; 39 | * random_layers: True if shuffle the layers; 40 | """ 41 | 42 | self.nb_vars = len(num_weights) 43 | self.num_weights = num_weights 44 | self.total_num_weights = sum(num_weights) 45 | # TODO: the design of states is experimental 46 | self.s_dims = self.nb_vars + 6 # one hot encoding for layer ids 47 | self.total_bits = total_bits 48 | self.w_bits_used = 0 49 | self.random_layers = random_layers 50 | self.layer_idxs = list(range(self.nb_vars)) 51 | self.num_weights_to_quantize = self.total_num_weights 52 | self.quantized_layers = 0 53 | 54 | self.var_shapes = [] 55 | self.states = np.zeros((self.nb_vars, self.s_dims)) 56 | 57 | for idx, var in enumerate(vars_list): 58 | var_shape = sess.run(tf.shape(var)) 59 | shape_len = var.shape.__len__() 60 | assert shape_len in [2, 4], \ 61 | "Unknown weight shape. Must be a 2 (fc) or 4 (conv) dimensional." 62 | if shape_len == 2: 63 | var_shape = np.hstack((np.ones(2), var_shape)) 64 | self.var_shapes += [var_shape] 65 | 66 | for idx in range(self.nb_vars): 67 | state = np.zeros(self.s_dims) 68 | state[idx] = 1.0 69 | state[self.nb_vars : self.nb_vars + 4] = self.var_shapes[idx] 70 | state[self.nb_vars + 4] = self.num_weights[idx] / np.max(self.num_weights) 71 | state[self.nb_vars + 5] = np.sum(self.num_weights[idx + 1 : ]) / self.total_num_weights 72 | self.states[idx] = state 73 | 74 | def calc_state(self, idx): 75 | """ return the rl state """ 76 | state = np.copy(self.states[idx]) 77 | return state[None, :] 78 | 79 | def calc_reward(self, accuracy): 80 | """ return the rl reward via reshaping """ 81 | return accuracy * np.ones((1, 1)) 82 | 83 | def reset(self): 84 | """ reset the helper params for each rollout """ 85 | self.w_bits_used = 0 86 | self.quantized_layers = 0 87 | if self.random_layers: 88 | random.shuffle(self.layer_idxs) 89 | self.num_weights_to_quantize = self.total_num_weights 90 | 91 | def __calc_w_duty(self, idx): 92 | """ Compute the maximum bits used for layer idx """ 93 | duty = self.total_bits - self.w_bits_used - self.num_weights_to_quantize*FLAGS.uql_w_bit_min 94 | assert duty >= 0, "Not enough budget for layer {}".format(idx) 95 | return duty 96 | 97 | def calc_w(self, action, idx): 98 | """ Generatee the action that satisfies the resource constraint 99 | 100 | Args: 101 | * action: an ndarray with shape (1,1), the output of actor network; 102 | * idx: A scalar, the id of layer; 103 | * num_weights: A list of scalars, the number of weights in each layer; 104 | 105 | Return: 106 | * An ndarray with the same shape with 'action' 107 | """ 108 | 109 | duty = self.__calc_w_duty(idx) 110 | if self.quantized_layers != self.nb_vars - 1: 111 | action = np.round(action) + FLAGS.uql_w_bit_min 112 | # duty will be automatically reshaped as the same as action 113 | action = np.minimum(action, FLAGS.uql_w_bit_min + \ 114 | np.floor(duty*1.0/self.num_weights[idx])) 115 | else: 116 | action = np.floor((self.total_bits - self.w_bits_used) / \ 117 | self.num_weights[idx]) * np.ones((1, 1)) 118 | action = np.minimum(action, FLAGS.uql_w_bit_max) 119 | self.w_bits_used += action[0][0] * self.num_weights[idx] 120 | self.num_weights_to_quantize -= self.num_weights[idx] 121 | self.quantized_layers += 1 122 | return action 123 | -------------------------------------------------------------------------------- /learners/uniform_quantization_tf/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/learners/uniform_quantization_tf/__init__.py -------------------------------------------------------------------------------- /learners/weight_sparsification/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/learners/weight_sparsification/__init__.py -------------------------------------------------------------------------------- /learners/weight_sparsification/utils.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Utility functions for the weight sparsification learner.""" 18 | 19 | def get_maskable_vars(trainable_vars): 20 | """Get a list of maskable variables. 21 | 22 | Args: 23 | * trainable_vars: list of trainable variables 24 | 25 | Returns: 26 | * maskable_vars: list of maskable variables 27 | 28 | Kernels in the following layer types will be matched: 29 | * tf.layers.conv2d 30 | * tf.layers.dense 31 | * Pointwise convolutional layer in slim.separable_conv2d 32 | """ 33 | 34 | vars_kernel = [var for var in trainable_vars if 'kernel' in var.name] 35 | vars_ptconv = [var for var in trainable_vars if 'pointwise/weights' in var.name] 36 | vars_fnconv = [var for var in trainable_vars if 'Conv2d_1c_1x1/weights' in var.name] 37 | maskable_vars = vars_kernel + vars_ptconv + vars_fnconv 38 | 39 | return maskable_vars 40 | -------------------------------------------------------------------------------- /main.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # configure pip to use internal source 4 | unset http_proxy 5 | unset https_proxy 6 | mkdir -p ~/.pip/ \ 7 | && echo "[global]" > ~/.pip/pip.conf \ 8 | && echo "index-url = http://mirror-sng.oa.com/pypi/web/simple/" >> ~/.pip/pip.conf \ 9 | && echo "trusted-host = mirror-sng.oa.com" >> ~/.pip/pip.conf 10 | cat ~/.pip/pip.conf 11 | 12 | # install Python packages with Internet access 13 | pip install tensorflow-gpu==1.12.0 14 | pip install horovod 15 | pip install docopt 16 | pip install hdfs 17 | pip install scipy 18 | pip install sklearn 19 | pip install pandas 20 | pip install mpi4py 21 | 22 | # add the current directory to PYTHONPATH 23 | export PYTHONPATH=${PYTHONPATH}:`pwd` 24 | export LD_LIBRARY_PATH=/opt/ml/disk/local/cuda/lib64:$LD_LIBRARY_PATH 25 | 26 | # start TensorBoard 27 | LOG_DIR=/opt/ml/log 28 | mkdir -p ${LOG_DIR} 29 | nohup tensorboard \ 30 | --port=${SEVEN_HTTP_FORWARD_PORT} \ 31 | --host=127.0.0.1 \ 32 | --logdir=${LOG_DIR} \ 33 | >/dev/null 2>&1 & 34 | 35 | # execute the main script 36 | mkdir models 37 | EXTRA_ARGS=`cat ./extra_args` 38 | if [ ${NB_GPUS} -eq 1 ]; then 39 | echo "multi-GPU training disabled" 40 | python main.py --log_dir ${LOG_DIR} ${EXTRA_ARGS} 41 | elif [ ${NB_GPUS} -le 8 ]; then 42 | echo "multi-GPU training enabled" 43 | options="-np ${NB_GPUS} -H localhost:${NB_GPUS} -bind-to none -map-by slot 44 | -x NCCL_DEBUG=INFO -x NCCL_SOCKET_IFNAME=eth1 -x NCCL_IB_DISABLE=1 45 | -x LD_LIBRARY_PATH --mca btl_tcp_if_include eth1" 46 | mpirun ${options} python main.py --enbl_multi_gpu --log_dir ${LOG_DIR} ${EXTRA_ARGS} 47 | fi 48 | 49 | # archive model files to HDFS 50 | mv models* /opt/ml/model 51 | 52 | # remove *.pyc files 53 | find . -name "*.pyc" -exec rm -f {} \; 54 | -------------------------------------------------------------------------------- /nets/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/nets/__init__.py -------------------------------------------------------------------------------- /nets/abstract_model_helper.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Abstract class for model helpers.""" 18 | 19 | from abc import ABC 20 | from abc import abstractmethod 21 | 22 | class AbstractModelHelper(ABC): 23 | """Abstract class for model helpers. 24 | 25 | A model helper should define the following function interface: 26 | 1. Data input pipeline for training and evaluation subsets. 27 | 2. Network's forward pass during training and evaluation. 28 | 3. Loss function (and some extra evaluation metrics). 29 | 30 | All functions marked with "@abstractmethod" must be explicitly implemented in the sub-class. 31 | """ 32 | 33 | def __init__(self, data_format, forward_w_labels=False): 34 | """Constructor function. 35 | 36 | Note: DO NOT create any TF operations here!!! 37 | 38 | Args: 39 | * data_format: data format ('channels_last' OR 'channels_first') 40 | """ 41 | 42 | self.data_format = data_format 43 | self.forward_w_labels = forward_w_labels 44 | 45 | @abstractmethod 46 | def build_dataset_train(self, enbl_trn_val_split): 47 | """Build the data subset for training, usually with data augmentation. 48 | 49 | Args: 50 | * enbl_trn_val_split: enable the training & validation splitting 51 | 52 | Returns: 53 | * iterator_trn: iterator for the training subset 54 | * iterator_val: iterator for the validation subset 55 | OR 56 | * iterator: iterator for the training subset 57 | """ 58 | pass 59 | 60 | @abstractmethod 61 | def build_dataset_eval(self): 62 | """Build the data subset for evaluation, usually without data augmentation. 63 | 64 | Returns: 65 | * iterator: iterator over the evaluation subset 66 | """ 67 | pass 68 | 69 | @abstractmethod 70 | def forward_train(self, inputs, labels=None): 71 | """Forward computation at training. 72 | 73 | Args: 74 | * inputs: inputs to the network's forward pass 75 | * labels: ground-truth labels 76 | 77 | Returns: 78 | * outputs: outputs from the network's forward pass 79 | """ 80 | pass 81 | 82 | @abstractmethod 83 | def forward_eval(self, inputs): 84 | """Forward computation at evaluation. 85 | 86 | Args: 87 | * inputs: inputs to the network's forward pass 88 | 89 | Returns: 90 | * outputs: outputs from the network's forward pass 91 | """ 92 | pass 93 | 94 | @abstractmethod 95 | def calc_loss(self, labels, outputs, trainable_vars): 96 | """Calculate loss (and some extra evaluation metrics). 97 | 98 | Args: 99 | * labels: ground-truth labels 100 | * outputs: outputs from the network's forward pass 101 | * trainable_vars: list of trainable variables 102 | 103 | Returns: 104 | * loss: loss function's value 105 | * metrics: dictionary of extra evaluation metrics 106 | """ 107 | pass 108 | 109 | @abstractmethod 110 | def setup_lrn_rate(self, global_step): 111 | """Setup the learning rate (and number of training iterations). 112 | 113 | Args: 114 | * global_step: training iteration counter 115 | 116 | Returns: 117 | * lrn_rate: learning rate 118 | * nb_iters: number of training iterations 119 | """ 120 | pass 121 | 122 | def warm_start(self, sess): 123 | """Initialize the model for warm-start. 124 | 125 | Args: 126 | * sess: TensorFlow session 127 | """ 128 | pass 129 | 130 | def dump_n_eval(self, outputs, action): 131 | """Dump the model's outputs to files and evaluate. 132 | 133 | Args: 134 | * outputs: outputs from the network's forward pass 135 | * action: 'init' | 'dump' | 'eval' 136 | """ 137 | pass 138 | 139 | @property 140 | @abstractmethod 141 | def model_name(self): 142 | """Model's name.""" 143 | pass 144 | 145 | @property 146 | @abstractmethod 147 | def dataset_name(self): 148 | """Dataset's name.""" 149 | pass 150 | -------------------------------------------------------------------------------- /nets/faster_rcnn_at_pascalvoc_run.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Execution script for VGG models on the Pascal VOC dataset.""" 18 | 19 | import traceback 20 | import tensorflow as tf 21 | 22 | from nets.faster_rcnn_at_pascalvoc import ModelHelper 23 | from learners.learner_utils import create_learner 24 | 25 | FLAGS = tf.app.flags.FLAGS 26 | 27 | tf.app.flags.DEFINE_string('log_dir', './logs', 'logging directory') 28 | tf.app.flags.DEFINE_boolean('enbl_multi_gpu', False, 'enable multi-GPU training') 29 | tf.app.flags.DEFINE_string('learner', 'full-prec', 'learner\'s name') 30 | tf.app.flags.DEFINE_string('exec_mode', 'train', 'execution mode: train / eval') 31 | tf.app.flags.DEFINE_boolean('debug', False, 'debugging information') 32 | 33 | def main(unused_argv): 34 | """Main entry.""" 35 | 36 | try: 37 | # setup the TF logging routine 38 | if FLAGS.debug: 39 | tf.logging.set_verbosity(tf.logging.DEBUG) 40 | else: 41 | tf.logging.set_verbosity(tf.logging.INFO) 42 | sm_writer = tf.summary.FileWriter(FLAGS.log_dir) 43 | 44 | # display FLAGS's values 45 | tf.logging.info('FLAGS:') 46 | for key, value in FLAGS.flag_values_dict().items(): 47 | tf.logging.info('{}: {}'.format(key, value)) 48 | 49 | # build the model helper & learner 50 | model_helper = ModelHelper() 51 | learner = create_learner(sm_writer, model_helper) 52 | 53 | # execute the learner 54 | if FLAGS.exec_mode == 'train': 55 | learner.train() 56 | elif FLAGS.exec_mode == 'eval': 57 | learner.download_model() 58 | learner.evaluate() 59 | else: 60 | raise ValueError('unrecognized execution mode: ' + FLAGS.exec_mode) 61 | 62 | # exit normally 63 | return 0 64 | except ValueError: 65 | traceback.print_exc() 66 | return 1 # exit with errors 67 | 68 | if __name__ == '__main__': 69 | tf.app.run() 70 | -------------------------------------------------------------------------------- /nets/lenet_at_cifar10.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Model helper for creating a LeNet-like model for the CIFAR-10 dataset.""" 18 | 19 | import tensorflow as tf 20 | 21 | from nets.abstract_model_helper import AbstractModelHelper 22 | from datasets.cifar10_dataset import Cifar10Dataset 23 | from utils.lrn_rate_utils import setup_lrn_rate_piecewise_constant 24 | from utils.multi_gpu_wrapper import MultiGpuWrapper as mgw 25 | 26 | FLAGS = tf.app.flags.FLAGS 27 | 28 | tf.app.flags.DEFINE_float('nb_epochs_rat', 1.0, '# of training epochs\'s ratio') 29 | tf.app.flags.DEFINE_float('lrn_rate_init', 1e-2, 'initial learning rate') 30 | tf.app.flags.DEFINE_float('batch_size_norm', 128, 'normalization factor of batch size') 31 | tf.app.flags.DEFINE_float('momentum', 0.9, 'momentum coefficient') 32 | tf.app.flags.DEFINE_float('loss_w_dcy', 5e-4, 'weight decaying loss\'s coefficient') 33 | 34 | def forward_fn(inputs, data_format): 35 | """Forward pass function. 36 | 37 | Args: 38 | * inputs: inputs to the network's forward pass 39 | * data_format: data format ('channels_last' OR 'channels_first') 40 | 41 | Returns: 42 | * inputs: outputs from the network's forward pass 43 | """ 44 | 45 | # convert inputs from channels_last (NHWC) to channels_first (NCHW) 46 | if data_format == 'channels_first': 47 | inputs = tf.transpose(inputs, [0, 3, 1, 2]) 48 | 49 | # conv1 50 | inputs = tf.layers.conv2d(inputs, 32, [5, 5], data_format=data_format, name='conv1') 51 | inputs = tf.nn.relu(inputs, name='relu1') 52 | inputs = tf.layers.max_pooling2d(inputs, [2, 2], 2, data_format=data_format, name='pool1') 53 | 54 | # conv2 55 | inputs = tf.layers.conv2d(inputs, 64, [5, 5], data_format=data_format, name='conv2') 56 | inputs = tf.nn.relu(inputs, name='relu2') 57 | inputs = tf.layers.max_pooling2d(inputs, [2, 2], 2, data_format=data_format, name='pool2') 58 | 59 | # fc3 60 | inputs = tf.layers.flatten(inputs, name='flatten') 61 | inputs = tf.layers.dense(inputs, 256, name='fc3') 62 | inputs = tf.nn.relu(inputs, name='relu3') 63 | 64 | # fc4 65 | inputs = tf.layers.dense(inputs, FLAGS.nb_classes, name='fc4') 66 | inputs = tf.nn.softmax(inputs, name='softmax') 67 | 68 | return inputs 69 | 70 | class ModelHelper(AbstractModelHelper): 71 | """Model helper for creating a LeNet-like model for the CIFAR-10 dataset.""" 72 | 73 | def __init__(self, data_format='channels_last'): 74 | """Constructor function.""" 75 | 76 | # class-independent initialization 77 | super(ModelHelper, self).__init__(data_format) 78 | 79 | # initialize training & evaluation subsets 80 | self.dataset_train = Cifar10Dataset(is_train=True) 81 | self.dataset_eval = Cifar10Dataset(is_train=False) 82 | 83 | def build_dataset_train(self, enbl_trn_val_split=False): 84 | """Build the data subset for training, usually with data augmentation.""" 85 | 86 | return self.dataset_train.build(enbl_trn_val_split) 87 | 88 | def build_dataset_eval(self): 89 | """Build the data subset for evaluation, usually without data augmentation.""" 90 | 91 | return self.dataset_eval.build() 92 | 93 | def forward_train(self, inputs): 94 | """Forward computation at training.""" 95 | 96 | return forward_fn(inputs, self.data_format) 97 | 98 | def forward_eval(self, inputs): 99 | """Forward computation at evaluation.""" 100 | 101 | return forward_fn(inputs, self.data_format) 102 | 103 | def calc_loss(self, labels, outputs, trainable_vars): 104 | """Calculate loss (and some extra evaluation metrics).""" 105 | 106 | loss = tf.losses.softmax_cross_entropy(labels, outputs) 107 | loss += FLAGS.loss_w_dcy * tf.add_n([tf.nn.l2_loss(var) for var in trainable_vars]) 108 | accuracy = tf.reduce_mean( 109 | tf.cast(tf.equal(tf.argmax(labels, axis=1), tf.argmax(outputs, axis=1)), tf.float32)) 110 | metrics = {'accuracy': accuracy} 111 | 112 | return loss, metrics 113 | 114 | def setup_lrn_rate(self, global_step): 115 | """Setup the learning rate (and number of training iterations).""" 116 | 117 | nb_epochs = 250 118 | idxs_epoch = [100, 150, 200] 119 | decay_rates = [1.0, 0.1, 0.01, 0.001] 120 | batch_size = FLAGS.batch_size * (1 if not FLAGS.enbl_multi_gpu else mgw.size()) 121 | lrn_rate = setup_lrn_rate_piecewise_constant(global_step, batch_size, idxs_epoch, decay_rates) 122 | nb_iters = int(FLAGS.nb_smpls_train * nb_epochs * FLAGS.nb_epochs_rat / batch_size) 123 | 124 | return lrn_rate, nb_iters 125 | 126 | @property 127 | def model_name(self): 128 | """Model's name.""" 129 | 130 | return 'lenet' 131 | 132 | @property 133 | def dataset_name(self): 134 | """Dataset's name.""" 135 | 136 | return 'cifar_10' 137 | -------------------------------------------------------------------------------- /nets/lenet_at_cifar10_run.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Execution script for LeNet-like models on the CIFAR-10 dataset.""" 18 | 19 | import traceback 20 | import tensorflow as tf 21 | 22 | from nets.lenet_at_cifar10 import ModelHelper 23 | from learners.learner_utils import create_learner 24 | 25 | FLAGS = tf.app.flags.FLAGS 26 | 27 | tf.app.flags.DEFINE_string('log_dir', './logs', 'logging directory') 28 | tf.app.flags.DEFINE_boolean('enbl_multi_gpu', False, 'enable multi-GPU training') 29 | tf.app.flags.DEFINE_string('learner', 'full-prec', 'learner\'s name') 30 | tf.app.flags.DEFINE_string('exec_mode', 'train', 'execution mode: train / eval') 31 | tf.app.flags.DEFINE_boolean('debug', False, 'debugging information') 32 | 33 | def main(unused_argv): 34 | """Main entry.""" 35 | 36 | try: 37 | # setup the TF logging routine 38 | if FLAGS.debug: 39 | tf.logging.set_verbosity(tf.logging.DEBUG) 40 | else: 41 | tf.logging.set_verbosity(tf.logging.INFO) 42 | sm_writer = tf.summary.FileWriter(FLAGS.log_dir) 43 | 44 | # display FLAGS's values 45 | tf.logging.info('FLAGS:') 46 | for key, value in FLAGS.flag_values_dict().items(): 47 | tf.logging.info('{}: {}'.format(key, value)) 48 | 49 | # build the model helper & learner 50 | model_helper = ModelHelper() 51 | learner = create_learner(sm_writer, model_helper) 52 | 53 | # execute the learner 54 | if FLAGS.exec_mode == 'train': 55 | learner.train() 56 | elif FLAGS.exec_mode == 'eval': 57 | learner.download_model() 58 | learner.evaluate() 59 | else: 60 | raise ValueError('unrecognized execution mode: ' + FLAGS.exec_mode) 61 | 62 | # exit normally 63 | return 0 64 | except ValueError: 65 | traceback.print_exc() 66 | return 1 # exit with errors 67 | 68 | if __name__ == '__main__': 69 | tf.app.run() 70 | -------------------------------------------------------------------------------- /nets/mobilenet_at_ilsvrc12_run.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Execution script for MobileNet models on the ILSVRC-12 dataset.""" 18 | 19 | import traceback 20 | import tensorflow as tf 21 | 22 | from nets.mobilenet_at_ilsvrc12 import ModelHelper 23 | from learners.learner_utils import create_learner 24 | 25 | FLAGS = tf.app.flags.FLAGS 26 | 27 | tf.app.flags.DEFINE_string('log_dir', './logs', 'logging directory') 28 | tf.app.flags.DEFINE_boolean('enbl_multi_gpu', False, 'enable multi-GPU training') 29 | tf.app.flags.DEFINE_string('learner', 'full-prec', 'learner\'s name') 30 | tf.app.flags.DEFINE_string('exec_mode', 'train', 'execution mode: train / eval') 31 | tf.app.flags.DEFINE_boolean('debug', False, 'debugging information') 32 | 33 | def main(unused_argv): 34 | """Main entry.""" 35 | 36 | try: 37 | # setup the TF logging routine 38 | if FLAGS.debug: 39 | tf.logging.set_verbosity(tf.logging.DEBUG) 40 | else: 41 | tf.logging.set_verbosity(tf.logging.INFO) 42 | sm_writer = tf.summary.FileWriter(FLAGS.log_dir) 43 | 44 | # display FLAGS's values 45 | tf.logging.info('FLAGS:') 46 | for key, value in FLAGS.flag_values_dict().items(): 47 | tf.logging.info('{}: {}'.format(key, value)) 48 | 49 | # build the model helper & learner 50 | model_helper = ModelHelper() 51 | learner = create_learner(sm_writer, model_helper) 52 | 53 | # execute the learner 54 | if FLAGS.exec_mode == 'train': 55 | learner.train() 56 | elif FLAGS.exec_mode == 'eval': 57 | learner.download_model() 58 | learner.evaluate() 59 | else: 60 | raise ValueError('unrecognized execution mode: ' + FLAGS.exec_mode) 61 | 62 | # exit normally 63 | return 0 64 | except ValueError: 65 | traceback.print_exc() 66 | return 1 # exit with errors 67 | 68 | if __name__ == '__main__': 69 | tf.app.run() 70 | -------------------------------------------------------------------------------- /nets/resnet_at_cifar10_run.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Execution script for ResNet models on the CIFAR-10 dataset.""" 18 | 19 | import traceback 20 | import tensorflow as tf 21 | 22 | from nets.resnet_at_cifar10 import ModelHelper 23 | from learners.learner_utils import create_learner 24 | 25 | FLAGS = tf.app.flags.FLAGS 26 | 27 | tf.app.flags.DEFINE_string('log_dir', './logs', 'logging directory') 28 | tf.app.flags.DEFINE_boolean('enbl_multi_gpu', False, 'enable multi-GPU training') 29 | tf.app.flags.DEFINE_string('learner', 'full-prec', 'learner\'s name') 30 | tf.app.flags.DEFINE_string('exec_mode', 'train', 'execution mode: train / eval') 31 | tf.app.flags.DEFINE_boolean('debug', False, 'debugging information') 32 | 33 | def main(unused_argv): 34 | """Main entry.""" 35 | 36 | try: 37 | # setup the TF logging routine 38 | if FLAGS.debug: 39 | tf.logging.set_verbosity(tf.logging.DEBUG) 40 | else: 41 | tf.logging.set_verbosity(tf.logging.INFO) 42 | sm_writer = tf.summary.FileWriter(FLAGS.log_dir) 43 | 44 | # display FLAGS's values 45 | tf.logging.info('FLAGS:') 46 | for key, value in FLAGS.flag_values_dict().items(): 47 | tf.logging.info('{}: {}'.format(key, value)) 48 | 49 | # build the model helper & learner 50 | model_helper = ModelHelper() 51 | learner = create_learner(sm_writer, model_helper) 52 | 53 | # execute the learner 54 | if FLAGS.exec_mode == 'train': 55 | learner.train() 56 | elif FLAGS.exec_mode == 'eval': 57 | learner.download_model() 58 | learner.evaluate() 59 | else: 60 | raise ValueError('unrecognized execution mode: ' + FLAGS.exec_mode) 61 | 62 | # exit normally 63 | return 0 64 | except ValueError: 65 | traceback.print_exc() 66 | return 1 # exit with errors 67 | 68 | if __name__ == '__main__': 69 | tf.app.run() 70 | -------------------------------------------------------------------------------- /nets/resnet_at_ilsvrc12_run.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Execution script for ResNet models on the ILSVRC-12 dataset.""" 18 | 19 | import traceback 20 | import tensorflow as tf 21 | 22 | from nets.resnet_at_ilsvrc12 import ModelHelper 23 | from learners.learner_utils import create_learner 24 | 25 | FLAGS = tf.app.flags.FLAGS 26 | 27 | tf.app.flags.DEFINE_string('log_dir', './logs', 'logging directory') 28 | tf.app.flags.DEFINE_boolean('enbl_multi_gpu', False, 'enable multi-GPU training') 29 | tf.app.flags.DEFINE_string('learner', 'full-prec', 'learner\'s name') 30 | tf.app.flags.DEFINE_string('exec_mode', 'train', 'execution mode: train / eval') 31 | tf.app.flags.DEFINE_boolean('debug', False, 'debugging information') 32 | 33 | def main(unused_argv): 34 | """Main entry.""" 35 | 36 | try: 37 | # setup the TF logging routine 38 | if FLAGS.debug: 39 | tf.logging.set_verbosity(tf.logging.DEBUG) 40 | else: 41 | tf.logging.set_verbosity(tf.logging.INFO) 42 | sm_writer = tf.summary.FileWriter(FLAGS.log_dir) 43 | 44 | # display FLAGS's values 45 | tf.logging.info('FLAGS:') 46 | for key, value in FLAGS.flag_values_dict().items(): 47 | tf.logging.info('{}: {}'.format(key, value)) 48 | 49 | # build the model helper & learner 50 | model_helper = ModelHelper() 51 | learner = create_learner(sm_writer, model_helper) 52 | 53 | # execute the learner 54 | if FLAGS.exec_mode == 'train': 55 | learner.train() 56 | elif FLAGS.exec_mode == 'eval': 57 | learner.download_model() 58 | learner.evaluate() 59 | else: 60 | raise ValueError('unrecognized execution mode: ' + FLAGS.exec_mode) 61 | 62 | # exit normally 63 | return 0 64 | except ValueError: 65 | traceback.print_exc() 66 | return 1 # exit with errors 67 | 68 | if __name__ == '__main__': 69 | tf.app.run() 70 | -------------------------------------------------------------------------------- /nets/vgg_at_pascalvoc_run.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Execution script for VGG models on the Pascal VOC dataset.""" 18 | 19 | import traceback 20 | import tensorflow as tf 21 | 22 | from nets.vgg_at_pascalvoc import ModelHelper 23 | from learners.learner_utils import create_learner 24 | 25 | FLAGS = tf.app.flags.FLAGS 26 | 27 | tf.app.flags.DEFINE_string('log_dir', './logs', 'logging directory') 28 | tf.app.flags.DEFINE_boolean('enbl_multi_gpu', False, 'enable multi-GPU training') 29 | tf.app.flags.DEFINE_string('learner', 'full-prec', 'learner\'s name') 30 | tf.app.flags.DEFINE_string('exec_mode', 'train', 'execution mode: train / eval') 31 | tf.app.flags.DEFINE_boolean('debug', False, 'debugging information') 32 | 33 | def main(unused_argv): 34 | """Main entry.""" 35 | 36 | try: 37 | # setup the TF logging routine 38 | if FLAGS.debug: 39 | tf.logging.set_verbosity(tf.logging.DEBUG) 40 | else: 41 | tf.logging.set_verbosity(tf.logging.INFO) 42 | sm_writer = tf.summary.FileWriter(FLAGS.log_dir) 43 | 44 | # display FLAGS's values 45 | tf.logging.info('FLAGS:') 46 | for key, value in FLAGS.flag_values_dict().items(): 47 | tf.logging.info('{}: {}'.format(key, value)) 48 | 49 | # build the model helper & learner 50 | model_helper = ModelHelper() 51 | learner = create_learner(sm_writer, model_helper) 52 | 53 | # execute the learner 54 | if FLAGS.exec_mode == 'train': 55 | learner.train() 56 | elif FLAGS.exec_mode == 'eval': 57 | learner.download_model() 58 | learner.evaluate() 59 | else: 60 | raise ValueError('unrecognized execution mode: ' + FLAGS.exec_mode) 61 | 62 | # exit normally 63 | return 0 64 | except ValueError: 65 | traceback.print_exc() 66 | return 1 # exit with errors 67 | 68 | if __name__ == '__main__': 69 | tf.app.run() 70 | -------------------------------------------------------------------------------- /path.conf.template: -------------------------------------------------------------------------------- 1 | # Notes on HDFS-related configurations: 2 | # * data_hdfs_host: 3 | # HDFS host for remote datasets. 4 | # This will be used by TensorFlow's data input pipeline. 5 | # Format: hdfs://IP_Address:Port 6 | # Leave 'None' unchanged if HDFS is not availabe. 7 | # * model_http_url: 8 | # HTTP/HTTPS url for remote model files. 9 | # This will be used to download pre-trained models via . 10 | # Format: http://aaa.bbb.ccc OR https://aaa.bbb.ccc 11 | # Set it to 'None' if such connection is not available. 12 | # 13 | # How to setup the directory path: 14 | # * If a data directory path exists, then replace 'None' with the actual path, e.g.: 15 | # data_dir_local_cifar10 = /home/user_name/datasets/cifar-10/cifar-10-batches-bin 16 | # which setups the path to the CIFAR-10 dataset on the local machine. 17 | # * Otherwise, leave 'None' unchanged. 18 | 19 | # data files 20 | data_hdfs_host = None 21 | data_dir_local_cifar10 = None 22 | data_dir_hdfs_cifar10 = None 23 | data_dir_seven_cifar10 = None 24 | data_dir_docker_cifar10 = /opt/ml/data # DO NOT EDIT 25 | data_dir_local_ilsvrc12 = None 26 | data_dir_hdfs_ilsvrc12 = None 27 | data_dir_seven_ilsvrc12 = None 28 | data_dir_docker_ilsvrc12 = /opt/ml/data # DO NOT EDIT 29 | 30 | # model files 31 | model_http_url = https://api.ai.tencent.com/pocketflow 32 | -------------------------------------------------------------------------------- /requirement.txt: -------------------------------------------------------------------------------- 1 | docopt>=0.6.2 2 | hdfs>=2.1.0 3 | numpy>=1.14.0 4 | scipy>=1.0.0 5 | sklearn>=0.19.1 6 | pandas>=0.22.0 7 | mpi4py>=3.0.0 8 | tensorflow>=1.10.0 9 | -------------------------------------------------------------------------------- /rl_agents/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/rl_agents/__init__.py -------------------------------------------------------------------------------- /rl_agents/ddpg/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/rl_agents/ddpg/__init__.py -------------------------------------------------------------------------------- /rl_agents/ddpg/actor_critic.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Actor & critic networks' definitions.""" 18 | 19 | import tensorflow as tf 20 | 21 | FLAGS = tf.app.flags.FLAGS 22 | 23 | tf.app.flags.DEFINE_integer('ddpg_actor_depth', 2, 'DDPG: actor network\'s depth') 24 | tf.app.flags.DEFINE_integer('ddpg_actor_width', 64, 'DDPG: actor network\'s width') 25 | tf.app.flags.DEFINE_integer('ddpg_critic_depth', 2, 'DDPG: critic network\'s depth') 26 | tf.app.flags.DEFINE_integer('ddpg_critic_width', 64, 'DDPG: critic network\'s width') 27 | 28 | ENBL_LAYER_NORM = True 29 | 30 | def dense_block(inputs, units): 31 | """A block of densely connected layers. 32 | 33 | Args: 34 | * inputs: input tensor to the block 35 | * units: number of neurons in each layer 36 | 37 | Returns: 38 | * inputs: output tensor from the block 39 | """ 40 | 41 | inputs = tf.layers.dense(inputs, units) 42 | if ENBL_LAYER_NORM: 43 | inputs = tf.contrib.layers.layer_norm(inputs) 44 | inputs = tf.nn.relu(inputs) 45 | 46 | return inputs 47 | 48 | class Model(object): 49 | """Abstract model for actor & critic networks.""" 50 | 51 | def __init__(self, scope): 52 | """Constructor function. 53 | 54 | Args: 55 | * scope: name scope in which the model is defined 56 | """ 57 | 58 | self.scope = scope 59 | 60 | @property 61 | def vars(self): 62 | """List of all global variables.""" 63 | 64 | return tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=self.scope) 65 | 66 | @property 67 | def trainable_vars(self): 68 | """List of all trainable variables.""" 69 | 70 | return tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=self.scope) 71 | 72 | @property 73 | def perturbable_vars(self): 74 | """List of all perturbable variables.""" 75 | 76 | return [var for var in self.trainable_vars if 'LayerNorm' not in var.name] 77 | 78 | class Actor(Model): 79 | """Actor network.""" 80 | 81 | def __init__(self, a_dims, a_min, a_max, scope='actor'): 82 | """Constructor function. 83 | 84 | Args: 85 | * a_dims: number of action vector's dimensions 86 | * a_min: minimal value in the action vector 87 | * a_max: maximal value in the action vector 88 | * scope: name scope in which the model is defined 89 | """ 90 | 91 | super(Actor, self).__init__(scope) 92 | 93 | self.a_dims = a_dims 94 | self.a_min = a_min 95 | self.a_max = a_max 96 | 97 | def __call__(self, states, reuse=False): 98 | """Create the actor network. 99 | 100 | Args: 101 | * states: state vectors (inputs to the actor network) 102 | * reuse: whether to reuse previously defined variables 103 | 104 | Returns: 105 | * inputs: action vectors (outputs from the actor network) 106 | """ 107 | 108 | with tf.variable_scope(self.scope) as scope: 109 | if reuse: 110 | scope.reuse_variables() 111 | 112 | inputs = states 113 | for __ in range(FLAGS.ddpg_actor_depth): 114 | inputs = dense_block(inputs, FLAGS.ddpg_actor_width) 115 | inputs = tf.layers.dense(inputs, self.a_dims) 116 | inputs = tf.sigmoid(inputs) * (self.a_max - self.a_min) + self.a_min 117 | 118 | return inputs 119 | 120 | class Critic(Model): 121 | """Critic network.""" 122 | 123 | def __init__(self, scope='critic'): 124 | """Constructor function. 125 | 126 | Args: 127 | * scope: name scope in which the model is defined 128 | """ 129 | 130 | super(Critic, self).__init__(scope) 131 | 132 | def __call__(self, states, actions, reuse=False): 133 | """Create the critic network. 134 | 135 | Args: 136 | * states: state vectors (inputs to the critic network) 137 | * actions: action vectors (inputs to the critic network) 138 | * reuse: whether to reuse previously defined variables 139 | 140 | Returns: 141 | * inputs: reward scalars (outputs from the critic network) 142 | """ 143 | 144 | with tf.variable_scope(self.scope) as scope: 145 | if reuse: 146 | scope.reuse_variables() 147 | 148 | inputs = dense_block(states, FLAGS.ddpg_critic_width) 149 | inputs = tf.concat([inputs, actions], axis=1) 150 | for __ in range(FLAGS.ddpg_critic_depth): 151 | inputs = dense_block(inputs, FLAGS.ddpg_critic_width) 152 | inputs = tf.layers.dense(inputs, 1) 153 | 154 | return inputs 155 | -------------------------------------------------------------------------------- /rl_agents/ddpg/noise.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Parameter & action noise's specifications.""" 18 | 19 | import tensorflow as tf 20 | 21 | FLAGS = tf.app.flags.FLAGS 22 | 23 | tf.app.flags.DEFINE_string('ddpg_noise_type', 'param', 24 | 'DDPG: noise type (\'action\' OR \'param\')') 25 | tf.app.flags.DEFINE_string('ddpg_noise_prtl', 'tdecy', 26 | 'DDPG: noise adjustment protocol (\'adapt\' OR \'tdecy\')') 27 | tf.app.flags.DEFINE_float('ddpg_noise_std_init', 1e+0, 28 | 'DDPG: parameter / action noise\'s initial stdev.') 29 | 30 | # for only 31 | tf.app.flags.DEFINE_float('ddpg_noise_dst_finl', 1e-2, 'DDPG: action noise\'s final distance') 32 | tf.app.flags.DEFINE_float('ddpg_noise_adpt_rat', 1.03, 'DDPG: parameter noise\'s adaption rate') 33 | 34 | # for only 35 | tf.app.flags.DEFINE_float('ddpg_noise_std_finl', 1e-5, 36 | 'DDPG: parameter / action noise\'s final stdev.') 37 | 38 | class AdaptiveNoiseSpec(object): 39 | """Adaptive parameter noise's specifications. 40 | 41 | To enable, set to 'param' and to 'adapt'. 42 | """ 43 | 44 | def __init__(self): 45 | """Constructor function.""" 46 | 47 | self.stdev_curr = FLAGS.ddpg_noise_std_init 48 | 49 | def reset(self): 50 | """Reset the standard deviation.""" 51 | 52 | self.stdev_curr = FLAGS.ddpg_noise_std_init 53 | 54 | def adapt(self, dst_curr): 55 | """Adjust the standard deviation to meet the distance requirement between actions. 56 | 57 | Args: 58 | * dst_curr: current distance between clean & distorted actions 59 | """ 60 | 61 | if dst_curr > FLAGS.ddpg_noise_dst_finl: 62 | self.stdev_curr /= FLAGS.ddpg_noise_adpt_rat 63 | else: 64 | self.stdev_curr *= FLAGS.ddpg_noise_adpt_rat 65 | 66 | class TimeDecayNoiseSpec(object): 67 | """Time-decaying action / parameter noise's specifications. 68 | 69 | To enable, set to 'action' / 'param' and to 'tdecy'. 70 | """ 71 | 72 | def __init__(self, nb_rlouts): 73 | """Constructor function.""" 74 | 75 | self.stdev_curr = FLAGS.ddpg_noise_std_init 76 | self.decy_rat = (FLAGS.ddpg_noise_std_finl / FLAGS.ddpg_noise_std_init) ** (1.0 / nb_rlouts) 77 | 78 | def reset(self): 79 | """Reset the standard deviation.""" 80 | 81 | self.stdev_curr = FLAGS.ddpg_noise_std_init 82 | 83 | def adapt(self): 84 | """Adjust the standard deviation by multiplying with the decaying ratio.""" 85 | 86 | self.stdev_curr *= self.decy_rat 87 | -------------------------------------------------------------------------------- /rl_agents/ddpg/replay_buffer.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Replay buffer for storing state-action-reward transitions.""" 18 | 19 | import numpy as np 20 | 21 | class ReplayBuffer(object): 22 | """Replay buffer for storing state-action-reward transitions. 23 | 24 | Each transition consists of five components: 25 | 1. state: current state 26 | 2. action: action under the current state 27 | 3. reward: reward for the chosen action under the current state 28 | 4. terminal: whether the terminal state is reached after the chosen action 29 | 5. state_next: next state 30 | """ 31 | 32 | def __init__(self, s_dims, a_dims, buf_size): 33 | """Constructor function. 34 | 35 | Args: 36 | * s_dims: number of state vector's dimensions 37 | * a_dims: number of action vector's dimensions 38 | * buf_size: maximal number of transitions to be stored 39 | """ 40 | 41 | self.s_dims = s_dims 42 | self.a_dims = a_dims 43 | self.buf_size = buf_size 44 | 45 | # initialize the buffer & counters 46 | self.idx_smpl = 0 47 | self.nb_smpls = 0 # number of valid samples in the buffer 48 | self.buffers = { 49 | 'states': np.zeros((self.buf_size, self.s_dims), dtype=np.float32), 50 | 'actions': np.zeros((self.buf_size, self.a_dims), dtype=np.float32), 51 | 'rewards': np.zeros((self.buf_size, 1), dtype=np.float32), 52 | 'terminals': np.zeros((self.buf_size, 1), dtype=np.float32), 53 | 'states_next': np.zeros((self.buf_size, self.s_dims), dtype=np.float32), 54 | } 55 | 56 | def reset(self): 57 | """Reset the replay buffer (all transitions are removed).""" 58 | 59 | self.idx_smpl = 0 60 | self.nb_smpls = 0 61 | 62 | def is_ready(self): 63 | """Check whether the replay buffer is ready for sampling transitions. 64 | 65 | Returns: 66 | * is_ready: True if the replay buffer is ready 67 | """ 68 | 69 | return self.nb_smpls == self.buf_size 70 | 71 | def append(self, states, actions, rewards, terminals, states_next): 72 | """Append multiple transitions into the replay buffer. 73 | 74 | Args: 75 | * states: np.array of current state vectors 76 | * actions: np.array of action vectors 77 | * rewards: np.array of reward scalars 78 | * terminals: np.array of terminal scalars (1: terminal; 0: non-terminal) 79 | * states_next: np.array of next state vectors 80 | """ 81 | 82 | # pack transactions into a mini-batch 83 | nb_smpls = states.shape[0] 84 | mbatch = { 85 | 'states': states, 86 | 'actions': actions, 87 | 'rewards': rewards, 88 | 'terminals': terminals, 89 | 'states_next': states_next, 90 | } 91 | 92 | # insert samples into the buffer 93 | if self.idx_smpl + nb_smpls <= self.buf_size: 94 | idxs = np.arange(self.idx_smpl, self.idx_smpl + nb_smpls) 95 | for key in self.buffers: 96 | self.buffers[key][idxs] = mbatch[key] 97 | else: 98 | nb_smpls_tail = self.buf_size - self.idx_smpl # samples to be added to the buffer tail 99 | nb_smpls_head = nb_smpls - nb_smpls_tail # samples to be added to the buffer head 100 | for key in self.buffers: 101 | self.buffers[key][self.idx_smpl:] = mbatch[key][:nb_smpls_tail] 102 | self.buffers[key][:nb_smpls_head] = mbatch[key][nb_smpls_tail:] 103 | 104 | # update counters 105 | self.idx_smpl = (self.idx_smpl + nb_smpls) % self.buf_size 106 | self.nb_smpls = min(self.nb_smpls + nb_smpls, self.buf_size) 107 | 108 | def sample(self, batch_size): 109 | """Sample a mini-batch of trasitions. 110 | 111 | Args: 112 | * batch_size: number of transitions in the mini-batch 113 | 114 | Returns: 115 | * mbatch: a mini-batch of transitions 116 | """ 117 | 118 | idxs_smpl = np.random.randint(0, self.nb_smpls, batch_size) 119 | mbatch = {key: self.buffers[key][idxs_smpl] for key in self.buffers} 120 | 121 | return mbatch 122 | -------------------------------------------------------------------------------- /rl_agents/ddpg/running_mean_std.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Running averages of mean value & standard deviation.""" 18 | 19 | import tensorflow as tf 20 | 21 | FLAGS = tf.app.flags.FLAGS 22 | 23 | tf.app.flags.DEFINE_float('ddpg_rms_eps', 1e-4, 'DDPG: running standard deviation\'s epsilon') 24 | 25 | class RunningMeanStd(object): 26 | """Running averages of mean value & standard deviation.""" 27 | 28 | def __init__(self, sess, nb_dims): 29 | """Constructor function. 30 | 31 | Args: 32 | * sess: TensorFlow session 33 | * nb_dims: number of sample's dimensions 34 | """ 35 | 36 | self.sess = sess 37 | 38 | # statistics for computing running mean & standard deviation 39 | x_sum = tf.get_variable( 40 | 'x_sum', shape=[nb_dims], initializer=tf.zeros_initializer(), trainable=False) 41 | x_sum_sq = tf.get_variable( 42 | 'x_sum_sq', shape=[nb_dims], initializer=tf.zeros_initializer(), trainable=False) 43 | x_cnt = tf.get_variable( 44 | 'x_cnt', shape=[], initializer=tf.zeros_initializer(), trainable=False) 45 | 46 | # update statistics with a mini-batch of samples 47 | self.x_new = tf.placeholder(tf.float32, shape=[None, nb_dims], name='x_new') 48 | self.updt_ops = [ 49 | x_sum.assign_add(tf.reduce_sum(self.x_new, axis=0)), 50 | x_sum_sq.assign_add(tf.reduce_sum(tf.square(self.x_new), axis=0)), 51 | x_cnt.assign_add(tf.cast(tf.shape(self.x_new)[0], tf.float32)) 52 | ] 53 | tf.summary.scalar('x_cnt', x_cnt) 54 | 55 | # compute running mean & standard deviation from statistics 56 | # Note: use default values if no samples have been added 57 | self.mean = tf.cond(x_cnt < 0.5, lambda: tf.zeros(shape=[nb_dims]), lambda: x_sum / x_cnt) 58 | self.std = tf.cond( 59 | x_cnt < 0.5, lambda: tf.ones(shape=[nb_dims]), 60 | lambda: tf.sqrt(tf.maximum(x_sum_sq / x_cnt - tf.square(self.mean), FLAGS.ddpg_rms_eps))) 61 | 62 | def updt(self, x_new): 63 | """Update running averages with a list of samples. 64 | 65 | Args: 66 | * x_new: np.array of list of samples (N x D) 67 | """ 68 | 69 | self.sess.run(self.updt_ops, feed_dict={self.x_new: x_new}) 70 | -------------------------------------------------------------------------------- /rl_agents/unit_tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/rl_agents/unit_tests/__init__.py -------------------------------------------------------------------------------- /rl_agents/unit_tests/pendulum_v0.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Unit-test for the Pendulum-v0 problem.""" 18 | 19 | import traceback 20 | import numpy as np 21 | import tensorflow as tf 22 | import gym 23 | 24 | from rl_agents.ddpg.agent import Agent as DdpgAgent 25 | 26 | FLAGS = tf.app.flags.FLAGS 27 | 28 | tf.app.flags.DEFINE_string('log_dir', './logs', 'logging directory') 29 | tf.app.flags.DEFINE_integer('nb_rlouts', 200, '# of roll-outs') 30 | tf.app.flags.DEFINE_integer('nb_rlouts_eval', 100, '# of roll-outs for evaluation') 31 | tf.app.flags.DEFINE_integer('rlout_len', 200, 'roll-out\'s length') 32 | 33 | def build_env_n_agent(sess): 34 | """Build the environment and an RL agent to solve it. 35 | 36 | Args: 37 | * sess: TensorFlow session 38 | 39 | Returns: 40 | * env: environment 41 | * agent: RL agent 42 | """ 43 | 44 | env = gym.make('Pendulum-v0') 45 | s_dims = env.observation_space.shape[-1] 46 | a_dims = env.action_space.shape[-1] 47 | buf_size = int(FLAGS.rlout_len * FLAGS.nb_rlouts * 0.25) 48 | a_lbnd = env.action_space.low[0] 49 | a_ubnd = env.action_space.high[0] 50 | agent = DdpgAgent(sess, s_dims, a_dims, FLAGS.nb_rlouts, buf_size, a_lbnd, a_ubnd) 51 | tf.logging.info('s_dims = %d, a_dims = %d' % (s_dims, a_dims)) 52 | tf.logging.info('a_lbnd = %f, a_ubnd = %f' % (a_lbnd, a_ubnd)) 53 | 54 | return env, agent 55 | 56 | def train_agent(sess, env, agent): 57 | """Train the RL agent through multiple roll-outs. 58 | 59 | Args: 60 | * sess: TensorFlow session 61 | * env: environment 62 | * agent: RL agent 63 | """ 64 | 65 | agent.init() 66 | for idx_rlout in range(FLAGS.nb_rlouts): 67 | agent.init_rlout() 68 | state = env.reset() 69 | rewards = np.zeros(FLAGS.rlout_len) 70 | for idx_iter in range(FLAGS.rlout_len): 71 | action = sess.run(agent.actions_noisy, feed_dict={agent.states: state[None, :]}) 72 | state_next, reward, __, __ = env.step(action.ravel()) 73 | terminal = np.ones((1, 1)) if idx_iter == FLAGS.rlout_len - 1 else np.zeros((1, 1)) 74 | agent.record(state[None, :], action, 75 | reward * np.ones((1, 1)), terminal, state_next[None, :]) 76 | actor_loss, critic_loss, noise_std = agent.train() 77 | state = state_next 78 | rewards[idx_iter] = reward 79 | agent.finalize_rlout(rewards) 80 | tf.logging.info('roll-out #%d: reward (ave.): %.2e' % (idx_rlout, np.mean(rewards))) 81 | tf.logging.info('roll-out #%d: a-loss = %.2e | c-loss = %.2e | noise std. = %.2e' 82 | % (idx_rlout, actor_loss, critic_loss, noise_std)) 83 | 84 | def eval_agent(sess, env, agent): 85 | """Evaluate the RL agent through multiple roll-outs. 86 | 87 | Args: 88 | * sess: TensorFlow session 89 | * env: environment 90 | * agent: RL agent 91 | """ 92 | 93 | reward_ave_list = [] 94 | for idx_rlout in range(FLAGS.nb_rlouts_eval): 95 | state = env.reset() 96 | rewards = np.zeros(FLAGS.rlout_len) 97 | for idx_iter in range(FLAGS.rlout_len): 98 | action = sess.run(agent.actions_clean, feed_dict={agent.states: state[None, :]}) 99 | state, reward, __, __ = env.step(action.ravel()) 100 | rewards[idx_iter] = reward 101 | tf.logging.info('roll-out #%d: reward (ave.): %.2e' % (idx_rlout, np.mean(rewards))) 102 | reward_ave_list += [np.mean(rewards)] 103 | tf.logging.info('[EVAL] reward (ave.): %.4e' % np.mean(np.array(reward_ave_list))) 104 | 105 | def main(unused_argv): 106 | """Main entry.""" 107 | 108 | try: 109 | # setup the TF logging routine 110 | tf.logging.set_verbosity(tf.logging.INFO) 111 | 112 | # display FLAGS's values 113 | tf.logging.info('FLAGS:') 114 | for key, value in FLAGS.flag_values_dict().items(): 115 | tf.logging.info('{}: {}'.format(key, value)) 116 | 117 | # build the environment & agent 118 | sess = tf.Session() 119 | env, agent = build_env_n_agent(sess) 120 | 121 | # train and evaluate the RL agent 122 | train_agent(sess, env, agent) 123 | eval_agent(sess, env, agent) 124 | 125 | # exit normally 126 | return 0 127 | except ValueError: 128 | traceback.print_exc() 129 | return 1 # exit with errors 130 | 131 | if __name__ == '__main__': 132 | tf.app.run() 133 | -------------------------------------------------------------------------------- /run.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # convert AutoML-generated hyper-parameter file to PocketFlow-compatible format 4 | python automl/cvt_hparam_file.py automl/automl_hparam.conf >> extra_args_from_automl 5 | 6 | # run seven job and get job id 7 | jobid=$(seven create --conf automl/seven.yaml --code `pwd` 2> seven.log | tee -a seven.log | \ 8 | python -c "import sys, json; print json.load(sys.stdin)['JobId']" 2>> seven.log) 9 | 10 | # save job id to file 11 | echo "seven_id="${jobid} > pid_file 12 | 13 | # wait until job finish 14 | seven wait --jobId ${jobid} >> seven.log 2>&1 15 | 16 | # get job log 17 | seven log --jobId ${jobid} > results 2>> seven.log 18 | 19 | # parse results to desired format 20 | python automl/parse_results.py results >> result_file 21 | -------------------------------------------------------------------------------- /run_pylint.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | check_dirs="" 4 | check_dirs=${check_dirs}" ./automl" 5 | check_dirs=${check_dirs}" ./datasets" 6 | check_dirs=${check_dirs}" ./learners" 7 | check_dirs=${check_dirs}" ./nets" 8 | check_dirs=${check_dirs}" ./rl_agents" 9 | check_dirs=${check_dirs}" ./tools" 10 | check_dirs=${check_dirs}" ./utils" 11 | echo "Folders to be checked:" 12 | echo " "${check_dirs} 13 | pylint --jobs=4 --reports=y ${check_dirs} --ignore=utils/external | tee .pylint_results 14 | 15 | echo 16 | echo "***************************************************" 17 | echo "*********** head -n50 .pylint_results ***********" 18 | echo "***************************************************" 19 | echo 20 | head -n50 .pylint_results 21 | -------------------------------------------------------------------------------- /scripts/create_minimal.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # obtain directory paths 4 | if [ "$#" -eq 2 ]; then 5 | dir_curr="$1" 6 | dir_temp="$2" 7 | else 8 | dir_curr=`pwd` 9 | dir_temp="${dir_curr}-minimal" 10 | fi 11 | echo "creating a minimal code directory: ${dir_curr} -> ${dir_temp}" 12 | 13 | # copy all files to a temporary directory 14 | rm -rf ${dir_temp} 15 | cp -r ${dir_curr} ${dir_temp} 16 | cd ${dir_temp} 17 | 18 | # remove redundant files 19 | #git clean -xdf # all files ignored by git 20 | rm -rf ./models* 21 | rm -rf ./logs 22 | rm -rf .git .gitignore 23 | cp ${dir_curr}/path.conf . 24 | 25 | # return to the original directory 26 | cd ${dir_curr} 27 | -------------------------------------------------------------------------------- /scripts/download_minimal.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # obtain directory paths 4 | dir_curr=`pwd` 5 | dir_temp="${dir_curr}-minimal" 6 | 7 | # create a minimal code directory 8 | ./scripts/create_minimal.sh ${dir_curr} ${dir_temp} 9 | 10 | # create a *.tar.gz archive and then download it 11 | dir_name_temp=`basename ${dir_temp}` 12 | tgz_name="${dir_name_temp}.tar.gz" 13 | cd .. 14 | tar -czvf ${tgz_name} ${dir_name_temp} 15 | sz -b ${tgz_name} 16 | cd ${dir_curr} 17 | -------------------------------------------------------------------------------- /scripts/run_docker.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # obtain directory paths 4 | dir_curr=`pwd` 5 | dir_temp="${dir_curr}-minimal" 6 | 7 | # create a minimal code directory 8 | ./scripts/create_minimal.sh ${dir_curr} ${dir_temp} 9 | cd ${dir_temp} 10 | 11 | # default arguments 12 | nb_gpus=1 13 | 14 | # parse arguments passed from the command line 15 | py_script="$1" 16 | shift 17 | extra_args="" 18 | for i in "$@" 19 | do 20 | case "$i" in 21 | -n=*|--nb_gpus=*) 22 | nb_gpus="${i#*=}" 23 | shift 24 | ;; 25 | *) 26 | # unknown option 27 | extra_args="${extra_args} ${i}" 28 | shift 29 | ;; 30 | esac 31 | done 32 | extra_args_path=`python utils/get_path_args.py docker ${py_script} path.conf` 33 | extra_args="${extra_args} ${extra_args_path}" 34 | echo ${extra_args} > extra_args 35 | echo "Python script: ${py_script}" 36 | echo "Data directory: ${dir_data}" 37 | echo "# of GPUs: ${nb_gpus}" 38 | echo "extra arguments: ${extra_args}" 39 | 40 | # create temporary directory for log & model 41 | dir_temp_log=`mktemp -d` 42 | dir_temp_model=`mktemp -d` 43 | dir_temp_data=`python utils/get_data_dir.py ${py_script} path.conf` 44 | 45 | # enter the docker environment 46 | cp -v ${py_script} main.py 47 | idle_gpus=`python utils/get_idle_gpus.py ${nb_gpus}` 48 | NV_GPU="${idle_gpus}" nvidia-docker run -it --rm --hostname=local-env \ 49 | --env SEVEN_HTTP_FORWARD_PORT= --env NB_GPUS=${nb_gpus} --network=host \ 50 | -v ${dir_temp}:/opt/ml/env \ 51 | -v ${dir_temp_log}:/opt/ml/log \ 52 | -v ${dir_temp_model}:/opt/ml/model \ 53 | -v ${dir_temp_data}:/opt/ml/data \ 54 | -w /opt/ml/env \ 55 | docker.oa.com/g_tfplus/horovod:python3.5 bash 56 | 57 | # return to the main directory 58 | rm -rf ${dir_temp_log} ${dir_temp_model} 59 | cd ${dir_curr} 60 | -------------------------------------------------------------------------------- /scripts/run_local.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # default arguments 4 | nb_gpus=1 5 | 6 | # parse arguments passed from the command line 7 | py_script="$1" 8 | shift 9 | extra_args="" 10 | for i in "$@" 11 | do 12 | case "$i" in 13 | -n=*|--nb_gpus=*) 14 | nb_gpus="${i#*=}" 15 | shift 16 | ;; 17 | *) 18 | # unknown option 19 | extra_args="${extra_args} ${i}" 20 | shift 21 | ;; 22 | esac 23 | done 24 | extra_args_path=`python utils/get_path_args.py local ${py_script} path.conf` 25 | extra_args="${extra_args} ${extra_args_path}" 26 | echo "Python script: ${py_script}" 27 | echo "# of GPUs: ${nb_gpus}" 28 | echo "extra arguments: ${extra_args}" 29 | 30 | # obtain list of idle GPUs 31 | idle_gpus=`python utils/get_idle_gpus.py ${nb_gpus}` 32 | export CUDA_VISIBLE_DEVICES=${idle_gpus} 33 | 34 | # re-create the logging directory 35 | rm -rf logs && mkdir logs 36 | 37 | # execute the specified Python script with one or more GPUs 38 | cp -v ${py_script} main.py 39 | if [ ${nb_gpus} -eq 1 ]; then 40 | echo "multi-GPU training disabled" 41 | python main.py ${extra_args} 42 | elif [ ${nb_gpus} -le 8 ]; then 43 | echo "multi-GPU training enabled" 44 | options="-np ${nb_gpus} -H localhost:${nb_gpus} -bind-to none -map-by slot 45 | -x NCCL_DEBUG=INFO -x NCCL_SOCKET_IFNAME=eth1 -x NCCL_IB_DISABLE=1 46 | -x LD_LIBRARY_PATH --mca btl_tcp_if_include eth1" 47 | mpirun ${options} python main.py --enbl_multi_gpu ${extra_args} 48 | fi 49 | -------------------------------------------------------------------------------- /scripts/run_seven.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # obtain directory paths 4 | dir_curr=`pwd` 5 | dir_temp="${dir_curr}-minimal" 6 | 7 | # create a minimal code directory 8 | ./scripts/create_minimal.sh ${dir_curr} ${dir_temp} 9 | cd ${dir_temp} 10 | 11 | # default arguments 12 | nb_gpus=1 13 | job_name="pocket-flow" 14 | 15 | # parse arguments passed from the command line 16 | py_script="$1" 17 | shift 18 | extra_args="" 19 | for i in "$@" 20 | do 21 | case "$i" in 22 | -n=*|--nb_gpus=*) 23 | nb_gpus="${i#*=}" 24 | shift 25 | ;; 26 | -j=*|--job_name=*) 27 | job_name="${i#*=}" 28 | shift 29 | ;; 30 | *) 31 | # unknown option 32 | extra_args="${extra_args} ${i}" 33 | shift 34 | ;; 35 | esac 36 | done 37 | extra_args_path=`python utils/get_path_args.py seven ${py_script} path.conf` 38 | extra_args="${extra_args} ${extra_args_path}" 39 | echo ${extra_args} > extra_args 40 | echo "Python script: ${py_script}" 41 | echo "Job name: ${job_name}" 42 | echo "# of GPUs: ${nb_gpus}" 43 | echo "extra arguments: ${extra_args}" 44 | 45 | # update the configuration file 46 | sed -i "s/nvidia.com\/gpu:\ 1/nvidia.com\/gpu:\ ${nb_gpus}/g" seven.yaml 47 | sed -i "s/name:\ NB_GPUS/name:\ NB_GPUS\n\ \ \ \ value: ${nb_gpus}/g" seven.yaml 48 | 49 | # start the seven job 50 | cp -v ${py_script} main.py 51 | seven create -cluster dev -conf seven.yaml -code `pwd` -name ${job_name} 52 | 53 | # return to the main directory 54 | cd ${dir_curr} 55 | -------------------------------------------------------------------------------- /seven.yaml: -------------------------------------------------------------------------------- 1 | version: 1 2 | kind: standalone 3 | jobname: pocket-flow 4 | container: 5 | image: 6 | docker.oa.com/g_tfplus/tfplus:tensorflow1.8-python3.6-cuda9.0-cudnn7.0.4.31-ubuntu16.04-tfplus-v3 7 | #docker.oa.com/g_tfplus/horovod:python3.5 8 | resources: 9 | nvidia.com/gpu: 1 10 | env: 11 | - name: SEVEN_HTTP_FORWARD_PORT 12 | - name: NB_GPUS 13 | -------------------------------------------------------------------------------- /tools/benchmark/calc_inference_time.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Measure the time consumption of *.pb and *.tflite models.""" 18 | 19 | import traceback 20 | from timeit import default_timer as timer 21 | import numpy as np 22 | import tensorflow as tf 23 | 24 | FLAGS = tf.app.flags.FLAGS 25 | 26 | tf.app.flags.DEFINE_string('model_file', None, 'model file path') 27 | tf.app.flags.DEFINE_string('input_name', 'net_input', 'input tensor\'s name in the *.pb model') 28 | tf.app.flags.DEFINE_string('output_name', 'net_output', 'output tensor\'s name in the *.pb model') 29 | tf.app.flags.DEFINE_string('input_dtype', 'float32', 30 | 'input tensor\'s data type in the *.tflite model') 31 | tf.app.flags.DEFINE_integer('batch_size', 1, 'batch size for run-time benchmark') 32 | tf.app.flags.DEFINE_integer('nb_repts_warmup', 100, '# of repeated runs for warm-up') 33 | tf.app.flags.DEFINE_integer('nb_repts', 100, '# of repeated runs for elapsed time measurement') 34 | 35 | def test_pb_model(): 36 | """Test the *.pb model.""" 37 | 38 | with tf.Graph().as_default() as graph: 39 | sess = tf.Session() 40 | 41 | # restore the model 42 | graph_def = tf.GraphDef() 43 | with tf.gfile.GFile(FLAGS.model_file, 'rb') as i_file: 44 | graph_def.ParseFromString(i_file.read()) 45 | tf.import_graph_def(graph_def) 46 | 47 | # obtain input & output nodes and then test the model 48 | net_input = graph.get_tensor_by_name('import/' + FLAGS.input_name + ':0') 49 | net_output = graph.get_tensor_by_name('import/' + FLAGS.output_name + ':0') 50 | net_input_data = np.zeros(tuple([FLAGS.batch_size] + list(net_input.shape[1:]))) 51 | for idx in range(FLAGS.nb_repts_warmup + FLAGS.nb_repts): 52 | if idx == FLAGS.nb_repts_warmup: 53 | time_beg = timer() 54 | sess.run(net_output, feed_dict={net_input: net_input_data}) 55 | time_elapsed = (timer() - time_beg) / FLAGS.nb_repts / FLAGS.batch_size 56 | tf.logging.info('time consumption of *.pb model: %.2f ms' % (time_elapsed * 1000)) 57 | 58 | def test_tflite_model(): 59 | """Test the *.tflite model.""" 60 | 61 | # restore the model and allocate tensors 62 | interpreter = tf.contrib.lite.Interpreter(model_path=FLAGS.model_file) 63 | interpreter.allocate_tensors() 64 | 65 | # get input & output tensors 66 | input_details = interpreter.get_input_details() 67 | output_details = interpreter.get_output_details() 68 | assert len(input_details) == 1, ' should contain only one element' 69 | if FLAGS.input_dtype == 'uint8': 70 | net_input_data = np.zeros(input_details[0]['shape'], dtype=np.uint8) 71 | elif FLAGS.input_dtype == 'float32': 72 | net_input_data = np.zeros(input_details[0]['shape'], dtype=np.float32) 73 | else: 74 | raise ValueError('unrecognized input data type: ' + FLAGS.input_dtype) 75 | 76 | # test the model with given inputs 77 | for idx in range(FLAGS.nb_repts_warmup + FLAGS.nb_repts): 78 | if idx == FLAGS.nb_repts_warmup: 79 | time_beg = timer() 80 | interpreter.set_tensor(input_details[0]['index'], net_input_data) 81 | interpreter.invoke() 82 | interpreter.get_tensor(output_details[0]['index']) 83 | time_elapsed = (timer() - time_beg) / FLAGS.nb_repts 84 | tf.logging.info('time consumption of *.tflite model: %.2f ms' % (time_elapsed * 1000)) 85 | 86 | def main(unused_argv): 87 | """Main entry. 88 | 89 | Args: 90 | * unused_argv: unused arguments (after FLAGS is parsed) 91 | """ 92 | 93 | try: 94 | # setup the TF logging routine 95 | tf.logging.set_verbosity(tf.logging.INFO) 96 | 97 | # call benchmark routines for *.pb / *.tflite models 98 | if FLAGS.model_file is None: 99 | raise ValueError(' must not be None') 100 | elif FLAGS.model_file.endswith('.pb'): 101 | test_pb_model() 102 | elif FLAGS.model_file.endswith('.tflite'): 103 | test_tflite_model() 104 | else: 105 | raise ValueError('unrecognized model file path: ' + FLAGS.model_file) 106 | 107 | # exit normally 108 | return 0 109 | except ValueError: 110 | traceback.print_exc() 111 | return 1 # exit with errors 112 | 113 | if __name__ == '__main__': 114 | tf.app.run() 115 | -------------------------------------------------------------------------------- /tools/conversion/convert_data_format.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Convert the data format from channels_last (NHWC) to channels_first (NCHW), or vice versa.""" 18 | 19 | import os 20 | import traceback 21 | import tensorflow as tf 22 | 23 | # NOTE: un-comment the corresponding before conversion 24 | #from nets.lenet_at_cifar10 import ModelHelper 25 | #from nets.resnet_at_cifar10 import ModelHelper 26 | from nets.resnet_at_ilsvrc12 import ModelHelper 27 | #from nets.mobilenet_at_ilsvrc12 import ModelHelper 28 | 29 | FLAGS = tf.app.flags.FLAGS 30 | 31 | tf.app.flags.DEFINE_string('log_dir', './logs', 'logging directory') 32 | tf.app.flags.DEFINE_boolean('enbl_multi_gpu', False, 'enable multi-GPU training') 33 | tf.app.flags.DEFINE_string('model_dir_in', './models', 'input model directory') 34 | tf.app.flags.DEFINE_string('model_dir_out', './models_out', 'output model directory') 35 | tf.app.flags.DEFINE_string('model_scope', 'model', 'model\'s variable scope name') 36 | tf.app.flags.DEFINE_string('data_format', 'channels_last', 'data format in the output model') 37 | 38 | def main(unused_argv): 39 | """Main entry. 40 | 41 | Args: 42 | * unused_argv: unused arguments (after FLAGS is parsed) 43 | """ 44 | 45 | try: 46 | # setup the TF logging routine 47 | tf.logging.set_verbosity(tf.logging.INFO) 48 | 49 | # create a TensorFlow session 50 | #sess = tf.Session() 51 | 52 | # create the model helper 53 | model_helper = ModelHelper(FLAGS.data_format) 54 | data_scope = 'data' 55 | model_scope = FLAGS.model_scope 56 | 57 | # bulid a graph with the target data format and rewrite checkpoint files 58 | with tf.Graph().as_default(): 59 | # data input pipeline 60 | with tf.variable_scope(data_scope): 61 | iterator = model_helper.build_dataset_eval() 62 | images, __ = iterator.get_next() 63 | 64 | # model definition 65 | with tf.variable_scope(model_scope): 66 | logits = model_helper.forward_eval(images) 67 | 68 | # add input & output tensors to certain collections 69 | tf.add_to_collection('images_final', images) 70 | tf.add_to_collection('logits_final', logits) 71 | 72 | # restore variables from checkpoint files 73 | vars_all = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope=model_scope) 74 | saver = tf.train.Saver(vars_all) 75 | save_path = tf.train.latest_checkpoint(FLAGS.model_dir_in) 76 | sess = tf.Session() 77 | saver.restore(sess, save_path) 78 | saver.save(sess, os.path.join(FLAGS.model_dir_out, 'model.ckpt')) 79 | 80 | # exit normally 81 | return 0 82 | except ValueError: 83 | traceback.print_exc() 84 | return 1 # exit with errors 85 | 86 | if __name__ == '__main__': 87 | tf.app.run() 88 | -------------------------------------------------------------------------------- /tools/graph_tools/add_to_collection.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Add a list of tensors to specified collections (useful when exporting *.pb & *.tflite models).""" 18 | 19 | import os 20 | import re 21 | import traceback 22 | import tensorflow as tf 23 | 24 | FLAGS = tf.app.flags.FLAGS 25 | 26 | tf.app.flags.DEFINE_string('model_dir_in', './models_in', 'input model directory') 27 | tf.app.flags.DEFINE_string('model_dir_out', './models_out', 'output model directory') 28 | tf.app.flags.DEFINE_string('tensor_names', None, 'list of tensors names (comma-separated)') 29 | tf.app.flags.DEFINE_string('coll_names', None, 'list of collection names (comma-separated)') 30 | 31 | ''' 32 | Example: SSD (VGG-16) @ Pascal VOC 33 | 34 | Input: 35 | * data/IteratorGetNext:1 / (?, 300, 300, 3) / images 36 | Output: 37 | * quant_model/ssd300/multibox_head/cls_5/Conv2D:0 / (?, 1, 1, 84) / cls_preds 38 | * quant_model/ssd300/multibox_head/loc_5/Conv2D:0 / (?, 1, 1, 16) / loc_preds 39 | ''' 40 | 41 | def main(unused_argv): 42 | """Main entry. 43 | 44 | Args: 45 | * unused_argv: unused arguments (after FLAGS is parsed) 46 | """ 47 | 48 | try: 49 | # setup the TF logging routine 50 | tf.logging.set_verbosity(tf.logging.INFO) 51 | 52 | # add a list of tensors to specified collections 53 | with tf.Graph().as_default() as graph: 54 | # create a TensorFlow session 55 | config = tf.ConfigProto() 56 | config.gpu_options.allow_growth = True 57 | sess = tf.Session(config=config) 58 | 59 | # restore a model from *.ckpt files 60 | ckpt_path = tf.train.latest_checkpoint(FLAGS.model_dir_in) 61 | meta_path = ckpt_path + '.meta' 62 | saver = tf.train.import_meta_graph(meta_path) 63 | saver.restore(sess, ckpt_path) 64 | 65 | # parse tensor & collection names 66 | tensor_names = [sub_str.strip() for sub_str in FLAGS.tensor_names.split(',')] 67 | coll_names = [sub_str.strip() for sub_str in FLAGS.coll_names.split(',')] 68 | assert len(tensor_names) == len(coll_names), \ 69 | '# of tensors and collections does not match: %d (tensor) vs. %d (collection)' \ 70 | % (len(tensor_names), len(coll_names)) 71 | 72 | # obtain the full list of tensors in the graph 73 | tensors = set() 74 | for op in graph.get_operations(): 75 | tensors |= set(op.inputs) | set(op.outputs) 76 | tensors = list(tensors) 77 | tensors.sort(key=lambda x: x.name) 78 | 79 | # find tensors and add them to corresponding collections 80 | for tensor in tensors: 81 | if tensor.name in tensor_names: 82 | tf.logging.info('tensor: {} / {}'.format(tensor.name, tensor.shape)) 83 | coll_name = coll_names[tensor_names.index(tensor.name)] 84 | tf.add_to_collection(coll_name, tensor) 85 | tf.logging.info('added tensor <{}> to collection <{}>'.format(tensor.name, coll_name)) 86 | 87 | # save the modified model 88 | vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) 89 | saver_new = tf.train.Saver(vars_list) 90 | save_path = saver_new.save(sess, os.path.join(FLAGS.model_dir_out, 'model.ckpt')) 91 | tf.logging.info('model saved to ' + save_path) 92 | 93 | # exit normally 94 | return 0 95 | except ValueError: 96 | traceback.print_exc() 97 | return 1 # exit with errors 98 | 99 | if __name__ == '__main__': 100 | tf.app.run() 101 | -------------------------------------------------------------------------------- /utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/utils/__init__.py -------------------------------------------------------------------------------- /utils/external/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/utils/external/__init__.py -------------------------------------------------------------------------------- /utils/external/faster_rcnn_tensorflow/configs/cfgs.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | from __future__ import division, print_function, absolute_import 3 | import os 4 | import tensorflow as tf 5 | 6 | # ------------------------------------------------ 7 | VERSION = 'MobileNetV2' 8 | NET_NAME = 'MobilenetV2' #'MobilenetV2' 9 | ADD_BOX_IN_TENSORBOARD = True 10 | 11 | if NET_NAME.startswith("resnet"): 12 | weights_name = NET_NAME 13 | elif NET_NAME.startswith("MobilenetV2"): 14 | weights_name = "mobilenet/mobilenet_v2_1.0_224" 15 | else: 16 | raise Exception('net name must in [resnet_v1_101, resnet_v1_50, MobilenetV2]') 17 | 18 | # ------------------------------------------ Train config 19 | RESTORE_FROM_RPN = False 20 | IS_FILTER_OUTSIDE_BOXES = True 21 | FIXED_BLOCKS = 1 # allow 0~3 22 | 23 | RPN_LOCATION_LOSS_WEIGHT = 1. 24 | RPN_CLASSIFICATION_LOSS_WEIGHT = 1.0 25 | 26 | FAST_RCNN_LOCATION_LOSS_WEIGHT = 1.0 27 | FAST_RCNN_CLASSIFICATION_LOSS_WEIGHT = 1.0 28 | RPN_SIGMA = 3.0 29 | FASTRCNN_SIGMA = 1.0 30 | 31 | MUTILPY_BIAS_GRADIENT = None # 2.0 # if None, will not multipy 32 | GRADIENT_CLIPPING_BY_NORM = None # 10.0 if None, will not clip 33 | 34 | EPSILON = 1e-5 35 | # LR = 0.001 # ResNet 36 | # DECAY_STEP = [50000, 70000] # ResNet 37 | LR = 0.0003 # MobileNet\ 38 | DECAY_STEP = [50000, 100000] # MobileNet 39 | MAX_ITERATION = 200000 40 | 41 | # -------------------------------------------- Data_preprocess_config 42 | DATASET_NAME = 'pascal' # 'ship', 'spacenet', 'pascal', 'coco' 43 | PIXEL_MEAN = [123.68, 116.779, 103.939] # R, G, B. In tf, channel is RGB. In openCV, channel is BGR 44 | IMG_SHORT_SIDE_LEN = 600 45 | IMG_MAX_LENGTH = 1000 46 | CLASS_NUM = 20 47 | 48 | # --------------------------------------------- Network_config 49 | BATCH_SIZE = 1 50 | INITIALIZER = tf.random_normal_initializer(mean=0.0, stddev=0.01) 51 | BBOX_INITIALIZER = tf.random_normal_initializer(mean=0.0, stddev=0.001) 52 | WEIGHT_DECAY = 0.00004 if NET_NAME.startswith('Mobilenet') else 0.0001 53 | 54 | # ---------------------------------------------Anchor config 55 | BASE_ANCHOR_SIZE_LIST = [256] # can be modified 56 | ANCHOR_STRIDE = [16] # can not be modified in most situations 57 | ANCHOR_SCALES = [0.5, 1., 2.0] # [4, 8, 16, 32] 58 | ANCHOR_RATIOS = [0.5, 1., 2.0] 59 | ROI_SCALE_FACTORS = [10., 10., 5.0, 5.0] 60 | ANCHOR_SCALE_FACTORS = None 61 | 62 | 63 | # --------------------------------------------RPN config 64 | KERNEL_SIZE = 3 65 | RPN_IOU_POSITIVE_THRESHOLD = 0.7 66 | RPN_IOU_NEGATIVE_THRESHOLD = 0.3 67 | TRAIN_RPN_CLOOBER_POSITIVES = False 68 | 69 | RPN_MINIBATCH_SIZE = 256 70 | RPN_POSITIVE_RATE = 0.5 71 | RPN_NMS_IOU_THRESHOLD = 0.7 72 | RPN_TOP_K_NMS_TRAIN = 12000 73 | RPN_MAXIMUM_PROPOSAL_TARIN = 2000 74 | 75 | RPN_TOP_K_NMS_TEST = 6000 # 5000 76 | RPN_MAXIMUM_PROPOSAL_TEST = 300 # 300 77 | 78 | 79 | # -------------------------------------------Fast-RCNN config 80 | ROI_SIZE = 14 81 | ROI_POOL_KERNEL_SIZE = 2 82 | USE_DROPOUT = False 83 | KEEP_PROB = 1.0 84 | SHOW_SCORE_THRSHOLD = 0.5 # only show in tensorboard 85 | 86 | FAST_RCNN_NMS_IOU_THRESHOLD = 0.3 # 0.6 87 | FAST_RCNN_NMS_MAX_BOXES_PER_CLASS = 100 88 | FAST_RCNN_IOU_POSITIVE_THRESHOLD = 0.5 89 | FAST_RCNN_IOU_NEGATIVE_THRESHOLD = 0.0 # 0.1 < IOU < 0.5 is negative 90 | FAST_RCNN_MINIBATCH_SIZE = 256 # if is -1, that is train with OHEM 91 | FAST_RCNN_POSITIVE_RATE = 0.25 92 | 93 | ADD_GTBOXES_TO_TRAIN = False 94 | -------------------------------------------------------------------------------- /utils/external/faster_rcnn_tensorflow/utility/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/utils/external/faster_rcnn_tensorflow/utility/__init__.py -------------------------------------------------------------------------------- /utils/external/faster_rcnn_tensorflow/utility/anchor_utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | from __future__ import absolute_import, print_function, division 3 | 4 | import tensorflow as tf 5 | 6 | def make_anchors(base_anchor_size, anchor_scales, anchor_ratios, 7 | featuremap_height, featuremap_width, 8 | stride, name='make_anchors'): 9 | ''' 10 | :param base_anchor_size:256 11 | :param anchor_scales: 12 | :param anchor_ratios: 13 | :param featuremap_height: 14 | :param featuremap_width: 15 | :param stride: 16 | :return: 17 | ''' 18 | with tf.variable_scope(name): 19 | base_anchor = tf.constant([0, 0, base_anchor_size, base_anchor_size], tf.float32) # [x_center, y_center, w, h] 20 | 21 | ws, hs = enum_ratios(enum_scales(base_anchor, anchor_scales), 22 | anchor_ratios) # per locations ws and hs 23 | 24 | x_centers = tf.range(featuremap_width, dtype=tf.float32) * stride 25 | y_centers = tf.range(featuremap_height, dtype=tf.float32) * stride 26 | 27 | x_centers, y_centers = tf.meshgrid(x_centers, y_centers) 28 | 29 | ws, x_centers = tf.meshgrid(ws, x_centers) 30 | hs, y_centers = tf.meshgrid(hs, y_centers) 31 | 32 | anchor_centers = tf.stack([x_centers, y_centers], 2) 33 | anchor_centers = tf.reshape(anchor_centers, [-1, 2]) 34 | 35 | box_sizes = tf.stack([ws, hs], axis=2) 36 | box_sizes = tf.reshape(box_sizes, [-1, 2]) 37 | # anchors = tf.concat([anchor_centers, box_sizes], axis=1) 38 | anchors = tf.concat([anchor_centers - 0.5*box_sizes, 39 | anchor_centers + 0.5*box_sizes], axis=1) 40 | return anchors 41 | 42 | 43 | def enum_scales(base_anchor, anchor_scales): 44 | 45 | anchor_scales = base_anchor * tf.constant(anchor_scales, dtype=tf.float32, shape=(len(anchor_scales), 1)) 46 | 47 | return anchor_scales 48 | 49 | 50 | def enum_ratios(anchors, anchor_ratios): 51 | ''' 52 | ratio = h /w 53 | :param anchors: 54 | :param anchor_ratios: 55 | :return: 56 | ''' 57 | ws = anchors[:, 2] # for base anchor: w == h 58 | hs = anchors[:, 3] 59 | sqrt_ratios = tf.sqrt(tf.constant(anchor_ratios)) 60 | 61 | ws = tf.reshape(ws / sqrt_ratios[:, tf.newaxis], [-1, 1]) 62 | hs = tf.reshape(hs * sqrt_ratios[:, tf.newaxis], [-1, 1]) 63 | 64 | return hs, ws 65 | 66 | 67 | -------------------------------------------------------------------------------- /utils/external/faster_rcnn_tensorflow/utility/boxes_utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | from __future__ import absolute_import 4 | from __future__ import division 5 | from __future__ import print_function 6 | 7 | import tensorflow as tf 8 | 9 | def ious_calu(boxes_1, boxes_2): 10 | ''' 11 | 12 | :param boxes_1: [N, 4] [xmin, ymin, xmax, ymax] 13 | :param boxes_2: [M, 4] [xmin, ymin. xmax, ymax] 14 | :return: 15 | ''' 16 | boxes_1 = tf.cast(boxes_1, tf.float32) 17 | boxes_2 = tf.cast(boxes_2, tf.float32) 18 | xmin_1, ymin_1, xmax_1, ymax_1 = tf.split(boxes_1, 4, axis=1) # xmin_1 shape is [N, 1].. 19 | xmin_2, ymin_2, xmax_2, ymax_2 = tf.unstack(boxes_2, axis=1) # xmin_2 shape is [M, ].. 20 | 21 | max_xmin = tf.maximum(xmin_1, xmin_2) 22 | min_xmax = tf.minimum(xmax_1, xmax_2) 23 | 24 | max_ymin = tf.maximum(ymin_1, ymin_2) 25 | min_ymax = tf.minimum(ymax_1, ymax_2) 26 | 27 | overlap_h = tf.maximum(0., min_ymax - max_ymin) # avoid h < 0 28 | overlap_w = tf.maximum(0., min_xmax - max_xmin) 29 | 30 | overlaps = overlap_h * overlap_w 31 | 32 | area_1 = (xmax_1 - xmin_1) * (ymax_1 - ymin_1) # [N, 1] 33 | area_2 = (xmax_2 - xmin_2) * (ymax_2 - ymin_2) # [M, ] 34 | 35 | ious = overlaps / (area_1 + area_2 - overlaps) 36 | 37 | return ious 38 | 39 | 40 | def clip_boxes_to_img_boundaries(decode_boxes, img_shape): 41 | ''' 42 | 43 | :param decode_boxes: 44 | :return: decode boxes, and already clip to boundaries 45 | ''' 46 | 47 | with tf.name_scope('clip_boxes_to_img_boundaries'): 48 | 49 | # xmin, ymin, xmax, ymax = tf.unstack(decode_boxes, axis=1) 50 | xmin = decode_boxes[:, 0] 51 | ymin = decode_boxes[:, 1] 52 | xmax = decode_boxes[:, 2] 53 | ymax = decode_boxes[:, 3] 54 | img_h, img_w = img_shape[1], img_shape[2] 55 | 56 | img_h, img_w = tf.cast(img_h, tf.float32), tf.cast(img_w, tf.float32) 57 | 58 | xmin = tf.maximum(tf.minimum(xmin, img_w-1.), 0.) 59 | ymin = tf.maximum(tf.minimum(ymin, img_h-1.), 0.) 60 | 61 | xmax = tf.maximum(tf.minimum(xmax, img_w-1.), 0.) 62 | ymax = tf.maximum(tf.minimum(ymax, img_h-1.), 0.) 63 | 64 | return tf.transpose(tf.stack([xmin, ymin, xmax, ymax])) 65 | 66 | 67 | def filter_outside_boxes(boxes, img_h, img_w): 68 | ''' 69 | :param anchors:boxes with format [xmin, ymin, xmax, ymax] 70 | :param img_h: height of image 71 | :param img_w: width of image 72 | :return: indices of anchors that inside the image boundary 73 | ''' 74 | 75 | with tf.name_scope('filter_outside_boxes'): 76 | xmin, ymin, xmax, ymax = tf.unstack(boxes, axis=1) 77 | 78 | xmin_index = tf.greater_equal(xmin, 0) 79 | ymin_index = tf.greater_equal(ymin, 0) 80 | xmax_index = tf.less_equal(xmax, tf.cast(img_w, tf.float32)) 81 | ymax_index = tf.less_equal(ymax, tf.cast(img_h, tf.float32)) 82 | 83 | indices = tf.transpose(tf.stack([xmin_index, ymin_index, xmax_index, ymax_index])) 84 | indices = tf.cast(indices, dtype=tf.int32) 85 | indices = tf.reduce_sum(indices, axis=1) 86 | indices = tf.where(tf.equal(indices, 4)) 87 | # indices = tf.equal(indices, 4) 88 | return tf.reshape(indices, [-1]) 89 | 90 | 91 | def padd_boxes_with_zeros(boxes, scores, max_num_of_boxes): 92 | 93 | ''' 94 | num of boxes less than max num of boxes, so it need to pad with zeros[0, 0, 0, 0] 95 | :param boxes: 96 | :param scores: [-1] 97 | :param max_num_of_boxes: 98 | :return: 99 | ''' 100 | 101 | pad_num = tf.cast(max_num_of_boxes, tf.int32) - tf.shape(boxes)[0] 102 | 103 | zero_boxes = tf.zeros(shape=[pad_num, 4], dtype=boxes.dtype) 104 | zero_scores = tf.zeros(shape=[pad_num], dtype=scores.dtype) 105 | 106 | final_boxes = tf.concat([boxes, zero_boxes], axis=0) 107 | 108 | final_scores = tf.concat([scores, zero_scores], axis=0) 109 | 110 | return final_boxes, final_scores -------------------------------------------------------------------------------- /utils/external/faster_rcnn_tensorflow/utility/coco_dict.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | from __future__ import absolute_import, print_function, division 4 | 5 | class_names = [ 6 | 'back_ground', 'person', 'bicycle', 'car', 'motorcycle', 7 | 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 8 | 'fire hydrant', 'stop sign', 'parking meter', 'bench', 9 | 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 10 | 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 11 | 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 12 | 'sports ball', 'kite', 'baseball bat', 'baseball glove', 13 | 'skateboard', 'surfboard', 'tennis racket', 'bottle', 14 | 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 15 | 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 16 | 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 17 | 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 18 | 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 19 | 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 20 | 'book', 'clock', 'vase', 'scissors', 'teddy bear', 21 | 'hair drier', 'toothbrush'] 22 | 23 | 24 | classes_originID = { 25 | 'person': 1, 'bicycle': 2, 'car': 3, 'motorcycle': 4, 26 | 'airplane': 5, 'bus': 6, 'train': 7, 'truck': 8, 'boat': 9, 27 | 'traffic light': 10, 'fire hydrant': 11, 'stop sign': 13, 28 | 'parking meter': 14, 'bench': 15, 'bird': 16, 'cat': 17, 29 | 'dog': 18, 'horse': 19, 'sheep': 20, 'cow': 21, 'elephant': 22, 30 | 'bear': 23, 'zebra': 24, 'giraffe': 25, 'backpack': 27, 31 | 'umbrella': 28, 'handbag': 31, 'tie': 32, 'suitcase': 33, 32 | 'frisbee': 34, 'skis': 35, 'snowboard': 36, 'sports ball': 37, 33 | 'kite': 38, 'baseball bat': 39, 'baseball glove': 40, 34 | 'skateboard': 41, 'surfboard': 42, 'tennis racket': 43, 35 | 'bottle': 44, 'wine glass': 46, 'cup': 47, 'fork': 48, 36 | 'knife': 49, 'spoon': 50, 'bowl': 51, 'banana': 52, 'apple': 53, 37 | 'sandwich': 54, 'orange': 55, 'broccoli': 56, 'carrot': 57, 38 | 'hot dog': 58, 'pizza': 59, 'donut': 60, 'cake': 61, 39 | 'chair': 62, 'couch': 63, 'potted plant': 64, 'bed': 65, 40 | 'dining table': 67, 'toilet': 70, 'tv': 72, 'laptop': 73, 41 | 'mouse': 74, 'remote': 75, 'keyboard': 76, 'cell phone': 77, 42 | 'microwave': 78, 'oven': 79, 'toaster': 80, 'sink': 81, 43 | 'refrigerator': 82, 'book': 84, 'clock': 85, 'vase': 86, 44 | 'scissors': 87, 'teddy bear': 88, 'hair drier': 89, 45 | 'toothbrush': 90} 46 | 47 | originID_classes = {item: key for key, item in classes_originID.items()} 48 | NAME_LABEL_MAP = dict(zip(class_names, range(len(class_names)))) 49 | LABEL_NAME_MAP = dict(zip(range(len(class_names)), class_names)) 50 | 51 | # print (originID_classes) 52 | 53 | 54 | 55 | -------------------------------------------------------------------------------- /utils/external/faster_rcnn_tensorflow/utility/encode_and_decode.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | from __future__ import absolute_import 4 | from __future__ import print_function 5 | from __future__ import division 6 | 7 | import tensorflow as tf 8 | import numpy as np 9 | 10 | 11 | def decode_boxes(encoded_boxes, reference_boxes, scale_factors=None): 12 | ''' 13 | 14 | :param encoded_boxes:[N, 4] 15 | :param reference_boxes: [N, 4] . 16 | :param scale_factors: use for scale. 17 | 18 | in the first stage, reference_boxes are anchors 19 | in the second stage, reference boxes are proposals(decode) produced by first stage 20 | :return:decode boxes [N, 4] 21 | ''' 22 | 23 | t_xcenter, t_ycenter, t_w, t_h = tf.unstack(encoded_boxes, axis=1) 24 | if scale_factors: 25 | t_xcenter /= scale_factors[0] 26 | t_ycenter /= scale_factors[1] 27 | t_w /= scale_factors[2] 28 | t_h /= scale_factors[3] 29 | 30 | reference_xmin, reference_ymin, reference_xmax, reference_ymax = tf.unstack(reference_boxes, axis=1) 31 | # reference boxes are anchors in the first stage 32 | 33 | # reference_xcenter = (reference_xmin + reference_xmax) / 2. 34 | # reference_ycenter = (reference_ymin + reference_ymax) / 2. 35 | reference_w = reference_xmax - reference_xmin 36 | reference_h = reference_ymax - reference_ymin 37 | reference_xcenter = reference_xmin + reference_w/2.0 38 | reference_ycenter = reference_ymin + reference_h/2.0 39 | 40 | predict_xcenter = t_xcenter * reference_w + reference_xcenter 41 | predict_ycenter = t_ycenter * reference_h + reference_ycenter 42 | predict_w = tf.exp(t_w) * reference_w 43 | predict_h = tf.exp(t_h) * reference_h 44 | 45 | predict_xmin = predict_xcenter - predict_w / 2. 46 | predict_xmax = predict_xcenter + predict_w / 2. 47 | predict_ymin = predict_ycenter - predict_h / 2. 48 | predict_ymax = predict_ycenter + predict_h / 2. 49 | 50 | return tf.transpose(tf.stack([predict_xmin, predict_ymin, 51 | predict_xmax, predict_ymax])) 52 | 53 | 54 | def encode_boxes(unencode_boxes, reference_boxes, scale_factors=None): 55 | ''' 56 | 57 | :param unencode_boxes: [-1, 4] 58 | :param reference_boxes: [-1, 4] 59 | :return: encode_boxes [-1, 4] 60 | ''' 61 | 62 | xmin, ymin, xmax, ymax = unencode_boxes[:, 0], unencode_boxes[:, 1], unencode_boxes[:, 2], unencode_boxes[:, 3] 63 | 64 | reference_xmin, reference_ymin, reference_xmax, reference_ymax = reference_boxes[:, 0], reference_boxes[:, 1], \ 65 | reference_boxes[:, 2], reference_boxes[:, 3] 66 | 67 | # x_center = (xmin + xmax) / 2. 68 | # y_center = (ymin + ymax) / 2. 69 | w = xmax - xmin + 1e-8 70 | h = ymax - ymin + 1e-8 71 | x_center = xmin + w/2.0 72 | y_center = ymin + h/2.0 73 | 74 | # reference_xcenter = (reference_xmin + reference_xmax) / 2. 75 | # reference_ycenter = (reference_ymin + reference_ymax) / 2. 76 | reference_w = reference_xmax - reference_xmin + 1e-8 77 | reference_h = reference_ymax - reference_ymin + 1e-8 78 | reference_xcenter = reference_xmin + reference_w/2.0 79 | reference_ycenter = reference_ymin + reference_h/2.0 80 | # w + 1e-8 to avoid NaN in division and log below 81 | 82 | t_xcenter = (x_center - reference_xcenter) / reference_w 83 | t_ycenter = (y_center - reference_ycenter) / reference_h 84 | t_w = np.log(w/reference_w) 85 | t_h = np.log(h/reference_h) 86 | 87 | if scale_factors: 88 | t_xcenter *= scale_factors[0] 89 | t_ycenter *= scale_factors[1] 90 | t_w *= scale_factors[2] 91 | t_h *= scale_factors[3] 92 | 93 | return np.transpose(np.stack([t_xcenter, t_ycenter, t_w, t_h], axis=0)) 94 | -------------------------------------------------------------------------------- /utils/external/faster_rcnn_tensorflow/utility/label_dict.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | from __future__ import division, print_function, absolute_import 3 | 4 | from utils.external.faster_rcnn_tensorflow.configs import cfgs 5 | 6 | if cfgs.DATASET_NAME == 'ship': 7 | NAME_LABEL_MAP = { 8 | 'back_ground': 0, 9 | 'ship': 1 10 | } 11 | elif cfgs.DATASET_NAME == 'FDDB': 12 | NAME_LABEL_MAP = { 13 | 'back_ground': 0, 14 | 'face': 1 15 | } 16 | elif cfgs.DATASET_NAME == 'icdar': 17 | NAME_LABEL_MAP = { 18 | 'back_ground': 0, 19 | 'text': 1 20 | } 21 | elif cfgs.DATASET_NAME.startswith('DOTA'): 22 | NAME_LABEL_MAP = { 23 | 'back_ground': 0, 24 | 'roundabout': 1, 25 | 'tennis-court': 2, 26 | 'swimming-pool': 3, 27 | 'storage-tank': 4, 28 | 'soccer-ball-field': 5, 29 | 'small-vehicle': 6, 30 | 'ship': 7, 31 | 'plane': 8, 32 | 'large-vehicle': 9, 33 | 'helicopter': 10, 34 | 'harbor': 11, 35 | 'ground-track-field': 12, 36 | 'bridge': 13, 37 | 'basketball-court': 14, 38 | 'baseball-diamond': 15 39 | } 40 | elif cfgs.DATASET_NAME == 'pascal': 41 | NAME_LABEL_MAP = { 42 | 'back_ground': 0, 43 | 'aeroplane': 1, 44 | 'bicycle': 2, 45 | 'bird': 3, 46 | 'boat': 4, 47 | 'bottle': 5, 48 | 'bus': 6, 49 | 'car': 7, 50 | 'cat': 8, 51 | 'chair': 9, 52 | 'cow': 10, 53 | 'diningtable': 11, 54 | 'dog': 12, 55 | 'horse': 13, 56 | 'motorbike': 14, 57 | 'person': 15, 58 | 'pottedplant': 16, 59 | 'sheep': 17, 60 | 'sofa': 18, 61 | 'train': 19, 62 | 'tvmonitor': 20 63 | } 64 | else: 65 | assert 'please set label dict!' 66 | 67 | 68 | def get_label_name_map(): 69 | reverse_dict = {} 70 | for name, label in NAME_LABEL_MAP.items(): 71 | reverse_dict[label] = name 72 | return reverse_dict 73 | 74 | LABEL_NAME_MAP = get_label_name_map() -------------------------------------------------------------------------------- /utils/external/faster_rcnn_tensorflow/utility/loss_utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | @author: jemmy li 4 | @contact: zengarden2009@gmail.com 5 | """ 6 | from __future__ import absolute_import 7 | from __future__ import division 8 | from __future__ import print_function 9 | 10 | import tensorflow as tf 11 | 12 | 13 | def _smooth_l1_loss_base(bbox_pred, bbox_targets, sigma=1.0): 14 | ''' 15 | 16 | :param bbox_pred: [-1, 4] in RPN. [-1, cls_num+1, 4] in Fast-rcnn 17 | :param bbox_targets: shape is same as bbox_pred 18 | :param sigma: 19 | :return: 20 | ''' 21 | sigma_2 = sigma**2 22 | 23 | box_diff = bbox_pred - bbox_targets 24 | 25 | abs_box_diff = tf.abs(box_diff) 26 | 27 | smoothL1_sign = tf.stop_gradient( 28 | tf.to_float(tf.less(abs_box_diff, 1. / sigma_2))) 29 | loss_box = tf.pow(box_diff, 2) * (sigma_2 / 2.0) * smoothL1_sign \ 30 | + (abs_box_diff - (0.5 / sigma_2)) * (1.0 - smoothL1_sign) 31 | return loss_box 32 | 33 | def smooth_l1_loss_rpn(bbox_pred, bbox_targets, label, sigma=1.0): 34 | ''' 35 | 36 | :param bbox_pred: [-1, 4] 37 | :param bbox_targets: [-1, 4] 38 | :param label: [-1] 39 | :param sigma: 40 | :return: 41 | ''' 42 | value = _smooth_l1_loss_base(bbox_pred, bbox_targets, sigma=sigma) 43 | value = tf.reduce_sum(value, axis=1) # to sum in axis 1 44 | rpn_select = tf.where(tf.greater(label, 0)) 45 | 46 | # rpn_select = tf.stop_gradient(rpn_select) # to avoid 47 | selected_value = tf.gather(value, rpn_select) 48 | non_ignored_mask = tf.stop_gradient( 49 | 1.0 - tf.to_float(tf.equal(label, -1))) # positve is 1.0 others is 0.0 50 | 51 | bbox_loss = tf.reduce_sum(selected_value) / tf.maximum(1.0, tf.reduce_sum(non_ignored_mask)) 52 | 53 | return bbox_loss 54 | 55 | 56 | 57 | def smooth_l1_loss_rcnn(bbox_pred, bbox_targets, label, num_classes, sigma=1.0): 58 | ''' 59 | 60 | :param bbox_pred: [-1, (cfgs.CLS_NUM +1) * 4] 61 | :param bbox_targets:[-1, (cfgs.CLS_NUM +1) * 4] 62 | :param label:[-1] 63 | :param num_classes: 64 | :param sigma: 65 | :return: 66 | ''' 67 | 68 | outside_mask = tf.stop_gradient(tf.to_float(tf.greater(label, 0))) 69 | 70 | bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 4]) 71 | bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 4]) 72 | 73 | value = _smooth_l1_loss_base(bbox_pred, 74 | bbox_targets, 75 | sigma=sigma) 76 | value = tf.reduce_sum(value, 2) 77 | value = tf.reshape(value, [-1, num_classes]) 78 | 79 | inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]), 80 | depth=num_classes, axis=1) 81 | 82 | inside_mask = tf.stop_gradient( 83 | tf.to_float(tf.reshape(inside_mask, [-1, num_classes]))) 84 | 85 | normalizer = tf.to_float(tf.shape(bbox_pred)[0]) 86 | bbox_loss = tf.reduce_sum( 87 | tf.reduce_sum(value * inside_mask, 1)*outside_mask) / normalizer 88 | 89 | return bbox_loss 90 | 91 | 92 | def sum_ohem_loss(cls_score, label, bbox_pred, bbox_targets, 93 | num_classes, num_ohem_samples=256, sigma=1.0): 94 | ''' 95 | 96 | :param cls_score: [-1, cls_num+1] 97 | :param label: [-1] 98 | :param bbox_pred: [-1, 4*(cls_num+1)] 99 | :param bbox_targets: [-1, 4*(cls_num+1)] 100 | :param num_ohem_samples: 256 by default 101 | :param num_classes: cls_num+1 102 | :param sigma: 103 | :return: 104 | ''' 105 | 106 | cls_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=cls_score, labels=label) # [-1, ] 107 | # cls_loss = tf.Print(cls_loss, [tf.shape(cls_loss)], summarize=10, message='CLS losss shape ****') 108 | 109 | outside_mask = tf.stop_gradient(tf.to_float(tf.greater(label, 0))) 110 | bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 4]) 111 | bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 4]) 112 | 113 | value = _smooth_l1_loss_base(bbox_pred, 114 | bbox_targets, 115 | sigma=sigma) 116 | value = tf.reduce_sum(value, 2) 117 | value = tf.reshape(value, [-1, num_classes]) 118 | 119 | inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]), 120 | depth=num_classes, axis=1) 121 | 122 | inside_mask = tf.stop_gradient( 123 | tf.to_float(tf.reshape(inside_mask, [-1, num_classes]))) 124 | loc_loss = tf.reduce_sum(value * inside_mask, 1)*outside_mask 125 | # loc_loss = tf.Print(loc_loss, [tf.shape(loc_loss)], summarize=10, message='loc_loss shape***') 126 | 127 | sum_loss = cls_loss + loc_loss 128 | 129 | num_ohem_samples = tf.stop_gradient(tf.minimum(num_ohem_samples, tf.shape(sum_loss)[0])) 130 | _, top_k_indices = tf.nn.top_k(sum_loss, k=num_ohem_samples) 131 | 132 | cls_loss_ohem = tf.gather(cls_loss, top_k_indices) 133 | cls_loss_ohem = tf.reduce_mean(cls_loss_ohem) 134 | 135 | loc_loss_ohem = tf.gather(loc_loss, top_k_indices) 136 | normalizer = tf.to_float(num_ohem_samples) 137 | loc_loss_ohem = tf.reduce_sum(loc_loss_ohem) / normalizer 138 | 139 | return cls_loss_ohem, loc_loss_ohem 140 | 141 | -------------------------------------------------------------------------------- /utils/external/faster_rcnn_tensorflow/utility/proposal_opr.py: -------------------------------------------------------------------------------- 1 | # encoding: utf-8 2 | """ 3 | @author: zeming li 4 | @contact: zengarden2009@gmail.com 5 | """ 6 | 7 | from utils.external.faster_rcnn_tensorflow.configs import cfgs 8 | from utils.external.faster_rcnn_tensorflow.utility import encode_and_decode 9 | from utils.external.faster_rcnn_tensorflow.utility import boxes_utils 10 | import tensorflow as tf 11 | import numpy as np 12 | 13 | 14 | def postprocess_rpn_proposals(rpn_bbox_pred, rpn_cls_prob, img_shape, anchors, is_training): 15 | ''' 16 | 17 | :param rpn_bbox_pred: [-1, 4] 18 | :param rpn_cls_prob: [-1, 2] 19 | :param img_shape: 20 | :param anchors:[-1, 4] 21 | :param is_training: 22 | :return: 23 | ''' 24 | 25 | if is_training: 26 | pre_nms_topN = cfgs.RPN_TOP_K_NMS_TRAIN 27 | post_nms_topN = cfgs.RPN_MAXIMUM_PROPOSAL_TARIN 28 | nms_thresh = cfgs.RPN_NMS_IOU_THRESHOLD 29 | else: 30 | pre_nms_topN = cfgs.RPN_TOP_K_NMS_TEST 31 | post_nms_topN = cfgs.RPN_MAXIMUM_PROPOSAL_TEST 32 | nms_thresh = cfgs.RPN_NMS_IOU_THRESHOLD 33 | 34 | cls_prob = rpn_cls_prob[:, 1] 35 | 36 | # 1. decode boxes 37 | decode_boxes = encode_and_decode.decode_boxes(encoded_boxes=rpn_bbox_pred, 38 | reference_boxes=anchors, 39 | scale_factors=cfgs.ANCHOR_SCALE_FACTORS) 40 | 41 | # decode_boxes = encode_and_decode.decode_boxes(boxes=anchors, 42 | # deltas=rpn_bbox_pred, 43 | # scale_factor=None) 44 | 45 | # 2. clip to img boundaries 46 | decode_boxes = boxes_utils.clip_boxes_to_img_boundaries(decode_boxes=decode_boxes, 47 | img_shape=img_shape) 48 | 49 | # 3. get top N to NMS 50 | if pre_nms_topN > 0: 51 | pre_nms_topN = tf.minimum(pre_nms_topN, tf.shape(decode_boxes)[0], name='avoid_unenough_boxes') 52 | cls_prob, top_k_indices = tf.nn.top_k(cls_prob, k=pre_nms_topN) 53 | decode_boxes = tf.gather(decode_boxes, top_k_indices) 54 | 55 | # 4. NMS 56 | keep = tf.image.non_max_suppression( 57 | boxes=decode_boxes, 58 | scores=cls_prob, 59 | max_output_size=post_nms_topN, 60 | iou_threshold=nms_thresh) 61 | 62 | final_boxes = tf.gather(decode_boxes, keep) 63 | final_probs = tf.gather(cls_prob, keep) 64 | 65 | return final_boxes, final_probs 66 | 67 | -------------------------------------------------------------------------------- /utils/external/faster_rcnn_tensorflow/utility/remote_sensing_dict.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | NAME_LABEL_MAP = { 4 | 'back_ground': 0, 5 | 'building': 1 6 | } 7 | 8 | 9 | def get_label_name_map(): 10 | reverse_dict = {} 11 | for name, label in NAME_LABEL_MAP.items(): 12 | reverse_dict[label] = name 13 | return reverse_dict 14 | 15 | LABEL_NAME_MAP = get_label_name_map() -------------------------------------------------------------------------------- /utils/external/faster_rcnn_tensorflow/utility/show_box_in_tensor.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | from __future__ import absolute_import 4 | from __future__ import division 5 | from __future__ import print_function 6 | 7 | import tensorflow as tf 8 | 9 | from utils.external.faster_rcnn_tensorflow.utility import draw_box_in_img 10 | 11 | def only_draw_boxes(img_batch, boxes): 12 | 13 | boxes = tf.stop_gradient(boxes) 14 | img_tensor = tf.squeeze(img_batch, 0) 15 | img_tensor = tf.cast(img_tensor, tf.float32) 16 | labels = tf.ones(shape=(tf.shape(boxes)[0], ), dtype=tf.int32) * draw_box_in_img.ONLY_DRAW_BOXES 17 | scores = tf.zeros_like(labels, dtype=tf.float32) 18 | img_tensor_with_boxes = tf.py_func(draw_box_in_img.draw_boxes_with_label_and_scores, 19 | inp=[img_tensor, boxes, labels, scores], 20 | Tout=tf.uint8) 21 | img_tensor_with_boxes = tf.reshape(img_tensor_with_boxes, tf.shape(img_batch)) # [batch_size, h, w, c] 22 | 23 | return img_tensor_with_boxes 24 | 25 | def draw_boxes_with_scores(img_batch, boxes, scores): 26 | 27 | boxes = tf.stop_gradient(boxes) 28 | scores = tf.stop_gradient(scores) 29 | 30 | img_tensor = tf.squeeze(img_batch, 0) 31 | img_tensor = tf.cast(img_tensor, tf.float32) 32 | labels = tf.ones(shape=(tf.shape(boxes)[0],), dtype=tf.int32) * draw_box_in_img.ONLY_DRAW_BOXES_WITH_SCORES 33 | img_tensor_with_boxes = tf.py_func(draw_box_in_img.draw_boxes_with_label_and_scores, 34 | inp=[img_tensor, boxes, labels, scores], 35 | Tout=[tf.uint8]) 36 | img_tensor_with_boxes = tf.reshape(img_tensor_with_boxes, tf.shape(img_batch)) 37 | return img_tensor_with_boxes 38 | 39 | def draw_boxes_with_categories(img_batch, boxes, labels): 40 | boxes = tf.stop_gradient(boxes) 41 | 42 | img_tensor = tf.squeeze(img_batch, 0) 43 | img_tensor = tf.cast(img_tensor, tf.float32) 44 | scores = tf.ones(shape=(tf.shape(boxes)[0],), dtype=tf.float32) 45 | img_tensor_with_boxes = tf.py_func(draw_box_in_img.draw_boxes_with_label_and_scores, 46 | inp=[img_tensor, boxes, labels, scores], 47 | Tout=[tf.uint8]) 48 | img_tensor_with_boxes = tf.reshape(img_tensor_with_boxes, tf.shape(img_batch)) 49 | return img_tensor_with_boxes 50 | 51 | def draw_boxes_with_categories_and_scores(img_batch, boxes, labels, scores): 52 | boxes = tf.stop_gradient(boxes) 53 | scores = tf.stop_gradient(scores) 54 | 55 | img_tensor = tf.squeeze(img_batch, 0) 56 | img_tensor = tf.cast(img_tensor, tf.float32) 57 | img_tensor_with_boxes = tf.py_func(draw_box_in_img.draw_boxes_with_label_and_scores, 58 | inp=[img_tensor, boxes, labels, scores], 59 | Tout=[tf.uint8]) 60 | img_tensor_with_boxes = tf.reshape(img_tensor_with_boxes, tf.shape(img_batch)) 61 | return img_tensor_with_boxes 62 | 63 | if __name__ == "__main__": 64 | print (1) 65 | 66 | -------------------------------------------------------------------------------- /utils/external/faster_rcnn_tensorflow/utility/tf_ops.py: -------------------------------------------------------------------------------- 1 | # -*- coding:utf-8 -*- 2 | 3 | from __future__ import absolute_import, print_function, division 4 | 5 | import tensorflow as tf 6 | 7 | ''' 8 | all of these ops are derived from tenosrflow Object Detection API 9 | ''' 10 | def indices_to_dense_vector(indices, 11 | size, 12 | indices_value=1., 13 | default_value=0, 14 | dtype=tf.float32): 15 | """Creates dense vector with indices set to specific (the para "indices_value" ) and rest to zeros. 16 | 17 | This function exists because it is unclear if it is safe to use 18 | tf.sparse_to_dense(indices, [size], 1, validate_indices=False) 19 | with indices which are not ordered. 20 | This function accepts a dynamic size (e.g. tf.shape(tensor)[0]) 21 | 22 | Args: 23 | indices: 1d Tensor with integer indices which are to be set to 24 | indices_values. 25 | size: scalar with size (integer) of output Tensor. 26 | indices_value: values of elements specified by indices in the output vector 27 | default_value: values of other elements in the output vector. 28 | dtype: data type. 29 | 30 | Returns: 31 | dense 1D Tensor of shape [size] with indices set to indices_values and the 32 | rest set to default_value. 33 | """ 34 | size = tf.to_int32(size) 35 | zeros = tf.ones([size], dtype=dtype) * default_value 36 | values = tf.ones_like(indices, dtype=dtype) * indices_value 37 | 38 | return tf.dynamic_stitch([tf.range(size), tf.to_int32(indices)], 39 | [zeros, values]) -------------------------------------------------------------------------------- /utils/external/ssd_tensorflow/dataset/dataset_inspect.py: -------------------------------------------------------------------------------- 1 | # Copyright 2018 Changan Wang 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 | from __future__ import absolute_import 16 | from __future__ import division 17 | from __future__ import print_function 18 | 19 | import os 20 | 21 | import tensorflow as tf 22 | 23 | def count_split_examples(split_path, file_prefix='.tfrecord'): 24 | # Count the total number of examples in all of these shard 25 | num_samples = 0 26 | tfrecords_to_count = tf.gfile.Glob(os.path.join(split_path, file_prefix)) 27 | opts = tf.python_io.TFRecordOptions(tf.python_io.TFRecordCompressionType.ZLIB) 28 | for tfrecord_file in tfrecords_to_count: 29 | for record in tf.python_io.tf_record_iterator(tfrecord_file):#, options = opts): 30 | num_samples += 1 31 | return num_samples 32 | 33 | if __name__ == '__main__': 34 | print('train:', count_split_examples('/media/rs/7A0EE8880EE83EAF/Detections/SSD/dataset/tfrecords', 'train-?????-of-?????')) 35 | print('val:', count_split_examples('/media/rs/7A0EE8880EE83EAF/Detections/SSD/dataset/tfrecords', 'val-?????-of-?????')) 36 | -------------------------------------------------------------------------------- /utils/external/ssd_tensorflow/demo/demo1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/utils/external/ssd_tensorflow/demo/demo1.jpg -------------------------------------------------------------------------------- /utils/external/ssd_tensorflow/demo/demo2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/utils/external/ssd_tensorflow/demo/demo2.jpg -------------------------------------------------------------------------------- /utils/external/ssd_tensorflow/demo/demo3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tencent/PocketFlow/53b82cba5a34834400619e7c335a23995d45c2a6/utils/external/ssd_tensorflow/demo/demo3.jpg -------------------------------------------------------------------------------- /utils/external/ssd_tensorflow/utility/checkpint_inspect.py: -------------------------------------------------------------------------------- 1 | # Copyright 2018 Changan Wang 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 | from __future__ import absolute_import 16 | from __future__ import division 17 | from __future__ import print_function 18 | 19 | import numpy as np 20 | 21 | from tensorflow.python import pywrap_tensorflow 22 | 23 | def print_tensors_in_checkpoint_file(file_name, tensor_name, all_tensors): 24 | try: 25 | reader = pywrap_tensorflow.NewCheckpointReader(file_name) 26 | if all_tensors: 27 | var_to_shape_map = reader.get_variable_to_shape_map() 28 | for key in var_to_shape_map: 29 | print("tensor_name: ", key) 30 | print(reader.get_tensor(key)) 31 | elif not tensor_name: 32 | print(reader.debug_string().decode("utf-8")) 33 | else: 34 | print("tensor_name: ", tensor_name) 35 | print(reader.get_tensor(tensor_name)) 36 | except Exception as e: # pylint: disable=broad-except 37 | print(str(e)) 38 | if "corrupted compressed block contents" in str(e): 39 | print("It's likely that your checkpoint file has been compressed " 40 | "with SNAPPY.") 41 | 42 | def print_all_tensors_name(file_name): 43 | try: 44 | reader = pywrap_tensorflow.NewCheckpointReader(file_name) 45 | var_to_shape_map = reader.get_variable_to_shape_map() 46 | for key in var_to_shape_map: 47 | print(key) 48 | except Exception as e: # pylint: disable=broad-except 49 | print(str(e)) 50 | if "corrupted compressed block contents" in str(e): 51 | print("It's likely that your checkpoint file has been compressed " 52 | "with SNAPPY.") 53 | 54 | if __name__ == "__main__": 55 | print_all_tensors_name('./model/vgg16_reducedfc.ckpt') 56 | -------------------------------------------------------------------------------- /utils/external/ssd_tensorflow/utility/draw_toolbox.py: -------------------------------------------------------------------------------- 1 | # Copyright 2018 Changan Wang 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 cv2 16 | import matplotlib.cm as mpcm 17 | 18 | from dataset import dataset_common 19 | 20 | def gain_translate_table(): 21 | label2name_table = {} 22 | for class_name, labels_pair in dataset_common.VOC_LABELS.items(): 23 | label2name_table[labels_pair[0]] = class_name 24 | return label2name_table 25 | 26 | label2name_table = gain_translate_table() 27 | 28 | def colors_subselect(colors, num_classes=21): 29 | dt = len(colors) // num_classes 30 | sub_colors = [] 31 | for i in range(num_classes): 32 | color = colors[i*dt] 33 | if isinstance(color[0], float): 34 | sub_colors.append([int(c * 255) for c in color]) 35 | else: 36 | sub_colors.append([c for c in color]) 37 | return sub_colors 38 | 39 | colors = colors_subselect(mpcm.plasma.colors, num_classes=21) 40 | colors_tableau = [(255, 255, 255), (31, 119, 180), (174, 199, 232), (255, 127, 14), (255, 187, 120), 41 | (44, 160, 44), (152, 223, 138), (214, 39, 40), (255, 152, 150), 42 | (148, 103, 189), (197, 176, 213), (140, 86, 75), (196, 156, 148), 43 | (227, 119, 194), (247, 182, 210), (127, 127, 127), (199, 199, 199), 44 | (188, 189, 34), (219, 219, 141), (23, 190, 207), (158, 218, 229)] 45 | 46 | def bboxes_draw_on_img(img, classes, scores, bboxes, thickness=2): 47 | shape = img.shape 48 | scale = 0.4 49 | text_thickness = 1 50 | line_type = 8 51 | for i in range(bboxes.shape[0]): 52 | if classes[i] < 1: continue 53 | bbox = bboxes[i] 54 | color = colors_tableau[classes[i]] 55 | # Draw bounding boxes 56 | p1 = (int(bbox[0] * shape[0]), int(bbox[1] * shape[1])) 57 | p2 = (int(bbox[2] * shape[0]), int(bbox[3] * shape[1])) 58 | if (p2[0] - p1[0] < 1) or (p2[1] - p1[1] < 1): 59 | continue 60 | 61 | cv2.rectangle(img, p1[::-1], p2[::-1], color, thickness) 62 | # Draw text 63 | s = '%s/%.1f%%' % (label2name_table[classes[i]], scores[i]*100) 64 | # text_size is (width, height) 65 | text_size, baseline = cv2.getTextSize(s, cv2.FONT_HERSHEY_SIMPLEX, scale, text_thickness) 66 | p1 = (p1[0] - text_size[1], p1[1]) 67 | 68 | cv2.rectangle(img, (p1[1] - thickness//2, p1[0] - thickness - baseline), (p1[1] + text_size[0], p1[0] + text_size[1]), color, -1) 69 | 70 | cv2.putText(img, s, (p1[1], p1[0] + baseline), cv2.FONT_HERSHEY_SIMPLEX, scale, (255,255,255), text_thickness, line_type) 71 | 72 | return img 73 | 74 | -------------------------------------------------------------------------------- /utils/external/ssd_tensorflow/utility/scaffolds.py: -------------------------------------------------------------------------------- 1 | # Copyright 2018 Changan Wang 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 | from __future__ import absolute_import 16 | from __future__ import division 17 | from __future__ import print_function 18 | 19 | import os 20 | import sys 21 | 22 | import tensorflow as tf 23 | 24 | def get_init_fn_for_scaffold(model_dir, checkpoint_path, model_scope, checkpoint_model_scope, checkpoint_exclude_scopes, ignore_missing_vars, name_remap=None): 25 | if tf.train.latest_checkpoint(model_dir): 26 | tf.logging.info('Ignoring --checkpoint_path because a checkpoint already exists in %s.' % model_dir) 27 | return None 28 | exclusion_scopes = [] 29 | if checkpoint_exclude_scopes: 30 | exclusion_scopes = [scope.strip() for scope in checkpoint_exclude_scopes.split(',')] 31 | 32 | variables_to_restore = [] 33 | for var in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES): 34 | excluded = False 35 | for exclusion in exclusion_scopes: 36 | if exclusion in var.op.name:#.startswith(exclusion): 37 | excluded = True 38 | break 39 | if not excluded: 40 | variables_to_restore.append(var) 41 | if checkpoint_model_scope is not None: 42 | if checkpoint_model_scope.strip() == '': 43 | variables_to_restore = {var.op.name.replace(model_scope + '/', ''): var for var in variables_to_restore} 44 | else: 45 | variables_to_restore = {var.op.name.replace(model_scope, checkpoint_model_scope.strip()): var for var in variables_to_restore} 46 | if name_remap is not None: 47 | renamed_variables_to_restore = dict() 48 | for var_name, var in variables_to_restore.items(): 49 | found = False 50 | for k, v in name_remap.items(): 51 | if k in var_name: 52 | renamed_variables_to_restore[var_name.replace(k, v)] = var 53 | found = True 54 | break 55 | if not found: 56 | renamed_variables_to_restore[var_name] = var 57 | variables_to_restore = renamed_variables_to_restore 58 | 59 | checkpoint_path = tf.train.latest_checkpoint(checkpoint_path) if tf.gfile.IsDirectory(checkpoint_path) else checkpoint_path 60 | 61 | tf.logging.info('Fine-tuning from %s. Ignoring missing vars: %s.' % (checkpoint_path, ignore_missing_vars)) 62 | 63 | if not variables_to_restore: 64 | raise ValueError('variables_to_restore cannot be empty') 65 | if ignore_missing_vars: 66 | reader = tf.train.NewCheckpointReader(checkpoint_path) 67 | if isinstance(variables_to_restore, dict): 68 | var_dict = variables_to_restore 69 | else: 70 | var_dict = {var.op.name: var for var in variables_to_restore} 71 | available_vars = {} 72 | for var in var_dict: 73 | if reader.has_tensor(var): 74 | available_vars[var] = var_dict[var] 75 | else: 76 | tf.logging.warning('Variable %s missing in checkpoint %s.', var, checkpoint_path) 77 | variables_to_restore = available_vars 78 | if variables_to_restore: 79 | saver = tf.train.Saver(variables_to_restore, reshape=False) 80 | saver.build() 81 | def callback(scaffold, session): 82 | saver.restore(session, checkpoint_path) 83 | return callback 84 | else: 85 | tf.logging.warning('No Variables to restore.') 86 | return None 87 | -------------------------------------------------------------------------------- /utils/get_data_dir.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Get data directory on the local machine.""" 18 | 19 | import re 20 | import sys 21 | 22 | # get the Python script's & path configuration file's path 23 | assert len(sys.argv) == 3 24 | py_file = sys.argv[1] 25 | conf_file = sys.argv[2] 26 | 27 | # obtain the dataset's name 28 | pattern = re.compile(r'at_[0-9A-Za-z]+_run.py$') 29 | match = re.search(pattern, py_file) 30 | assert match is not None, 'unable to match pattern in ' + py_file 31 | dataset_name = match.group(0).split('_')[1] 32 | 33 | # extract local directory path to the dataset 34 | pattern = re.compile(r'^([^#]*)#(.*)$') 35 | with open(conf_file, 'r') as i_file: 36 | for i_line in i_file: 37 | # remove comments and whitespaces 38 | match = re.match(pattern, i_line) 39 | if match: 40 | i_line = match.group(1) 41 | i_line = i_line.strip() # remote whitespaces 42 | if i_line == '': 43 | continue 44 | 45 | # extract the (key, value) pair 46 | key_n_value = i_line.split(' = ') 47 | assert len(key_n_value) == 2, 'each line must contains exactly one \' = \'' 48 | key = key_n_value[0].strip() 49 | value = key_n_value[1].strip() 50 | if key == 'data_dir_local_%s' % dataset_name: 51 | print(value) 52 | break 53 | -------------------------------------------------------------------------------- /utils/get_idle_gpus.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Get a list of idle GPUs. 18 | 19 | This script sorts GPUs in the ascending order of memory usage, and return the top-k ones. 20 | """ 21 | 22 | import os 23 | import sys 24 | import subprocess 25 | 26 | # get the required number of idle GPUs 27 | assert len(sys.argv) == 2 28 | nb_idle_gpus = int(sys.argv[1]) 29 | 30 | # assume: idle gpu has no more than 50% of total card memory used 31 | mem_usage_ulimit = .5 32 | 33 | # command to execute to get gpu id and corresponding memory used 34 | # and total memory. It gives output in the format 35 | # gpu id, memory used, total memory 36 | cmd = 'nvidia-smi --query-gpu=index,memory.used,memory.total ' \ 37 | '--format=csv,noheader,nounits' 38 | gpu_smi_output = subprocess.check_output(cmd, shell=True) 39 | gpu_smi_output = gpu_smi_output.decode('utf-8') 40 | 41 | idle_gpus = [] 42 | for gpu in gpu_smi_output.split(sep='\n')[:-1]: 43 | (gpu_id, mem_used, mem_total) = [int(value) for value in gpu.split(sep=',')] 44 | mem_usage = float(mem_used) / mem_total 45 | if mem_usage < mem_usage_ulimit: 46 | idle_gpus += [(gpu_id, mem_usage)] 47 | idle_gpus.sort(key=lambda x: x[1]) 48 | idle_gpus = [x[0] for x in idle_gpus] # only keep GPU ids 49 | 50 | if len(idle_gpus) < nb_idle_gpus: 51 | raise ValueError('not enough idle GPUs; idle GPUs are: {}'.format(idle_gpus)) 52 | idle_gpus = idle_gpus[:nb_idle_gpus] 53 | idle_gpus_str = ','.join([str(idle_gpu) for idle_gpu in idle_gpus]) 54 | print(idle_gpus_str) 55 | -------------------------------------------------------------------------------- /utils/get_path_args.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Get path-related arguments.""" 18 | 19 | import re 20 | import sys 21 | 22 | # parse input arguments 23 | assert len(sys.argv) == 4 24 | exec_mode = sys.argv[1] 25 | py_file = sys.argv[2] 26 | conf_file = sys.argv[3] 27 | 28 | # obtain the dataset's name 29 | pattern = re.compile(r'at_[0-9A-Za-z]+_run.py$') 30 | match = re.search(pattern, py_file) 31 | assert match is not None, 'unable to match pattern in ' + py_file 32 | dataset_name = match.group(0).split('_')[1] 33 | 34 | # extract path-related arguments from path.conf 35 | arg_list = [] 36 | pattern_comment = re.compile(r'^([^#]*)#(.*)$') 37 | pattern_data_dir = re.compile(r'^data_dir_[a-z]+_%s$' % dataset_name) 38 | data_dirs = {'local': None, 'docker': None, 'seven': None, 'hdfs': None} 39 | with open(conf_file, 'r') as i_file: 40 | for i_line in i_file: 41 | # remove comments and whitespaces 42 | match = re.match(pattern_comment, i_line) 43 | if match: 44 | i_line = match.group(1) 45 | i_line = i_line.strip() # remote whitespaces 46 | if i_line == '': 47 | continue 48 | 49 | # extract the (key, value) pair 50 | key_n_value = i_line.split(' = ') 51 | assert len(key_n_value) == 2, 'each line must contains exactly one \' = \'' 52 | key = key_n_value[0].strip() 53 | value = key_n_value[1].strip() 54 | if value == 'None': 55 | continue 56 | 57 | # extract arguments 58 | if not key.startswith('data_dir_'): 59 | arg_list += ['--%s %s' % (key, value)] 60 | elif re.match(pattern_data_dir, key): 61 | data_disk = key.split('_')[2] 62 | data_dirs[data_disk] = value 63 | 64 | # append path-related arguments 65 | if exec_mode in ['local', 'seven'] and data_dirs[exec_mode] is not None: 66 | arg_list += ['--data_dir_local %s' % data_dirs[exec_mode]] 67 | elif data_dirs['local'] is not None: 68 | arg_list += ['--data_dir_local %s' % data_dirs['docker']] 69 | if data_dirs['hdfs'] is not None: 70 | arg_list += ['--data_dir_hdfs %s' % data_dirs['hdfs']] 71 | 72 | # concatenate all arguments into one string 73 | arg_list_str = ' '.join(arg_list) 74 | print(arg_list_str) 75 | -------------------------------------------------------------------------------- /utils/lrn_rate_utils.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Utility functions for learning rates.""" 18 | 19 | import tensorflow as tf 20 | 21 | FLAGS = tf.app.flags.FLAGS 22 | 23 | def setup_lrn_rate_piecewise_constant(global_step, batch_size, idxs_epoch, decay_rates): 24 | """Setup the learning rate with piecewise constant strategy. 25 | 26 | Args: 27 | * global_step: training iteration counter 28 | * batch_size: number of samples in each mini-batch 29 | * idxs_epoch: indices of epoches to decay the learning rate 30 | * decay_rates: list of decaying rates 31 | 32 | Returns: 33 | * lrn_rate: learning rate 34 | """ 35 | 36 | # adjust interval endpoints w.r.t. FLAGS.nb_epochs_rat 37 | idxs_epoch = [idx_epoch * FLAGS.nb_epochs_rat for idx_epoch in idxs_epoch] 38 | 39 | # setup learning rate with the piecewise constant strategy 40 | lrn_rate_init = FLAGS.lrn_rate_init * batch_size / FLAGS.batch_size_norm 41 | nb_batches_per_epoch = float(FLAGS.nb_smpls_train) / batch_size 42 | bnds = [int(nb_batches_per_epoch * idx_epoch) for idx_epoch in idxs_epoch] 43 | vals = [lrn_rate_init * decay_rate for decay_rate in decay_rates] 44 | lrn_rate = tf.train.piecewise_constant(global_step, bnds, vals) 45 | 46 | return lrn_rate 47 | 48 | def setup_lrn_rate_exponential_decay(global_step, batch_size, epoch_step, decay_rate): 49 | """Setup the learning rate with exponential decaying strategy. 50 | 51 | Args: 52 | * global_step: training iteration counter 53 | * batch_size: number of samples in each mini-batch 54 | * epoch_step: epoch step-size for applying the decaying step 55 | * decay_rate: decaying rate 56 | 57 | Returns: 58 | * lrn_rate: learning rate 59 | """ 60 | 61 | # adjust the step size & decaying rate w.r.t. FLAGS.nb_epochs_rat 62 | epoch_step *= FLAGS.nb_epochs_rat 63 | 64 | # setup learning rate with the exponential decay strategy 65 | lrn_rate_init = FLAGS.lrn_rate_init * batch_size / FLAGS.batch_size_norm 66 | batch_step = int(FLAGS.nb_smpls_train * epoch_step / batch_size) 67 | lrn_rate = tf.train.exponential_decay( 68 | lrn_rate_init, tf.cast(global_step, tf.int32), batch_step, decay_rate, staircase=True) 69 | 70 | return lrn_rate 71 | -------------------------------------------------------------------------------- /utils/misc_utils.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Miscellaneous utility functions.""" 18 | 19 | import tensorflow as tf 20 | 21 | from utils.multi_gpu_wrapper import MultiGpuWrapper as mgw 22 | 23 | FLAGS = tf.app.flags.FLAGS 24 | 25 | def auto_barrier(mpi_comm=None): 26 | """Automatically insert a barrier for multi-GPU training, or pass for single-GPU training. 27 | 28 | Args: 29 | * mpi_comm: MPI communication object 30 | """ 31 | 32 | if FLAGS.enbl_multi_gpu: 33 | mpi_comm.Barrier() 34 | else: 35 | pass 36 | 37 | def is_primary_worker(scope='global'): 38 | """Check whether is the primary worker of all nodes (global) or the current node (local). 39 | 40 | Args: 41 | * scope: check scope ('global' OR 'local') 42 | 43 | Returns: 44 | * flag: whether is the primary worker 45 | """ 46 | 47 | if scope == 'global': 48 | return True if not FLAGS.enbl_multi_gpu else mgw.rank() == 0 49 | elif scope == 'local': 50 | return True if not FLAGS.enbl_multi_gpu else mgw.local_rank() == 0 51 | else: 52 | raise ValueError('unrecognized worker scope: ' + scope) 53 | -------------------------------------------------------------------------------- /utils/multi_gpu_wrapper.py: -------------------------------------------------------------------------------- 1 | # Tencent is pleased to support the open source community by making PocketFlow available. 2 | # 3 | # Copyright (C) 2018 THL A29 Limited, a Tencent company. All rights reserved. 4 | # 5 | # Licensed under the BSD 3-Clause License (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://opensource.org/licenses/BSD-3-Clause 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 | """Wrapper for multi-GPU training.""" 18 | 19 | # use Hovorod / TF-Plus for multi-GPU training 20 | try: 21 | import horovod.tensorflow as mgw 22 | print('using Horovod for multi-GPU training') 23 | except ImportError: 24 | try: 25 | import tfplus.tensorflow as mgw 26 | print('using TF-Plus for multi-GPU training') 27 | except ImportError: 28 | print('[WARNING] TF-Plus & Horovod cannot be imported; multi-GPU training is unsupported') 29 | 30 | class MultiGpuWrapper(object): 31 | """Wrapper for multi-GPU training.""" 32 | 33 | def __init__(self): 34 | """Constructor function.""" 35 | pass 36 | 37 | @classmethod 38 | def init(cls, *args): 39 | """Initialization.""" 40 | 41 | try: 42 | return mgw.init(*args) 43 | except NameError: 44 | raise NameError('module not imported') 45 | 46 | @classmethod 47 | def size(cls, *args): 48 | """Get the number of workers at all nodes.""" 49 | 50 | try: 51 | return mgw.size(*args) 52 | except NameError: 53 | raise NameError('module not imported') 54 | 55 | @classmethod 56 | def rank(cls, *args): 57 | """Get the rank of current worker at all nodes.""" 58 | 59 | try: 60 | return mgw.rank(*args) 61 | except NameError: 62 | raise NameError('module not imported') 63 | 64 | @classmethod 65 | def local_size(cls, *args): 66 | """Get the number of workers at the current node.""" 67 | 68 | try: 69 | return mgw.local_size(*args) 70 | except NameError: 71 | raise NameError('module not imported') 72 | 73 | @classmethod 74 | def local_rank(cls, *args): 75 | """Get the rank of current worker at the current node.""" 76 | 77 | try: 78 | return mgw.local_rank(*args) 79 | except NameError: 80 | raise NameError('module not imported') 81 | 82 | @classmethod 83 | def DistributedOptimizer(cls, *args): 84 | """Get a distributed optimizer from the base optimizer.""" 85 | 86 | try: 87 | return mgw.DistributedOptimizer(*args) 88 | except NameError: 89 | raise NameError('module not imported') 90 | 91 | @classmethod 92 | def broadcast_global_variables(cls, *args): 93 | """Get a TensorFlow operation to broadcast all the global variables.""" 94 | 95 | try: 96 | return mgw.broadcast_global_variables(*args) 97 | except NameError: 98 | raise NameError('module not imported') 99 | --------------------------------------------------------------------------------