├── src └── zen_nas │ ├── hotfix │ ├── __init__.py │ ├── vision.py │ ├── transforms.py │ └── folder.py │ ├── ModelLoader │ ├── geffnet │ │ ├── version.py │ │ ├── __init__.py │ │ ├── model_factory.py │ │ ├── activations │ │ │ ├── activations_jit.py │ │ │ ├── activations.py │ │ │ ├── __init__.py │ │ │ └── activations_me.py │ │ ├── helpers.py │ │ └── config.py │ └── __init__.py │ ├── misc │ └── ZenNet_speed.png │ ├── __init__.py │ ├── ZenNet │ ├── zennet_imagenet1k_flops600M_res224.txt │ ├── zennet_imagenet1k_flops400M_res224.txt │ ├── zennet_cifar_model_size05M_res32.txt │ ├── zennet_imagenet1k_flops900M_res224.txt │ ├── zennet_cifar_model_size1M_res32.txt │ ├── zennet_imagenet1k_latency01ms_res224.txt │ ├── zennet_imagenet1k_latency02ms_res224.txt │ ├── zennet_cifar_model_size2M_res32.txt │ ├── zennet_imagenet1k_latency03ms_res224.txt │ ├── zennet_imagenet1k_latency05ms_res224.txt │ ├── zennet_imagenet1k_latency08ms_res224.txt │ ├── zennet_imagenet1k_latency12ms_res224.txt │ └── masternet.py │ ├── scripts │ ├── Zen_NAS_ImageNet_latency0.1ms_train.sh │ ├── Zen_NAS_ImageNet_latency0.5ms_train.sh │ ├── Zen_NAS_ImageNet_latency0.5ms_train_apex.sh │ ├── Zen_NAS_ImageNet_latency1.2ms_train_apex.sh │ ├── Zen_NAS_ImageNet_latency0.1ms_train_apex.sh │ ├── Zen_NAS_ImageNet_flops400M.sh │ ├── Zen_NAS_ImageNet_flops600M.sh │ ├── Zen_NAS_ImageNet_flops800M.sh │ ├── Zen_NAS_ImageNet_latency0.1ms.sh │ ├── Zen_NAS_ImageNet_latency0.2ms.sh │ ├── Zen_NAS_ImageNet_latency0.3ms.sh │ ├── Zen_NAS_ImageNet_latency0.5ms.sh │ ├── Zen_NAS_ImageNet_latency0.8ms.sh │ ├── Zen_NAS_ImageNet_latency1.2ms.sh │ ├── Zen_NAS_cifar_params2M.sh │ ├── Flops_NAS_cifar_params1M.sh │ ├── Params_NAS_cifar_params1M.sh │ ├── Random_NAS_cifar_params1M.sh │ ├── TE_NAS_cifar_params1M.sh │ ├── Zen_NAS_cifar_params1M.sh │ ├── NASWOT_NAS_cifar_params1M.sh │ ├── Syncflow_NAS_cifar_params1M.sh │ └── GradNorm_NAS_cifar_params1M.sh │ ├── analyze_model.py │ ├── compute_score.py │ ├── ZeroShotProxy │ ├── compute_gradnorm_score.py │ ├── compute_syncflow_score.py │ ├── compute_NASWOT_score.py │ └── compute_zen_score.py │ ├── test.py │ ├── test_cifar.py │ ├── SearchSpace │ ├── search_space_XXBL.py │ ├── search_choice.py │ └── search_space_IDW_fixfc.py │ ├── benchmark_network_latency.py │ ├── Masternet.py │ └── PlainNet │ ├── SuperResKXKX.py │ ├── super_blocks.py │ └── SuperResK1KXK1.py ├── .vscode └── settings.json ├── tests ├── __init__.py └── test.py ├── .markdownlint.json ├── .flake8 ├── bors.toml ├── mypy.ini ├── tox.ini ├── .pylintrc ├── azure-pipelines.yml ├── setup.py ├── .gitignore └── README.md /src/zen_nas/hotfix/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | { 2 | "python.pythonPath": "/opt/conda/bin/python" 3 | } -------------------------------------------------------------------------------- /src/zen_nas/ModelLoader/geffnet/version.py: -------------------------------------------------------------------------------- 1 | """__version__""" 2 | 3 | __version__ = '1.0.0' 4 | -------------------------------------------------------------------------------- /src/zen_nas/misc/ZenNet_speed.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Adlik/zen_nas/HEAD/src/zen_nas/misc/ZenNet_speed.png -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 ZTE corporation. All Rights Reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | -------------------------------------------------------------------------------- /src/zen_nas/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 ZTE corporation. All Rights Reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | -------------------------------------------------------------------------------- /.markdownlint.json: -------------------------------------------------------------------------------- 1 | { 2 | "default": true, 3 | "line-length": { 4 | "line_length": 120, 5 | "tables": false 6 | } 7 | } -------------------------------------------------------------------------------- /.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | exclude = src/zen_nas/ModelLoader/geffnet/*.py,src/zen_nas/tf_train_image_classification.py 3 | ignore = E402 4 | max-line-length = 120 5 | -------------------------------------------------------------------------------- /bors.toml: -------------------------------------------------------------------------------- 1 | status = [ "Adlik.zen_nas" ] 2 | block_labels = [ "do not merge", "work in progress" ] 3 | timeout_sec = 21600 4 | delete_merged_branches = true 5 | use_squash_merge = true 6 | -------------------------------------------------------------------------------- /src/zen_nas/ZenNet/zennet_imagenet1k_flops600M_res224.txt: -------------------------------------------------------------------------------- 1 | SuperConvK3BNRELU(3,24,2,1)SuperResIDWE1K7(24,48,2,48,1)SuperResIDWE2K7(48,72,2,72,1)SuperResIDWE6K7(72,96,2,88,5)SuperResIDWE4K7(96,192,2,168,5)SuperConvK1BNRELU(192,2048,1,1) -------------------------------------------------------------------------------- /src/zen_nas/ZenNet/zennet_imagenet1k_flops400M_res224.txt: -------------------------------------------------------------------------------- 1 | SuperConvK3BNRELU(3,16,2,1)SuperResIDWE1K7(16,40,2,40,1)SuperResIDWE1K7(40,64,2,64,1)SuperResIDWE4K7(64,96,2,96,5)SuperResIDWE2K7(96,224,2,224,5)SuperConvK1BNRELU(224,2048,1,1) 2 | -------------------------------------------------------------------------------- /src/zen_nas/ZenNet/zennet_cifar_model_size05M_res32.txt: -------------------------------------------------------------------------------- 1 | SuperConvK3BNRELU(3,64,1,1)SuperResK1K5K1(64,168,1,16,3)SuperResK1K3K1(168,80,2,32,4)SuperResK1K5K1(80,112,2,16,3)SuperResK1K5K1(112,144,1,24,3)SuperResK1K3K1(144,32,2,40,1)SuperConvK1BNRELU(32,512,1,1) 2 | -------------------------------------------------------------------------------- /src/zen_nas/ZenNet/zennet_imagenet1k_flops900M_res224.txt: -------------------------------------------------------------------------------- 1 | SuperConvK3BNRELU(3,16,2,1)SuperResIDWE1K7(16,48,2,72,1)SuperResIDWE2K7(48,72,2,64,3)SuperResIDWE2K7(72,152,2,144,3)SuperResIDWE2K7(152,360,2,352,4)SuperResIDWE4K7(360,288,1,264,3)SuperConvK1BNRELU(288,2048,1,1) 2 | -------------------------------------------------------------------------------- /mypy.ini: -------------------------------------------------------------------------------- 1 | [mypy] 2 | ignore_missing_imports = True 3 | disable_error_code = no-redef 4 | mypy_path = src 5 | exclude = src/zen_nas/ModelLoader/geffnet/*.py,src/zen_nas/tf_train_image_classification.py 6 | 7 | [mypy-src.zen_nas.tf_train_image_classification] 8 | ignore_errors = True 9 | -------------------------------------------------------------------------------- /src/zen_nas/ZenNet/zennet_cifar_model_size1M_res32.txt: -------------------------------------------------------------------------------- 1 | SuperConvK3BNRELU(3,88,1,1)SuperResK1K7K1(88,120,1,16,1)SuperResK1K7K1(120,192,2,16,3)SuperResK1K5K1(192,224,1,24,4)SuperResK1K5K1(224,96,2,24,2)SuperResK1K3K1(96,168,2,40,3)SuperResK1K3K1(168,112,1,48,3)SuperConvK1BNRELU(112,512,1,1) 2 | -------------------------------------------------------------------------------- /src/zen_nas/ZenNet/zennet_imagenet1k_latency01ms_res224.txt: -------------------------------------------------------------------------------- 1 | SuperConvK3BNRELU(3,24,2,1)SuperResK3K3(24,32,2,64,1)SuperResK5K5(32,64,2,32,1)SuperResK5K5(64,168,2,96,1)SuperResK1K5K1(168,320,1,120,1)SuperResK1K5K1(320,640,2,304,3)SuperResK1K5K1(640,512,1,384,1)SuperConvK1BNRELU(512,2384,1,1) -------------------------------------------------------------------------------- /src/zen_nas/ZenNet/zennet_imagenet1k_latency02ms_res224.txt: -------------------------------------------------------------------------------- 1 | SuperConvK3BNRELU(3,24,2,1)SuperResK1K5K1(24,32,2,32,1)SuperResK1K7K1(32,104,2,64,1)SuperResK1K5K1(104,512,2,160,1)SuperResK1K5K1(512,344,1,192,1)SuperResK1K5K1(344,688,2,320,4)SuperResK1K5K1(688,680,1,304,3)SuperConvK1BNRELU(680,2552,1,1) -------------------------------------------------------------------------------- /src/zen_nas/ZenNet/zennet_cifar_model_size2M_res32.txt: -------------------------------------------------------------------------------- 1 | SuperConvK3BNRELU(3,32,1,1)SuperResK1K5K1(32,120,1,40,1)SuperResK1K5K1(120,176,2,32,3)SuperResK1K7K1(176,272,1,24,3)SuperResK1K3K1(272,176,1,56,3)SuperResK1K3K1(176,176,1,64,4)SuperResK1K5K1(176,216,2,40,2)SuperResK1K3K1(216,72,2,56,2)SuperConvK1BNRELU(72,512,1,1) 2 | -------------------------------------------------------------------------------- /src/zen_nas/ZenNet/zennet_imagenet1k_latency03ms_res224.txt: -------------------------------------------------------------------------------- 1 | SuperConvK3BNRELU(3,24,2,1)SuperResK1K5K1(24,64,2,32,1)SuperResK1K3K1(64,128,2,128,1)SuperResK1K7K1(128,432,2,128,1)SuperResK1K5K1(432,272,1,160,1)SuperResK1K5K1(272,848,2,384,4)SuperResK1K5K1(848,848,1,320,3)SuperResK1K5K1(848,456,1,320,3)SuperConvK1BNRELU(456,6704,1,1) -------------------------------------------------------------------------------- /src/zen_nas/ZenNet/zennet_imagenet1k_latency05ms_res224.txt: -------------------------------------------------------------------------------- 1 | SuperConvK3BNRELU(3,8,2,1)SuperResK1K7K1(8,64,2,32,1)SuperResK1K3K1(64,152,2,128,1)SuperResK1K5K1(152,640,2,192,4)SuperResK1K5K1(640,640,1,192,2)SuperResK1K5K1(640,1536,2,384,4)SuperResK1K5K1(1536,816,1,384,3)SuperResK1K5K1(816,816,1,384,3)SuperConvK1BNRELU(816,5304,1,1) -------------------------------------------------------------------------------- /src/zen_nas/ZenNet/zennet_imagenet1k_latency08ms_res224.txt: -------------------------------------------------------------------------------- 1 | SuperConvK3BNRELU(3,16,2,1)SuperResK1K5K1(16,64,2,64,1)SuperResK1K3K1(64,240,2,128,2)SuperResK1K7K1(240,640,2,160,3)SuperResK1K7K1(640,768,1,192,4)SuperResK1K5K1(768,1536,2,384,5)SuperResK1K5K1(1536,1536,1,384,3)SuperResK1K5K1(1536,2304,1,384,5)SuperConvK1BNRELU(2304,4912,1,1) -------------------------------------------------------------------------------- /src/zen_nas/ZenNet/zennet_imagenet1k_latency12ms_res224.txt: -------------------------------------------------------------------------------- 1 | SuperConvK3BNRELU(3,32,2,1)SuperResK1K5K1(32,80,2,32,1)SuperResK1K7K1(80,432,2,128,5)SuperResK1K7K1(432,640,2,192,3)SuperResK1K7K1(640,1008,1,160,5)SuperResK1K7K1(1008,976,1,160,4)SuperResK1K5K1(976,2304,2,384,5)SuperResK1K5K1(2304,2496,1,384,5)SuperConvK1BNRELU(2496,3072,1,1) -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [testenv] 2 | commands = bandit -r -s B101,B310,B311 src 3 | bandit -r -s B101,B310,B311 tests 4 | 5 | flake8 src tests 6 | pylint src tests 7 | 8 | pytest --flake8 \ 9 | --mypy \ 10 | -n auto \ 11 | --pylint 12 | 13 | deps = .[test] 14 | -------------------------------------------------------------------------------- /.pylintrc: -------------------------------------------------------------------------------- 1 | [MASTER] 2 | jobs = 0 3 | ignore = **/ModelLoader/geffnet/*.py,src/zen_nas/tf_train_image_classification.py 4 | 5 | [MESSAGES CONTROL] 6 | disable = missing-class-docstring, 7 | missing-function-docstring, 8 | missing-module-docstring, 9 | duplicate-code 10 | 11 | 12 | [FORMAT] 13 | max-line-length = 120 14 | 15 | [TYPECHECK] 16 | ignored-modules=torch 17 | ignored-classes=torch 18 | -------------------------------------------------------------------------------- /src/zen_nas/ModelLoader/geffnet/__init__.py: -------------------------------------------------------------------------------- 1 | ''' 2 | The geffnet module is modified from: 3 | https://github.com/rwightman/gen-efficientnet-pytorch 4 | ''' 5 | import os 6 | import sys 7 | sys.path.append(os.path.dirname(os.path.abspath(__file__))) 8 | try: 9 | from gen_efficientnet import * 10 | from mobilenetv3 import * 11 | from model_factory import create_model 12 | from config import is_exportable, is_scriptable, set_exportable, set_scriptable 13 | from activations import * 14 | except ImportError: 15 | print('fail to import zen_nas modules') 16 | -------------------------------------------------------------------------------- /tests/test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 ZTE corporation. All Rights Reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | 4 | """Tests using ModelLoader modules to build model """ 5 | import os 6 | import sys 7 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 8 | try: 9 | from src.zen_nas import global_utils, ModelLoader 10 | except ImportError: 11 | print('fail to import ModelLoader') 12 | 13 | 14 | if __name__ == '__main__': 15 | opt = global_utils.parse_cmd_options(sys.argv) 16 | 17 | model = ModelLoader.get_model(opt, sys.argv) 18 | 19 | model = model.cuda(opt.gpu) 20 | -------------------------------------------------------------------------------- /azure-pipelines.yml: -------------------------------------------------------------------------------- 1 | trigger: 2 | - staging 3 | - trying 4 | - main 5 | jobs: 6 | - job: Tox 7 | displayName: Tox 8 | pool: 9 | vmImage: ubuntu-18.04 10 | steps: 11 | - task: UsePythonVersion@0 12 | inputs: 13 | versionSpec: '3.8' 14 | - script: python3 -m pip install -U tox 15 | displayName: Install tox 16 | - script: tox 17 | displayName: Run tox 18 | 19 | - job: Markdownlint 20 | displayName: Markdownlint 21 | pool: 22 | vmImage: ubuntu-18.04 23 | steps: 24 | - script: sudo npm install -g markdownlint-cli 25 | displayName: Install markdownlint-cli 26 | - script: markdownlint '**/*.md' -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 ZTE corporation. All Rights Reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | 4 | """Setup""" 5 | 6 | import setuptools 7 | 8 | INSTALL_REQUIRES = [ 9 | 'numpy', 10 | 'torch >= 1.8.1', 11 | 'torchvision >= 0.9.1', 12 | 'ptflops', 13 | 'tensorboard >= 1.15', 14 | 'horovod >= 0.22.1', 15 | 'apex' 16 | ] 17 | 18 | TEST_REQUIRES = [ 19 | 'bandit', 20 | 'flake8', 21 | 'mypy==0.812', 22 | 'pylint', 23 | 'pytest-cov', 24 | 'pytest-flake8', 25 | 'pytest-mypy', 26 | 'pytest-pylint', 27 | 'pytest-xdist' 28 | ] 29 | 30 | setuptools.setup( 31 | name='zen_nas', 32 | packages=setuptools.find_packages('src'), 33 | package_dir={'': 'src'}, 34 | install_requires=INSTALL_REQUIRES, 35 | extras_require={ 36 | 'test': TEST_REQUIRES 37 | }, 38 | python_requires='>= 3.7' 39 | ) 40 | -------------------------------------------------------------------------------- /src/zen_nas/ModelLoader/geffnet/model_factory.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=W0401,import-error 2 | import os 3 | import sys 4 | from helpers import load_checkpoint 5 | sys.path.append(os.path.dirname(os.path.abspath(__file__))) 6 | try: 7 | from gen_efficientnet import * 8 | from mobilenetv3 import * 9 | except ImportError: 10 | print('fail to import zen_nas modules') 11 | 12 | 13 | def create_model( 14 | model_name='mnasnet_100', 15 | pretrained=None, 16 | num_classes=1000, 17 | in_chans=3, 18 | checkpoint_path='', 19 | **kwargs): 20 | """ create a model based on model name 21 | 22 | :param model_name (str): model name 23 | :param pretrained (bool): pretrained model 24 | :param num_classes (int): class number 25 | :param in_chans (int): input channels 26 | :param checkpoint_path (str): checkpoint file path 27 | :return model 28 | """ 29 | 30 | model_kwargs = dict(num_classes=num_classes, in_chans=in_chans, pretrained=pretrained, **kwargs) 31 | 32 | if model_name in globals(): 33 | create_fn = globals()[model_name] 34 | model = create_fn(**model_kwargs) 35 | else: 36 | raise RuntimeError(f'Unknown model ({model_name})') 37 | 38 | if checkpoint_path and not pretrained: 39 | load_checkpoint(model, checkpoint_path) 40 | 41 | return model 42 | -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_latency0.1ms_train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/Zen_NAS_ImageNet_latency0.1ms 8 | mkdir -p ${save_dir} 9 | 10 | 11 | 12 | resolution=224 13 | budget_latency=1e-4 14 | max_layers=10 15 | population_size=512 16 | epochs=480 17 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 18 | 19 | 20 | 21 | 22 | 23 | 24 | horovodrun -np 8 -H localhost:8 python ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 25 | --dist_mode horovod --workers_per_gpu 8 --sync_bn \ 26 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 27 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 28 | --label_smoothing --random_erase --mixup --auto_augment \ 29 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 30 | --arch Masternet.py:MasterNet \ 31 | --plainnet_struct_txt ${save_dir}/zennet_imagenet1k_latency01ms_res224.txt \ 32 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 33 | --teacher_pretrained \ 34 | --teacher_input_image_size 320 \ 35 | --teacher_feature_weight 1.0 \ 36 | --teacher_logit_weight 1.0 \ 37 | --ts_proj_no_relu \ 38 | --ts_proj_no_bn \ 39 | --target_downsample_ratio 16 \ 40 | --batch_size_per_gpu 64 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/analyze_model.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | 4 | get model FLOPs and parameters 5 | ''' 6 | import os 7 | import sys 8 | from ptflops import get_model_complexity_info 9 | sys.path.append(os.path.dirname(os.path.abspath(__file__))) 10 | try: 11 | import ModelLoader 12 | import global_utils 13 | import ZenNet 14 | except ImportError: 15 | print('fail to import zen_nas modules') 16 | 17 | 18 | def main(opt, argv): 19 | """get model flops and parameters""" 20 | model = ModelLoader.get_model(opt, argv) 21 | flops, params = get_model_complexity_info(model, (3, opt.input_image_size, opt.input_image_size), 22 | as_strings=False, 23 | print_per_layer_stat=True) 24 | print(f'Flops: {flops:4g}') 25 | print(f'Params: {params:4g}') 26 | 27 | 28 | def get_flops_params(opt): 29 | """get model flops and parameters""" 30 | model = ZenNet.get_ZenNet(opt.arch) 31 | flops, params = get_model_complexity_info(model, (3, opt.input_image_size, opt.input_image_size), 32 | as_strings=False, 33 | print_per_layer_stat=True) 34 | print(f'Flops: {flops:4g}') 35 | print(f'Params: {params:4g}') 36 | 37 | 38 | if __name__ == "__main__": 39 | option = global_utils.parse_cmd_options(sys.argv) 40 | 41 | # get_flops_params(opt) 42 | main(option, sys.argv) 43 | -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_latency0.5ms_train.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/Zen_NAS_ImageNet_latency0.5ms 8 | mkdir -p ${save_dir} 9 | 10 | 11 | resolution=224 12 | budget_latency=5e-4 13 | max_layers=20 14 | population_size=512 15 | epochs=480 16 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 17 | 18 | echo "SuperConvK3BNRELU(3,32,2,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,128,2,64,1)\ 19 | SuperResK3K3(128,256,2,128,1)SuperResK3K3(256,512,2,256,1)\ 20 | SuperConvK1BNRELU(256,512,1,1)" > ${save_dir}/init_plainnet.txt 21 | 22 | 23 | 24 | horovodrun -np 8 -H localhost:8 python ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 25 | --dist_mode horovod --workers_per_gpu 6 --sync_bn \ 26 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 27 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 28 | --label_smoothing --random_erase --mixup --auto_augment \ 29 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 30 | --arch Masternet.py:MasterNet \ 31 | --plainnet_struct_txt ${save_dir}/zennet_imagenet1k_latency05ms_res224.txt \ 32 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 33 | --teacher_pretrained \ 34 | --teacher_input_image_size 320 \ 35 | --teacher_feature_weight 1.0 \ 36 | --teacher_logit_weight 1.0 \ 37 | --ts_proj_no_relu \ 38 | --ts_proj_no_bn \ 39 | --target_downsample_ratio 16 \ 40 | --batch_size_per_gpu 128 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_latency0.5ms_train_apex.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/search_Zen_NAS_ImageNet_latency0.5ms 8 | mkdir -p ${save_dir} 9 | 10 | 11 | resolution=224 12 | budget_latency=5e-4 13 | max_layers=20 14 | population_size=512 15 | epochs=480 16 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 17 | 18 | echo "SuperConvK3BNRELU(3,32,2,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,128,2,64,1)\ 19 | SuperResK3K3(128,256,2,128,1)SuperResK3K3(256,512,2,256,1)\ 20 | SuperConvK1BNRELU(256,512,1,1)" > ${save_dir}/init_plainnet.txt 21 | 22 | 23 | 24 | python -m torch.distributed.launch --nproc_per_node=8 ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 25 | --dist_mode apex --workers_per_gpu 6 --apex \ 26 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 27 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 28 | --label_smoothing --random_erase --mixup --auto_augment \ 29 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 30 | --arch Masternet.py:MasterNet \ 31 | --plainnet_struct_txt ${save_dir}/search_imagenet1k_latency05ms_ress224.txt \ 32 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 33 | --teacher_pretrained \ 34 | --teacher_input_image_size 320 \ 35 | --teacher_feature_weight 1.0 \ 36 | --teacher_logit_weight 1.0 \ 37 | --ts_proj_no_relu \ 38 | --ts_proj_no_bn \ 39 | --target_downsample_ratio 16 \ 40 | --batch_size_per_gpu 128 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_latency1.2ms_train_apex.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/search_Zen_NAS_ImageNet_latency1.2ms 8 | mkdir -p ${save_dir} 9 | 10 | 11 | resolution=224 12 | budget_latency=12e-4 13 | max_layers=30 14 | population_size=512 15 | epochs=100 16 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 17 | 18 | echo "SuperConvK3BNRELU(3,32,2,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,128,2,64,1)\ 19 | SuperResK3K3(128,256,2,128,1)SuperResK3K3(256,512,2,256,1)\ 20 | SuperConvK1BNRELU(256,512,1,1)" > ${save_dir}/init_plainnet.txt 21 | 22 | 23 | 24 | python -m torch.distributed.launch --nproc_per_node=8 ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 25 | --dist_mode apex --workers_per_gpu 6 --apex \ 26 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 27 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 28 | --label_smoothing --random_erase --mixup --auto_augment \ 29 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 30 | --arch Masternet.py:MasterNet \ 31 | --plainnet_struct_txt ${save_dir}/search_imagenet1k_latency12ms_res224.txt \ 32 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 33 | --teacher_pretrained \ 34 | --teacher_input_image_size 320 \ 35 | --teacher_feature_weight 1.0 \ 36 | --teacher_logit_weight 1.0 \ 37 | --ts_proj_no_relu \ 38 | --ts_proj_no_bn \ 39 | --target_downsample_ratio 16 \ 40 | --batch_size_per_gpu 256 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_latency0.1ms_train_apex.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/searchv4_Zen_NAS_ImageNet_latency0.1ms 8 | mkdir -p ${save_dir} 9 | 10 | 11 | resolution=224 12 | budget_latency=1e-4 13 | max_layers=10 14 | population_size=512 15 | epochs=480 16 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 17 | 18 | echo "SuperConvK3BNRELU(3,32,2,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,128,2,64,1)\ 19 | SuperResK3K3(128,256,2,128,1)SuperResK3K3(256,512,2,256,1)\ 20 | SuperConvK1BNRELU(256,512,1,1)" > ${save_dir}/init_plainnet.txt 21 | 22 | 23 | 24 | python -m torch.distributed.launch --nproc_per_node=8 ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 25 | --dist_mode apex --workers_per_gpu 6 --apex \ 26 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 27 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 28 | --label_smoothing --random_erase --mixup --auto_augment \ 29 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 30 | --arch Masternet.py:MasterNet \ 31 | --plainnet_struct_txt ${save_dir}/searchv4_imagenet1k_latency01ms_ress224.txt \ 32 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 33 | --teacher_pretrained \ 34 | --teacher_input_image_size 320 \ 35 | --teacher_feature_weight 1.0 \ 36 | --teacher_logit_weight 1.0 \ 37 | --ts_proj_no_relu \ 38 | --ts_proj_no_bn \ 39 | --target_downsample_ratio 16 \ 40 | --batch_size_per_gpu 256 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/compute_score.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 ZTE corporation. All Rights Reserved. 2 | # SPDX-License-Identifier: Apache-2.0 3 | # pylint: disable=invalid-name 4 | import os 5 | import sys 6 | import argparse 7 | sys.path.append(os.path.dirname(os.path.abspath(__file__))) 8 | try: 9 | import Masternet 10 | from evolution_search import compute_nas_score 11 | except ImportError: 12 | print('fail to import Masternet, evolution_search') 13 | 14 | 15 | def parse_cmd_options(argv): 16 | parser = argparse.ArgumentParser() 17 | parser.add_argument('--gpu', type=int, default=0) 18 | parser.add_argument('--zero_shot_score', type=str, default='Zen', 19 | help='could be: Zen (for Zen-NAS), TE (for TE-NAS)') 20 | parser.add_argument('--batch_size', type=int, default=16, help='number of instances in one mini-batch.') 21 | parser.add_argument('--input_image_size', type=int, default=224, 22 | help='resolution of input image, usually 32 for CIFAR and 224 for ImageNet.') 23 | parser.add_argument('--gamma', type=float, default=1e-2, 24 | help='noise perturbation coefficient') 25 | parser.add_argument('--repeat_times', type=int, default=32) 26 | parser.add_argument('--num_classes', type=int, default=1000, 27 | help='number of classes') 28 | parser.add_argument('--plain_structure', type=str, default=None, 29 | help='the text file with model structure str') 30 | module_opt, _ = parser.parse_known_args(argv) 31 | return module_opt 32 | 33 | 34 | PATH = './ZenNet' 35 | 36 | if __name__ == '__main__': 37 | opt = parse_cmd_options(sys.argv) 38 | 39 | gpu = opt.gpu 40 | 41 | model_plainnet_str_txt = os.path.join(PATH, opt.plain_structure) 42 | with open(model_plainnet_str_txt, 'r', encoding='utf8') as fid: 43 | model_plainnet_str = fid.readline().strip() 44 | 45 | Any_Plain_Net = Masternet.MasterNet 46 | 47 | the_nas_score = compute_nas_score(Any_Plain_Net, model_plainnet_str, gpu, opt) 48 | 49 | print(f'zen-score={the_nas_score:.4g}') 50 | -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_flops400M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/Zen_NAS_ImageNet_flops400M 8 | mkdir -p ${save_dir} 9 | 10 | 11 | resolution=224 12 | budget_flops=400e6 13 | max_layers=14 14 | population_size=512 15 | epochs=480 16 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 17 | 18 | echo "SuperConvK3BNRELU(3,8,2,1)SuperResIDWE6K3(8,32,2,8,1)SuperResIDWE6K3(32,48,2,32,1)\ 19 | SuperResIDWE6K3(48,96,2,48,1)SuperResIDWE6K3(96,128,2,96,1)\ 20 | SuperConvK1BNRELU(128,2048,1,1)" > ${save_dir}/init_plainnet.txt 21 | 22 | python evolution_search.py --gpu 0 \ 23 | --zero_shot_score Zen \ 24 | --search_space SearchSpace/search_space_IDW_fixfc.py \ 25 | --budget_flops ${budget_flops} \ 26 | --max_layers ${max_layers} \ 27 | --batch_size 64 \ 28 | --input_image_size ${resolution} \ 29 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 30 | --num_classes 1000 \ 31 | --evolution_max_iter ${evolution_max_iter} \ 32 | --population_size ${population_size} \ 33 | --save_dir ${save_dir} 34 | 35 | 36 | python analyze_model.py \ 37 | --input_image_size 224 \ 38 | --num_classes 1000 \ 39 | --arch Masternet.py:MasterNet \ 40 | --plainnet_struct_txt ${save_dir}/best_structure.txt 41 | 42 | 43 | python ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 44 | --dist_mode single --workers_per_gpu 6 \ 45 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 46 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 47 | --label_smoothing --random_erase --mixup --auto_augment \ 48 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 49 | --arch Masternet.py:MasterNet \ 50 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 51 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 52 | --teacher_pretrained \ 53 | --teacher_input_image_size 320 \ 54 | --teacher_feature_weight 1.0 \ 55 | --teacher_logit_weight 1.0 \ 56 | --ts_proj_no_relu \ 57 | --ts_proj_no_bn \ 58 | --use_se \ 59 | --target_downsample_ratio 16 \ 60 | --batch_size_per_gpu 64 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_flops600M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/Zen_NAS_ImageNet_flops600M 8 | mkdir -p ${save_dir} 9 | 10 | 11 | resolution=224 12 | budget_flops=600e6 13 | max_layers=14 14 | population_size=512 15 | epochs=480 16 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 17 | 18 | echo "SuperConvK3BNRELU(3,8,2,1)SuperResIDWE6K3(8,32,2,8,1)SuperResIDWE6K3(32,48,2,32,1)\ 19 | SuperResIDWE6K3(48,96,2,48,1)SuperResIDWE6K3(96,128,2,96,1)\ 20 | SuperConvK1BNRELU(128,2048,1,1)" > ${save_dir}/init_plainnet.txt 21 | 22 | python evolution_search.py --gpu 0 \ 23 | --zero_shot_score Zen \ 24 | --search_space SearchSpace/search_space_IDW_fixfc.py \ 25 | --budget_flops ${budget_flops} \ 26 | --max_layers ${max_layers} \ 27 | --batch_size 64 \ 28 | --input_image_size ${resolution} \ 29 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 30 | --num_classes 1000 \ 31 | --evolution_max_iter ${evolution_max_iter} \ 32 | --population_size ${population_size} \ 33 | --save_dir ${save_dir} 34 | 35 | 36 | python analyze_model.py \ 37 | --input_image_size 224 \ 38 | --num_classes 1000 \ 39 | --arch Masternet.py:MasterNet \ 40 | --plainnet_struct_txt ${save_dir}/best_structure.txt 41 | 42 | 43 | python ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 44 | --dist_mode single --workers_per_gpu 6 \ 45 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 46 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 47 | --label_smoothing --random_erase --mixup --auto_augment \ 48 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 49 | --arch Masternet.py:MasterNet \ 50 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 51 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 52 | --teacher_pretrained \ 53 | --teacher_input_image_size 320 \ 54 | --teacher_feature_weight 1.0 \ 55 | --teacher_logit_weight 1.0 \ 56 | --ts_proj_no_relu \ 57 | --ts_proj_no_bn \ 58 | --use_se \ 59 | --target_downsample_ratio 16 \ 60 | --batch_size_per_gpu 64 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_flops800M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/Zen_NAS_ImageNet_flops800M 8 | mkdir -p ${save_dir} 9 | 10 | 11 | resolution=224 12 | budget_flops=800e6 13 | max_layers=16 14 | population_size=512 15 | epochs=480 16 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 17 | 18 | echo "SuperConvK3BNRELU(3,8,2,1)SuperResIDWE6K3(8,32,2,8,1)SuperResIDWE6K3(32,48,2,32,1)\ 19 | SuperResIDWE6K3(48,96,2,48,1)SuperResIDWE6K3(96,128,2,96,1)\ 20 | SuperConvK1BNRELU(128,2048,1,1)" > ${save_dir}/init_plainnet.txt 21 | 22 | python evolution_search.py --gpu 0 \ 23 | --zero_shot_score Zen \ 24 | --search_space SearchSpace/search_space_IDW_fixfc.py \ 25 | --budget_flops ${budget_flops} \ 26 | --max_layers ${max_layers} \ 27 | --batch_size 64 \ 28 | --input_image_size ${resolution} \ 29 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 30 | --num_classes 1000 \ 31 | --evolution_max_iter ${evolution_max_iter} \ 32 | --population_size ${population_size} \ 33 | --save_dir ${save_dir} 34 | 35 | 36 | python analyze_model.py \ 37 | --input_image_size 224 \ 38 | --num_classes 1000 \ 39 | --arch Masternet.py:MasterNet \ 40 | --plainnet_struct_txt ${save_dir}/best_structure.txt 41 | 42 | 43 | python ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 44 | --dist_mode single --workers_per_gpu 6 \ 45 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 46 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 47 | --label_smoothing --random_erase --mixup --auto_augment \ 48 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 49 | --arch Masternet.py:MasterNet \ 50 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 51 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 52 | --teacher_pretrained \ 53 | --teacher_input_image_size 320 \ 54 | --teacher_feature_weight 1.0 \ 55 | --teacher_logit_weight 1.0 \ 56 | --ts_proj_no_relu \ 57 | --ts_proj_no_bn \ 58 | --use_se \ 59 | --target_downsample_ratio 16 \ 60 | --batch_size_per_gpu 64 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_latency0.1ms.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/Zen_NAS_ImageNet_latency0.1ms 8 | mkdir -p ${save_dir} 9 | 10 | 11 | resolution=224 12 | budget_latency=1e-4 13 | max_layers=10 14 | population_size=512 15 | epochs=480 16 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 17 | 18 | echo "SuperConvK3BNRELU(3,32,2,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,128,2,64,1)\ 19 | SuperResK3K3(128,256,2,128,1)SuperResK3K3(256,512,2,256,1)\ 20 | SuperConvK1BNRELU(256,512,1,1)" > ${save_dir}/init_plainnet.txt 21 | 22 | python evolution_search.py --gpu 0 \ 23 | --zero_shot_score Zen \ 24 | --fix_initialize \ 25 | --origin \ 26 | --search_space SearchSpace/search_space_XXBL.py \ 27 | --budget_latency ${budget_latency} \ 28 | --max_layers ${max_layers} \ 29 | --batch_size 64 \ 30 | --input_image_size ${resolution} \ 31 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 32 | --num_classes 1000 \ 33 | --evolution_max_iter ${evolution_max_iter} \ 34 | --population_size ${population_size} \ 35 | --save_dir ${save_dir} 36 | 37 | 38 | python analyze_model.py \ 39 | --input_image_size 224 \ 40 | --num_classes 1000 \ 41 | --arch Masternet.py:MasterNet \ 42 | --plainnet_struct_txt ${save_dir}/best_structure.txt 43 | 44 | 45 | python ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 46 | --dist_mode single --workers_per_gpu 6 \ 47 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 48 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 49 | --label_smoothing --random_erase --mixup --auto_augment \ 50 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 51 | --arch Masternet.py:MasterNet \ 52 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 53 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 54 | --teacher_pretrained \ 55 | --teacher_input_image_size 320 \ 56 | --teacher_feature_weight 1.0 \ 57 | --teacher_logit_weight 1.0 \ 58 | --ts_proj_no_relu \ 59 | --ts_proj_no_bn \ 60 | --target_downsample_ratio 16 \ 61 | --batch_size_per_gpu 64 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_latency0.2ms.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/Zen_NAS_ImageNet_latency0.2ms 8 | mkdir -p ${save_dir} 9 | 10 | 11 | resolution=224 12 | budget_latency=2e-4 13 | max_layers=13 14 | population_size=512 15 | epochs=480 16 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 17 | 18 | echo "SuperConvK3BNRELU(3,32,2,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,128,2,64,1)\ 19 | SuperResK3K3(128,256,2,128,1)SuperResK3K3(256,512,2,256,1)\ 20 | SuperConvK1BNRELU(256,512,1,1)" > ${save_dir}/init_plainnet.txt 21 | 22 | python evolution_search.py --gpu 0 \ 23 | --zero_shot_score Zen \ 24 | --fix_initialize \ 25 | --origin \ 26 | --search_space SearchSpace/search_space_XXBL.py \ 27 | --budget_latency ${budget_latency} \ 28 | --max_layers ${max_layers} \ 29 | --batch_size 64 \ 30 | --input_image_size ${resolution} \ 31 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 32 | --num_classes 1000 \ 33 | --evolution_max_iter ${evolution_max_iter} \ 34 | --population_size ${population_size} \ 35 | --save_dir ${save_dir} 36 | 37 | 38 | python analyze_model.py \ 39 | --input_image_size 224 \ 40 | --num_classes 1000 \ 41 | --arch Masternet.py:MasterNet \ 42 | --plainnet_struct_txt ${save_dir}/best_structure.txt 43 | 44 | 45 | python ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 46 | --dist_mode single --workers_per_gpu 6 \ 47 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 48 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 49 | --label_smoothing --random_erase --mixup --auto_augment \ 50 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 51 | --arch Masternet.py:MasterNet \ 52 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 53 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 54 | --teacher_pretrained \ 55 | --teacher_input_image_size 320 \ 56 | --teacher_feature_weight 1.0 \ 57 | --teacher_logit_weight 1.0 \ 58 | --ts_proj_no_relu \ 59 | --ts_proj_no_bn \ 60 | --target_downsample_ratio 16 \ 61 | --batch_size_per_gpu 64 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_latency0.3ms.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/Zen_NAS_ImageNet_latency0.3ms 8 | mkdir -p ${save_dir} 9 | 10 | 11 | resolution=224 12 | budget_latency=3e-4 13 | max_layers=16 14 | population_size=512 15 | epochs=480 16 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 17 | 18 | echo "SuperConvK3BNRELU(3,32,2,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,128,2,64,1)\ 19 | SuperResK3K3(128,256,2,128,1)SuperResK3K3(256,512,2,256,1)\ 20 | SuperConvK1BNRELU(256,512,1,1)" > ${save_dir}/init_plainnet.txt 21 | 22 | python evolution_search.py --gpu 0 \ 23 | --zero_shot_score Zen \ 24 | --fix_initialize \ 25 | --origin \ 26 | --search_space SearchSpace/search_space_XXBL.py \ 27 | --budget_latency ${budget_latency} \ 28 | --max_layers ${max_layers} \ 29 | --batch_size 64 \ 30 | --input_image_size ${resolution} \ 31 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 32 | --num_classes 1000 \ 33 | --evolution_max_iter ${evolution_max_iter} \ 34 | --population_size ${population_size} \ 35 | --save_dir ${save_dir} 36 | 37 | 38 | python analyze_model.py \ 39 | --input_image_size 224 \ 40 | --num_classes 1000 \ 41 | --arch Masternet.py:MasterNet \ 42 | --plainnet_struct_txt ${save_dir}/best_structure.txt 43 | 44 | 45 | python ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 46 | --dist_mode single --workers_per_gpu 6 \ 47 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 48 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 49 | --label_smoothing --random_erase --mixup --auto_augment \ 50 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 51 | --arch Masternet.py:MasterNet \ 52 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 53 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 54 | --teacher_pretrained \ 55 | --teacher_input_image_size 320 \ 56 | --teacher_feature_weight 1.0 \ 57 | --teacher_logit_weight 1.0 \ 58 | --ts_proj_no_relu \ 59 | --ts_proj_no_bn \ 60 | --target_downsample_ratio 16 \ 61 | --batch_size_per_gpu 64 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_latency0.5ms.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/Zen_NAS_ImageNet_latency0.5ms 8 | mkdir -p ${save_dir} 9 | 10 | 11 | resolution=224 12 | budget_latency=5e-4 13 | max_layers=20 14 | population_size=512 15 | epochs=480 16 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 17 | 18 | echo "SuperConvK3BNRELU(3,32,2,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,128,2,64,1)\ 19 | SuperResK3K3(128,256,2,128,1)SuperResK3K3(256,512,2,256,1)\ 20 | SuperConvK1BNRELU(256,512,1,1)" > ${save_dir}/init_plainnet.txt 21 | 22 | python evolution_search.py --gpu 0 \ 23 | --zero_shot_score Zen \ 24 | --fix_initialize \ 25 | --origin \ 26 | --search_space SearchSpace/search_space_XXBL.py \ 27 | --budget_latency ${budget_latency} \ 28 | --max_layers ${max_layers} \ 29 | --batch_size 64 \ 30 | --input_image_size ${resolution} \ 31 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 32 | --num_classes 1000 \ 33 | --evolution_max_iter ${evolution_max_iter} \ 34 | --population_size ${population_size} \ 35 | --save_dir ${save_dir} 36 | 37 | 38 | python analyze_model.py \ 39 | --input_image_size 224 \ 40 | --num_classes 1000 \ 41 | --arch Masternet.py:MasterNet \ 42 | --plainnet_struct_txt ${save_dir}/best_structure.txt 43 | 44 | 45 | python ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 46 | --dist_mode single --workers_per_gpu 6 \ 47 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 48 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 49 | --label_smoothing --random_erase --mixup --auto_augment \ 50 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 51 | --arch Masternet.py:MasterNet \ 52 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 53 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 54 | --teacher_pretrained \ 55 | --teacher_input_image_size 320 \ 56 | --teacher_feature_weight 1.0 \ 57 | --teacher_logit_weight 1.0 \ 58 | --ts_proj_no_relu \ 59 | --ts_proj_no_bn \ 60 | --target_downsample_ratio 16 \ 61 | --batch_size_per_gpu 64 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_latency0.8ms.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/Zen_NAS_ImageNet_latency0.8ms 8 | mkdir -p ${save_dir} 9 | 10 | 11 | resolution=224 12 | budget_latency=8e-4 13 | max_layers=25 14 | population_size=512 15 | epochs=480 16 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 17 | 18 | echo "SuperConvK3BNRELU(3,32,2,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,128,2,64,1)\ 19 | SuperResK3K3(128,256,2,128,1)SuperResK3K3(256,512,2,256,1)\ 20 | SuperConvK1BNRELU(256,512,1,1)" > ${save_dir}/init_plainnet.txt 21 | 22 | python evolution_search.py --gpu 0 \ 23 | --zero_shot_score Zen \ 24 | --fix_initialize \ 25 | --origin \ 26 | --search_space SearchSpace/search_space_XXBL.py \ 27 | --budget_latency ${budget_latency} \ 28 | --max_layers ${max_layers} \ 29 | --batch_size 64 \ 30 | --input_image_size ${resolution} \ 31 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 32 | --num_classes 1000 \ 33 | --evolution_max_iter ${evolution_max_iter} \ 34 | --population_size ${population_size} \ 35 | --save_dir ${save_dir} 36 | 37 | 38 | python analyze_model.py \ 39 | --input_image_size 224 \ 40 | --num_classes 1000 \ 41 | --arch Masternet.py:MasterNet \ 42 | --plainnet_struct_txt ${save_dir}/best_structure.txt 43 | 44 | 45 | python ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 46 | --dist_mode single --workers_per_gpu 6 \ 47 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 48 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 49 | --label_smoothing --random_erase --mixup --auto_augment \ 50 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 51 | --arch Masternet.py:MasterNet \ 52 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 53 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 54 | --teacher_pretrained \ 55 | --teacher_input_image_size 320 \ 56 | --teacher_feature_weight 1.0 \ 57 | --teacher_logit_weight 1.0 \ 58 | --ts_proj_no_relu \ 59 | --ts_proj_no_bn \ 60 | --target_downsample_ratio 16 \ 61 | --batch_size_per_gpu 64 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_ImageNet_latency1.2ms.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | save_dir=../../save_dir/Zen_NAS_ImageNet_latency1.2ms 8 | mkdir -p ${save_dir} 9 | 10 | 11 | resolution=224 12 | budget_latency=12e-4 13 | max_layers=30 14 | population_size=512 15 | epochs=480 16 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for sufficient searching 17 | 18 | echo "SuperConvK3BNRELU(3,32,2,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,128,2,64,1)\ 19 | SuperResK3K3(128,256,2,128,1)SuperResK3K3(256,512,2,256,1)\ 20 | SuperConvK1BNRELU(256,512,1,1)" > ${save_dir}/init_plainnet.txt 21 | 22 | python evolution_search.py --gpu 0 \ 23 | --zero_shot_score Zen \ 24 | --fix_initialize \ 25 | --origin \ 26 | --search_space SearchSpace/search_space_XXBL.py \ 27 | --budget_latency ${budget_latency} \ 28 | --max_layers ${max_layers} \ 29 | --batch_size 64 \ 30 | --input_image_size ${resolution} \ 31 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 32 | --num_classes 1000 \ 33 | --evolution_max_iter ${evolution_max_iter} \ 34 | --population_size ${population_size} \ 35 | --save_dir ${save_dir} 36 | 37 | 38 | python analyze_model.py \ 39 | --input_image_size 224 \ 40 | --num_classes 1000 \ 41 | --arch Masternet.py:MasterNet \ 42 | --plainnet_struct_txt ${save_dir}/best_structure.txt 43 | 44 | 45 | python ts_train_image_classification.py --dataset imagenet --num_classes 1000 \ 46 | --dist_mode single --workers_per_gpu 6 \ 47 | --input_image_size ${resolution} --epochs ${epochs} --warmup 5 \ 48 | --optimizer sgd --bn_momentum 0.01 --wd 4e-5 --nesterov --weight_init custom \ 49 | --label_smoothing --random_erase --mixup --auto_augment \ 50 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 51 | --arch Masternet.py:MasterNet \ 52 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 53 | --teacher_arch geffnet_tf_efficientnet_b3_ns \ 54 | --teacher_pretrained \ 55 | --teacher_input_image_size 320 \ 56 | --teacher_feature_weight 1.0 \ 57 | --teacher_logit_weight 1.0 \ 58 | --ts_proj_no_relu \ 59 | --ts_proj_no_bn \ 60 | --target_downsample_ratio 16 \ 61 | --batch_size_per_gpu 64 --save_dir ${save_dir}/ts_effnet_b3ns_epochs${epochs} -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_cifar_params2M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | budget_model_size=2e6 8 | max_layers=20 9 | population_size=512 10 | evolution_max_iter=480000 11 | 12 | save_dir=../../save_dir/Zen_NAS_cifar_params2M 13 | mkdir -p ${save_dir} 14 | 15 | echo "SuperConvK3BNRELU(3,8,1,1)SuperResK3K3(8,16,1,8,1)SuperResK3K3(16,32,2,16,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,64,2,32,1)SuperConvK1BNRELU(64,128,1,1)" \ 16 | > ${save_dir}/init_plainnet.txt 17 | 18 | python evolution_search.py --gpu 0 \ 19 | --zero_shot_score Zen \ 20 | --search_space SearchSpace/search_space_XXBL.py \ 21 | --budget_model_size ${budget_model_size} \ 22 | --max_layers ${max_layers} \ 23 | --batch_size 64 \ 24 | --input_image_size 32 \ 25 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 26 | --num_classes 100 \ 27 | --evolution_max_iter ${evolution_max_iter} \ 28 | --population_size ${population_size} \ 29 | --save_dir ${save_dir} 30 | 31 | 32 | python analyze_model.py \ 33 | --input_image_size 32 \ 34 | --num_classes 100 \ 35 | --arch Masternet.py:MasterNet \ 36 | --plainnet_struct_txt ${save_dir}/best_structure.txt 37 | 38 | python train_image_classification.py --dataset cifar10 --num_classes 10 \ 39 | --dist_mode single --workers_per_gpu 6 \ 40 | --input_image_size 32 --epochs 1440 --warmup 5 \ 41 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 42 | --label_smoothing --random_erase --mixup --auto_augment \ 43 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 44 | --arch Masternet.py:MasterNet \ 45 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 46 | --batch_size_per_gpu 64 \ 47 | --save_dir ${save_dir}/cifar10_1440epochs 48 | 49 | python train_image_classification.py --dataset cifar100 --num_classes 100 \ 50 | --dist_mode single --workers_per_gpu 6 \ 51 | --input_image_size 32 --epochs 1440 --warmup 5 \ 52 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 53 | --label_smoothing --random_erase --mixup --auto_augment \ 54 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 55 | --arch Masternet.py:MasterNet \ 56 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 57 | --batch_size_per_gpu 64 \ 58 | --save_dir ${save_dir}/cifar100_1440epochs 59 | -------------------------------------------------------------------------------- /src/zen_nas/scripts/Flops_NAS_cifar_params1M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | budget_model_size=1e6 8 | max_layers=18 9 | population_size=512 10 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for 11 | 12 | save_dir=../../save_dir/Flops_NAS_cifar_params1M 13 | mkdir -p ${save_dir} 14 | 15 | echo "SuperConvK3BNRELU(3,8,1,1)SuperResK3K3(8,16,1,8,1)SuperResK3K3(16,32,2,16,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,64,2,32,1)SuperConvK1BNRELU(64,128,1,1)" \ 16 | > ${save_dir}/init_plainnet.txt 17 | 18 | python evolution_search.py --gpu 0 \ 19 | --zero_shot_score Flops \ 20 | --search_space SearchSpace/search_space_XXBL.py \ 21 | --budget_model_size ${budget_model_size} \ 22 | --max_layers ${max_layers} \ 23 | --batch_size 64 \ 24 | --input_image_size 32 \ 25 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 26 | --num_classes 100 \ 27 | --evolution_max_iter ${evolution_max_iter} \ 28 | --population_size ${population_size} \ 29 | --save_dir ${save_dir} 30 | 31 | 32 | python analyze_model.py \ 33 | --input_image_size 32 \ 34 | --num_classes 100 \ 35 | --arch Masternet.py:MasterNet \ 36 | --plainnet_struct_txt ${save_dir}/best_structure.txt 37 | 38 | python train_image_classification.py --dataset cifar10 --num_classes 10 \ 39 | --dist_mode single --workers_per_gpu 6 \ 40 | --input_image_size 32 --epochs 1440 --warmup 5 \ 41 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 42 | --label_smoothing --random_erase --mixup --auto_augment \ 43 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 44 | --arch Masternet.py:MasterNet \ 45 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 46 | --batch_size_per_gpu 64 \ 47 | --save_dir ${save_dir}/cifar10_1440epochs 48 | 49 | python train_image_classification.py --dataset cifar100 --num_classes 100 \ 50 | --dist_mode single --workers_per_gpu 6 \ 51 | --input_image_size 32 --epochs 1440 --warmup 5 \ 52 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 53 | --label_smoothing --random_erase --mixup --auto_augment \ 54 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 55 | --arch Masternet.py:MasterNet \ 56 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 57 | --batch_size_per_gpu 64 \ 58 | --save_dir ${save_dir}/cifar100_1440epochs 59 | -------------------------------------------------------------------------------- /src/zen_nas/scripts/Params_NAS_cifar_params1M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | budget_model_size=1e6 8 | max_layers=18 9 | population_size=512 10 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for 11 | 12 | 13 | save_dir=../../save_dir/Params_NAS_cifar_params1M 14 | mkdir -p ${save_dir} 15 | 16 | echo "SuperConvK3BNRELU(3,8,1,1)SuperResK3K3(8,16,1,8,1)SuperResK3K3(16,32,2,16,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,64,2,32,1)SuperConvK1BNRELU(64,128,1,1)" \ 17 | > ${save_dir}/init_plainnet.txt 18 | 19 | python evolution_search.py --gpu 0 \ 20 | --zero_shot_score Params \ 21 | --search_space SearchSpace/search_space_XXBL.py \ 22 | --budget_model_size ${budget_model_size} \ 23 | --max_layers ${max_layers} \ 24 | --batch_size 64 \ 25 | --input_image_size 32 \ 26 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 27 | --num_classes 100 \ 28 | --evolution_max_iter ${evolution_max_iter} \ 29 | --population_size ${population_size} \ 30 | --save_dir ${save_dir} 31 | 32 | 33 | python analyze_model.py \ 34 | --input_image_size 32 \ 35 | --num_classes 100 \ 36 | --arch Masternet.py:MasterNet \ 37 | --plainnet_struct_txt ${save_dir}/best_structure.txt 38 | 39 | python train_image_classification.py --dataset cifar10 --num_classes 10 \ 40 | --dist_mode single --workers_per_gpu 6 \ 41 | --input_image_size 32 --epochs 1440 --warmup 5 \ 42 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 43 | --label_smoothing --random_erase --mixup --auto_augment \ 44 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 45 | --arch Masternet.py:MasterNet \ 46 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 47 | --batch_size_per_gpu 64 \ 48 | --save_dir ${save_dir}/cifar10_1440epochs 49 | 50 | python train_image_classification.py --dataset cifar100 --num_classes 100 \ 51 | --dist_mode single --workers_per_gpu 6 \ 52 | --input_image_size 32 --epochs 1440 --warmup 5 \ 53 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 54 | --label_smoothing --random_erase --mixup --auto_augment \ 55 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 56 | --arch Masternet.py:MasterNet \ 57 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 58 | --batch_size_per_gpu 64 \ 59 | --save_dir ${save_dir}/cifar100_1440epochs 60 | -------------------------------------------------------------------------------- /src/zen_nas/scripts/Random_NAS_cifar_params1M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | budget_model_size=1e6 8 | max_layers=18 9 | population_size=512 10 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for 11 | 12 | 13 | save_dir=../../save_dir/Random_NAS_cifar_params1M 14 | mkdir -p ${save_dir} 15 | 16 | echo "SuperConvK3BNRELU(3,8,1,1)SuperResK3K3(8,16,1,8,1)SuperResK3K3(16,32,2,16,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,64,2,32,1)SuperConvK1BNRELU(64,128,1,1)" \ 17 | > ${save_dir}/init_plainnet.txt 18 | 19 | python evolution_search.py --gpu 0 \ 20 | --zero_shot_score Random \ 21 | --search_space SearchSpace/search_space_XXBL.py \ 22 | --budget_model_size ${budget_model_size} \ 23 | --max_layers ${max_layers} \ 24 | --batch_size 64 \ 25 | --input_image_size 32 \ 26 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 27 | --num_classes 100 \ 28 | --evolution_max_iter ${evolution_max_iter} \ 29 | --population_size ${population_size} \ 30 | --save_dir ${save_dir} 31 | 32 | 33 | python analyze_model.py \ 34 | --input_image_size 32 \ 35 | --num_classes 100 \ 36 | --arch Masternet.py:MasterNet \ 37 | --plainnet_struct_txt ${save_dir}/best_structure.txt 38 | 39 | python train_image_classification.py --dataset cifar10 --num_classes 10 \ 40 | --dist_mode single --workers_per_gpu 6 \ 41 | --input_image_size 32 --epochs 1440 --warmup 5 \ 42 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 43 | --label_smoothing --random_erase --mixup --auto_augment \ 44 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 45 | --arch Masternet.py:MasterNet \ 46 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 47 | --batch_size_per_gpu 64 \ 48 | --save_dir ${save_dir}/cifar10_1440epochs 49 | 50 | python train_image_classification.py --dataset cifar100 --num_classes 100 \ 51 | --dist_mode single --workers_per_gpu 6 \ 52 | --input_image_size 32 --epochs 1440 --warmup 5 \ 53 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 54 | --label_smoothing --random_erase --mixup --auto_augment \ 55 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 56 | --arch Masternet.py:MasterNet \ 57 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 58 | --batch_size_per_gpu 64 \ 59 | --save_dir ${save_dir}/cifar100_1440epochs 60 | -------------------------------------------------------------------------------- /src/zen_nas/scripts/TE_NAS_cifar_params1M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | budget_model_size=1e6 8 | max_layers=18 9 | population_size=512 10 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for 11 | 12 | 13 | save_dir=../../save_dir/TE_NAS_cifar_params1M 14 | mkdir -p ${save_dir} 15 | 16 | echo "SuperConvK3BNRELU(3,8,1,1)SuperResK3K3(8,16,1,8,1)SuperResK3K3(16,32,2,16,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,64,2,32,1)SuperConvK1BNRELU(64,128,1,1)" \ 17 | > ${save_dir}/init_plainnet.txt 18 | 19 | python evolution_search.py --gpu 0 \ 20 | --zero_shot_score TE-NAS \ 21 | --search_space SearchSpace/search_space_XXBL.py \ 22 | --budget_model_size ${budget_model_size} \ 23 | --max_layers ${max_layers} \ 24 | --batch_size 64 \ 25 | --input_image_size 32 \ 26 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 27 | --num_classes 100 \ 28 | --evolution_max_iter ${evolution_max_iter} \ 29 | --population_size ${population_size} \ 30 | --save_dir ${save_dir} 31 | 32 | 33 | python analyze_model.py \ 34 | --input_image_size 32 \ 35 | --num_classes 100 \ 36 | --arch Masternet.py:MasterNet \ 37 | --plainnet_struct_txt ${save_dir}/best_structure.txt 38 | 39 | python train_image_classification.py --dataset cifar10 --num_classes 10 \ 40 | --dist_mode single --workers_per_gpu 6 \ 41 | --input_image_size 32 --epochs 1440 --warmup 5 \ 42 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 43 | --label_smoothing --random_erase --mixup --auto_augment \ 44 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 45 | --arch Masternet.py:MasterNet \ 46 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 47 | --batch_size_per_gpu 64 \ 48 | --save_dir ${save_dir}/cifar10_1440epochs 49 | 50 | 51 | python train_image_classification.py --dataset cifar100 --num_classes 100 \ 52 | --dist_mode single --workers_per_gpu 6 \ 53 | --input_image_size 32 --epochs 1440 --warmup 5 \ 54 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 55 | --label_smoothing --random_erase --mixup --auto_augment \ 56 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 57 | --arch Masternet.py:MasterNet \ 58 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 59 | --batch_size_per_gpu 64 \ 60 | --save_dir ${save_dir}/cifar100_1440epochs 61 | -------------------------------------------------------------------------------- /src/zen_nas/scripts/Zen_NAS_cifar_params1M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | budget_model_size=1e6 8 | max_layers=18 9 | population_size=512 10 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for 11 | 12 | 13 | save_dir=../../save_dir/Zen_NAS_cifar_params1M 14 | mkdir -p ${save_dir} 15 | 16 | echo "SuperConvK3BNRELU(3,8,1,1)SuperResK3K3(8,16,1,8,1)SuperResK3K3(16,32,2,16,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,64,2,32,1)SuperConvK1BNRELU(64,128,1,1)" \ 17 | > ${save_dir}/init_plainnet.txt 18 | 19 | python evolution_search.py --gpu 0 \ 20 | --zero_shot_score Zen \ 21 | --search_space SearchSpace/search_space_XXBL.py \ 22 | --budget_model_size ${budget_model_size} \ 23 | --max_layers ${max_layers} \ 24 | --batch_size 64 \ 25 | --input_image_size 32 \ 26 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 27 | --num_classes 100 \ 28 | --evolution_max_iter ${evolution_max_iter} \ 29 | --population_size ${population_size} \ 30 | --save_dir ${save_dir} 31 | 32 | 33 | python analyze_model.py \ 34 | --input_image_size 32 \ 35 | --num_classes 100 \ 36 | --arch Masternet.py:MasterNet \ 37 | --plainnet_struct_txt ${save_dir}/best_structure.txt 38 | 39 | python train_image_classification.py --dataset cifar10 --num_classes 10 \ 40 | --dist_mode single --workers_per_gpu 6 \ 41 | --input_image_size 32 --epochs 1440 --warmup 5 \ 42 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 43 | --label_smoothing --random_erase --mixup --auto_augment \ 44 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 45 | --arch Masternet.py:MasterNet \ 46 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 47 | --batch_size_per_gpu 64 \ 48 | --save_dir ${save_dir}/cifar10_1440epochs 49 | 50 | 51 | python train_image_classification.py --dataset cifar100 --num_classes 100 \ 52 | --dist_mode single --workers_per_gpu 6 \ 53 | --input_image_size 32 --epochs 1440 --warmup 5 \ 54 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 55 | --label_smoothing --random_erase --mixup --auto_augment \ 56 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 57 | --arch Masternet.py:MasterNet \ 58 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 59 | --batch_size_per_gpu 64 \ 60 | --save_dir ${save_dir}/cifar100_1440epochs 61 | -------------------------------------------------------------------------------- /src/zen_nas/scripts/NASWOT_NAS_cifar_params1M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | budget_model_size=1e6 8 | max_layers=18 9 | population_size=512 10 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for 11 | 12 | 13 | save_dir=../../save_dir/NASWOT_NAS_cifar_params1M 14 | mkdir -p ${save_dir} 15 | 16 | echo "SuperConvK3BNRELU(3,8,1,1)SuperResK3K3(8,16,1,8,1)SuperResK3K3(16,32,2,16,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,64,2,32,1)SuperConvK1BNRELU(64,128,1,1)" \ 17 | > ${save_dir}/init_plainnet.txt 18 | 19 | 20 | python evolution_search.py --gpu 0 \ 21 | --zero_shot_score NASWOT \ 22 | --search_space SearchSpace/search_space_XXBL.py \ 23 | --budget_model_size ${budget_model_size} \ 24 | --max_layers ${max_layers} \ 25 | --batch_size 8 \ 26 | --input_image_size 32 \ 27 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 28 | --num_classes 100 \ 29 | --evolution_max_iter ${evolution_max_iter} \ 30 | --population_size ${population_size} \ 31 | --save_dir ${save_dir} 32 | 33 | 34 | python analyze_model.py \ 35 | --input_image_size 32 \ 36 | --num_classes 100 \ 37 | --arch Masternet.py:MasterNet \ 38 | --plainnet_struct_txt ${save_dir}/best_structure.txt 39 | 40 | python train_image_classification.py --dataset cifar10 --num_classes 10 \ 41 | --dist_mode single --workers_per_gpu 6 \ 42 | --input_image_size 32 --epochs 1440 --warmup 5 \ 43 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 44 | --label_smoothing --random_erase --mixup --auto_augment \ 45 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 46 | --arch Masternet.py:MasterNet \ 47 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 48 | --batch_size_per_gpu 64 \ 49 | --save_dir ${save_dir}/cifar10_1440epochs 50 | 51 | python train_image_classification.py --dataset cifar100 --num_classes 100 \ 52 | --dist_mode single --workers_per_gpu 6 \ 53 | --input_image_size 32 --epochs 1440 --warmup 5 \ 54 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 55 | --label_smoothing --random_erase --mixup --auto_augment \ 56 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 57 | --arch Masternet.py:MasterNet \ 58 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 59 | --batch_size_per_gpu 64 \ 60 | --save_dir ${save_dir}/cifar100_1440epochs 61 | -------------------------------------------------------------------------------- /src/zen_nas/scripts/Syncflow_NAS_cifar_params1M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | budget_model_size=1e6 8 | max_layers=18 9 | population_size=512 10 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for 11 | 12 | 13 | save_dir=../../save_dir/Syncflow_NAS_cifar_params1M 14 | mkdir -p ${save_dir} 15 | 16 | echo "SuperConvK3BNRELU(3,8,1,1)SuperResK3K3(8,16,1,8,1)SuperResK3K3(16,32,2,16,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,64,2,32,1)SuperConvK1BNRELU(64,128,1,1)" \ 17 | > ${save_dir}/init_plainnet.txt 18 | 19 | python evolution_search.py --gpu 0 \ 20 | --zero_shot_score Syncflow \ 21 | --search_space SearchSpace/search_space_XXBL.py \ 22 | --budget_model_size ${budget_model_size} \ 23 | --max_layers ${max_layers} \ 24 | --batch_size 64 \ 25 | --input_image_size 32 \ 26 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 27 | --num_classes 100 \ 28 | --evolution_max_iter ${evolution_max_iter} \ 29 | --population_size ${population_size} \ 30 | --save_dir ${save_dir} 31 | 32 | 33 | python analyze_model.py \ 34 | --input_image_size 32 \ 35 | --num_classes 100 \ 36 | --arch Masternet.py:MasterNet \ 37 | --plainnet_struct_txt ${save_dir}/best_structure.txt 38 | 39 | python train_image_classification.py --dataset cifar10 --num_classes 10 \ 40 | --dist_mode single --workers_per_gpu 6 \ 41 | --input_image_size 32 --epochs 1440 --warmup 5 \ 42 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 43 | --label_smoothing --random_erase --mixup --auto_augment \ 44 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 45 | --arch Masternet.py:MasterNet \ 46 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 47 | --batch_size_per_gpu 64 \ 48 | --save_dir ${save_dir}/cifar10_1440epochs 49 | 50 | python train_image_classification.py --dataset cifar100 --num_classes 100 \ 51 | --dist_mode single --workers_per_gpu 6 \ 52 | --input_image_size 32 --epochs 1440 --warmup 5 \ 53 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 54 | --label_smoothing --random_erase --mixup --auto_augment \ 55 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 56 | --arch Masternet.py:MasterNet \ 57 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 58 | --batch_size_per_gpu 64 \ 59 | --save_dir ${save_dir}/cifar100_1440epochs 60 | -------------------------------------------------------------------------------- /src/zen_nas/scripts/GradNorm_NAS_cifar_params1M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | cd "$(dirname "$0")" 3 | set -e 4 | 5 | cd ../ 6 | 7 | budget_model_size=1e6 8 | max_layers=18 9 | population_size=512 10 | evolution_max_iter=480000 # we suggest evolution_max_iter=480000 for 11 | 12 | 13 | save_dir=../../save_dir/GradNorm_NAS_cifar_params1M 14 | mkdir -p ${save_dir} 15 | 16 | echo "SuperConvK3BNRELU(3,8,1,1)SuperResK3K3(8,16,1,8,1)SuperResK3K3(16,32,2,16,1)SuperResK3K3(32,64,2,32,1)SuperResK3K3(64,64,2,32,1)SuperConvK1BNRELU(64,128,1,1)" \ 17 | > ${save_dir}/init_plainnet.txt 18 | 19 | python evolution_search.py --gpu 0 \ 20 | --zero_shot_score GradNorm \ 21 | --search_space SearchSpace/search_space_XXBL.py \ 22 | --budget_model_size ${budget_model_size} \ 23 | --max_layers ${max_layers} \ 24 | --batch_size 64 \ 25 | --input_image_size 32 \ 26 | --plainnet_struct_txt ${save_dir}/init_plainnet.txt \ 27 | --num_classes 100 \ 28 | --evolution_max_iter ${evolution_max_iter} \ 29 | --population_size ${population_size} \ 30 | --save_dir ${save_dir} 31 | 32 | 33 | python analyze_model.py \ 34 | --input_image_size 32 \ 35 | --num_classes 100 \ 36 | --arch Masternet.py:MasterNet \ 37 | --plainnet_struct_txt ${save_dir}/best_structure.txt 38 | 39 | 40 | python train_image_classification.py --dataset cifar10 --num_classes 10 \ 41 | --dist_mode single --workers_per_gpu 6 \ 42 | --input_image_size 32 --epochs 1440 --warmup 5 \ 43 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 44 | --label_smoothing --random_erase --mixup --auto_augment \ 45 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 46 | --arch Masternet.py:MasterNet \ 47 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 48 | --batch_size_per_gpu 64 \ 49 | --save_dir ${save_dir}/cifar10_1440epochs 50 | 51 | 52 | python train_image_classification.py --dataset cifar100 --num_classes 100 \ 53 | --dist_mode single --workers_per_gpu 6 \ 54 | --input_image_size 32 --epochs 1440 --warmup 5 \ 55 | --optimizer sgd --bn_momentum 0.01 --wd 5e-4 --nesterov --weight_init custom \ 56 | --label_smoothing --random_erase --mixup --auto_augment \ 57 | --lr_per_256 0.1 --target_lr_per_256 0.0 --lr_mode cosine \ 58 | --arch Masternet.py:MasterNet \ 59 | --plainnet_struct_txt ${save_dir}/best_structure.txt \ 60 | --batch_size_per_gpu 64 \ 61 | --save_dir ${save_dir}/cifar100_1440epochs 62 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | -------------------------------------------------------------------------------- /src/zen_nas/ZeroShotProxy/compute_gradnorm_score.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | 4 | compute gradient norm score 5 | ''' 6 | import os 7 | import sys 8 | import torch 9 | from torch import nn 10 | import torch.nn.functional as F 11 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 12 | 13 | 14 | def network_weight_gaussian_init(net: nn.Module): 15 | """gaussian initialization""" 16 | with torch.no_grad(): 17 | for module in net.modules(): 18 | if isinstance(module, nn.Conv2d): 19 | nn.init.normal_(module.weight) 20 | if hasattr(module, 'bias') and module.bias is not None: 21 | nn.init.zeros_(module.bias) 22 | elif isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)): 23 | nn.init.ones_(module.weight) 24 | nn.init.zeros_(module.bias) 25 | elif isinstance(module, nn.Linear): 26 | nn.init.normal_(module.weight) 27 | if hasattr(module, 'bias') and module.bias is not None: 28 | nn.init.zeros_(module.bias) 29 | else: 30 | continue 31 | 32 | return net 33 | 34 | 35 | def cross_entropy(logit, target): 36 | """compute cross entropy loss""" 37 | # target must be one-hot format!! 38 | prob_logit = F.log_softmax(logit, dim=1) 39 | loss = -(target * prob_logit).sum(dim=1).mean() 40 | return loss 41 | 42 | 43 | def compute_nas_score(gpu, model, resolution, batch_size): 44 | """compute gradient norm score""" 45 | model.train() 46 | model.requires_grad_(True) 47 | 48 | model.zero_grad() 49 | 50 | if gpu is not None: 51 | torch.cuda.set_device(gpu) 52 | model = model.cuda(gpu) 53 | 54 | network_weight_gaussian_init(model) 55 | input_ = torch.randn(size=[batch_size, 3, resolution, resolution]) 56 | if gpu is not None: 57 | input_ = input_.cuda(gpu) 58 | output = model(input_) 59 | # y_true = torch.rand(size=[batch_size, output.shape[1]], device=torch.device('cuda:{}'.format(gpu))) + 1e-10 60 | # y_true = y_true / torch.sum(y_true, dim=1, keepdim=True) 61 | 62 | num_classes = output.shape[1] 63 | label = torch.randint(low=0, high=num_classes, size=[batch_size]) 64 | 65 | one_hot_y = F.one_hot(label, num_classes).float() 66 | if gpu is not None: 67 | one_hot_y = one_hot_y.cuda(gpu) 68 | 69 | loss = cross_entropy(output, one_hot_y) 70 | loss.backward() 71 | norm2_sum = 0 72 | with torch.no_grad(): 73 | for parameter in model.parameters(): 74 | if hasattr(parameter, 'grad') and parameter.grad is not None: 75 | norm2_sum += torch.norm(parameter.grad) ** 2 76 | 77 | grad_norm = float(torch.sqrt(norm2_sum)) 78 | 79 | return grad_norm 80 | -------------------------------------------------------------------------------- /src/zen_nas/ModelLoader/geffnet/activations/activations_jit.py: -------------------------------------------------------------------------------- 1 | """ Activations (jit) 2 | 3 | A collection of jit-scripted activations fn and modules with a common interface so that they can 4 | easily be swapped. All have an `inplace` arg even if not used. 5 | 6 | All jit scripted activations are lacking in-place variations on purpose, scripted kernel fusion does not 7 | currently work across in-place op boundaries, thus performance is equal to or less than the non-scripted 8 | versions if they contain in-place ops. 9 | 10 | Copyright 2020 Ross Wightman 11 | """ 12 | # pylint: disable=W0613,no-self-use 13 | import torch 14 | from torch import nn 15 | from torch.nn import functional as F 16 | 17 | __all__ = ['swish_jit', 'SwishJit', 'mish_jit', 'MishJit', 18 | 'hard_sigmoid_jit', 'HardSigmoidJit', 'hard_swish_jit', 'HardSwishJit'] 19 | 20 | 21 | @torch.jit.script 22 | def swish_jit(input_, inplace: bool = False): 23 | """Swish - Described originally as SiLU (https://arxiv.org/abs/1702.03118v3) 24 | and also as Swish (https://arxiv.org/abs/1710.05941). 25 | 26 | TODO Rename to SiLU with addition to PyTorch 27 | """ 28 | return input_.mul(input_.sigmoid()) 29 | 30 | 31 | @torch.jit.script 32 | def mish_jit(input_, _inplace: bool = False): 33 | """Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681 34 | """ 35 | return input_.mul(F.softplus(input_).tanh()) 36 | 37 | 38 | class SwishJit(nn.Module): 39 | """jit-scripted swish module""" 40 | def __init__(self, inplace: bool = False): 41 | super().__init__() 42 | 43 | def forward(self, input_): 44 | return swish_jit(input_) 45 | 46 | 47 | class MishJit(nn.Module): 48 | """jit-scripted mish module""" 49 | def __init__(self, inplace: bool = False): 50 | super().__init__() 51 | 52 | def forward(self, input_): 53 | return mish_jit(input_) 54 | 55 | 56 | @torch.jit.script 57 | def hard_sigmoid_jit(input_, inplace: bool = False): 58 | """jit-scripted hard sigmoid function""" 59 | # return F.relu6(x + 3.) / 6. 60 | return (input_ + 3).clamp(min=0, max=6).div(6.) # clamp seems ever so slightly faster? 61 | 62 | 63 | class HardSigmoidJit(nn.Module): 64 | """jit-scripted hard sigmoid module""" 65 | def __init__(self, inplace: bool = False): 66 | super().__init__() 67 | 68 | def forward(self, input_): 69 | return hard_sigmoid_jit(input_) 70 | 71 | 72 | @torch.jit.script 73 | def hard_swish_jit(input_, inplace: bool = False): 74 | """jit-scripted hard swish function""" 75 | # return x * (F.relu6(x + 3.) / 6) 76 | return input_ * (input_ + 3).clamp(min=0, max=6).div(6.) # clamp seems ever so slightly faster? 77 | 78 | 79 | class HardSwishJit(nn.Module): 80 | """jit-scripted hard swish module""" 81 | def __init__(self, inplace: bool = False): 82 | super().__init__() 83 | 84 | def forward(self, input_): 85 | return hard_swish_jit(input_) 86 | -------------------------------------------------------------------------------- /src/zen_nas/ModelLoader/geffnet/helpers.py: -------------------------------------------------------------------------------- 1 | """ Checkpoint loading / state_dict helpers 2 | Copyright 2020 Ross Wightman 3 | """ 4 | import os 5 | from collections import OrderedDict 6 | import torch 7 | try: 8 | from torch.hub import load_state_dict_from_url 9 | except ImportError: 10 | from torch.utils.model_zoo import load_url as load_state_dict_from_url 11 | 12 | 13 | def load_checkpoint(model, checkpoint_path): 14 | """load checkpoint""" 15 | if checkpoint_path and os.path.isfile(checkpoint_path): 16 | print(f"=> Loading checkpoint '{checkpoint_path}'") 17 | checkpoint = torch.load(checkpoint_path) 18 | if isinstance(checkpoint, dict) and 'state_dict' in checkpoint: 19 | new_state_dict = OrderedDict() 20 | for k, value in checkpoint['state_dict'].items(): 21 | if k.startswith('module'): 22 | name = k[7:] # remove `module.` 23 | else: 24 | name = k 25 | new_state_dict[name] = value 26 | model.load_state_dict(new_state_dict) 27 | else: 28 | model.load_state_dict(checkpoint) 29 | print(f"=> Loaded checkpoint '{checkpoint_path}'") 30 | else: 31 | print(f"=> Error: No checkpoint found at '{checkpoint_path}'") 32 | raise FileNotFoundError() 33 | 34 | 35 | def load_pretrained(model, url, filter_fn=None, strict=True): 36 | """load pretrained model and process weights according to the model """ 37 | if not url: 38 | print("=> Warning: Pretrained model URL is empty, using random initialization.") 39 | return 40 | 41 | state_dict = load_state_dict_from_url(url, progress=False, map_location='cpu') 42 | 43 | input_conv = 'conv_stem' 44 | classifier = 'classifier' 45 | in_chans = getattr(model, input_conv).weight.shape[1] 46 | num_classes = getattr(model, classifier).weight.shape[0] 47 | 48 | input_conv_weight = input_conv + '.weight' 49 | pretrained_in_chans = state_dict[input_conv_weight].shape[1] 50 | if in_chans != pretrained_in_chans: 51 | if in_chans == 1: 52 | print(f'=> Converting pretrained input conv {input_conv_weight} from {pretrained_in_chans} to 1 channel') 53 | conv1_weight = state_dict[input_conv_weight] 54 | state_dict[input_conv_weight] = conv1_weight.sum(dim=1, keepdim=True) 55 | else: 56 | print('=> Discarding pretrained input conv {input_conv_weight} \ 57 | since input channel count != {pretrained_in_chans}') 58 | del state_dict[input_conv_weight] 59 | strict = False 60 | 61 | classifier_weight = classifier + '.weight' 62 | pretrained_num_classes = state_dict[classifier_weight].shape[0] 63 | if num_classes != pretrained_num_classes: 64 | print(f'=> Discarding pretrained classifier since num_classes != {pretrained_num_classes}') 65 | del state_dict[classifier_weight] 66 | del state_dict[classifier + '.bias'] 67 | strict = False 68 | 69 | if filter_fn is not None: 70 | state_dict = filter_fn(state_dict) 71 | 72 | model.load_state_dict(state_dict, strict=strict) 73 | -------------------------------------------------------------------------------- /src/zen_nas/hotfix/vision.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=no-self-use 2 | import os 3 | import torch 4 | from torch.utils import data 5 | 6 | 7 | class VisionDataset(data.Dataset): 8 | _repr_indent = 4 9 | 10 | def __init__(self, root, transforms=None, transform=None, target_transform=None): 11 | if isinstance(root, torch._six.string_classes): 12 | root = os.path.expanduser(root) 13 | self.root = root 14 | 15 | has_transforms = transforms is not None 16 | has_separate_transform = transform is not None or target_transform is not None 17 | if has_transforms and has_separate_transform: 18 | raise ValueError("Only transforms or transform/target_transform can " 19 | "be passed as argument") 20 | 21 | # for backwards-compatibility 22 | self.transform = transform 23 | self.target_transform = target_transform 24 | 25 | if has_separate_transform: 26 | transforms = StandardTransform(transform, target_transform) 27 | self.transforms = transforms 28 | 29 | def __getitem__(self, index): 30 | raise NotImplementedError 31 | 32 | def __len__(self): 33 | raise NotImplementedError 34 | 35 | def __repr__(self): 36 | head = "Dataset " + self.__class__.__name__ 37 | body = [f"Number of datapoints: {self.__len__()}"] 38 | if self.root is not None: 39 | body.append(f"Root location: {self.root}") 40 | body += self.extra_repr().splitlines() 41 | if self.transforms is not None: 42 | body += [repr(self.transforms)] 43 | lines = [head] + [" " * self._repr_indent + line for line in body] 44 | return '\n'.join(lines) 45 | 46 | def _format_transform_repr(self, transform, head): 47 | lines = transform.__repr__().splitlines() 48 | return ([f"{head}{lines[0]}"] + [f"{' ' * len(head)}{line}" for line in lines[1:]]) 49 | 50 | def extra_repr(self): 51 | return "" 52 | 53 | 54 | class StandardTransform(): 55 | """transform""" 56 | def __init__(self, transform=None, target_transform=None): 57 | self.transform = transform 58 | self.target_transform = target_transform 59 | 60 | def __call__(self, input_, target): 61 | if self.transform is not None: 62 | input_ = self.transform(input_) 63 | if self.target_transform is not None: 64 | target = self.target_transform(target) 65 | return input_, target 66 | 67 | def _format_transform_repr(self, transform, head): 68 | lines = transform.__repr__().splitlines() 69 | return ([f"{head}{lines[0]}"] + [f"{' ' * len(head)}{line}" for line in lines[1:]]) 70 | 71 | def __repr__(self): 72 | body = [self.__class__.__name__] 73 | if self.transform is not None: 74 | body += self._format_transform_repr(self.transform, 75 | "Transform: ") 76 | if self.target_transform is not None: 77 | body += self._format_transform_repr(self.target_transform, 78 | "Target transform: ") 79 | 80 | return '\n'.join(body) 81 | -------------------------------------------------------------------------------- /src/zen_nas/ModelLoader/geffnet/activations/activations.py: -------------------------------------------------------------------------------- 1 | """ Activations 2 | 3 | A collection of activations fn and modules with a common interface so that they can 4 | easily be swapped. All have an `inplace` arg even if not used. 5 | 6 | Copyright 2020 Ross Wightman 7 | """ 8 | # pylint: disable=W0613 9 | from torch import nn 10 | from torch.nn import functional as F 11 | 12 | 13 | def swish(input_, inplace: bool = False): 14 | """Swish - Described originally as SiLU (https://arxiv.org/abs/1702.03118v3) 15 | and also as Swish (https://arxiv.org/abs/1710.05941). 16 | 17 | TODO Rename to SiLU with addition to PyTorch 18 | """ 19 | return input_.mul_(input_.sigmoid()) if inplace else input_.mul(input_.sigmoid()) 20 | 21 | 22 | class Swish(nn.Module): 23 | """Swish module""" 24 | def __init__(self, inplace: bool = False): 25 | super().__init__() 26 | self.inplace = inplace 27 | 28 | def forward(self, input_): 29 | return swish(input_, self.inplace) 30 | 31 | 32 | def mish(input_, inplace: bool = False): 33 | """Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681 34 | """ 35 | return input_.mul(F.softplus(input_).tanh()) 36 | 37 | 38 | class Mish(nn.Module): 39 | """Mish module""" 40 | def __init__(self, inplace: bool = False): 41 | super().__init__() 42 | self.inplace = inplace 43 | 44 | def forward(self, input_): 45 | return mish(input_, self.inplace) 46 | 47 | 48 | def sigmoid(input_, inplace: bool = False): 49 | """sigmoid function""" 50 | return input_.sigmoid_() if inplace else input_.sigmoid() 51 | 52 | 53 | # PyTorch has this, but not with a consistent inplace argmument interface 54 | class Sigmoid(nn.Module): 55 | """sigmoid module""" 56 | def __init__(self, inplace: bool = False): 57 | super().__init__() 58 | self.inplace = inplace 59 | 60 | def forward(self, input_): 61 | return input_.sigmoid_() if self.inplace else input_.sigmoid() 62 | 63 | 64 | def tanh(input_, inplace: bool = False): 65 | """Tanh function""" 66 | return input_.tanh_() if inplace else input_.tanh() 67 | 68 | 69 | # PyTorch has this, but not with a consistent inplace argmument interface 70 | class Tanh(nn.Module): 71 | """Tanh module""" 72 | def __init__(self, inplace: bool = False): 73 | super().__init__() 74 | self.inplace = inplace 75 | 76 | def forward(self, input_): 77 | return input_.tanh_() if self.inplace else input_.tanh() 78 | 79 | 80 | def hard_swish(input_, inplace: bool = False): 81 | """hard swish function""" 82 | inner = F.relu6(input_ + 3.).div_(6.) 83 | return input_.mul_(inner) if inplace else input_.mul(inner) 84 | 85 | 86 | class HardSwish(nn.Module): 87 | """hard swish module""" 88 | def __init__(self, inplace: bool = False): 89 | super().__init__() 90 | self.inplace = inplace 91 | 92 | def forward(self, input_): 93 | return hard_swish(input_, self.inplace) 94 | 95 | 96 | def hard_sigmoid(input_, inplace: bool = False): 97 | """hard sigmoid function""" 98 | if inplace: 99 | return input_.add_(3.).clamp_(0., 6.).div_(6.) 100 | return F.relu6(input_ + 3.) / 6. 101 | 102 | 103 | class HardSigmoid(nn.Module): 104 | """hard sigmoid module""" 105 | def __init__(self, inplace: bool = False): 106 | super().__init__() 107 | self.inplace = inplace 108 | 109 | def forward(self, input_): 110 | return hard_sigmoid(input_, self.inplace) 111 | -------------------------------------------------------------------------------- /src/zen_nas/ModelLoader/__init__.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | 4 | The geffnet module is modified from: 5 | https://github.com/rwightman/gen-efficientnet-pytorch 6 | ''' 7 | # pylint: disable=invalid-name 8 | import os 9 | import sys 10 | import importlib.util 11 | import torchvision.models 12 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 13 | sys.path.append(os.path.dirname(os.path.abspath(__file__))) 14 | try: 15 | import PlainNet 16 | import geffnet 17 | import myresnet 18 | except ImportError: 19 | print('fail to import zen_nas modules') 20 | 21 | 22 | torchvision_model_name_list = sorted(name for name, value in torchvision.models.__dict__.items() 23 | if name.islower() and not name.startswith("__") and callable(value)) 24 | 25 | 26 | def _get_model_(arch, num_classes, pretrained=False, opt=None, argv=None): 27 | """create arch model and return 28 | 29 | :param arch (str): model 30 | :param num_classes (int): class number 31 | :param pretrained (bool): pretrained 32 | :return arch model 33 | """ 34 | 35 | # load torch vision model 36 | if arch in torchvision_model_name_list: 37 | if pretrained: 38 | print(f'Using pretrained model: {arch}') 39 | model = torchvision.models.__dict__[arch](pretrained=True, num_classes=num_classes) 40 | else: 41 | print(f'Create torchvision model: {arch}') 42 | model = torchvision.models.__dict__[arch](num_classes=num_classes) 43 | 44 | # my implementation of resnet 45 | elif arch == 'myresnet18': 46 | print(f'Create model: {arch}') 47 | model = myresnet.resnet18(pretrained=False, opt=opt, argv=argv) 48 | elif arch == 'myresnet34': 49 | print(f'Create model: {arch}') 50 | model = myresnet.resnet34(pretrained=False, opt=opt, argv=argv) 51 | elif arch == 'myresnet50': 52 | print(f'Create model: {arch}') 53 | model = myresnet.resnet50(pretrained=False, opt=opt, argv=argv) 54 | elif arch == 'myresnet101': 55 | print(f'Create model: {arch}') 56 | model = myresnet.resnet101(pretrained=False, opt=opt, argv=argv) 57 | elif arch == 'myresnet152': 58 | print(f'Create model: {arch}') 59 | model = myresnet.resnet152(pretrained=False, opt=opt, argv=argv) 60 | 61 | # geffnet 62 | elif arch.startswith('geffnet_'): 63 | model_name = arch[len('geffnet_'):] 64 | model = geffnet.create_model(model_name, pretrained=pretrained) 65 | 66 | # PlainNet 67 | elif arch == 'PlainNet': 68 | print(f'Create model: {arch}') 69 | model = PlainNet.PlainNet(num_classes=num_classes, opt=opt, argv=argv) 70 | 71 | # Any PlainNet 72 | elif arch.find('.py:MasterNet') >= 0: 73 | module_path = arch.split(':')[0] 74 | assert arch.split(':')[1] == 'MasterNet' 75 | my_working_dir = os.path.dirname(os.path.dirname(__file__)) 76 | module_full_path = os.path.join(my_working_dir, module_path) 77 | 78 | spec = importlib.util.spec_from_file_location('any_plain_net', module_full_path) 79 | any_plain_net = importlib.util.module_from_spec(spec) 80 | spec.loader.exec_module(any_plain_net) 81 | print(f'Create model: {arch}') 82 | model = any_plain_net.MasterNet(num_classes=num_classes, opt=opt, argv=argv) 83 | 84 | else: 85 | raise ValueError(f'Unknown model arch: {arch}') 86 | 87 | return model 88 | 89 | 90 | def get_model(opt, argv): 91 | """get arch model""" 92 | return _get_model_(arch=opt.arch, num_classes=opt.num_classes, pretrained=opt.pretrained, opt=opt, argv=argv) 93 | -------------------------------------------------------------------------------- /src/zen_nas/ModelLoader/geffnet/config.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=global-statement,invalid-name 2 | from typing import Any, Optional 3 | from typing_extensions import Literal 4 | 5 | __all__ = [ 6 | 'is_exportable', 'is_scriptable', 'is_no_jit', 'layer_config_kwargs', 7 | 'set_exportable', 'set_scriptable', 'set_no_jit', 'set_layer_config' 8 | ] 9 | 10 | # Set to True if prefer to have layers with no jit optimization (includes activations) 11 | _NO_JIT = False 12 | 13 | # Set to True if prefer to have activation layers with no jit optimization 14 | # NOTE not currently used as no difference between no_jit and no_activation jit as only layers obeying 15 | # the jit flags so far are activations. This will change as more layers are updated and/or added. 16 | _NO_ACTIVATION_JIT = False 17 | 18 | # Set to True if exporting a model with Same padding via ONNX 19 | _EXPORTABLE = False 20 | 21 | # Set to True if wanting to use torch.jit.script on a model 22 | _SCRIPTABLE = False 23 | 24 | 25 | def is_no_jit(): 26 | """get no_jit state""" 27 | return _NO_JIT 28 | 29 | 30 | class set_no_jit: 31 | def __init__(self, mode: bool) -> None: 32 | """set no_jit and save previous state""" 33 | global _NO_JIT 34 | self.prev = _NO_JIT 35 | _NO_JIT = mode 36 | 37 | def __enter__(self) -> None: 38 | pass 39 | 40 | def __exit__(self, *args: Any) -> Literal[False]: 41 | global _NO_JIT 42 | _NO_JIT = self.prev 43 | return False 44 | 45 | 46 | def is_exportable(): 47 | """get exportable state""" 48 | return _EXPORTABLE 49 | 50 | 51 | class set_exportable: 52 | """set exportable and save previous state""" 53 | def __init__(self, mode: bool) -> None: 54 | global _EXPORTABLE 55 | self.prev = _EXPORTABLE 56 | _EXPORTABLE = mode 57 | 58 | def __enter__(self) -> None: 59 | pass 60 | 61 | def __exit__(self, *args: Any) -> Literal[False]: 62 | global _EXPORTABLE 63 | _EXPORTABLE = self.prev 64 | return False 65 | 66 | 67 | def is_scriptable(): 68 | """get scriptable state""" 69 | return _SCRIPTABLE 70 | 71 | 72 | class set_scriptable: 73 | """set scriptable and save previous state""" 74 | def __init__(self, mode: bool) -> None: 75 | global _SCRIPTABLE 76 | self.prev = _SCRIPTABLE 77 | _SCRIPTABLE = mode 78 | 79 | def __enter__(self) -> None: 80 | pass 81 | 82 | def __exit__(self, *args: Any) -> Literal[False]: 83 | global _SCRIPTABLE 84 | _SCRIPTABLE = self.prev 85 | return False 86 | 87 | 88 | class set_layer_config: 89 | """ Layer config context manager that allows setting all layer config flags at once. 90 | If a flag arg is None, it will not change the current value. 91 | """ 92 | def __init__( 93 | self, 94 | scriptable: Optional[bool] = None, 95 | exportable: Optional[bool] = None, 96 | no_jit: Optional[bool] = None, 97 | no_activation_jit: Optional[bool] = None): 98 | global _SCRIPTABLE 99 | global _EXPORTABLE 100 | global _NO_JIT 101 | global _NO_ACTIVATION_JIT 102 | self.prev = _SCRIPTABLE, _EXPORTABLE, _NO_JIT, _NO_ACTIVATION_JIT 103 | if scriptable is not None: 104 | _SCRIPTABLE = scriptable 105 | if exportable is not None: 106 | _EXPORTABLE = exportable 107 | if no_jit is not None: 108 | _NO_JIT = no_jit 109 | if no_activation_jit is not None: 110 | _NO_ACTIVATION_JIT = no_activation_jit 111 | 112 | def __enter__(self) -> None: 113 | pass 114 | 115 | def __exit__(self, *args: Any) -> Literal[False]: 116 | global _SCRIPTABLE 117 | global _EXPORTABLE 118 | global _NO_JIT 119 | global _NO_ACTIVATION_JIT 120 | _SCRIPTABLE, _EXPORTABLE, _NO_JIT, _NO_ACTIVATION_JIT = self.prev 121 | return False 122 | 123 | 124 | def layer_config_kwargs(kwargs): 125 | """ Consume config kwargs and return contextmgr obj """ 126 | return set_layer_config( 127 | scriptable=kwargs.pop('scriptable', None), 128 | exportable=kwargs.pop('exportable', None), 129 | no_jit=kwargs.pop('no_jit', None)) 130 | -------------------------------------------------------------------------------- /src/zen_nas/ZeroShotProxy/compute_syncflow_score.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | This file is modified from 4 | https://github.com/SamsungLabs/zero-cost-nas 5 | ''' 6 | # Copyright 2021 Samsung Electronics Co., Ltd. 7 | # 8 | # Licensed under the Apache License, Version 2.0 (the "License"); 9 | # you may not use this file except in compliance with the License. 10 | # You may obtain a copy of the License at 11 | 12 | # http://www.apache.org/licenses/LICENSE-2.0 13 | 14 | # Unless required by applicable law or agreed to in writing, software 15 | # distributed under the License is distributed on an "AS IS" BASIS, 16 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 17 | # See the License for the specific language governing permissions and 18 | # limitations under the License. 19 | # ============================================================================= 20 | import os 21 | import sys 22 | import torch 23 | from torch import nn 24 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 25 | 26 | 27 | def network_weight_gaussian_init(net: nn.Module): 28 | with torch.no_grad(): 29 | for module in net.modules(): 30 | if isinstance(module, nn.Conv2d): 31 | nn.init.normal_(module.weight) 32 | if hasattr(module, 'bias') and module.bias is not None: 33 | nn.init.zeros_(module.bias) 34 | elif isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)): 35 | nn.init.ones_(module.weight) 36 | nn.init.zeros_(module.bias) 37 | elif isinstance(module, nn.Linear): 38 | nn.init.normal_(module.weight) 39 | if hasattr(module, 'bias') and module.bias is not None: 40 | nn.init.zeros_(module.bias) 41 | else: 42 | continue 43 | 44 | return net 45 | 46 | 47 | def get_layer_metric_array(net, metric, mode): 48 | metric_array = [] 49 | 50 | for layer in net.modules(): 51 | if mode == 'channel' and hasattr(layer, 'dont_ch_prune'): 52 | continue 53 | if isinstance(layer, (nn.Conv2d, nn.Linear)): 54 | metric_array.append(metric(layer)) 55 | 56 | return metric_array 57 | 58 | 59 | def compute_synflow_per_weight(net, inputs, mode): 60 | """compute synflow for each weight""" 61 | device = inputs.device 62 | 63 | # convert params to their abs. Keep sign for converting it back. 64 | @torch.no_grad() 65 | def linearize(net): 66 | signs = {} 67 | for name, param in net.state_dict().items(): 68 | signs[name] = torch.sign(param) 69 | param.abs_() 70 | return signs 71 | 72 | # convert to orig values 73 | @torch.no_grad() 74 | def nonlinearize(net, signs): 75 | for name, param in net.state_dict().items(): 76 | if 'weight_mask' not in name: 77 | param.mul_(signs[name]) 78 | 79 | # keep signs of all params 80 | signs = linearize(net) 81 | 82 | # Compute gradients with input of 1s 83 | net.zero_grad() 84 | net.double() 85 | input_dim = list(inputs[0, :].shape) 86 | inputs = torch.ones([1] + input_dim).double().to(device) 87 | output = net.forward(inputs) 88 | torch.sum(output).backward() 89 | 90 | # select the gradients that we want to use for search/prune 91 | def synflow(layer): 92 | if layer.weight.grad is not None: 93 | return torch.abs(layer.weight * layer.weight.grad) 94 | return torch.zeros_like(layer.weight) 95 | 96 | grads_abs = get_layer_metric_array(net, synflow, mode) 97 | 98 | # apply signs of all params 99 | nonlinearize(net, signs) 100 | 101 | return grads_abs 102 | 103 | 104 | def do_compute_nas_score(gpu, model, resolution, batch_size): 105 | """compute syncflow score""" 106 | model.train() 107 | model.requires_grad_(True) 108 | 109 | model.zero_grad() 110 | 111 | if gpu is not None: 112 | torch.cuda.set_device(gpu) 113 | model = model.cuda(gpu) 114 | 115 | network_weight_gaussian_init(model) 116 | input_ = torch.randn(size=[batch_size, 3, resolution, resolution]) 117 | if gpu is not None: 118 | input_ = input_.cuda(gpu) 119 | 120 | grads_abs_list = compute_synflow_per_weight(net=model, inputs=input_, mode='') 121 | score = 0 122 | for grad_abs in grads_abs_list: 123 | if len(grad_abs.shape) == 4: 124 | score += float(torch.mean(torch.sum(grad_abs, dim=[1, 2, 3]))) 125 | elif len(grad_abs.shape) == 2: 126 | score += float(torch.mean(torch.sum(grad_abs, dim=[1]))) 127 | else: 128 | raise RuntimeError('!!!') 129 | 130 | return -1 * score 131 | -------------------------------------------------------------------------------- /src/zen_nas/ModelLoader/geffnet/activations/__init__.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=global-variable-not-assigned 2 | import os 3 | import sys 4 | from torch import nn 5 | from torch.nn import functional as F 6 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 7 | try: 8 | from geffnet import config 9 | # from geffnet.activations.activations_me import * 10 | # from geffnet.activations.activations_jit import * 11 | # from geffnet.activations.activations import * 12 | from geffnet.activations.activations import swish, Swish, sigmoid, Sigmoid,\ 13 | mish, Mish, tanh, Tanh, hard_sigmoid, hard_swish, HardSigmoid, HardSwish 14 | from geffnet.activations.activations_jit import swish_jit, SwishJit, mish_jit, MishJit 15 | from geffnet.activations.activations_me import swish_me, mish_me, hard_swish_me,\ 16 | hard_sigmoid_me, HardSigmoidMe, HardSwishMe, SwishMe, MishMe 17 | except ImportError: 18 | print('fail to import zen_nas modules') 19 | 20 | _ACT_FN_DEFAULT = dict( 21 | swish=swish, 22 | mish=mish, 23 | relu=F.relu, 24 | relu6=F.relu6, 25 | sigmoid=sigmoid, 26 | tanh=tanh, 27 | hard_sigmoid=hard_sigmoid, 28 | hard_swish=hard_swish, 29 | ) 30 | 31 | _ACT_FN_JIT = dict( 32 | swish=swish_jit, 33 | mish=mish_jit, 34 | ) 35 | 36 | _ACT_FN_ME = dict( 37 | swish=swish_me, 38 | mish=mish_me, 39 | hard_swish=hard_swish_me, 40 | hard_sigmoid_jit=hard_sigmoid_me, 41 | ) 42 | 43 | _ACT_LAYER_DEFAULT = dict( 44 | swish=Swish, 45 | mish=Mish, 46 | relu=nn.ReLU, 47 | relu6=nn.ReLU6, 48 | sigmoid=Sigmoid, 49 | tanh=Tanh, 50 | hard_sigmoid=HardSigmoid, 51 | hard_swish=HardSwish, 52 | ) 53 | 54 | _ACT_LAYER_JIT = dict( 55 | swish=SwishJit, 56 | mish=MishJit, 57 | ) 58 | 59 | _ACT_LAYER_ME = dict( 60 | swish=SwishMe, 61 | mish=MishMe, 62 | hard_swish=HardSwishMe, 63 | hard_sigmoid=HardSigmoidMe 64 | ) 65 | 66 | _OVERRIDE_FN = {} 67 | _OVERRIDE_LAYER = {} 68 | 69 | 70 | def add_override_act_fn(name, function): 71 | """add function to dict""" 72 | global _OVERRIDE_FN 73 | _OVERRIDE_FN[name] = function 74 | 75 | 76 | def update_override_act_fn(overrides): 77 | """update function dict""" 78 | assert isinstance(overrides, dict) 79 | global _OVERRIDE_FN 80 | _OVERRIDE_FN.update(overrides) 81 | 82 | 83 | # pylint: disable=global-statement 84 | def clear_override_act_fn(): 85 | global _OVERRIDE_FN 86 | _OVERRIDE_FN = {} 87 | 88 | 89 | def add_override_act_layer(name, function): 90 | """add layer to dict""" 91 | _OVERRIDE_LAYER[name] = function 92 | 93 | 94 | def update_override_act_layer(overrides): 95 | assert isinstance(overrides, dict) 96 | global _OVERRIDE_LAYER 97 | _OVERRIDE_LAYER.update(overrides) 98 | 99 | 100 | # pylint: disable=global-statement 101 | def clear_override_act_layer(): 102 | global _OVERRIDE_LAYER 103 | _OVERRIDE_LAYER = {} 104 | 105 | 106 | def get_act_fn(name='relu'): 107 | """ Activation Function Factory 108 | Fetching activation fns by name with this function allows export or torch script friendly 109 | functions to be returned dynamically based on current config. 110 | """ 111 | if name in _OVERRIDE_FN: 112 | return _OVERRIDE_FN[name] 113 | no_me = config.is_exportable() or config.is_scriptable() or config.is_no_jit() 114 | if not no_me and name in _ACT_FN_ME: 115 | # If not exporting or scripting the model, first look for a memory optimized version 116 | # activation with custom autograd, then fallback to jit scripted, then a Python or Torch builtin 117 | return _ACT_FN_ME[name] 118 | no_jit = config.is_exportable() or config.is_no_jit() 119 | # NOTE: export tracing should work with jit scripted components, but I keep running into issues 120 | if no_jit and name in _ACT_FN_JIT: # jit scripted models should be okay for export/scripting 121 | return _ACT_FN_JIT[name] 122 | return _ACT_FN_DEFAULT[name] 123 | 124 | 125 | def get_act_layer(name='relu'): 126 | """ Activation Layer Factory 127 | Fetching activation layers by name with this function allows export or torch script friendly 128 | functions to be returned dynamically based on current config. 129 | """ 130 | if name in _OVERRIDE_LAYER: 131 | return _OVERRIDE_LAYER[name] 132 | no_me = config.is_exportable() or config.is_scriptable() or config.is_no_jit() 133 | if not no_me and name in _ACT_LAYER_ME: 134 | return _ACT_LAYER_ME[name] 135 | no_jit = config.is_exportable() or config.is_no_jit() 136 | if not no_jit and name in _ACT_LAYER_JIT: # jit scripted models should be okay for export/scripting 137 | return _ACT_LAYER_JIT[name] 138 | return _ACT_LAYER_DEFAULT[name] 139 | -------------------------------------------------------------------------------- /src/zen_nas/ZeroShotProxy/compute_NASWOT_score.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | 4 | The implementation of NASWOT score is modified from: 5 | https://github.com/BayesWatch/nas-without-training 6 | ''' 7 | # pylint: disable=W0613,invalid-name 8 | import os 9 | import sys 10 | import time 11 | import argparse 12 | import torch 13 | from torch import nn 14 | import numpy as np 15 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 16 | try: 17 | from PlainNet import basic_blocks 18 | import global_utils 19 | import ModelLoader 20 | except ImportError: 21 | print('fail to import zen_nas modules') 22 | 23 | 24 | def network_weight_gaussian_init(net: nn.Module): 25 | with torch.no_grad(): 26 | for module in net.modules(): 27 | if isinstance(module, nn.Conv2d): 28 | nn.init.normal_(module.weight) 29 | if hasattr(module, 'bias') and module.bias is not None: 30 | nn.init.zeros_(module.bias) 31 | elif isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)): 32 | nn.init.ones_(module.weight) 33 | nn.init.zeros_(module.bias) 34 | elif isinstance(module, nn.Linear): 35 | nn.init.normal_(module.weight) 36 | if hasattr(module, 'bias') and module.bias is not None: 37 | nn.init.zeros_(module.bias) 38 | else: 39 | continue 40 | 41 | return net 42 | 43 | 44 | def logdet(K): 45 | """get the natural log of the absolute value of the determinant""" 46 | _, ld = np.linalg.slogdet(K) 47 | return ld 48 | 49 | 50 | def get_batch_jacobian(net, x): 51 | """get jacobian and net's output""" 52 | net.zero_grad() 53 | x.requires_grad_(True) 54 | y = net(x) 55 | y.backward(torch.ones_like(y)) 56 | jacob = x.grad.detach() 57 | # return jacob, target.detach(), y.detach() 58 | return jacob, y.detach() 59 | 60 | 61 | def compute_nas_score(gpu, model, resolution, batch_size): 62 | """compute NASWOT score""" 63 | if gpu is not None: 64 | torch.cuda.set_device(gpu) 65 | model = model.cuda(gpu) 66 | 67 | network_weight_gaussian_init(model) 68 | input_ = torch.randn(size=[batch_size, 3, resolution, resolution]) 69 | if gpu is not None: 70 | input_ = input_.cuda(gpu) 71 | 72 | model.K = np.zeros((batch_size, batch_size)) 73 | 74 | def counting_forward_hook(module, inp, out): 75 | try: 76 | if not module.visited_backwards: 77 | return 78 | if isinstance(inp, tuple): 79 | inp = inp[0] 80 | inp = inp.view(inp.size(0), -1) 81 | x = (inp > 0).float() 82 | K = x @ x.t() 83 | K2 = (1. - x) @ (1. - x.t()) 84 | model.K = model.K + K.cpu().numpy() + K2.cpu().numpy() 85 | except Exception as err: 86 | print('---- error on model : ') 87 | print(model) 88 | raise err 89 | 90 | def counting_backward_hook(module, inp, out): 91 | module.visited_backwards = True 92 | 93 | for _, module in model.named_modules(): 94 | # if 'ReLU' in str(type(module)): 95 | if isinstance(module, basic_blocks.RELU): 96 | # hooks[name] = module.register_forward_hook(counting_hook) 97 | module.visited_backwards = True 98 | module.register_forward_hook(counting_forward_hook) 99 | module.register_backward_hook(counting_backward_hook) 100 | 101 | x = input_ 102 | # jacobs, y = get_batch_jacobian(model, x) 103 | _, _ = get_batch_jacobian(model, x) 104 | 105 | score = logdet(model.K) 106 | 107 | return float(score) 108 | 109 | 110 | def parse_cmd_options(argv): 111 | """parse command line""" 112 | parser = argparse.ArgumentParser() 113 | parser.add_argument('--batch_size', type=int, default=16, help='number of instances in one mini-batch.') 114 | parser.add_argument('--input_image_size', type=int, default=None, 115 | help='resolution of input image, usually 32 for CIFAR and 224 for ImageNet.') 116 | parser.add_argument('--repeat_times', type=int, default=32) 117 | parser.add_argument('--gpu', type=int, default=None) 118 | module_opt, _ = parser.parse_known_args(argv) 119 | return module_opt 120 | 121 | 122 | if __name__ == "__main__": 123 | opt = global_utils.parse_cmd_options(sys.argv) 124 | args = parse_cmd_options(sys.argv) 125 | the_model = ModelLoader.get_model(opt, sys.argv) 126 | if args.gpu is not None: 127 | the_model = the_model.cuda(args.gpu) 128 | 129 | start_timer = time.time() 130 | 131 | for repeat_count in range(args.repeat_times): 132 | the_score = compute_nas_score(gpu=args.gpu, model=the_model, 133 | resolution=args.input_image_size, batch_size=args.batch_size) 134 | 135 | time_cost = (time.time() - start_timer) / args.repeat_times 136 | 137 | print(f'NASWOT={the_score:.4g}, time cost={time_cost:.4g} second(s)') 138 | -------------------------------------------------------------------------------- /src/zen_nas/hotfix/transforms.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=invalid-name,not-callable,protected-access,too-many-arguments 2 | import math 3 | import random 4 | import numbers 5 | import warnings 6 | import torch 7 | 8 | 9 | def erase(img, i, j, h, w, v, inplace=False): 10 | """ Erase the input Tensor Image with given value. 11 | 12 | Args: 13 | img (Tensor Image): Tensor image of size (C, H, W) to be erased 14 | i (int): i in (i,j) i.e coordinates of the upper left corner. 15 | j (int): j in (i,j) i.e coordinates of the upper left corner. 16 | h (int): Height of the erased region. 17 | w (int): Width of the erased region. 18 | v: Erasing value. 19 | inplace(bool, optional): For in-place operations. By default is set False. 20 | 21 | Returns: 22 | Tensor Image: Erased image. 23 | """ 24 | if not isinstance(img, torch.Tensor): 25 | raise TypeError(f'img should be Tensor Image. Got {type(img)}') 26 | 27 | if not inplace: 28 | img = img.clone() 29 | 30 | img[:, i:i + h, j:j + w] = v 31 | return img 32 | 33 | 34 | # pylint: disable=too-many-locals 35 | class RandomErasing(): 36 | """ Randomly selects a rectangle region in an image and erases its pixels. 37 | 'Random Erasing Data Augmentation' by Zhong et al. 38 | See https://arxiv.org/pdf/1708.04896.pdf 39 | Args: 40 | p: probability that the random erasing operation will be performed. 41 | scale: range of proportion of erased area against input image. 42 | ratio: range of aspect ratio of erased area. 43 | value: erasing value. Default is 0. If a single int, it is used to 44 | erase all pixels. If a tuple of length 3, it is used to erase 45 | R, G, B channels respectively. 46 | If a str of 'random', erasing each pixel with random values. 47 | inplace: boolean to make this transform inplace. Default set to False. 48 | 49 | Returns: 50 | Erased Image. 51 | # Examples: 52 | >>> transform = transforms.Compose([ 53 | >>> transforms.RandomHorizontalFlip(), 54 | >>> transforms.ToTensor(), 55 | >>> transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), 56 | >>> transforms.RandomErasing(), 57 | >>> ]) 58 | """ 59 | 60 | def __init__(self, p=0.5, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=0, inplace=False): 61 | assert isinstance(value, (numbers.Number, str, tuple, list)) 62 | if (scale[0] > scale[1]) or (ratio[0] > ratio[1]): 63 | warnings.warn("range should be of kind (min, max)") 64 | if scale[0] < 0 or scale[1] > 1: 65 | raise ValueError("range of scale should be between 0 and 1") 66 | if p < 0 or p > 1: 67 | raise ValueError("range of random erasing probability should be between 0 and 1") 68 | 69 | self.p = p 70 | self.scale = scale 71 | self.ratio = ratio 72 | self.value = value 73 | self.inplace = inplace 74 | 75 | @staticmethod 76 | def get_params(img, scale, ratio, value=0): 77 | """Get parameters for ``erase`` for a random erasing. 78 | 79 | Args: 80 | img (Tensor): Tensor image of size (C, H, W) to be erased. 81 | scale: range of proportion of erased area against input image. 82 | ratio: range of aspect ratio of erased area. 83 | 84 | Returns: 85 | tuple: params (i, j, h, w, v) to be passed to ``erase`` for random erasing. 86 | """ 87 | img_c, img_h, img_w = img.shape 88 | area = img_h * img_w 89 | 90 | for _ in range(10): 91 | erase_area = random.uniform(scale[0], scale[1]) * area 92 | aspect_ratio = random.uniform(ratio[0], ratio[1]) 93 | 94 | h = int(round(math.sqrt(erase_area * aspect_ratio))) 95 | w = int(round(math.sqrt(erase_area / aspect_ratio))) 96 | 97 | if h < img_h and w < img_w: 98 | i = random.randint(0, img_h - h) 99 | j = random.randint(0, img_w - w) 100 | if isinstance(value, numbers.Number): 101 | v = value 102 | elif isinstance(value, torch._six.string_classes): 103 | v = torch.empty([img_c, h, w], dtype=torch.float32).normal_() 104 | elif isinstance(value, (list, tuple)): 105 | v = torch.tensor(value, dtype=torch.float32).view(-1, 1, 1).expand(-1, h, w) 106 | return i, j, h, w, v 107 | 108 | # Return original image 109 | return 0, 0, img_h, img_w, img 110 | 111 | def __call__(self, img): 112 | """ 113 | Args: 114 | img (Tensor): Tensor image of size (C, H, W) to be erased. 115 | 116 | Returns: 117 | img (Tensor): Erased Tensor image. 118 | """ 119 | if random.uniform(0, 1) < self.p: 120 | x, y, h, w, v = self.get_params(img, scale=self.scale, ratio=self.ratio, value=self.value) 121 | return erase(img, x, y, h, w, v, self.inplace) 122 | return img 123 | -------------------------------------------------------------------------------- /src/zen_nas/ZeroShotProxy/compute_zen_score.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | ''' 4 | import os 5 | import sys 6 | import time 7 | import argparse 8 | import torch 9 | from torch import nn 10 | import numpy as np 11 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 12 | try: 13 | import global_utils 14 | import ModelLoader 15 | except ImportError: 16 | print('fail to import zen_nas modules') 17 | 18 | 19 | def network_weight_gaussian_init(net: nn.Module): 20 | """gaussian initialization""" 21 | with torch.no_grad(): 22 | for module in net.modules(): 23 | if isinstance(module, nn.Conv2d): 24 | nn.init.normal_(module.weight) 25 | if hasattr(module, 'bias') and module.bias is not None: 26 | nn.init.zeros_(module.bias) 27 | elif isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)): 28 | nn.init.ones_(module.weight) 29 | nn.init.zeros_(module.bias) 30 | elif isinstance(module, nn.Linear): 31 | nn.init.normal_(module.weight) 32 | if hasattr(module, 'bias') and module.bias is not None: 33 | nn.init.zeros_(module.bias) 34 | else: 35 | continue 36 | 37 | return net 38 | 39 | 40 | # pylint: disable=too-many-locals,too-many-arguments 41 | def compute_nas_score(gpu, model, mixup_gamma, resolution, batch_size, repeat, fp16=False): 42 | """compute zen score 43 | 44 | :param gpu (int): gpu index 45 | :param model: model 46 | :param mixup_gamma (double): mixing coefficient 47 | :param resolution (int): input image resolution 48 | :param batch_size (int): batch size 49 | :param repeat (int): repeat time 50 | :oaram fp16 (bool): whether using half precision 51 | :return dict(score) 52 | """ 53 | 54 | score_info = {} 55 | nas_score_list = [] 56 | if gpu is not None: 57 | device = torch.device(f'cuda:{gpu}') 58 | else: 59 | device = torch.device('cpu') 60 | 61 | if fp16: 62 | dtype = torch.half 63 | model = model.half() 64 | else: 65 | dtype = torch.float32 66 | 67 | with torch.no_grad(): 68 | for _ in range(repeat): 69 | network_weight_gaussian_init(model) 70 | input1 = torch.randn(size=[batch_size, 3, resolution, resolution], device=device, dtype=dtype) 71 | input2 = torch.randn(size=[batch_size, 3, resolution, resolution], device=device, dtype=dtype) 72 | mixup_input = input1 + mixup_gamma * input2 73 | output = model.forward_pre_GAP(input1) 74 | mixup_output = model.forward_pre_GAP(mixup_input) 75 | 76 | nas_score = torch.sum(torch.abs(output - mixup_output), dim=[1, 2, 3]) 77 | nas_score = torch.mean(nas_score) 78 | 79 | # compute BN scaling 80 | log_bn_scaling_factor = 0.0 81 | for module in model.modules(): 82 | if isinstance(module, nn.BatchNorm2d): 83 | bn_scaling_factor = torch.sqrt(torch.mean(module.running_var)) 84 | log_bn_scaling_factor += torch.log(bn_scaling_factor) 85 | 86 | nas_score = torch.log(nas_score) + log_bn_scaling_factor 87 | nas_score_list.append(float(nas_score)) 88 | 89 | std_nas_score = np.std(nas_score_list) 90 | avg_precision = 1.96 * std_nas_score / np.sqrt(len(nas_score_list)) 91 | avg_nas_score = np.mean(nas_score_list) 92 | 93 | score_info['avg_nas_score'] = float(avg_nas_score) 94 | score_info['std_nas_score'] = float(std_nas_score) 95 | score_info['avg_precision'] = float(avg_precision) 96 | return score_info 97 | 98 | 99 | def parse_cmd_options(argv): 100 | """parse command line""" 101 | parser = argparse.ArgumentParser() 102 | parser.add_argument('--batch_size', type=int, default=16, help='number of instances in one mini-batch.') 103 | parser.add_argument('--input_image_size', type=int, default=None, 104 | help='resolution of input image, usually 32 for CIFAR and 224 for ImageNet.') 105 | parser.add_argument('--repeat_times', type=int, default=32) 106 | parser.add_argument('--gpu', type=int, default=None) 107 | parser.add_argument('--mixup_gamma', type=float, default=1e-2) 108 | module_opt, _ = parser.parse_known_args(argv) 109 | return module_opt 110 | 111 | 112 | if __name__ == "__main__": 113 | opt = global_utils.parse_cmd_options(sys.argv) 114 | args = parse_cmd_options(sys.argv) 115 | the_model = ModelLoader.get_model(opt, sys.argv) 116 | if args.gpu is not None: 117 | the_model = the_model.cuda(args.gpu) 118 | 119 | start_timer = time.time() 120 | info = compute_nas_score(gpu=args.gpu, model=the_model, mixup_gamma=args.mixup_gamma, 121 | resolution=args.input_image_size, batch_size=args.batch_size, 122 | repeat=args.repeat_times, fp16=False) 123 | time_cost = (time.time() - start_timer) / args.repeat_times 124 | zen_score = info['avg_nas_score'] 125 | print(f'zen-score={zen_score:.4g}, time cost={time_cost:.4g} second(s)') 126 | -------------------------------------------------------------------------------- /src/zen_nas/test.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | 4 | Usage: 5 | python val.py --gpu 0 --arch zennet_imagenet1k_latency02ms_res192 6 | ''' 7 | # pylint: disable=invalid-name 8 | import os 9 | import sys 10 | import argparse 11 | import math 12 | import PIL 13 | import torch 14 | from torchvision import transforms, datasets 15 | sys.path.append(os.path.dirname(os.path.abspath(__file__))) 16 | try: 17 | import ZenNet 18 | except ImportError: 19 | print('fail to import ZenNet') 20 | 21 | imagenet_data_dir = os.path.expanduser('~/data/imagenet/imagenet-torch/') 22 | 23 | 24 | def accuracy(output_, target_, topk=(1, )): 25 | """Computes the accuracy over the k top predictions for the specified values of k""" 26 | with torch.no_grad(): 27 | maxk = max(topk) 28 | batch_size = target_.size(0) 29 | 30 | _, pred = output_.topk(maxk, 1, True, True) 31 | pred = pred.t() 32 | correct = pred.eq(target_.view(1, -1).expand_as(pred)) 33 | 34 | res = [] 35 | for k in topk: 36 | correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True) 37 | res.append(correct_k.mul_(100.0 / batch_size)) 38 | return res 39 | 40 | 41 | if __name__ == '__main__': 42 | 43 | parser = argparse.ArgumentParser() 44 | parser.add_argument('--batch_size', type=int, default=128, 45 | help='batch size for evaluation.') 46 | parser.add_argument('--workers', type=int, default=6, 47 | help='number of workers to load dataset.') 48 | parser.add_argument('--gpu', type=int, default=None, 49 | help='GPU device ID. None for CPU.') 50 | parser.add_argument('--data', type=str, default=imagenet_data_dir, 51 | help='ImageNet data directory.') 52 | parser.add_argument('--arch', type=str, default=None, 53 | help='model to be evaluated.') 54 | parser.add_argument('--fp16', action='store_true') 55 | parser.add_argument('--apex', action='store_true', 56 | help='Use NVIDIA Apex (float16 precision).') 57 | 58 | opt, _ = parser.parse_known_args(sys.argv) 59 | 60 | if opt.apex: 61 | # pylint: disable=no-name-in-module 62 | from apex import amp 63 | # else: 64 | # print('Warning!!! The GENets are trained by NVIDIA Apex, ' 65 | # 'it is suggested to turn on --apex in the evaluation. ' 66 | # 'Otherwise the model accuracy might be harmed.') 67 | 68 | input_image_size = ZenNet.zennet_model_zoo[opt.arch]['resolution'] 69 | crop_image_size = ZenNet.zennet_model_zoo[opt.arch]['crop_image_size'] 70 | 71 | print(f'Evaluate {opt.arch} at {input_image_size}x{input_image_size} resolution.') 72 | 73 | # load dataset 74 | val_dir = os.path.join(opt.data, 'val') 75 | input_image_crop = 0.875 76 | resize_image_size = int(math.ceil(crop_image_size / input_image_crop)) 77 | transforms_normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) 78 | transform_list = [transforms.Resize(resize_image_size, interpolation=PIL.Image.BICUBIC), 79 | transforms.CenterCrop(crop_image_size), transforms.ToTensor(), transforms_normalize] 80 | transformer = transforms.Compose(transform_list) 81 | val_dataset = datasets.ImageFolder(val_dir, transformer) 82 | val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=opt.batch_size, shuffle=False, 83 | num_workers=opt.workers, pin_memory=True, sampler=None) 84 | 85 | # load model 86 | model = ZenNet.get_ZenNet(opt.arch, pretrained=True) 87 | if opt.gpu is not None: 88 | torch.cuda.set_device(opt.gpu) 89 | torch.backends.cudnn.benchmark = True 90 | model = model.cuda(opt.gpu) 91 | print(f'Using GPU {opt.gpu}.') 92 | if opt.apex: 93 | model = amp.initialize(model, opt_level="O1") 94 | elif opt.fp16: 95 | model = model.half() 96 | 97 | model.eval() 98 | acc1_sum = 0 99 | acc5_sum = 0 100 | num = 0 101 | with torch.no_grad(): 102 | for i, (input_, target) in enumerate(val_loader): 103 | 104 | if opt.gpu is not None: 105 | input_ = input_.cuda(opt.gpu, non_blocking=True) 106 | target = target.cuda(opt.gpu, non_blocking=True) 107 | if opt.fp16: 108 | input_ = input_.half() 109 | 110 | input_ = torch.nn.functional.interpolate(input_, input_image_size, mode='bilinear') 111 | output = model(input_) 112 | # pylint: disable=unbalanced-tuple-unpacking 113 | acc1, acc5 = accuracy(output, target, topk=(1, 5)) 114 | 115 | acc1_sum += acc1[0] * input_.shape[0] 116 | acc5_sum += acc5[0] * input_.shape[0] 117 | num += input_.shape[0] 118 | 119 | if i % 100 == 0: 120 | print(f'mini_batch {i}, top-1 acc={acc1[0]:4g}%, top-5 acc={acc5[0]:4g}%, \ 121 | number of evaluated images={num}') 122 | 123 | acc1_avg = acc1_sum / num 124 | acc5_avg = acc5_sum / num 125 | 126 | print(f'*** arch={opt.arch}, validation top-1 acc={acc1_avg}%, top-5 acc={acc5_avg}%, \ 127 | number of evaluated images={num}') 128 | -------------------------------------------------------------------------------- /src/zen_nas/test_cifar.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | 4 | Usage: 5 | python val_cifar.py --dataset cifar10 --gpu 0 --arch zennet_cifar10_model_size05M_res32 6 | ''' 7 | # pylint: disable=invalid-name 8 | import os 9 | import sys 10 | import argparse 11 | import torch 12 | from torchvision import transforms, datasets 13 | sys.path.append(os.path.dirname(os.path.abspath(__file__))) 14 | try: 15 | import ZenNet 16 | except ImportError: 17 | print('fail to import ZenNet') 18 | 19 | cifar10_data_dir = '~/data/pytorch_cifar10' 20 | cifar100_data_dir = '~/data/pytorch_cifar100' 21 | 22 | 23 | def accuracy(output_, target_, topk=(1, )): 24 | """Computes the accuracy over the k top predictions for the specified values of k""" 25 | with torch.no_grad(): 26 | maxk = max(topk) 27 | batch_size = target_.size(0) 28 | 29 | _, pred = output_.topk(maxk, 1, True, True) 30 | pred = pred.t() 31 | correct = pred.eq(target_.view(1, -1).expand_as(pred)) 32 | 33 | res = [] 34 | for k in topk: 35 | correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True) 36 | res.append(correct_k.mul_(100.0 / batch_size)) 37 | return res 38 | 39 | 40 | if __name__ == '__main__': 41 | 42 | parser = argparse.ArgumentParser() 43 | parser.add_argument('--batch_size', type=int, default=512, 44 | help='batch size for evaluation.') 45 | parser.add_argument('--dataset', type=str, default=None, 46 | help='cifar10 or cifar100') 47 | parser.add_argument('--workers', type=int, default=6, 48 | help='number of workers to load dataset.') 49 | parser.add_argument('--gpu', type=int, default=None, 50 | help='GPU device ID. None for CPU.') 51 | parser.add_argument('--arch', type=str, default=None, 52 | help='model to be evaluated.') 53 | parser.add_argument('--fp16', action='store_true') 54 | parser.add_argument('--apex', action='store_true', 55 | help='Use NVIDIA Apex (float16 precision).') 56 | 57 | opt, _ = parser.parse_known_args(sys.argv) 58 | 59 | if opt.apex: 60 | # pylint: disable=no-name-in-module 61 | from apex import amp 62 | else: 63 | print('Warning!!! The GENets are trained by NVIDIA Apex, ' 64 | 'it is suggested to turn on --apex in the evaluation. ' 65 | 'Otherwise the model accuracy might be harmed.') 66 | 67 | input_image_size = ZenNet.zennet_model_zoo[opt.arch]['resolution'] 68 | crop_image_size = ZenNet.zennet_model_zoo[opt.arch]['crop_image_size'] 69 | 70 | print(f'Evaluate {opt.arch} at {input_image_size}x{input_image_size} resolution.') 71 | 72 | # load dataset 73 | transforms_normalize = transforms.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010]) 74 | transform_list = [transforms.ToTensor(), transforms_normalize] 75 | transformer = transforms.Compose(transform_list) 76 | 77 | if opt.dataset == 'cifar10': 78 | the_dataset = datasets.CIFAR10(root=cifar10_data_dir, train=False, download=True, transform=transformer) 79 | elif opt.dataset == 'cifar100': 80 | the_dataset = datasets.CIFAR100(root=cifar100_data_dir, train=False, download=True, transform=transformer) 81 | else: 82 | raise ValueError('Unknown dataset_name=' + opt.dataset) 83 | val_loader = torch.utils.data.DataLoader(the_dataset, batch_size=opt.batch_size, shuffle=False, 84 | num_workers=opt.workers, pin_memory=True, sampler=None) 85 | 86 | # load model 87 | model = ZenNet.get_ZenNet(opt.arch, pretrained=True) 88 | if opt.gpu is not None: 89 | torch.cuda.set_device(opt.gpu) 90 | torch.backends.cudnn.benchmark = True 91 | model = model.cuda(opt.gpu) 92 | print(f'Using GPU {opt.gpu}.') 93 | if opt.apex: 94 | model = amp.initialize(model, opt_level="O1") 95 | elif opt.fp16: 96 | model = model.half() 97 | 98 | model.eval() 99 | acc1_sum = 0 100 | acc5_sum = 0 101 | num = 0 102 | with torch.no_grad(): 103 | for i, (input_, target) in enumerate(val_loader): 104 | 105 | if opt.gpu is not None: 106 | input_ = input_.cuda(opt.gpu, non_blocking=True) 107 | target = target.cuda(opt.gpu, non_blocking=True) 108 | if opt.fp16: 109 | input_ = input_.half() 110 | 111 | input_ = torch.nn.functional.interpolate(input_, input_image_size, mode='bilinear') 112 | 113 | output = model(input_) 114 | # pylint: disable=unbalanced-tuple-unpacking 115 | acc1, acc5 = accuracy(output, target, topk=(1, 5)) 116 | 117 | acc1_sum += acc1[0] * input_.shape[0] 118 | acc5_sum += acc5[0] * input_.shape[0] 119 | num += input_.shape[0] 120 | 121 | if i % 100 == 0: 122 | print(f'mini_batch {i}, top-1 acc={acc1[0]:4g}%, top-5 acc={acc5[0]:4g}%, \ 123 | number of evaluated images={num}') 124 | 125 | acc1_avg = acc1_sum / num 126 | acc5_avg = acc5_sum / num 127 | 128 | print(f'*** arch={opt.arch}, validation top-1 acc={acc1_avg}%, top-5 acc={acc5_avg}%, \ 129 | number of evaluated images={num}') 130 | -------------------------------------------------------------------------------- /src/zen_nas/SearchSpace/search_space_XXBL.py: -------------------------------------------------------------------------------- 1 | """define SuperRes Block search space""" 2 | # pylint: disable=invalid-name 3 | import os 4 | import sys 5 | import itertools 6 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 7 | try: 8 | import global_utils 9 | from PlainNet import super_blocks, SuperResKXKX, SuperResK1KXK1 10 | except ImportError: 11 | print('fail to import zen_nas modules') 12 | 13 | seach_space_block_type_list_list = [ 14 | [SuperResK1KXK1.SuperResK1K3K1, SuperResK1KXK1.SuperResK1K5K1, SuperResK1KXK1.SuperResK1K7K1], 15 | [SuperResKXKX.SuperResK3K3, SuperResKXKX.SuperResK5K5, SuperResKXKX.SuperResK7K7], 16 | ] 17 | 18 | __block_type_round_channels_base_dict__ = { 19 | SuperResKXKX.SuperResK3K3: 8, 20 | SuperResKXKX.SuperResK5K5: 8, 21 | SuperResKXKX.SuperResK7K7: 8, 22 | SuperResK1KXK1.SuperResK1K3K1: 8, SuperResK1KXK1.SuperResK1K5K1: 8, SuperResK1KXK1.SuperResK1K7K1: 8, 23 | } 24 | 25 | __block_type_min_channels_base_dict__ = { 26 | SuperResKXKX.SuperResK3K3: 8, 27 | SuperResKXKX.SuperResK5K5: 8, 28 | SuperResKXKX.SuperResK7K7: 8, 29 | SuperResK1KXK1.SuperResK1K3K1: 8, 30 | SuperResK1KXK1.SuperResK1K5K1: 8, 31 | SuperResK1KXK1.SuperResK1K7K1: 8, 32 | } 33 | 34 | 35 | def get_select_student_channels_list(out_channels): 36 | """generate all possible channels""" 37 | the_list = [out_channels * 2.5, out_channels * 2, out_channels * 1.5, out_channels * 1.25, 38 | out_channels, 39 | out_channels / 1.25, out_channels / 1.5, out_channels / 2, out_channels / 2.5] 40 | the_list = [min(2048, max(8, x)) for x in the_list] 41 | the_list = [global_utils.smart_round(x, base=8) for x in the_list] 42 | the_list = list(set(the_list)) 43 | the_list.sort(reverse=True) 44 | return the_list 45 | 46 | 47 | def get_select_student_sublayers_list(sub_layers): 48 | """generate all possible sublayers""" 49 | the_list = [sub_layers, 50 | sub_layers + 1, sub_layers + 2, 51 | sub_layers - 1, sub_layers - 2, ] 52 | the_list = [max(0, round(x)) for x in the_list] 53 | the_list = list(set(the_list)) 54 | the_list.sort(reverse=True) 55 | return the_list 56 | 57 | 58 | # pylint: disable=too-many-locals 59 | def gen_search_space(block_list, block_id): 60 | """ generate all possible blocks 61 | 62 | :param block_list (list): block list 63 | :param block_id (int): block index 64 | :return list(list()) 65 | """ 66 | 67 | the_block = block_list[block_id] 68 | student_blocks_list_list = [] 69 | 70 | if isinstance(the_block, super_blocks.SuperConvKXBNRELU): 71 | student_blocks_list = [] 72 | student_out_channels_list = get_select_student_channels_list(the_block.out_channels) 73 | for student_out_channels in student_out_channels_list: 74 | tmp_block_str = type(the_block).__name__ + f'({the_block.in_channels},\ 75 | {student_out_channels},\ 76 | {the_block.stride},1)' 77 | student_blocks_list.append(tmp_block_str) 78 | student_blocks_list = list(set(student_blocks_list)) 79 | assert len(student_blocks_list) >= 1 80 | student_blocks_list_list.append(student_blocks_list) 81 | else: 82 | for student_block_type_list in seach_space_block_type_list_list: 83 | student_blocks_list = [] 84 | student_out_channels_list = get_select_student_channels_list(the_block.out_channels) 85 | student_sublayers_list = get_select_student_sublayers_list(sub_layers=the_block.sub_layers) 86 | student_bottleneck_channels_list = get_select_student_channels_list(the_block.bottleneck_channels) 87 | for student_block_type in student_block_type_list: 88 | for student_out_channels, student_sublayers, student_bottleneck_channels in itertools.product( 89 | student_out_channels_list, student_sublayers_list, student_bottleneck_channels_list): 90 | 91 | # filter smallest possible channel for this block type 92 | min_possible_channels = __block_type_round_channels_base_dict__[student_block_type] 93 | channel_round_base = __block_type_round_channels_base_dict__[student_block_type] 94 | student_out_channels = global_utils.smart_round(student_out_channels, channel_round_base) 95 | student_bottleneck_channels = global_utils.smart_round(student_bottleneck_channels, 96 | channel_round_base) 97 | 98 | if student_out_channels < min_possible_channels or \ 99 | student_bottleneck_channels < min_possible_channels: 100 | continue 101 | if student_sublayers <= 0: # no empty layer 102 | continue 103 | tmp_block_str = student_block_type.__name__ + f'({the_block.in_channels},\ 104 | {student_out_channels},\ 105 | {the_block.stride},\ 106 | {student_bottleneck_channels},\ 107 | {student_sublayers})' 108 | student_blocks_list.append(tmp_block_str) 109 | student_blocks_list = list(set(student_blocks_list)) 110 | assert len(student_blocks_list) >= 1 111 | student_blocks_list_list.append(student_blocks_list) 112 | return student_blocks_list_list 113 | -------------------------------------------------------------------------------- /src/zen_nas/benchmark_network_latency.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | 4 | ''' 5 | # pylint: disable=too-many-arguments 6 | import os 7 | import sys 8 | import argparse 9 | import time 10 | import numpy as np 11 | import torch 12 | sys.path.append(os.path.dirname(os.path.abspath(__file__))) 13 | try: 14 | import global_utils 15 | import ModelLoader 16 | import train_image_classification as tic 17 | import ZenNet 18 | except ImportError: 19 | print('fail to import zen_nas modules') 20 | 21 | 22 | def __get_latency__(model, batch_size, resolution, channel, gpu, benchmark_repeat_times, fp16): 23 | """compute model latency""" 24 | device = torch.device(f'cuda:{gpu}') 25 | torch.backends.cudnn.benchmark = True 26 | 27 | torch.cuda.set_device(gpu) 28 | model = model.cuda(gpu) 29 | if fp16: 30 | model = model.half() 31 | dtype = torch.float16 32 | else: 33 | dtype = torch.float32 34 | 35 | the_image = torch.randn(batch_size, channel, resolution, resolution, dtype=dtype, 36 | device=device) 37 | model.eval() 38 | warmup_times = 3 39 | with torch.no_grad(): 40 | for _ in range(warmup_times): 41 | _ = model(the_image) 42 | start_timer = time.time() 43 | for _ in range(benchmark_repeat_times): 44 | _ = model(the_image) 45 | 46 | end_timer = time.time() 47 | the_latency = (end_timer - start_timer) / float(benchmark_repeat_times) / batch_size 48 | return the_latency 49 | 50 | 51 | def get_robust_latency_mean_std(model, batch_size, resolution, channel, gpu, benchmark_repeat_times=30, fp16=False): 52 | """get robust latency""" 53 | robust_repeat_times = 10 54 | latency_list = [] 55 | model = model.cuda(gpu) 56 | for _ in range(robust_repeat_times): 57 | try: 58 | the_latency = __get_latency__(model, batch_size, resolution, channel, gpu, benchmark_repeat_times, fp16) 59 | except Exception as error: # pylint: disable=broad-except 60 | print(error) 61 | the_latency = np.inf 62 | 63 | latency_list.append(the_latency) 64 | 65 | latency_list.sort() 66 | avg_latency = np.mean(latency_list[2:8]) 67 | std_latency = np.std(latency_list[2:8]) 68 | return avg_latency, std_latency 69 | 70 | 71 | def main(opt, argv): 72 | """get model and compute latency""" 73 | global_utils.create_logging() 74 | 75 | batch_size_list = [int(x) for x in opt.batch_size_list.split(',')] 76 | opt.batch_size = 1 77 | opt = tic.config_dist_env_and_opt(opt) 78 | 79 | # create model 80 | model = ModelLoader.get_model(opt, argv) 81 | 82 | print('batch_size, latency_per_image') 83 | 84 | for the_batch_size_per_gpu in batch_size_list: 85 | 86 | the_latency, _ = get_robust_latency_mean_std(model=model, batch_size=the_batch_size_per_gpu, 87 | resolution=opt.input_image_size, channel=3, gpu=opt.gpu, 88 | benchmark_repeat_times=opt.repeat_times, 89 | fp16=opt.fp16) 90 | print(f'{the_batch_size_per_gpu},{the_latency:4g}') 91 | 92 | 93 | def get_model_latency(model, batch_size, resolution, in_channels, gpu, repeat_times, fp16): 94 | """get model latency 95 | 96 | :param model: model 97 | :param batch_size (int): batch size 98 | :param resolution (int): input image size 99 | :param in_channels (int): input channels 100 | :param gpu (int): gpu index 101 | :param repeat_times (int): repeat times 102 | :param fp16 (bool): whether using half precision 103 | :return latency 104 | """ 105 | if gpu is not None: 106 | device = torch.device(f'cuda:{gpu}') 107 | else: 108 | device = torch.device('cpu') 109 | 110 | if fp16: 111 | model = model.half() 112 | dtype = torch.float16 113 | else: 114 | dtype = torch.float32 115 | 116 | the_image = torch.randn(batch_size, in_channels, resolution, resolution, dtype=dtype, 117 | device=device) 118 | 119 | model.eval() 120 | warmup_times = 3 121 | with torch.no_grad(): 122 | for _ in range(warmup_times): 123 | _ = model(the_image) 124 | start_timer = time.time() 125 | for _ in range(repeat_times): 126 | _ = model(the_image) 127 | # torch.cuda.synchronize(device=device) 128 | end_timer = time.time() 129 | the_latency = (end_timer - start_timer) / float(repeat_times) / batch_size 130 | return the_latency 131 | 132 | 133 | def parse_cmd_options(argv): 134 | parser = argparse.ArgumentParser() 135 | parser.add_argument('--batch_size', type=int, default=None, help='number of instances in one mini-batch.') 136 | parser.add_argument('--input_image_size', type=int, default=None, 137 | help='resolution of input image, usually 32 for CIFAR and 224 for ImageNet.') 138 | parser.add_argument('--save_dir', type=str, default=None, 139 | help='output directory') 140 | parser.add_argument('--repeat_times', type=int, default=1) 141 | parser.add_argument('--gpu', type=int, default=None) 142 | parser.add_argument('--fp16', action='store_true') 143 | module_opt, _ = parser.parse_known_args(argv) 144 | return module_opt 145 | 146 | 147 | if __name__ == "__main__": 148 | option = global_utils.parse_cmd_options(sys.argv) 149 | args = parse_cmd_options(sys.argv) 150 | # the_model = ModelLoader.get_model(opt, sys.argv) 151 | the_model = ZenNet.get_ZenNet(option.arch) 152 | if args.gpu is not None: 153 | the_model = the_model.cuda(args.gpu) 154 | 155 | latency = get_model_latency(model=the_model, batch_size=args.batch_size, 156 | resolution=args.input_image_size, 157 | in_channels=3, gpu=args.gpu, repeat_times=args.repeat_times, 158 | fp16=args.fp16) 159 | print(f'{(latency * 1000):.4g} millisecond(s) per image, or {1.0/latency:.4g} image(s) per second.') 160 | -------------------------------------------------------------------------------- /src/zen_nas/SearchSpace/search_choice.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright 2021 ZTE corporation. All Rights Reserved. 3 | SPDX-License-Identifier: Apache-2. 4 | 5 | generate random network and mutated network 6 | ''' 7 | # pylint: disable=global-statement 8 | import os 9 | import sys 10 | import random 11 | import numpy as np 12 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 13 | try: 14 | import global_utils 15 | import PlainNet 16 | from PlainNet import super_blocks, SuperResKXKX, SuperResK1KXK1 17 | except ImportError: 18 | print('fail to import zen_nas modules') 19 | 20 | SEARCH_SPACE = [SuperResKXKX, SuperResK1KXK1] 21 | 22 | 23 | def get_block_from_module(module, block_dict: dict): 24 | """collect all optional block from module""" 25 | assert hasattr(module, 'register_netblocks_dict') 26 | block_dict = module.register_netblocks_dict(block_dict) 27 | return block_dict 28 | 29 | 30 | def register_serach_space(candidate_blocks): 31 | """register search space""" 32 | global SEARCH_SPACE 33 | SEARCH_SPACE.clear() 34 | SEARCH_SPACE = candidate_blocks.copy() 35 | 36 | 37 | def get_random_block(): 38 | """get random block from serach space""" 39 | search_space_block_dict = {} 40 | for module in SEARCH_SPACE: 41 | search_space_block_dict = get_block_from_module(module, search_space_block_dict) 42 | 43 | block_list = list(search_space_block_dict.values()) 44 | return np.random.choice(block_list) 45 | 46 | 47 | def get_random_channels_ratio(search_range: list = None): 48 | """determin mutated channel ratio""" 49 | if search_range is not None: 50 | assert len(search_range) == 2, 'need two number to limit to choose number in ranges' 51 | search_range.sort() 52 | return random.uniform(search_range[0], search_range[1]) 53 | candidate = [2.5, 2, 1.5, 1.25, 1, 1 / 1.25, 1 / 1.5, 1 / 2, 1 / 2.5] 54 | return np.random.choice(candidate) 55 | 56 | 57 | def get_random_sublayers_change(): 58 | """determin mutated sublayer """ 59 | candidate = [-2, -1, 0, 1, 2] 60 | return np.random.choice(candidate) 61 | 62 | 63 | def verify_channels(channel, ratio): 64 | """restrict range""" 65 | channel *= ratio 66 | new_channel = global_utils.smart_round(channel, base=8) 67 | return min(new_channel, 2048) 68 | 69 | 70 | def verify_sublayers(sublayers, change: int): 71 | """restrict range""" 72 | sublayers += change 73 | return max(1, sublayers) 74 | 75 | 76 | def mutated_block(block_list: list, block_id: int): 77 | """construct mutated block objec""" 78 | chosen_block = block_list[block_id] 79 | in_channel = chosen_block.in_channels 80 | out_channel = chosen_block.out_channels 81 | stride = chosen_block.stride 82 | sub_layers = chosen_block.sub_layers 83 | ratio = get_random_channels_ratio() 84 | 85 | if isinstance(chosen_block, super_blocks.SuperConvKXBNRELU): 86 | out_channel = verify_channels(out_channel, ratio) 87 | mutated_block_obj = chosen_block.__class__(in_channel, out_channel, stride, 1) 88 | else: 89 | bottleneck_channel = chosen_block.bottleneck_channels 90 | new_block = get_random_block() 91 | out_channel = verify_channels(out_channel, ratio) 92 | bottleneck_channel = verify_channels(bottleneck_channel, ratio) 93 | # out_channel and bottleneck size too big, CUDNN_STATUS_NOT_SUPPORTED may occur in model training 94 | bottleneck_channel = min(bottleneck_channel, 512) 95 | change = get_random_sublayers_change() 96 | sub_layers = verify_sublayers(sub_layers, change) 97 | mutated_block_obj = new_block(in_channel, out_channel, stride, bottleneck_channel, sub_layers) 98 | 99 | return mutated_block_obj 100 | 101 | 102 | def get_mutated_structure_str(any_plain_net, structure_str, num_classes, num_replaces=1): 103 | """generate mutated network string""" 104 | the_net = any_plain_net(num_classes=num_classes, plainnet_struct=structure_str, no_create=True) 105 | assert isinstance(the_net, PlainNet.PlainNet) 106 | selected_random_id_set = set() 107 | for _ in range(num_replaces): 108 | random_id = random.randint(0, len(the_net.block_list) - 1) 109 | if random_id in selected_random_id_set: 110 | continue 111 | selected_random_id_set.add(random_id) 112 | 113 | new_block_obj = mutated_block(the_net.block_list, random_id) 114 | if random_id < len(the_net.block_list) - 1: 115 | out_channel = new_block_obj.out_channels 116 | the_net.block_list[random_id + 1].set_in_channels(out_channel) 117 | the_net.block_list[random_id] = new_block_obj 118 | 119 | assert hasattr(the_net, 'split') 120 | new_random_structure_str = the_net.split(split_layer_threshold=6) 121 | return new_random_structure_str 122 | 123 | 124 | def get_random_initialized_structure_str(): 125 | """generate random network from serach space""" 126 | struct_str = "" 127 | input_channel = 3 128 | out_channel = 32 129 | str_len = random.randint(6, 10) 130 | channel_choice = [8 * i for i in range(1, 16)] 131 | # sub_layer_choice = [1, 2] 132 | for i in range(str_len): 133 | if i <= 3 or i == round((3 + str_len - 1) / 2): 134 | stride = 2 135 | else: 136 | stride = 1 137 | 138 | if i == 0: 139 | struct_str += f'SuperConvK3BNRELU({input_channel},{out_channel},{stride},{1})' 140 | input_channel = out_channel 141 | elif i == str_len - 1: 142 | out_channel = np.random.choice(channel_choice) 143 | struct_str += f'SuperConvK1BNRELU({input_channel},{out_channel},{stride},{1})' 144 | else: 145 | out_channel = np.random.choice(channel_choice) 146 | bottleneck_channel = np.random.choice(channel_choice) 147 | # sub_layer = np.random.choice(sub_layer_choice) 148 | sub_layer = 1 149 | 150 | block = get_random_block() 151 | struct_str += block.__name__ + f'({input_channel},{out_channel},\ 152 | {stride},{bottleneck_channel},{sub_layer})' 153 | input_channel = out_channel 154 | return struct_str 155 | -------------------------------------------------------------------------------- /src/zen_nas/ModelLoader/geffnet/activations/activations_me.py: -------------------------------------------------------------------------------- 1 | """ Activations (memory-efficient w/ custom autograd) 2 | 3 | A collection of activations fn and modules with a common interface so that they can 4 | easily be swapped. All have an `inplace` arg even if not used. 5 | 6 | These activations are not compatible with jit scripting or ONNX export of the model, please use either 7 | the JIT or basic versions of the activations. 8 | 9 | Copyright 2020 Ross Wightman 10 | """ 11 | # pylint: disable=W0613,arguments-differ,no-self-use,abstract-method 12 | import torch 13 | from torch import nn 14 | from torch.nn import functional as F 15 | 16 | 17 | __all__ = ['swish_me', 'SwishMe', 'mish_me', 'MishMe', 18 | 'hard_sigmoid_me', 'HardSigmoidMe', 'hard_swish_me', 'HardSwishMe'] 19 | 20 | 21 | @torch.jit.script 22 | def swish_jit_fwd(input_): 23 | """jit-scripted swish forward""" 24 | return input_.mul(torch.sigmoid(input_)) 25 | 26 | 27 | @torch.jit.script 28 | def swish_jit_bwd(input_, grad_output): 29 | """jit-scripted swish backward""" 30 | input_sigmoid = torch.sigmoid(input_) 31 | return grad_output * (input_sigmoid * (1 + input_ * (1 - input_sigmoid))) 32 | 33 | 34 | class SwishJitAutoFn(torch.autograd.Function): 35 | """ torch.jit.script optimised Swish w/ memory-efficient checkpoint 36 | Inspired by conversation btw Jeremy Howard & Adam Pazske 37 | https://twitter.com/jeremyphoward/status/1188251041835315200 38 | 39 | Swish - Described originally as SiLU (https://arxiv.org/abs/1702.03118v3) 40 | and also as Swish (https://arxiv.org/abs/1710.05941). 41 | 42 | TODO Rename to SiLU with addition to PyTorch 43 | """ 44 | 45 | @staticmethod 46 | def forward(ctx, x): 47 | ctx.save_for_backward(x) 48 | return swish_jit_fwd(x) 49 | 50 | @staticmethod 51 | def backward(ctx, grad_output): 52 | output = ctx.saved_tensors[0] 53 | return swish_jit_bwd(output, grad_output) 54 | 55 | 56 | def swish_me(input_, inplace=False): 57 | """jit-scripted swish_me function""" 58 | return SwishJitAutoFn.apply(input_) 59 | 60 | 61 | class SwishMe(nn.Module): 62 | """jit-scripted swish_me module""" 63 | def __init__(self, inplace: bool = False): 64 | super().__init__() 65 | 66 | def forward(self, input_): 67 | return SwishJitAutoFn.apply(input_) 68 | 69 | 70 | @torch.jit.script 71 | def mish_jit_fwd(input_): 72 | """jit-scripted mish forward""" 73 | return input_.mul(torch.tanh(F.softplus(input_))) 74 | 75 | 76 | @torch.jit.script 77 | def mish_jit_bwd(output, grad_output): 78 | """jit-scripted mish backward""" 79 | output_sigmoid = torch.sigmoid(output) 80 | output_tanh_sp = F.softplus(output).tanh() 81 | return grad_output.mul(output_tanh_sp + output * output_sigmoid * (1 - output_tanh_sp * output_tanh_sp)) 82 | 83 | 84 | class MishJitAutoFn(torch.autograd.Function): 85 | """ Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681 86 | A memory efficient, jit scripted variant of Mish 87 | """ 88 | @staticmethod 89 | def forward(ctx, x): 90 | ctx.save_for_backward(x) 91 | return mish_jit_fwd(x) 92 | 93 | @staticmethod 94 | def backward(ctx, grad_output): 95 | output = ctx.saved_tensors[0] 96 | return mish_jit_bwd(output, grad_output) 97 | 98 | 99 | def mish_me(input_, inplace=False): 100 | """jit-scripted mish_me function""" 101 | return MishJitAutoFn.apply(input_) 102 | 103 | 104 | class MishMe(nn.Module): 105 | """jit-scripted mish_me module""" 106 | def __init__(self, inplace: bool = False): 107 | super().__init__() 108 | 109 | def forward(self, input_): 110 | return MishJitAutoFn.apply(input_) 111 | 112 | 113 | @torch.jit.script 114 | def hard_sigmoid_jit_fwd(input_, inplace: bool = False): 115 | """jit-scripted hard sigmoid forward""" 116 | return (input_ + 3).clamp(min=0, max=6).div(6.) 117 | 118 | 119 | @torch.jit.script 120 | def hard_sigmoid_jit_bwd(output, grad_output): 121 | """git-scripted hard sigmoid backward""" 122 | tmp = torch.ones_like(output) * ((output >= -3.) & (output <= 3.)) / 6. 123 | return grad_output * tmp 124 | 125 | 126 | class HardSigmoidJitAutoFn(torch.autograd.Function): 127 | """jit-scripted hard sigmoid module""" 128 | @staticmethod 129 | def forward(ctx, x): 130 | ctx.save_for_backward(x) 131 | return hard_sigmoid_jit_fwd(x) 132 | 133 | @staticmethod 134 | def backward(ctx, grad_output): 135 | output = ctx.saved_tensors[0] 136 | return hard_sigmoid_jit_bwd(output, grad_output) 137 | 138 | 139 | def hard_sigmoid_me(input_, inplace: bool = False): 140 | """jit-scripted hard_sigmoid_me function""" 141 | return HardSigmoidJitAutoFn.apply(input_) 142 | 143 | 144 | class HardSigmoidMe(nn.Module): 145 | """jit-scripted hard_sigmoid_me module""" 146 | def __init__(self, inplace: bool = False): 147 | super().__init__() 148 | 149 | def forward(self, input_): 150 | return HardSigmoidJitAutoFn.apply(input_) 151 | 152 | 153 | @torch.jit.script 154 | def hard_swish_jit_fwd(input_): 155 | """jit-scripted hard swish forward""" 156 | return input_ * (input_ + 3).clamp(min=0, max=6).div(6.) 157 | 158 | 159 | @torch.jit.script 160 | def hard_swish_jit_bwd(output, grad_output): 161 | """jit-scripted hard swish backward""" 162 | tmp = torch.ones_like(output) * (output >= 3.) 163 | tmp = torch.where((output >= -3.) & (output <= 3.), output / 3. + .5, tmp) 164 | return grad_output * tmp 165 | 166 | 167 | class HardSwishJitAutoFn(torch.autograd.Function): 168 | """A memory efficient, jit-scripted HardSwish activation""" 169 | @staticmethod 170 | def forward(ctx, x): 171 | ctx.save_for_backward(x) 172 | return hard_swish_jit_fwd(x) 173 | 174 | @staticmethod 175 | def backward(ctx, grad_output): 176 | output = ctx.saved_tensors[0] 177 | return hard_swish_jit_bwd(output, grad_output) 178 | 179 | 180 | def hard_swish_me(input_, inplace=False): 181 | """jit-scripted hard_swish_me function""" 182 | return HardSwishJitAutoFn.apply(input_) 183 | 184 | 185 | class HardSwishMe(nn.Module): 186 | """jit-scripted hard_swish_me module""" 187 | 188 | def __init__(self, inplace: bool = False): 189 | super().__init__() 190 | 191 | def forward(self, input_): 192 | return HardSwishJitAutoFn.apply(input_) 193 | -------------------------------------------------------------------------------- /src/zen_nas/SearchSpace/search_space_IDW_fixfc.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | 4 | define SuperResIDW Block search space 5 | ''' 6 | # pylint: disable=invalid-name 7 | import os 8 | import sys 9 | import itertools 10 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 11 | try: 12 | import global_utils 13 | from PlainNet import super_blocks, SuperResIDWEXKX 14 | except ImportError: 15 | print('fail to import zen_nas modules') 16 | 17 | seach_space_block_type_list_list = [ 18 | [SuperResIDWEXKX.SuperResIDWE1K3, SuperResIDWEXKX.SuperResIDWE2K3, SuperResIDWEXKX.SuperResIDWE4K3, 19 | SuperResIDWEXKX.SuperResIDWE6K3, 20 | SuperResIDWEXKX.SuperResIDWE1K5, SuperResIDWEXKX.SuperResIDWE2K5, SuperResIDWEXKX.SuperResIDWE4K5, 21 | SuperResIDWEXKX.SuperResIDWE6K5, 22 | SuperResIDWEXKX.SuperResIDWE1K7, SuperResIDWEXKX.SuperResIDWE2K7, SuperResIDWEXKX.SuperResIDWE4K7, 23 | SuperResIDWEXKX.SuperResIDWE6K7], 24 | ] 25 | 26 | __block_type_round_channels_base_dict__ = { 27 | SuperResIDWEXKX.SuperResIDWE1K3: 8, 28 | SuperResIDWEXKX.SuperResIDWE2K3: 8, 29 | SuperResIDWEXKX.SuperResIDWE4K3: 8, 30 | SuperResIDWEXKX.SuperResIDWE6K3: 8, 31 | SuperResIDWEXKX.SuperResIDWE1K5: 8, 32 | SuperResIDWEXKX.SuperResIDWE2K5: 8, 33 | SuperResIDWEXKX.SuperResIDWE4K5: 8, 34 | SuperResIDWEXKX.SuperResIDWE6K5: 8, 35 | SuperResIDWEXKX.SuperResIDWE1K7: 8, 36 | SuperResIDWEXKX.SuperResIDWE2K7: 8, 37 | SuperResIDWEXKX.SuperResIDWE4K7: 8, 38 | SuperResIDWEXKX.SuperResIDWE6K7: 8, 39 | } 40 | 41 | __block_type_min_channels_base_dict__ = { 42 | SuperResIDWEXKX.SuperResIDWE1K3: 8, 43 | SuperResIDWEXKX.SuperResIDWE2K3: 8, 44 | SuperResIDWEXKX.SuperResIDWE4K3: 8, 45 | SuperResIDWEXKX.SuperResIDWE6K3: 8, 46 | SuperResIDWEXKX.SuperResIDWE1K5: 8, 47 | SuperResIDWEXKX.SuperResIDWE2K5: 8, 48 | SuperResIDWEXKX.SuperResIDWE4K5: 8, 49 | SuperResIDWEXKX.SuperResIDWE6K5: 8, 50 | SuperResIDWEXKX.SuperResIDWE1K7: 8, 51 | SuperResIDWEXKX.SuperResIDWE2K7: 8, 52 | SuperResIDWEXKX.SuperResIDWE4K7: 8, 53 | SuperResIDWEXKX.SuperResIDWE6K7: 8, 54 | } 55 | 56 | 57 | def get_select_student_channels_list(out_channels): 58 | """generate all possible channels""" 59 | the_list = [out_channels * 2.5, out_channels * 2, out_channels * 1.5, out_channels * 1.25, 60 | out_channels, 61 | out_channels / 1.25, out_channels / 1.5, out_channels / 2, out_channels / 2.5] 62 | the_list = [max(8, x) for x in the_list] 63 | the_list = [global_utils.smart_round(x, base=8) for x in the_list] 64 | the_list = list(set(the_list)) 65 | the_list.sort(reverse=True) 66 | return the_list 67 | 68 | 69 | def get_select_student_sublayers_list(sub_layers): 70 | """generate all possible sublayers""" 71 | the_list = [sub_layers, 72 | sub_layers + 1, sub_layers + 2, 73 | sub_layers - 1, sub_layers - 2, ] 74 | the_list = [max(0, round(x)) for x in the_list] 75 | the_list = list(set(the_list)) 76 | the_list.sort(reverse=True) 77 | return the_list 78 | 79 | 80 | # pylint: disable=too-many-locals 81 | def gen_search_space(block_list, block_id): 82 | """ generate all possible blocks 83 | 84 | :param block_list (list): block list 85 | :param block_id (int): block index 86 | :return list(list()) 87 | """ 88 | 89 | the_block = block_list[block_id] 90 | student_blocks_list_list = [] 91 | 92 | if isinstance(the_block, super_blocks.SuperConvKXBNRELU): 93 | student_blocks_list = [] 94 | 95 | if the_block.kernel_size == 1: # last fc layer, never change fc 96 | student_out_channels_list = [the_block.out_channels] 97 | else: 98 | student_out_channels_list = get_select_student_channels_list(the_block.out_channels) 99 | 100 | for student_out_channels in student_out_channels_list: 101 | tmp_block_str = type(the_block).__name__ + f'({the_block.in_channels},\ 102 | {student_out_channels},\ 103 | {the_block.stride},1)' 104 | student_blocks_list.append(tmp_block_str) 105 | student_blocks_list = list(set(student_blocks_list)) 106 | assert len(student_blocks_list) >= 1 107 | student_blocks_list_list.append(student_blocks_list) 108 | else: 109 | for student_block_type_list in seach_space_block_type_list_list: 110 | student_blocks_list = [] 111 | student_out_channels_list = get_select_student_channels_list(the_block.out_channels) 112 | student_sublayers_list = get_select_student_sublayers_list(sub_layers=the_block.sub_layers) 113 | student_bottleneck_channels_list = get_select_student_channels_list(the_block.bottleneck_channels) 114 | for student_block_type in student_block_type_list: 115 | for student_out_channels, student_sublayers, student_bottleneck_channels in itertools.product( 116 | student_out_channels_list, student_sublayers_list, student_bottleneck_channels_list): 117 | 118 | # filter smallest possible channel for this block type 119 | min_possible_channels = __block_type_round_channels_base_dict__[student_block_type] 120 | channel_round_base = __block_type_round_channels_base_dict__[student_block_type] 121 | student_out_channels = global_utils.smart_round(student_out_channels, channel_round_base) 122 | student_bottleneck_channels = global_utils.smart_round(student_bottleneck_channels, 123 | channel_round_base) 124 | 125 | if student_out_channels < min_possible_channels or \ 126 | student_bottleneck_channels < min_possible_channels: 127 | continue 128 | if student_sublayers <= 0: # no empty layer 129 | continue 130 | tmp_block_str = student_block_type.__name__ + f'({the_block.in_channels},{student_out_channels},\ 131 | {the_block.stride},{student_bottleneck_channels},\ 132 | {student_sublayers})' 133 | student_blocks_list.append(tmp_block_str) 134 | student_blocks_list = list(set(student_blocks_list)) 135 | assert len(student_blocks_list) >= 1 136 | student_blocks_list_list.append(student_blocks_list) 137 | return student_blocks_list_list 138 | -------------------------------------------------------------------------------- /src/zen_nas/Masternet.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | 4 | defin MasterNet class 5 | ''' 6 | # pylint: disable=invalid-name,too-many-branches 7 | import os 8 | import sys 9 | import argparse 10 | import numpy as np 11 | import torch 12 | from torch import nn 13 | import torch.nn.functional as F 14 | sys.path.append(os.path.dirname(os.path.abspath(__file__))) 15 | try: 16 | import PlainNet 17 | from PlainNet import basic_blocks, super_blocks 18 | except ImportError: 19 | print('fail to import PlainNet') 20 | 21 | 22 | def parse_cmd_options(argv): 23 | parser = argparse.ArgumentParser() 24 | parser.add_argument('--no_BN', action='store_true') 25 | parser.add_argument('--no_reslink', action='store_true') 26 | parser.add_argument('--use_se', action='store_true') 27 | module_opt, _ = parser.parse_known_args(argv) 28 | return module_opt 29 | 30 | 31 | class MasterNet(PlainNet.PlainNet): 32 | """model calss""" 33 | 34 | # pylint: disable=too-many-arguments 35 | def __init__(self, argv=None, opt=None, num_classes=None, plainnet_struct=None, no_create=False, 36 | no_reslink=None, no_BN=None, use_se=None): 37 | 38 | if argv is not None: 39 | module_opt = parse_cmd_options(argv) 40 | else: 41 | module_opt = None 42 | 43 | if no_BN is None: 44 | if module_opt is not None: 45 | no_BN = module_opt.no_BN 46 | else: 47 | no_BN = False 48 | 49 | if no_reslink is None: 50 | if module_opt is not None: 51 | no_reslink = module_opt.no_reslink 52 | else: 53 | no_reslink = False 54 | 55 | if use_se is None: 56 | if module_opt is not None: 57 | use_se = module_opt.use_se 58 | else: 59 | use_se = False 60 | 61 | super().__init__(argv=argv, opt=opt, num_classes=num_classes, plainnet_struct=plainnet_struct, 62 | no_create=no_create, no_reslink=no_reslink, no_BN=no_BN, use_se=use_se) 63 | self.last_channels = self.block_list[-1].out_channels 64 | self.fc_linear = basic_blocks.Linear(in_channels=self.last_channels, 65 | out_channels=self.num_classes, no_create=no_create) 66 | 67 | self.no_create = no_create 68 | self.no_reslink = no_reslink 69 | self.no_BN = no_BN 70 | self.use_se = use_se 71 | self.module_list = None 72 | 73 | # bn eps 74 | for layer in self.modules(): 75 | if isinstance(layer, nn.BatchNorm2d): 76 | layer.eps = 1e-3 77 | 78 | def extract_stage_features_and_logit(self, input_, target_downsample_ratio=None): 79 | """get stages according to downsample ratio""" 80 | stage_features_list = [] 81 | image_size = input_.shape[2] 82 | output = input_ 83 | 84 | for _, the_block in enumerate(self.block_list): 85 | output = the_block(output) 86 | dowsample_ratio = round(image_size / output.shape[2]) 87 | if dowsample_ratio == target_downsample_ratio: 88 | stage_features_list.append(output) 89 | target_downsample_ratio *= 2 90 | 91 | output = F.adaptive_avg_pool2d(output, output_size=1) 92 | output = torch.flatten(output, 1) 93 | logit = self.fc_linear(output) 94 | return stage_features_list, logit 95 | 96 | def forward(self, input_): 97 | output = input_ 98 | for _, the_block in enumerate(self.block_list): 99 | output = the_block(output) 100 | 101 | output = F.adaptive_avg_pool2d(output, output_size=1) 102 | 103 | output = torch.flatten(output, 1) 104 | output = self.fc_linear(output) 105 | return output 106 | 107 | def forward_pre_GAP(self, input_): 108 | """compute output before global averag pool""" 109 | output = input_ 110 | for the_block in self.block_list: 111 | output = the_block(output) 112 | return output 113 | 114 | def get_FLOPs(self, input_resolution): 115 | the_res = input_resolution 116 | the_flops = 0 117 | for the_block in self.block_list: 118 | the_flops += the_block.get_FLOPs(the_res) 119 | the_res = the_block.get_output_resolution(the_res) 120 | 121 | the_flops += self.fc_linear.get_FLOPs(the_res) 122 | 123 | return the_flops 124 | 125 | def get_model_size(self): 126 | the_size = 0 127 | for the_block in self.block_list: 128 | the_size += the_block.get_model_size() 129 | 130 | the_size += self.fc_linear.get_model_size() 131 | 132 | return the_size 133 | 134 | def get_num_layers(self): 135 | num_layers = 0 136 | for block in self.block_list: 137 | assert isinstance(block, super_blocks.PlainNetSuperBlockClass) 138 | num_layers += block.sub_layers 139 | return num_layers 140 | 141 | def replace_block(self, block_id, new_block): 142 | """replace block_list[block_id] with new_block""" 143 | self.block_list[block_id] = new_block 144 | 145 | if block_id < len(self.block_list) - 1: 146 | if self.block_list[block_id + 1].in_channels != new_block.out_channels: 147 | self.block_list[block_id + 1].set_in_channels(new_block.out_channels) 148 | else: 149 | assert block_id == len(self.block_list) - 1 150 | self.last_channels = self.block_list[-1].out_channels 151 | if self.fc_linear.in_channels != self.last_channels: 152 | self.fc_linear.set_in_channels(self.last_channels) 153 | 154 | self.module_list = nn.ModuleList(self.block_list) 155 | 156 | def split(self, split_layer_threshold): 157 | """split block when block layers exceed threshold""" 158 | new_str = '' 159 | for block in self.block_list: 160 | new_str += block.split(split_layer_threshold=split_layer_threshold) 161 | return new_str 162 | 163 | def init_parameters(self): 164 | for module in self.modules(): 165 | if isinstance(module, nn.Conv2d): 166 | nn.init.xavier_normal_(module.weight.data, gain=3.26033) 167 | if hasattr(module, 'bias') and module.bias is not None: 168 | nn.init.zeros_(module.bias) 169 | elif isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)): 170 | nn.init.ones_(module.weight) 171 | nn.init.zeros_(module.bias) 172 | elif isinstance(module, nn.Linear): 173 | nn.init.normal_(module.weight, 0, 174 | 3.26033 * np.sqrt(2 / (module.weight.shape[0] + module.weight.shape[1]))) 175 | if hasattr(module, 'bias') and module.bias is not None: 176 | nn.init.zeros_(module.bias) 177 | else: 178 | pass 179 | 180 | for superblock in self.block_list: 181 | if not isinstance(superblock, super_blocks.PlainNetSuperBlockClass): 182 | continue 183 | for block in superblock.block_list: 184 | if not isinstance(block, (basic_blocks.ResBlock, basic_blocks.ResBlockProj)): 185 | continue 186 | # print('---debug set bn weight zero in resblock {}:{}'.format(superblock, block)) 187 | last_bn_block = None 188 | for inner_resblock in block.block_list: 189 | if isinstance(inner_resblock, basic_blocks.BN): 190 | last_bn_block = inner_resblock 191 | assert last_bn_block is not None 192 | # print('-------- last_bn_block={}'.format(last_bn_block)) 193 | nn.init.zeros_(last_bn_block.netblock.weight) 194 | -------------------------------------------------------------------------------- /src/zen_nas/hotfix/folder.py: -------------------------------------------------------------------------------- 1 | 2 | """load image and get dataset""" 3 | # pylint: disable=import-error,function-redefined,too-many-arguments 4 | import os 5 | import os.path 6 | import sys 7 | from PIL import Image 8 | from torchvision import get_image_backend 9 | import accimage 10 | from .vision import VisionDataset 11 | 12 | 13 | def has_file_allowed_extension(filename, extensions): 14 | """Checks if a file is an allowed extension. 15 | 16 | Args: 17 | filename (string): path to a file 18 | extensions (tuple of strings): extensions to consider (lowercase) 19 | 20 | Returns: 21 | bool: True if the filename ends with one of given extensions 22 | """ 23 | return filename.lower().endswith(extensions) 24 | 25 | 26 | def is_image_file(filename): 27 | """Checks if a file is an allowed image extension. 28 | 29 | Args: 30 | filename (string): path to a file 31 | 32 | Returns: 33 | bool: True if the filename ends with a known image extension 34 | """ 35 | return has_file_allowed_extension(filename, IMG_EXTENSIONS) 36 | 37 | 38 | def make_dataset(dir_path, class_to_idx, extensions=None, is_valid_file=None): 39 | images = [] 40 | dir_path = os.path.expanduser(dir_path) 41 | if not (extensions is None) ^ (is_valid_file is None): 42 | raise ValueError("Both extensions and is_valid_file cannot be None or not None at the same time") 43 | if extensions is not None: 44 | def is_valid_file(file_path): 45 | return has_file_allowed_extension(file_path, extensions) 46 | for target in sorted(class_to_idx.keys()): 47 | target_path = os.path.join(dir_path, target) 48 | if not os.path.isdir(target_path): 49 | continue 50 | for root, _, fnames in sorted(os.walk(target_path)): 51 | for fname in sorted(fnames): 52 | path = os.path.join(root, fname) 53 | if is_valid_file(path): 54 | item = (path, class_to_idx[target]) 55 | images.append(item) 56 | 57 | return images 58 | 59 | 60 | # pylint: disable=too-many-instance-attributes 61 | class DatasetFolder(VisionDataset): 62 | """A generic data loader where the samples are arranged in this way: :: 63 | 64 | root/class_x/xxx.ext 65 | root/class_x/xxy.ext 66 | root/class_x/xxz.ext 67 | 68 | root/class_y/123.ext 69 | root/class_y/nsdf3.ext 70 | root/class_y/asd932_.ext 71 | 72 | Args: 73 | root (string): Root directory path. 74 | loader (callable): A function to load a sample given its path. 75 | extensions (tuple[string]): A list of allowed extensions. 76 | both extensions and is_valid_file should not be passed. 77 | transform (callable, optional): A function/transform that takes in 78 | a sample and returns a transformed version. 79 | E.g, ``transforms.RandomCrop`` for images. 80 | target_transform (callable, optional): A function/transform that takes 81 | in the target and transforms it. 82 | is_valid_file (callable, optional): A function that takes path of an Image file 83 | and check if the file is a valid_file (used to check of corrupt files) 84 | both extensions and is_valid_file should not be passed. 85 | 86 | Attributes: 87 | classes (list): List of the class names. 88 | class_to_idx (dict): Dict with items (class_name, class_index). 89 | samples (list): List of (sample path, class_index) tuples 90 | targets (list): The class_index value for each image in the dataset 91 | """ 92 | 93 | def __init__(self, root, loader, extensions=None, transform=None, target_transform=None, is_valid_file=None): 94 | super().__init__(root) 95 | self.transform = transform 96 | self.target_transform = target_transform 97 | classes, class_to_idx = self._find_classes(self.root) 98 | samples = make_dataset(self.root, class_to_idx, extensions, is_valid_file) 99 | if len(samples) == 0: 100 | raise (RuntimeError("Found 0 files in subfolders of: " + self.root + "\n" 101 | "Supported extensions are: " + ",".join(extensions))) 102 | 103 | self.loader = loader 104 | self.extensions = extensions 105 | 106 | self.classes = classes 107 | self.class_to_idx = class_to_idx 108 | self.samples = samples 109 | self.targets = [s[1] for s in samples] 110 | 111 | # pylint: disable=no-self-use 112 | def _find_classes(self, dir_path): 113 | """ 114 | Finds the class folders in a dataset. 115 | 116 | Args: 117 | dir_path (string): Root directory path. 118 | 119 | Returns: 120 | tuple: (classes, class_to_idx) where classes are relative to (dir_path), and class_to_idx is a dictionary. 121 | 122 | Ensures: 123 | No class is a subdirectory of another. 124 | """ 125 | if sys.version_info >= (3, 5): 126 | # Faster and available in Python 3.5 and above 127 | classes = [d.name for d in os.scandir(dir_path) if d.is_dir()] 128 | else: 129 | classes = [d for d in os.listdir(dir_path) if os.path.isdir(os.path.join(dir_path, d))] 130 | classes.sort() 131 | class_to_idx = {classes[i]: i for i in range(len(classes))} 132 | return classes, class_to_idx 133 | 134 | def __getitem__(self, index): 135 | """ 136 | Args: 137 | index (int): Index 138 | 139 | Returns: 140 | tuple: (sample, target) where target is class_index of the target class. 141 | """ 142 | path, target = self.samples[index] 143 | sample = self.loader(path) 144 | if self.transform is not None: 145 | sample = self.transform(sample) 146 | if self.target_transform is not None: 147 | target = self.target_transform(target) 148 | 149 | return sample, target 150 | 151 | def __len__(self): 152 | return len(self.samples) 153 | 154 | 155 | IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp') 156 | 157 | 158 | def pil_loader(path): 159 | """return RGB img""" 160 | # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835) 161 | with open(path, 'rb') as image: 162 | img = Image.open(image) 163 | return img.convert('RGB') 164 | 165 | 166 | def accimage_loader(path): 167 | """load img from path""" 168 | try: 169 | return accimage.Image(path) 170 | except IOError: 171 | # Potentially a decoding problem, fall back to PIL.Image 172 | return pil_loader(path) 173 | 174 | 175 | def default_loader(path): 176 | """default load img from path""" 177 | if get_image_backend() == 'accimage': 178 | return accimage_loader(path) 179 | return pil_loader(path) 180 | 181 | 182 | # pylint: disable=too-few-public-methods 183 | class ImageFolder(DatasetFolder): 184 | """A generic data loader where the images are arranged in this way: :: 185 | 186 | root/dog/xxx.png 187 | root/dog/xxy.png 188 | root/dog/xxz.png 189 | 190 | root/cat/123.png 191 | root/cat/nsdf3.png 192 | root/cat/asd932_.png 193 | 194 | Args: 195 | root (string): Root directory path. 196 | transform (callable, optional): A function/transform that takes in an PIL image 197 | and returns a transformed version. E.g, ``transforms.RandomCrop`` 198 | target_transform (callable, optional): A function/transform that takes in the 199 | target and transforms it. 200 | loader (callable, optional): A function to load an image given its path. 201 | is_valid_file (callable, optional): A function that takes path of an Image file 202 | and check if the file is a valid_file (used to check of corrupt files) 203 | 204 | Attributes: 205 | classes (list): List of the class names. 206 | class_to_idx (dict): Dict with items (class_name, class_index). 207 | imgs (list): List of (image path, class_index) tuples 208 | """ 209 | 210 | def __init__(self, root, transform=None, target_transform=None, 211 | loader=default_loader, is_valid_file=None): 212 | super().__init__(root, loader, IMG_EXTENSIONS if is_valid_file is None else None, 213 | transform=transform, 214 | target_transform=target_transform, 215 | is_valid_file=is_valid_file) 216 | self.imgs = self.samples 217 | -------------------------------------------------------------------------------- /src/zen_nas/PlainNet/SuperResKXKX.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | ''' 4 | # pylint: disable=invalid-name 5 | import os 6 | import sys 7 | import uuid 8 | from torch import nn 9 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 10 | try: 11 | import PlainNet 12 | from PlainNet import _get_right_parentheses_index_ 13 | from PlainNet.super_blocks import PlainNetSuperBlockClass 14 | import global_utils 15 | except ImportError: 16 | print('fail to import zen_nas modules') 17 | 18 | 19 | # pylint: disable=too-many-instance-attributes,too-many-locals,too-many-arguments 20 | class SuperResKXKX(PlainNetSuperBlockClass): 21 | """Resnet BasicBlock-like block""" 22 | 23 | def __init__(self, in_channels=None, out_channels=None, stride=None, bottleneck_channels=None, 24 | sub_layers=None, kernel_size=None, no_create=False, no_reslink=False, 25 | no_BN=False, use_se=False, **kwargs): 26 | super().__init__(**kwargs) 27 | self.in_channels = in_channels 28 | self.out_channels = out_channels 29 | self.stride = stride 30 | self.bottleneck_channels = bottleneck_channels 31 | self.sub_layers = sub_layers 32 | self.kernel_size = kernel_size 33 | self.no_create = no_create 34 | self.no_reslink = no_reslink 35 | self.no_BN = no_BN 36 | self.use_se = use_se 37 | if self.use_se: 38 | print('---debug use_se in ' + str(self)) 39 | 40 | full_str = '' 41 | last_channels = in_channels 42 | current_stride = stride 43 | for i in range(self.sub_layers): 44 | inner_str = '' 45 | 46 | inner_str += f'ConvKX({last_channels},{self.bottleneck_channels},{self.kernel_size},{current_stride})' 47 | if not self.no_BN: 48 | inner_str += f'BN({self.bottleneck_channels})' 49 | inner_str += f'RELU({self.bottleneck_channels})' 50 | if self.use_se: 51 | inner_str += f'SE({bottleneck_channels})' 52 | 53 | inner_str += f'ConvKX({self.bottleneck_channels},{self.out_channels},{self.kernel_size},{1})' 54 | if not self.no_BN: 55 | inner_str += f'BN({self.out_channels})' 56 | 57 | if not self.no_reslink: 58 | if i == 0: 59 | res_str = f'ResBlockProj({inner_str})RELU({out_channels})' 60 | else: 61 | res_str = f'ResBlock({inner_str})RELU({out_channels})' 62 | else: 63 | res_str = f'{inner_str}RELU({out_channels})' 64 | 65 | full_str += res_str 66 | 67 | last_channels = out_channels 68 | current_stride = 1 69 | 70 | self.block_list = PlainNet.create_netblock_list_from_str(full_str, no_create=no_create, 71 | no_reslink=no_reslink, no_BN=no_BN, **kwargs) 72 | if not no_create: 73 | self.module_list = nn.ModuleList(self.block_list) 74 | else: 75 | self.module_list = None 76 | 77 | def forward_pre_relu(self, input_): 78 | """calculate the value before relu""" 79 | 80 | output = input_ 81 | for block in self.block_list[0:-1]: 82 | output = block(output) 83 | return output 84 | 85 | def __str__(self): 86 | return type(self).__name__ + f'({self.in_channels},{self.out_channels},{self.stride},\ 87 | {self.bottleneck_channels},{self.sub_layers})' 88 | 89 | def __repr__(self): 90 | return type(self).__name__ + f'({self.block_name}|in={self.in_channels},out={self.out_channels},\ 91 | stride={self.stride},btl_channels={self.bottleneck_channels},\ 92 | sub_layers={self.sub_layers},kernel_size={self.kernel_size})' 93 | 94 | def encode_structure(self): 95 | """pack channels and sub_layers to a list""" 96 | 97 | return [self.out_channels, self.sub_layers, self.bottleneck_channels] 98 | 99 | def split(self, split_layer_threshold): 100 | """split the layer when exceeding threshold""" 101 | 102 | if self.sub_layers >= split_layer_threshold: 103 | new_sublayers_1 = split_layer_threshold // 2 104 | new_sublayers_2 = self.sub_layers - new_sublayers_1 105 | new_block_str1 = type(self).__name__ + f'({self.in_channels},{self.out_channels},{self.stride},\ 106 | {self.bottleneck_channels},{new_sublayers_1})' 107 | 108 | new_block_str2 = type(self).__name__ + f'({self.out_channels},{self.out_channels},{1},\ 109 | {self.bottleneck_channels},{new_sublayers_2})' 110 | return new_block_str1 + new_block_str2 111 | return str(self) 112 | 113 | def structure_scale(self, scale=1.0, channel_scale=None, sub_layer_scale=None): 114 | """ adjust the number to a specific multiple or range""" 115 | 116 | if channel_scale is None: 117 | channel_scale = scale 118 | if sub_layer_scale is None: 119 | sub_layer_scale = scale 120 | 121 | new_out_channels = global_utils.smart_round(self.out_channels * channel_scale) 122 | new_bottleneck_channels = global_utils.smart_round(self.bottleneck_channels * channel_scale) 123 | new_sub_layers = max(1, round(self.sub_layers * sub_layer_scale)) 124 | 125 | return type(self).__name__ + f'({self.in_channels},{new_out_channels},{self.stride},\ 126 | {new_bottleneck_channels},{new_sub_layers})' 127 | 128 | # pylint: disable=arguments-differ 129 | @classmethod 130 | def create_from_str(cls, struct_str, **kwargs): 131 | """ class method 132 | 133 | :param s (str): SuperRes block str 134 | :return cls instance 135 | """ 136 | 137 | assert cls.is_instance_from_str(struct_str) 138 | idx = _get_right_parentheses_index_(struct_str) 139 | assert idx is not None 140 | param_str = struct_str[len(cls.__name__ + '('):idx] 141 | 142 | # find block_name 143 | tmp_idx = param_str.find('|') 144 | if tmp_idx < 0: 145 | tmp_block_name = f'uuid{uuid.uuid4().hex}' 146 | else: 147 | tmp_block_name = param_str[0:tmp_idx] 148 | param_str = param_str[tmp_idx + 1:] 149 | 150 | param_str_split = param_str.split(',') 151 | in_channels = int(param_str_split[0]) 152 | out_channels = int(param_str_split[1]) 153 | stride = int(param_str_split[2]) 154 | bottleneck_channels = int(param_str_split[3]) 155 | sub_layers = int(param_str_split[4]) 156 | return cls(in_channels=in_channels, out_channels=out_channels, stride=stride, 157 | bottleneck_channels=bottleneck_channels, sub_layers=sub_layers, 158 | block_name=tmp_block_name, **kwargs), struct_str[idx + 1:] 159 | 160 | 161 | class SuperResK3K3(SuperResKXKX): 162 | """ kernel size 3x3""" 163 | 164 | def __init__(self, in_channels=None, out_channels=None, stride=None, bottleneck_channels=None, 165 | sub_layers=None, no_create=False, **kwargs): 166 | super().__init__(in_channels=in_channels, out_channels=out_channels, stride=stride, 167 | bottleneck_channels=bottleneck_channels, sub_layers=sub_layers, 168 | kernel_size=3, 169 | no_create=no_create, **kwargs) 170 | 171 | 172 | class SuperResK5K5(SuperResKXKX): 173 | """kernel size 5x5""" 174 | 175 | def __init__(self, in_channels=None, out_channels=None, stride=None, bottleneck_channels=None, 176 | sub_layers=None, no_create=False, **kwargs): 177 | super().__init__(in_channels=in_channels, out_channels=out_channels, stride=stride, 178 | bottleneck_channels=bottleneck_channels, sub_layers=sub_layers, 179 | kernel_size=5, 180 | no_create=no_create, **kwargs) 181 | 182 | 183 | class SuperResK7K7(SuperResKXKX): 184 | """kernel size 7x7""" 185 | 186 | def __init__(self, in_channels=None, out_channels=None, stride=None, bottleneck_channels=None, 187 | sub_layers=None, no_create=False, **kwargs): 188 | super().__init__(in_channels=in_channels, out_channels=out_channels, stride=stride, 189 | bottleneck_channels=bottleneck_channels, sub_layers=sub_layers, 190 | kernel_size=7, 191 | no_create=no_create, **kwargs) 192 | 193 | 194 | def register_netblocks_dict(netblocks_dict: dict): 195 | """add different kernel size block to block dict""" 196 | 197 | this_py_file_netblocks_dict = { 198 | 'SuperResK3K3': SuperResK3K3, 199 | 'SuperResK5K5': SuperResK5K5, 200 | 'SuperResK7K7': SuperResK7K7, 201 | 202 | } 203 | netblocks_dict.update(this_py_file_netblocks_dict) 204 | return netblocks_dict 205 | -------------------------------------------------------------------------------- /src/zen_nas/ZenNet/masternet.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | ''' 4 | # pylint: disable=invalid-name,function-redefined,too-many-branches 5 | import os 6 | import sys 7 | import argparse 8 | import numpy as np 9 | import torch 10 | from torch import nn 11 | import torch.nn.functional as F 12 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 13 | try: 14 | import PlainNet 15 | from PlainNet import basic_blocks, super_blocks 16 | except ImportError: 17 | print('fail to import zen_nas modules') 18 | 19 | 20 | def parse_cmd_options(argv): 21 | parser = argparse.ArgumentParser() 22 | parser.add_argument('--no_BN', action='store_true') 23 | parser.add_argument('--no_reslink', action='store_true') 24 | parser.add_argument('--use_se', action='store_true') 25 | parser.add_argument('--dropout', type=float, default=None) 26 | module_opt, _ = parser.parse_known_args(argv) 27 | return module_opt 28 | 29 | 30 | # pylint: disable=too-many-instance-attributes,too-many-arguments 31 | class PlainNet(PlainNet.PlainNet): 32 | """model class""" 33 | 34 | def __init__(self, argv=None, opt=None, num_classes=None, plainnet_struct=None, no_create=False, 35 | no_reslink=None, no_BN=None, use_se=None, dropout=None, 36 | **kwargs): 37 | 38 | if argv is not None: 39 | module_opt = parse_cmd_options(argv) 40 | else: 41 | module_opt = None 42 | 43 | if no_BN is None: 44 | if module_opt is not None: 45 | no_BN = module_opt.no_BN 46 | else: 47 | no_BN = False 48 | 49 | if no_reslink is None: 50 | if module_opt is not None: 51 | no_reslink = module_opt.no_reslink 52 | else: 53 | no_reslink = False 54 | 55 | if use_se is None: 56 | if module_opt is not None: 57 | use_se = module_opt.use_se 58 | else: 59 | use_se = False 60 | 61 | if dropout is None: 62 | if module_opt is not None: 63 | self.dropout = module_opt.dropout 64 | else: 65 | self.dropout = None 66 | else: 67 | self.dropout = dropout 68 | 69 | if self.dropout is not None: 70 | print(f'--- using dropout={self.dropout:4g}') 71 | 72 | super().__init__(argv=argv, opt=opt, num_classes=num_classes, plainnet_struct=plainnet_struct, 73 | no_create=no_create, no_reslink=no_reslink, no_BN=no_BN, use_se=use_se, **kwargs) 74 | self.last_channels = self.block_list[-1].out_channels 75 | self.fc_linear = basic_blocks.Linear(in_channels=self.last_channels, 76 | out_channels=self.num_classes, no_create=no_create) 77 | 78 | self.no_create = no_create 79 | self.no_reslink = no_reslink 80 | self.no_BN = no_BN 81 | self.use_se = use_se 82 | self.module_list = None 83 | 84 | # bn eps 85 | for layer in self.modules(): 86 | if isinstance(layer, nn.BatchNorm2d): 87 | layer.eps = 1e-3 88 | 89 | def extract_stage_features_and_logit(self, input_, target_downsample_ratio=None): 90 | """split model into several stages given downsample_ratio""" 91 | stage_features_list = [] 92 | image_size = input_.shape[2] 93 | output = input_ 94 | 95 | for block_id, the_block in enumerate(self.block_list): 96 | output = the_block(output) 97 | if self.dropout is not None: 98 | dropout_p = float(block_id) / len(self.block_list) * self.dropout 99 | output = F.dropout(output, dropout_p, training=self.training, inplace=True) 100 | dowsample_ratio = round(image_size / output.shape[2]) 101 | if dowsample_ratio == target_downsample_ratio: 102 | stage_features_list.append(output) 103 | target_downsample_ratio *= 2 104 | 105 | output = F.adaptive_avg_pool2d(output, output_size=1) 106 | # if self.dropout is not None: 107 | # output = F.dropout(output, self.dropout, training=self.training, inplace=True) 108 | output = torch.flatten(output, 1) 109 | logit = self.fc_linear(output) 110 | return stage_features_list, logit 111 | 112 | def forward(self, input_): 113 | output = input_ 114 | for block_id, the_block in enumerate(self.block_list): 115 | output = the_block(output) 116 | if self.dropout is not None: 117 | dropout_p = float(block_id) / len(self.block_list) * self.dropout 118 | output = F.dropout(output, dropout_p, training=self.training, inplace=True) 119 | 120 | output = F.adaptive_avg_pool2d(output, output_size=1) 121 | if self.dropout is not None: 122 | output = F.dropout(output, self.dropout, training=self.training, inplace=True) 123 | output = torch.flatten(output, 1) 124 | output = self.fc_linear(output) 125 | return output 126 | 127 | def forward_pre_GAP(self, input_): 128 | """compute result before the Global Average Pool""" 129 | output = input_ 130 | for the_block in self.block_list: 131 | output = the_block(output) 132 | return output 133 | 134 | def get_FLOPs(self, input_resolution): 135 | """model FLOPs""" 136 | the_res = input_resolution 137 | the_flops = 0 138 | for the_block in self.block_list: 139 | the_flops += the_block.get_FLOPs(the_res) 140 | the_res = the_block.get_output_resolution(the_res) 141 | 142 | the_flops += self.fc_linear.get_FLOPs(the_res) 143 | 144 | return the_flops 145 | 146 | def get_model_size(self): 147 | """model parameters""" 148 | the_size = 0 149 | for the_block in self.block_list: 150 | the_size += the_block.get_model_size() 151 | 152 | the_size += self.fc_linear.get_model_size() 153 | 154 | return the_size 155 | 156 | def get_num_layers(self): 157 | """total layers""" 158 | num_layers = 0 159 | for block in self.block_list: 160 | assert isinstance(block, super_blocks.PlainNetSuperBlockClass) 161 | num_layers += block.sub_layers 162 | return num_layers 163 | 164 | def replace_block(self, block_id, new_block): 165 | """replace block_list[block_id] with new_block""" 166 | self.block_list[block_id] = new_block 167 | 168 | if block_id < len(self.block_list) - 1: 169 | if self.block_list[block_id + 1].in_channels != new_block.out_channels: 170 | self.block_list[block_id + 1].set_in_channels(new_block.out_channels) 171 | else: 172 | assert block_id == len(self.block_list) - 1 173 | self.last_channels = self.block_list[-1].out_channels 174 | if self.fc_linear.in_channels != self.last_channels: 175 | self.fc_linear.set_in_channels(self.last_channels) 176 | 177 | self.module_list = nn.ModuleList(self.block_list) 178 | 179 | def split(self, split_layer_threshold): 180 | """split block when exceeding threshold""" 181 | new_str = '' 182 | for block in self.block_list: 183 | new_str += block.split(split_layer_threshold=split_layer_threshold) 184 | return new_str 185 | 186 | def init_parameters(self): 187 | """initilize model""" 188 | for module in self.modules(): 189 | if isinstance(module, nn.Conv2d): 190 | nn.init.xavier_normal_(module.weight.data, gain=3.26033) 191 | if hasattr(module, 'bias') and module.bias is not None: 192 | nn.init.zeros_(module.bias) 193 | elif isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)): 194 | nn.init.ones_(module.weight) 195 | nn.init.zeros_(module.bias) 196 | elif isinstance(module, nn.Linear): 197 | nn.init.normal_(module.weight, 0, 198 | 3.26033 * np.sqrt(2 / (module.weight.shape[0] + module.weight.shape[1]))) 199 | if hasattr(module, 'bias') and module.bias is not None: 200 | nn.init.zeros_(module.bias) 201 | else: 202 | pass 203 | 204 | for superblock in self.block_list: 205 | if not isinstance(superblock, super_blocks.PlainNetSuperBlockClass): 206 | continue 207 | for block in superblock.block_list: 208 | if not isinstance(block, basic_blocks.ResBlock): 209 | continue 210 | # print('---debug set bn weight zero in resblock {}:{}'.format(superblock, block)) 211 | last_bn_block = None 212 | for inner_resblock in block.block_list: 213 | if isinstance(inner_resblock, basic_blocks.BN): 214 | last_bn_block = inner_resblock 215 | assert last_bn_block is not None 216 | # print('-------- last_bn_block={}'.format(last_bn_block)) 217 | nn.init.zeros_(last_bn_block.netblock.weight) 218 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # zen_nas 2 | 3 | [![Build Status](https://dev.azure.com/Adlik/GitHub/_apis/build/status/Adlik.zen_nas?branchName=main)](https://dev.azure.com/Adlik/GitHub/_build/latest?definitionId=5&branchName=main) 4 | [![Bors enabled](https://bors.tech/images/badge_small.svg)](https://app.bors.tech/repositories/38905) 5 | [![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0) 6 | 7 | ## Introduction 8 | 9 | Our based code is forked from [ZenNAS](https://github.com/idstcv/ZenNAS). 10 | We modify the code to make it easier to use, and according to paper 11 | [Zen-NAS: A Zero-Shot NAS for High-Performance Image Recognition](https://arxiv.org/abs/2102.01063), 12 | we add some features that ZenNAS does not have. 13 | 14 | We mainly made the following changes: 15 | 16 | - support Horovod, PyTorch distributed training 17 | - code refactoring for evolutionary search, speed up searching 18 | - repair THCudaTensor sizes too large problem when search latency model 19 | 20 | Besides, as some functions can not work correctly under distributed training, we provide a distributed version. 21 | 22 | ## Experimental results 23 | 24 | We tested the modified code and verified its correctness. The results are as follows: 25 | 26 | We used apex with mixed precision to complete the training within 5 days on 8 tesla V100 GPUs, 27 | and the results are consistent with the paper. 28 | 29 | | | paper model accuracy | distributed training accuracy | 30 | | :----------: | :------------------: | :---------------------------: | 31 | | ZenNet-0.1ms | 77.8% | 77.922% | 32 | | ZenNet-0.8ms | 83.0% | 83.214% | 33 | 34 | Note that the same results can be obtained with Horovod, but we took more time to complete the training. 35 | So we recommend using apex for distributed training. 36 | 37 | Before the code was released, we reproduced the paper algorithm and searched the model according to the paper's conditions. 38 | The following table shows the comparison results between the search model and the paper model. 39 | 40 | | | paper accuracy | searched model accuracy | 41 | | :---------: | :------------: | :---------------------: | 42 | | latency01ms | 77.8% | 78.622% | 43 | | latency05ms | 82.7% | 82.752% | 44 | | latency12ms | 83.6% | 83.466% | 45 | 46 | For proving the effectiveness of the algorithm, we experimented with several different model searches 47 | and get the following result. 48 | We use a single Tesla V100 GPU to evolve the population 50000 times. 49 | | method | model | search time(hours) | model score | 50 | | :-----------: | :---------: | :----------------: | :---------: | 51 | | ZenNAS | latency01ms | 98.4274 | 126.038 | 52 | | \ | latency05ms | 22.0189 | 243.101 | 53 | | \ | latency08ms | 28.5952 | 304.323 | 54 | | \ | latency12ms | 44.6237 | 375.027 | 55 | | modify-ZenNAS | latency01ms | 64.988 | 134.896 | 56 | | \ | latency05ms | 20.9895 | 245.712 | 57 | | \ | latency08ms | 25.0358 | 310.629 | 58 | | \ | latency12ms | 43.239 | 386.669 | 59 | 60 | ## Reproduce Paper Experiments 61 | 62 | ### System Requirements 63 | 64 | - PyTorch >= 1.6, Python >= 3.7 65 | - By default, ImageNet dataset is stored under \~/data/imagenet; 66 | CIFAR-10/CIFAR-100 is stored under \~/data/pytorch\_cifar10 or \~/data/pytorch\_cifar100 67 | - Pre-trained parameters are cached under \~/.cache/pytorch/checkpoints/zennet\_pretrained 68 | 69 | ### Package Requirements 70 | 71 | - ptflops 72 | - tensorboard >= 1.15 (optional) 73 | - apex 74 | 75 | ### Pre-trained model download 76 | 77 | If you want to evaluate pre-trained models, 78 | please go to [ZenNAS](https://github.com/idstcv/ZenNAS) to download the pre-trained model. 79 | 80 | ### Evaluate pre-trained models on ImageNet and CIFAR-10/100 81 | 82 | To evaluate the pre-trained model on ImageNet using GPU 0: 83 | 84 | ``` bash 85 | python val.py --fp16 --gpu 0 --arch ${zennet_model_name} 86 | ``` 87 | 88 | where ${zennet\_model\_name} should be replaced by a valid ZenNet model name. 89 | The complete list of model names can be found in the 'Pre-trained Models' section. 90 | 91 | To evaluate the pre-trained model on CIFAR-10 or CIFAR-100 using GPU 0: 92 | 93 | ``` bash 94 | python val_cifar.py --dataset cifar10 --gpu 0 --arch ${zennet_model_name} 95 | ``` 96 | 97 | To create a ZenNet in your python code: 98 | 99 | ``` python 100 | gpu=0 101 | model = ZenNet.get_ZenNet(opt.arch, pretrained=True) 102 | torch.cuda.set_device(gpu) 103 | torch.backends.cudnn.benchmark = True 104 | model = model.cuda(gpu) 105 | model = model.half() 106 | model.eval() 107 | ``` 108 | 109 | ### usage 110 | 111 | We supply apex and Horovod distributed training scripts, you can modify other original scripts based on these scripts. 112 | apex script: 113 | 114 | ```bash 115 | scripts/Zen_NAS_ImageNet_latency0.1ms_train_apex.sh 116 | ``` 117 | 118 | Horovod script: 119 | 120 | ```bash 121 | scripts/Zen_NAS_ImageNet_latency0.1ms_train.sh 122 | ``` 123 | 124 | If you want to search model, please notice the choices "--fix_initialize" and "--origin". 125 | "--fix_initialize" decides how to initialize population, the algorithm default choice is random initialization. 126 | "--origin" determines how the mutation model is generated. 127 | When specified "--origin", the mutated model will be produced using the original method. 128 | 129 | ### Searching on CIFAR-10/100 130 | 131 | Searching for CIFAR-10/100 models with budget params < 1M, using different zero-shot proxies: 132 | 133 | ```bash 134 | scripts/Flops_NAS_cifar_params1M.sh 135 | scripts/GradNorm_NAS_cifar_params1M.sh 136 | scripts/NASWOT_NAS_cifar_params1M.sh 137 | scripts/Params_NAS_cifar_params1M.sh 138 | scripts/Random_NAS_cifar_params1M.sh 139 | scripts/Syncflow_NAS_cifar_params1M.sh 140 | scripts/TE_NAS_cifar_params1M.sh 141 | scripts/Zen_NAS_cifar_params1M.sh 142 | ``` 143 | 144 | ### Searching on ImageNet 145 | 146 | Searching for ImageNet models, with latency budget on NVIDIA V100 from 0.1 ms/image to 1.2 ms/image at batch size 64 FP16: 147 | 148 | ```bash 149 | scripts/Zen_NAS_ImageNet_latency0.1ms.sh 150 | scripts/Zen_NAS_ImageNet_latency0.2ms.sh 151 | scripts/Zen_NAS_ImageNet_latency0.3ms.sh 152 | scripts/Zen_NAS_ImageNet_latency0.5ms.sh 153 | scripts/Zen_NAS_ImageNet_latency0.8ms.sh 154 | scripts/Zen_NAS_ImageNet_latency1.2ms.sh 155 | ``` 156 | 157 | Searching for ImageNet models, with FLOPs budget from 400M to 800M: 158 | 159 | ``` bash 160 | scripts/Zen_NAS_ImageNet_flops400M.sh 161 | scripts/Zen_NAS_ImageNet_flops600M.sh 162 | scripts/Zen_NAS_ImageNet_flops800M.sh 163 | ``` 164 | 165 | ## Customize Your Own Search Space and Zero-Shot Proxy 166 | 167 | The masternet definition is stored in "Masternet.py". 168 | The masternet takes in a structure string and parses it into a PyTorch nn.Module object. 169 | The structure string defines the layer structure which is implemented in "PlainNet/*.py" files. 170 | For example, in "PlainNet/SuperResK1KXK1.py", 171 | we defined SuperResK1K3K1 block, which consists of multiple layers of ResNet blocks. 172 | To define your block, e.g. ABC_Block, first, implement "PlainNet/ABC_Block.py". 173 | Then in "PlainNet/\_\_init\_\_.py", after the last line, append the following lines to register the new block definition: 174 | 175 | ```python 176 | from PlainNet import ABC_Block 177 | _all_netblocks_dict_ = ABC_Block.register_netblocks_dict(_all_netblocks_dict_) 178 | ``` 179 | 180 | After the above registration call, the PlainNet module can parse your customized block from the structure string. 181 | 182 | The search space definitions are stored in SearchSpace/*.py. The important function is 183 | 184 | ```python 185 | gen_search_space(block_list, block_id) 186 | ``` 187 | 188 | block_list is a list of super-blocks parsed by the masternet. 189 | block_id is the index of the block in block_list which will be replaced later by a mutated block 190 | This function must return a list of mutated blocks. 191 | 192 | ### Direct specify search space 193 | 194 | "PlainNet/AABC_Block.py" has defined the candidate blocks, 195 | you can directly specify candidate blocks in the search spaces by passing parameters "--search_space_list". 196 | So you have two methods to specify search spaces. 197 | Taking ResNet-like search space as an example, you can use "--search_space SearchSpace/search_space_XXBL.py" or 198 | "--search_space_list PlainNet/SuperResK1KXK1.py PlainNet/SuperResKXKX.py" to specify search space. Both of them are equivalent. 199 | 200 | In scripts, when you choose to use the first method to specify search space, 201 | **you should also add other two parameters "--fix_initialize" and"--origin"**, 202 | so the algorithm will initialize with a fixed model. 203 | 204 | The zero-shot proxies are implemented in "ZeroShotProxy/*.py". The evolutionary algorithm is implemented in "evolution_search.py". 205 | "analyze_model.py" prints the FLOPs and model size of the given network. 206 | "benchmark_network_latency.py" measures the network inference latency. 207 | "train_image_classification.py" implements SGD gradient training 208 | and "ts_train_image_classification.py" implements teacher-student distillation. 209 | 210 | ## Open Source 211 | 212 | A few files in this repository are modified from the following open-source implementations: 213 | 214 | ```text 215 | https://github.com/DeepVoltaire/AutoAugment/blob/master/autoaugment.py 216 | https://github.com/VITA-Group/TENAS 217 | https://github.com/SamsungLabs/zero-cost-nas 218 | https://github.com/BayesWatch/nas-without-training 219 | https://github.com/rwightman/gen-efficientnet-pytorch 220 | https://pytorch.org/vision/0.8/_modules/torchvision/models/resnet.html 221 | ``` 222 | 223 | ## Copyright 224 | 225 | Copyright 2021 ZTE corporation. All Rights Reserved. 226 | -------------------------------------------------------------------------------- /src/zen_nas/PlainNet/super_blocks.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | 4 | define SuperConvBlock class with different kernel size 5 | ''' 6 | # pylint: disable=W0613,not-an-iterable,unsubscriptable-object,too-many-arguments 7 | import os 8 | import sys 9 | import uuid 10 | from torch import nn 11 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 12 | try: 13 | import global_utils 14 | import PlainNet 15 | from PlainNet import _get_right_parentheses_index_, basic_blocks 16 | except ImportError: 17 | print('fail to import zen_nas modules') 18 | 19 | 20 | class PlainNetSuperBlockClass(basic_blocks.PlainNetBasicBlockClass): 21 | """SuperBlock base class""" 22 | def __init__(self, in_channels=None, out_channels=None, stride=None, sub_layers=None, no_create=False, **kwargs): 23 | super().__init__() 24 | self.in_channels = in_channels 25 | self.out_channels = out_channels 26 | self.stride = stride 27 | self.sub_layers = sub_layers 28 | self.no_create = no_create 29 | self.block_list = None 30 | self.module_list = None 31 | 32 | def forward(self, input_): 33 | output = input_ 34 | for block in self.block_list: 35 | output = block(output) 36 | return output 37 | 38 | def __str__(self): 39 | return type(self).__name__ + f'({self.in_channels},{self.out_channels},\ 40 | {self.stride},{self.sub_layers})' 41 | 42 | def __repr__(self): 43 | return type(self).__name__ + f'({self.block_name}|{self.in_channels},\ 44 | {self.out_channels},{self.stride},{self.sub_layers})' 45 | 46 | def get_output_resolution(self, input_resolution): 47 | resolution = input_resolution 48 | for block in self.block_list: 49 | resolution = block.get_output_resolution(resolution) 50 | return resolution 51 | 52 | # pylint: disable=invalid-name 53 | def get_FLOPs(self, input_resolution): 54 | resolution = input_resolution 55 | flops = 0.0 56 | for block in self.block_list: 57 | flops += block.get_FLOPs(resolution) 58 | resolution = block.get_output_resolution(resolution) 59 | return flops 60 | 61 | def get_model_size(self): 62 | model_size = 0.0 63 | for block in self.block_list: 64 | model_size += block.get_model_size() 65 | return model_size 66 | 67 | def set_in_channels(self, channels): 68 | self.in_channels = channels 69 | if len(self.block_list) == 0: 70 | self.out_channels = channels 71 | return 72 | 73 | self.block_list[0].set_in_channels(channels) 74 | last_channels = self.block_list[0].out_channels 75 | if len(self.block_list) >= 2 and \ 76 | (isinstance(self.block_list[0], (basic_blocks.ConvKX, basic_blocks.ConvDW))) and \ 77 | isinstance(self.block_list[1], basic_blocks.BN): 78 | self.block_list[1].set_in_channels(last_channels) 79 | 80 | def encode_structure(self): 81 | """pack channels and sub_layers with list""" 82 | 83 | return [self.out_channels, self.sub_layers] 84 | 85 | @classmethod 86 | def create_from_str(cls, struct_str, no_create=False, **kwargs): 87 | assert cls.is_instance_from_str(struct_str) 88 | idx = _get_right_parentheses_index_(struct_str) 89 | assert idx is not None 90 | param_str = struct_str[len(cls.__name__ + '('):idx] 91 | 92 | # find block_name 93 | tmp_idx = param_str.find('|') 94 | if tmp_idx < 0: 95 | tmp_block_name = f'uuid{uuid.uuid4().hex}' 96 | else: 97 | tmp_block_name = param_str[0:tmp_idx] 98 | param_str = param_str[tmp_idx + 1:] 99 | 100 | param_str_split = param_str.split(',') 101 | in_channels = int(param_str_split[0]) 102 | out_channels = int(param_str_split[1]) 103 | stride = int(param_str_split[2]) 104 | sub_layers = int(param_str_split[3]) 105 | return cls(in_channels=in_channels, out_channels=out_channels, stride=stride, 106 | sub_layers=sub_layers, block_name=tmp_block_name, no_create=no_create, 107 | **kwargs), struct_str[idx + 1:] 108 | 109 | 110 | # pylint: disable=invalid-name,too-many-instance-attributes 111 | class SuperConvKXBNRELU(PlainNetSuperBlockClass): 112 | """SuperConv block""" 113 | 114 | def __init__(self, in_channels=None, out_channels=None, stride=None, sub_layers=None, kernel_size=None, 115 | no_create=False, no_reslink=False, no_BN=False, **kwargs): 116 | super().__init__(**kwargs) 117 | self.in_channels = in_channels 118 | self.out_channels = out_channels 119 | self.stride = stride 120 | self.sub_layers = sub_layers 121 | self.kernel_size = kernel_size 122 | self.no_create = no_create 123 | self.no_reslink = no_reslink 124 | self.no_BN = no_BN 125 | 126 | # if self.no_reslink: 127 | # print('Warning! {} use no_reslink'.format(str(self))) 128 | # if self.no_BN: 129 | # print('Warning! {} use no_BN'.format(str(self))) 130 | 131 | full_str = '' 132 | last_channels = in_channels 133 | current_stride = stride 134 | for _ in range(self.sub_layers): 135 | if not self.no_BN: 136 | inner_str = f'ConvKX({last_channels},{self.out_channels},\ 137 | {self.kernel_size},{current_stride})BN({self.out_channels})RELU({self.out_channels})' 138 | else: 139 | inner_str = f'ConvKX({last_channels},{self.out_channels},\ 140 | {self.kernel_size},{current_stride})RELU({self.out_channels})' 141 | full_str += inner_str 142 | 143 | last_channels = out_channels 144 | current_stride = 1 145 | 146 | self.block_list = PlainNet.create_netblock_list_from_str(full_str, no_create=no_create, 147 | no_reslink=no_reslink, no_BN=no_BN) 148 | if not no_create: 149 | self.module_list = nn.ModuleList(self.block_list) 150 | else: 151 | self.module_list = None 152 | 153 | def forward_pre_relu(self, input_): 154 | output = input_ 155 | for block in self.block_list[0:-1]: 156 | output = block(output) 157 | return output 158 | 159 | def __str__(self): 160 | return type(self).__name__ + f'({self.in_channels},{self.out_channels},\ 161 | {self.stride},{self.sub_layers})' 162 | 163 | def __repr__(self): 164 | return type(self).__name__ + f'({self.block_name}|in={self.in_channels},out={self.out_channels},\ 165 | stride={self.stride},sub_layers={self.sub_layers},\ 166 | kernel_size={self.kernel_size})' 167 | 168 | def split(self, split_layer_threshold): 169 | """return str(self)""" 170 | 171 | return str(self) 172 | 173 | def structure_scale(self, scale=1.0, channel_scale=None, sub_layer_scale=None): 174 | """ adjust the number to a specific multiple or range""" 175 | 176 | if channel_scale is None: 177 | channel_scale = scale 178 | if sub_layer_scale is None: 179 | sub_layer_scale = scale 180 | 181 | new_out_channels = global_utils.smart_round(self.out_channels * channel_scale) 182 | new_sub_layers = max(1, round(self.sub_layers * sub_layer_scale)) 183 | 184 | return type(self).__name__ + f'({self.in_channels},{new_out_channels},{self.stride},{new_sub_layers})' 185 | 186 | 187 | class SuperConvK1BNRELU(SuperConvKXBNRELU): 188 | """kernel size 1x1""" 189 | 190 | def __init__(self, in_channels=None, out_channels=None, stride=None, sub_layers=None, no_create=False, **kwargs): 191 | super().__init__(in_channels=in_channels, out_channels=out_channels, stride=stride, 192 | sub_layers=sub_layers, 193 | kernel_size=1, 194 | no_create=no_create, **kwargs) 195 | 196 | 197 | class SuperConvK3BNRELU(SuperConvKXBNRELU): 198 | """kernel size 3x3""" 199 | 200 | def __init__(self, in_channels=None, out_channels=None, stride=None, sub_layers=None, no_create=False, **kwargs): 201 | super().__init__(in_channels=in_channels, out_channels=out_channels, stride=stride, 202 | sub_layers=sub_layers, 203 | kernel_size=3, 204 | no_create=no_create, **kwargs) 205 | 206 | 207 | class SuperConvK5BNRELU(SuperConvKXBNRELU): 208 | """kernel size 5x5""" 209 | 210 | def __init__(self, in_channels=None, out_channels=None, stride=None, sub_layers=None, no_create=False, **kwargs): 211 | super().__init__(in_channels=in_channels, out_channels=out_channels, stride=stride, 212 | sub_layers=sub_layers, 213 | kernel_size=5, 214 | no_create=no_create, **kwargs) 215 | 216 | 217 | class SuperConvK7BNRELU(SuperConvKXBNRELU): 218 | """"kernel size 7x7""" 219 | 220 | def __init__(self, in_channels=None, out_channels=None, stride=None, sub_layers=None, no_create=False, **kwargs): 221 | super().__init__(in_channels=in_channels, out_channels=out_channels, stride=stride, 222 | sub_layers=sub_layers, 223 | kernel_size=7, 224 | no_create=no_create, **kwargs) 225 | 226 | 227 | def register_netblocks_dict(netblocks_dict: dict): 228 | """add different kernel size block to block dict""" 229 | 230 | this_py_file_netblocks_dict = { 231 | 'SuperConvK1BNRELU': SuperConvK1BNRELU, 232 | 'SuperConvK3BNRELU': SuperConvK3BNRELU, 233 | 'SuperConvK5BNRELU': SuperConvK5BNRELU, 234 | 'SuperConvK7BNRELU': SuperConvK7BNRELU, 235 | 236 | } 237 | netblocks_dict.update(this_py_file_netblocks_dict) 238 | return netblocks_dict 239 | -------------------------------------------------------------------------------- /src/zen_nas/PlainNet/SuperResK1KXK1.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Copyright (C) 2010-2021 Alibaba Group Holding Limited. 3 | ''' 4 | # pylint: disable=invalid-name,too-many-branches,too-many-statements 5 | import os 6 | import sys 7 | import uuid 8 | from torch import nn 9 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) 10 | try: 11 | import PlainNet 12 | from PlainNet import _get_right_parentheses_index_ 13 | from PlainNet.super_blocks import PlainNetSuperBlockClass 14 | import global_utils 15 | except ImportError: 16 | print('fail to import zen_nas modules') 17 | 18 | 19 | # pylint: disable=too-many-instance-attributes,too-many-locals,too-many-arguments 20 | class SuperResK1KXK1(PlainNetSuperBlockClass): 21 | """Resnet BottleneckBlock-like block""" 22 | 23 | def __init__(self, in_channels=None, out_channels=None, stride=None, bottleneck_channels=None, 24 | sub_layers=None, kernel_size=None, no_create=False, no_reslink=False, 25 | no_BN=False, use_se=False, **kwargs): 26 | super().__init__(**kwargs) 27 | self.in_channels = in_channels 28 | self.out_channels = out_channels 29 | self.stride = stride 30 | self.bottleneck_channels = bottleneck_channels 31 | self.sub_layers = sub_layers 32 | self.kernel_size = kernel_size 33 | self.no_create = no_create 34 | self.no_reslink = no_reslink 35 | self.no_BN = no_BN 36 | self.use_se = use_se 37 | if self.use_se: 38 | print('---debug use_se in ' + str(self)) 39 | 40 | full_str = '' 41 | last_channels = in_channels 42 | current_stride = stride 43 | for i in range(self.sub_layers): 44 | inner_str = '' 45 | 46 | # first bl-block with reslink 47 | inner_str += f'ConvKX({last_channels},{self.bottleneck_channels},{1},{1})' 48 | if not self.no_BN: 49 | inner_str += f'BN({self.bottleneck_channels})' 50 | inner_str += f'RELU({self.bottleneck_channels})' 51 | inner_str += f'ConvKX({self.bottleneck_channels},{self.bottleneck_channels},\ 52 | {self.kernel_size},{current_stride})' 53 | if not self.no_BN: 54 | inner_str += f'BN({self.bottleneck_channels})' 55 | inner_str += f'RELU({self.bottleneck_channels})' 56 | if self.use_se: 57 | inner_str += f'SE({self.bottleneck_channels})' 58 | 59 | inner_str += f'ConvKX({self.bottleneck_channels},{self.out_channels},{1},{1})' 60 | if not self.no_BN: 61 | inner_str += f'BN({self.out_channels})'.format() 62 | 63 | if not self.no_reslink: 64 | if i == 0: 65 | res_str = f'ResBlockProj({inner_str})RELU({out_channels})' 66 | else: 67 | res_str = f'ResBlock({inner_str})RELU({out_channels})' 68 | else: 69 | res_str = f'{inner_str}RELU({out_channels})' 70 | 71 | full_str += res_str 72 | 73 | # second bl-block with reslink 74 | inner_str = '' 75 | inner_str += f'ConvKX({self.out_channels},{self.bottleneck_channels},{1},{1})' 76 | if not self.no_BN: 77 | inner_str += f'BN({self.bottleneck_channels})' 78 | inner_str += f'RELU({self.bottleneck_channels})' 79 | 80 | inner_str += f'ConvKX({self.bottleneck_channels},{self.bottleneck_channels},\ 81 | {self.kernel_size},{1})' 82 | if not self.no_BN: 83 | inner_str += f'BN({self.bottleneck_channels})' 84 | inner_str += f'RELU({self.bottleneck_channels})' 85 | if self.use_se: 86 | inner_str += f'SE({self.bottleneck_channels})' 87 | 88 | inner_str += f'ConvKX({self.bottleneck_channels},{self.out_channels},{1},{1})' 89 | if not self.no_BN: 90 | inner_str += f'BN({self.out_channels})' 91 | 92 | if not self.no_reslink: 93 | res_str = f'ResBlock({inner_str})RELU({self.out_channels})' 94 | else: 95 | res_str = f'{inner_str}RELU({self.out_channels})' 96 | 97 | full_str += res_str 98 | 99 | last_channels = out_channels 100 | current_stride = 1 101 | 102 | self.block_list = PlainNet.create_netblock_list_from_str(full_str, no_create=no_create, 103 | no_reslink=no_reslink, no_BN=no_BN, **kwargs) 104 | if not no_create: 105 | self.module_list = nn.ModuleList(self.block_list) 106 | else: 107 | self.module_list = None 108 | 109 | def __str__(self): 110 | return type(self).__name__ + f'({self.in_channels},{self.out_channels},\ 111 | {self.stride},{self.bottleneck_channels},{self.sub_layers})' 112 | 113 | def __repr__(self): 114 | return type(self).__name__ + f'({self.block_name}|in={self.in_channels},out={self.out_channels},\ 115 | stride={self.stride},btl_channels={self.bottleneck_channels},\ 116 | sub_layers={self.sub_layers},kernel_size={self.kernel_size})' 117 | 118 | def encode_structure(self): 119 | """pack channels and sub_layers to a list""" 120 | 121 | return [self.out_channels, self.sub_layers, self.bottleneck_channels] 122 | 123 | def split(self, split_layer_threshold): 124 | """split the layer when exceeding threshold""" 125 | 126 | if self.sub_layers >= split_layer_threshold: 127 | new_sublayers_1 = split_layer_threshold // 2 128 | new_sublayers_2 = self.sub_layers - new_sublayers_1 129 | new_block_str1 = type(self).__name__ + f'({self.in_channels},{self.out_channels},\ 130 | {self.stride},{self.bottleneck_channels},{new_sublayers_1})' 131 | 132 | new_block_str2 = type(self).__name__ + f'({self.out_channels},{self.out_channels},\ 133 | {1},{self.bottleneck_channels},{new_sublayers_2})' 134 | return new_block_str1 + new_block_str2 135 | return str(self) 136 | 137 | def structure_scale(self, scale=1.0, channel_scale=None, sub_layer_scale=None): 138 | """ adjust the number to a specific multiple or range""" 139 | 140 | if channel_scale is None: 141 | channel_scale = scale 142 | if sub_layer_scale is None: 143 | sub_layer_scale = scale 144 | 145 | new_out_channels = global_utils.smart_round(self.out_channels * channel_scale) 146 | new_bottleneck_channels = global_utils.smart_round(self.bottleneck_channels * channel_scale) 147 | new_sub_layers = max(1, round(self.sub_layers * sub_layer_scale)) 148 | 149 | return type(self).__name__ + f'({self.in_channels},{new_out_channels}, \ 150 | {self.stride},{new_bottleneck_channels},{new_sub_layers})' 151 | 152 | # pylint: disable=arguments-differ 153 | @classmethod 154 | def create_from_str(cls, struct_str, **kwargs): 155 | assert cls.is_instance_from_str(struct_str) 156 | idx = _get_right_parentheses_index_(struct_str) 157 | assert idx is not None 158 | param_str = struct_str[len(cls.__name__ + '('):idx] 159 | 160 | # find block_name 161 | tmp_idx = param_str.find('|') 162 | if tmp_idx < 0: 163 | tmp_block_name = f'uuid{uuid.uuid4().hex}' 164 | else: 165 | tmp_block_name = param_str[0:tmp_idx] 166 | param_str = param_str[tmp_idx + 1:] 167 | 168 | param_str_split = param_str.split(',') 169 | in_channels = int(param_str_split[0]) 170 | out_channels = int(param_str_split[1]) 171 | stride = int(param_str_split[2]) 172 | bottleneck_channels = int(param_str_split[3]) 173 | sub_layers = int(param_str_split[4]) 174 | return cls(in_channels=in_channels, out_channels=out_channels, stride=stride, 175 | bottleneck_channels=bottleneck_channels, sub_layers=sub_layers, 176 | block_name=tmp_block_name, **kwargs), struct_str[idx + 1:] 177 | 178 | 179 | class SuperResK1K3K1(SuperResK1KXK1): 180 | """ kernel size 3x3""" 181 | 182 | def __init__(self, in_channels=None, out_channels=None, stride=None, bottleneck_channels=None, 183 | sub_layers=None, no_create=False, **kwargs): 184 | super().__init__(in_channels=in_channels, out_channels=out_channels, stride=stride, 185 | bottleneck_channels=bottleneck_channels, sub_layers=sub_layers, 186 | kernel_size=3, 187 | no_create=no_create, **kwargs) 188 | 189 | 190 | class SuperResK1K5K1(SuperResK1KXK1): 191 | """ kernel size 5x5""" 192 | 193 | def __init__(self, in_channels=None, out_channels=None, stride=None, bottleneck_channels=None, 194 | sub_layers=None, no_create=False, **kwargs): 195 | super().__init__(in_channels=in_channels, out_channels=out_channels, stride=stride, 196 | bottleneck_channels=bottleneck_channels, sub_layers=sub_layers, 197 | kernel_size=5, 198 | no_create=no_create, **kwargs) 199 | 200 | 201 | class SuperResK1K7K1(SuperResK1KXK1): 202 | """ kernel size 7x7""" 203 | 204 | def __init__(self, in_channels=None, out_channels=None, stride=None, bottleneck_channels=None, 205 | sub_layers=None, no_create=False, **kwargs): 206 | super().__init__(in_channels=in_channels, out_channels=out_channels, stride=stride, 207 | bottleneck_channels=bottleneck_channels, sub_layers=sub_layers, 208 | kernel_size=7, 209 | no_create=no_create, **kwargs) 210 | 211 | 212 | def register_netblocks_dict(netblocks_dict: dict): 213 | """add different kernel size block to block dict""" 214 | 215 | this_py_file_netblocks_dict = { 216 | 'SuperResK1K3K1': SuperResK1K3K1, 217 | 'SuperResK1K5K1': SuperResK1K5K1, 218 | 'SuperResK1K7K1': SuperResK1K7K1, 219 | } 220 | netblocks_dict.update(this_py_file_netblocks_dict) 221 | return netblocks_dict 222 | --------------------------------------------------------------------------------